Autor AI agentes em código

Este artigo mostra como criar um agente AI em Python usando o Mosaic AI Agent Framework e uma biblioteca popular de autoria de agentes, como LangGraph, PyFunc e OpenAI.

Requisitos

A Databricks recomenda instalar a versão mais recente do cliente MLflow Python ao desenvolver agentes.

Para criar e implantar agentes usando a abordagem deste artigo, o senhor deve ter as seguintes versões mínimas de pacote:

• databricks-agents versão 0.16.0 e acima
• mlflow versão 2.20.2 e acima
• Python 3,10 ou acima. O senhor pode usar serverless compute ou Databricks Runtime 13.3 LTS e acima para atender a esse requisito.

%pip install -U -qqqq databricks-agents>=0.16.0 mlflow>=2.20.2

Databricks também recomenda a instalação do pacote de integração Databricks AI Bridge ao criar agentes. Esses pacotes de integração (como databricks-langchain, databricks-openai) fornecem uma camada compartilhada de APIs para interagir com Databricks AI recurso, como Databricks AI/BI Genie e Vector Search, em estruturas de criação de agentes e SDKs.

LangChain/LangGraph

%pip install -U -qqqq databricks-langchain

OpenAI

%pip install -U -qqqq databricks-openai

Pure Python agents

%pip install -U -qqqq databricks-ai-bridge

Use ChatAgent para criar agentes

A Databricks recomenda usar a interface ChatAgent do MLflow para criar agentes de nível de produção. Essa especificação de esquema de bate-papo foi projetada para cenários de agentes e é semelhante, mas não estritamente compatível com, o esquema OpenAI ChatCompletion. ChatAgent também adiciona funcionalidade para agentes de vários turnos que usam ferramentas.

Criar seu agente usando ChatAgent oferece os seguintes benefícios:

• Capacidades avançadas de agente

  • transmissão de saída : Possibilite experiências interativas para o usuário por meio da transmissão da saída em partes menores.
  • Histórico abrangente de mensagens de tool-calling : Retorna várias mensagens, inclusive mensagens intermediárias de chamadas de ferramentas, para melhorar a qualidade e o gerenciamento de conversas.
  • Suporte para confirmação de chamadas de ferramentas
  • Suporte a sistemas multiagentes

• Desenvolvimento, implantação e monitoramento simplificados

  • Integração de recurso Databricks independente de estrutura : Escreva seu agente em qualquer estrutura de sua escolha e obtenha compatibilidade imediata com AI Playground, Agent Evaluation e Agent monitoring.
  • Interfaces de criação digitadas : Escreva o código do agente usando classes Python digitadas, beneficiando-se do preenchimento automático do IDE e do Notebook.
  • Inferência automática de assinaturas : O MLflow infere automaticamente assinaturas ChatAgent ao registrar o agente, simplificando o registro e a implantação. Consulte Inferir assinatura do modelo durante o registro.
  • AI Tabelas de inferência aprimoradas pelo gateway : AI As tabelas de inferência do gateway são ativadas automaticamente para agentes implantados, fornecendo acesso a solicitações detalhadas log metadados.

Para saber como criar um ChatAgent, consulte os exemplos na seção a seguir e a documentação do MLflow - O que é a interface do ChatAgent.

Exemplos deChatAgent

O Notebook a seguir mostra ao senhor como criar transmissão e não-transmissão ChatAgents usando a popular biblioteca OpenAI e LangGraph.

Agente de chamada de ferramentas LangGraph

Open notebook in new tab

Agente de chamada de ferramentas OpenAI

Open notebook in new tab

Ferramenta da API de respostas da OpenAI - agente de chamadas

Open notebook in new tab

Agente somente de bate-papo OpenAI

Open notebook in new tab

Agente somente de bate-papo DSPy

Open notebook in new tab

Para saber como expandir os recursos desses agentes adicionando ferramentas, consulte AI agent tools.

Exemplo do sistema multiagente ChatAgent

Para saber como criar um sistema multiagentes usando o Genie, consulte Usar o Genie em sistemas multiagentes.

Autor agentes de transmissão de resultados

Os agentes de transmissão fornecem respostas em uma transmissão contínua de partes menores e incrementais. A saída de transmissão permite que os usuários finais leiam a saída do agente à medida que ela é gerada, reduzindo a latência percebida e melhorando a experiência geral do usuário para agentes de conversação.

Para criar uma transmissão ChatAgent, defina um método predict_stream que retorne um gerador que produza objetos ChatAgentChunk, cada um contendo uma parte da resposta. Leia mais sobre o comportamento ideal de transmissão ChatAgent nos documentos do siteMLflow.

O código a seguir mostra um exemplo de função predict_stream. Para obter exemplos completos de agentes de transmissão, consulte Exemplos de ChatAgent:

Python

def predict_stream(
  self,
  messages: list[ChatAgentMessage],
  context: Optional[ChatContext] = None,
  custom_inputs: Optional[dict[str, Any]] = None,
) -> Generator[ChatAgentChunk, None, None]:
  # Convert messages to a format suitable for your agent
  request = {"messages": self._convert_messages_to_dict(messages)}

  # Stream the response from your agent
  for event in self.agent.stream(request, stream_mode="updates"):
    for node_data in event.values():
      # Yield each chunk of the response
      yield from (
        ChatAgentChunk(**{"delta": msg}) for msg in node_data["messages"]
      )

Author deployment-ready ChatAgents for Databricks servindo modelo

Databricks implantado ChatAgents em um ambiente distribuído em Databricks servindo modelo, o que significa que, durante uma conversa com vários turnos, a mesma réplica de serviço pode não atender a todas as solicitações. Preste atenção às seguintes implicações para gerenciar o estado do agente:

• Evite o armazenamento em cache local : ao implantar um ChatAgent, não presuma que a mesma réplica atenderá a todas as solicitações em uma conversa com vários turnos. Reconstrua o estado interno usando um esquema ChatAgentRequest do dicionário para cada turno.

• Estado de thread-safe: projete o estado do agente para que seja seguro para thread-safe, evitando conflitos em ambientes com vários segmentos.

• Inicializar estado na função predict : Inicialize o estado sempre que a função predict for chamada, não durante a inicialização ChatAgent. O armazenamento de estado no nível ChatAgent poderia vazar informações entre as conversas e causar conflitos, pois uma única réplica ChatAgent poderia lidar com solicitações de várias conversas.

Entradas e saídas personalizadas

Alguns cenários podem exigir entradas adicionais do agente, como client_type e session_id, ou saídas como links de fontes de recuperação que não devem ser incluídos no histórico do chat para interações futuras.

Para esses cenários, o MLflow ChatAgent suporta nativamente os campos custom_inputs e custom_outputs.

atenção

Atualmente, o aplicativo Agent Evaluation Review não oferece suporte à renderização de rastreamentos para agentes com campos de entrada adicionais.

Veja os exemplos a seguir para saber como definir entradas e saídas personalizadas para agentes OpenAI/PyFunc e LangGraph.

Agente de esquema personalizado OpenAI + PyFunc Notebook

Open notebook in new tab

Agente de esquema personalizado LangGraph Notebook

Open notebook in new tab

Fornecer custom_inputs no AI Playground e no aplicativo de avaliação de agentes

Se o seu agente aceitar entradas adicionais usando o campo custom_inputs, o senhor poderá fornecer manualmente essas entradas no AI Playground e no aplicativo de avaliação do agente.

1. No AI Playground ou no aplicativo Agent Review, selecione o ícone de engrenagem .

2. Habilite custom_inputs .

3. Forneça um objeto JSON que corresponda ao esquema de entrada definido pelo seu agente.

   

Especifique esquemas de recuperação personalizados

AI Os agentes geralmente usam recuperadores para localizar e consultar dados não estruturados de índices de pesquisa de vetores. Para obter exemplos de ferramentas de recuperação, consulte Recuperação não estruturada AI ferramentas de agente.

Rastreie esses recuperadores dentro de seu agente com MLflow RETRIEVER spans para permitir Databricks produto recurso, incluindo:

• Exibição automática de links para documentos de origem recuperados na interface do usuário do AI Playground
• Executando automaticamente juízes de fundamentação e relevância de recuperação na Avaliação de Agentes

nota

Databricks recomenda o uso de ferramentas retriever fornecidas por Databricks AI Bridge pacote como databricks_langchain.VectorSearchRetrieverTool e databricks_openai.VectorSearchRetrieverTool porque elas já estão em conformidade com o esquema retriever MLflow. Consulte Desenvolver localmente ferramentas de recuperação do Vector Search com o AI Bridge.

Se seu agente incluir extensões de retriever com um esquema personalizado, chame mlflow.models.set_retriever_schema quando você define seu agente em código. Isso mapeia as colunas de saída do seu retriever para os campos esperados do MLflow (primary_key, text_column, doc_uri).

Python

import mlflow
# Define the retriever's schema by providing your column names
# For example, the following call specifies the schema of a retriever that returns a list of objects like
# [
#     {
#         'document_id': '9a8292da3a9d4005a988bf0bfdd0024c',
#         'chunk_text': 'MLflow is an open-source platform, purpose-built to assist machine learning practitioners...',
#         'doc_uri': 'https://mlflow.org/docs/latest/index.html',
#         'title': 'MLflow: A Tool for Managing the Machine Learning Lifecycle'
#     },
#     {
#         'document_id': '7537fe93c97f4fdb9867412e9c1f9e5b',
#         'chunk_text': 'A great way to get started with MLflow is to use the autologging feature. Autologging automatically logs your model...',
#         'doc_uri': 'https://mlflow.org/docs/latest/getting-started/',
#         'title': 'Getting Started with MLflow'
#     },
# ...
# ]
mlflow.models.set_retriever_schema(
    # Specify the name of your retriever span
    name="mlflow_docs_vector_search",
    # Specify the output column name to treat as the primary key (ID) of each retrieved document
    primary_key="document_id",
    # Specify the output column name to treat as the text content (page content) of each retrieved document
    text_column="chunk_text",
    # Specify the output column name to treat as the document URI of each retrieved document
    doc_uri="doc_uri",
    # Specify any other columns returned by the retriever
    other_columns=["title"],
)

nota

A coluna doc_uri é especialmente importante ao avaliar o desempenho do retriever. doc_uri é o identificador principal dos documentos retornados pelo recuperador, permitindo compará-los com os conjuntos de avaliação da verdade básica. Consulte Conjuntos de avaliação.

Parametrize o código do agente para implantação em todos os ambientes

Você pode parametrizar o código do agente para reutilizar o mesmo código do agente em diferentes ambientes.

Os parâmetros são pares chave-valor que você define em um dicionário no Python ou em um arquivo .yaml.

Para configurar o código, crie um ModelConfig usando um dicionário Python ou um arquivo .yaml. ModelConfig é um conjunto de parâmetros key-value que permite o gerenciamento flexível da configuração. Por exemplo, o senhor pode usar um dicionário durante o desenvolvimento e, em seguida, convertê-lo em um arquivo .yaml para implantação de produção e CI/CD.

Para obter detalhes sobre ModelConfig, consulte a documentação do MLflow.

Um exemplo ModelConfig é mostrado abaixo:

YAML

llm_parameters:
  max_tokens: 500
  temperature: 0.01
model_serving_endpoint: databricks-dbrx-instruct
vector_search_index: ml.docs.databricks_docs_index
prompt_template: 'You are a hello world bot. Respond with a reply to the user''s
  question that indicates your prompt template came from a YAML file. Your response
  must use the word "YAML" somewhere. User''s question: {question}'
prompt_template_input_vars:
  - question

No código do agente, o senhor pode fazer referência a uma configuração default (desenvolvimento) do arquivo ou dicionário .yaml:

Python

import mlflow
# Example for loading from a .yml file
config_file = "configs/hello_world_config.yml"
model_config = mlflow.models.ModelConfig(development_config=config_file)

# Example of using a dictionary
config_dict = {
    "prompt_template": "You are a hello world bot. Respond with a reply to the user's question that is fun and interesting to the user. User's question: {question}",
    "prompt_template_input_vars": ["question"],
    "model_serving_endpoint": "databricks-dbrx-instruct",
    "llm_parameters": {"temperature": 0.01, "max_tokens": 500},
}

model_config = mlflow.models.ModelConfig(development_config=config_dict)

# Use model_config.get() to retrieve a parameter value
# You can also use model_config.to_dict() to convert the loaded config object
# into a dictionary
value = model_config.get('sample_param')

Em seguida, ao registrar o agente, especifique o parâmetro model_config como log_model para especificar um conjunto personalizado de parâmetros a serem usados ao carregar o agente de registros. Consulte Documentação do MLflow - ModelConfig.

transmissão propagação de erros

Mosaic AI propagação de quaisquer erros encontrados durante a transmissão com os últimos tokens em databricks_output.error. Cabe ao cliente chamador tratar e revelar adequadamente esse erro.

Bash

{
  "delta": …,
  "databricks_output": {
    "trace": {...},
    "error": {
      "error_code": BAD_REQUEST,
      "message": "TimeoutException: Tool XYZ failed to execute."
    }
  }
}

Autenticação em nome do usuário

info

Beta

Esse recurso está na versão beta.

atenção

Embora a autenticação do usuário no local seja uma ferramenta avançada para impor o acesso seguro a dados confidenciais, ela permite que os usuários do workspace criem agentes que agem em nome de outros usuários no Databricks. Dessa forma, ele é desativado pelo site default durante a versão beta e deve ser ativado por um administrador do site workspace. Leia as considerações de segurança para a autenticação no local do usuário antes de ativar esse recurso.

Com a autenticação no local do usuário, os agentes implantados via Mosaic AI servindo modelo podem acessar Databricks recurso usando a identidade do Databricks usuário final que consultou o agente. Isso permite o acesso a informações confidenciais por usuário, com aplicação refinada do controle de acesso a dados em Unity Catalog.

A autenticação no local do usuário restringe ainda mais os tokens de usuário recebidos por meio de downscoping, garantindo que os tokens expostos ao código do agente sejam limitados a acessar apenas o APIs específico definido pelo autor do agente. Isso aumenta a segurança, impedindo ações não autorizadas e reduzindo o risco de uso indevido de tokens.

Ao criar seu agente, o senhor pode continuar a usar os SDKs existentes para acessar Databricks recurso, como índices de pesquisa vetorial. Para executar em nome do usuário o acesso ao recurso, há duas etapas:

1. No código do agente, atualize as chamadas SDK para indicar que o recurso deve ser acessado em nome do usuário final do agente
2. No momento do registro do agente (antes da implementação do agente), especifique os escopos da API REST do usuário final exigidos pelo seu agente. Para obter mais detalhes, consulte Autenticação em nome do usuário

Para ver um exemplo de ponta a ponta de autenticação em nome do usuário, consulte Exemplo de ponta a ponta.

Os recursos a seguir são compatíveis com a autenticação no local do usuário com agentes.

Databricks recurso	Clientes compatíveis
Índice de pesquisa vetorial	databricks_langchain.VectorSearchRetrieverTool, databricks_openai.VectorSearchRetrieverTool ou VectorSearchClient
Endpoint do serviço de modelos	databricks.sdk.WorkspaceClient
SQL Warehouse	databricks.sdk.WorkspaceClient
Conexões UC	databricks.sdk.WorkspaceClient
Tabelas e funções de UC	Atualmente, não oferecemos suporte a clientes diretos para acessar tabelas ou funções de UC com autenticação em nome do usuário. Em vez disso, incentivamos os usuários a usar o Genie para acessar dados estruturados com autenticação no local do usuário
Espaço Genie	databricks_langchain.GenieAgent ou databricks_openai.GenieAgent

Ao inicializar ferramentas com o cliente em nome do usuário, você pode agrupar a inicialização da ferramenta em um bloco try-except ou permitir que os erros sejam expostos aos usuários. Ao lidar com os erros, o agente ainda pode dar a melhor resposta, mesmo que o usuário final não tenha acesso a todas as ferramentas necessárias. No entanto, se o senhor optar por não tratar os erros durante a inicialização da ferramenta, o agente emitirá um erro se o usuário não tiver algum recurso necessário.

Configurar SDKs

Os snippets abaixo demonstram como configurar, em nome do usuário final, o acesso a diferentes Databricks recurso usando vários SDKs

Vector Search Retriever Tool

Python

from databricks.sdk import WorkspaceClient
from databricks.sdk.credentials_provider import ModelServingUserCredentials
from databricks_langchain import VectorSearchRetrieverTool

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy = ModelServingUserCredentials())

vector_search_tools = []
# Exclude exception handling if you want the agent to fail
# when users lack access to all required Databricks resources
try:
    tool = VectorSearchRetrieverTool(
       index_name="<index_name>",
       description="...",
       tool_name="...",
       workspace_client=user_client # Specify the user authenticated client
     )
     vector_search_tools.append(tool)
except Exception as e:
     _logger.debug("Skipping adding tool as user does not have permissions)

Vector Search Client

Python

from databricks.vector_search.client import VectorSearchClient
from databricks.vector_search.utils import CredentialStrategy

# Configure a VectorSearch Client to use on behalf of end
# user authentication
user_authenticated_vsc = VectorSearchClient(credential_strategy=CredentialStrategy.MODEL_SERVING_USER_CREDENTIALS)
# Exclude exception handling if you want the agent to fail when
# users lack access to all required Databricks resources
try:
  vs_index = user_authenticated_vsc.get_index(endpoint_name="endpoint_name", index_name="index_name")
  ...
except Exception as e:
  _logger.debug("Skipping Vector Index because user does not have permissions)

Model Serving Endpoint

Python

from databricks.sdk import WorkspaceClient
from databricks.sdk.credentials_provider import ModelServingUserCredentials

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

# Exclude exception handling if you want the agent to fail
# when users lack access to all required Databricks resources
try:
    user_client.serving_endpoints.query("endpoint_name", input="")
except Exception as e:
    _logger.debug("Skipping Model Serving Endpoint due to no permissions")

UC Connections

Python

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

user_client.serving_endpoints.http_request(
  conn="connection_name",
  method=ExternalFunctionRequestHttpMethod.POST,
  path="/api/v1/resource",
  json={"key": "value"},
  headers={"extra_header_key": "extra_header_value"},
)

Genie Spaces

Python

from databricks.sdk import WorkspaceClient
from databricks.sdk.credentials_provider import ModelServingUserCredentials
from databricks_langchain.genie import GenieAgent

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

genie_agent = GenieAgent(
      genie_space_id="<genie_space_id>",
      genie_agent_name="Genie",
      description="Genie_description",
      client=user_client, # Specify the user client here
  )

Inicializando o agente

a autenticação em nome do usuário é compatível com a interface do ChatAgent. Ao usar a autenticação no local do usuário, a identidade do usuário final só é conhecida quando o agente implantado é consultado, ou seja, nas funções predict e predict_stream da interface do ChatAgent. Como resultado, o senhor deve executar qualquer acesso do usuário ao recurso em nome do usuário (por exemplo liste os índices de pesquisa vetorial aos quais o usuário final tem acesso) a partir desses métodos, em vez do método __init__ da sua implementação do ChatAgent. Isso garante que os recursos sejam isolados entre as invocações

Python

from mlflow.pyfunc import ChatAgent


class LangGraphChatAgent(ChatAgent):
    def initialize_agent():
        user_client = WorkspaceClient(
            credentials_strategy=ModelServingUserCredentials()
        )
        system_authorized_client = WorkspaceClient()
        ### Use the clients above to access resources with either system or user authorization

    def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> ChatAgentResponse:
        agent = initialize_agent() # Initialize the Agent in Predict
        request = {"messages": self._convert_messages_to_dict(messages)}

        messages = []
        for event in self.agent.stream(request, stream_mode="updates"):
            for node_data in event.values():
                messages.extend(
                    ChatAgentMessage(**msg) for msg in node_data.get("messages", [])
                )
        return ChatAgentResponse(messages=messages)

Considerações de segurança

Há algumas implicações de segurança a serem consideradas antes de ativar a autenticação em nome do usuário com agentes:

1. Acesso a recursos confidenciais Databricks: a ativação da autenticação na metade do usuário permite que os agentes acessem recursos confidenciais Databricks. Embora tenhamos implementado os escopos do API para restringir os recursos que os desenvolvedores podem acessar e reduzir o risco de uso indevido de tokens, alguns riscos ainda permanecem. Por exemplo, o escopo da API serving.serving-endpoints concede a um agente permissão para executar um endpoint de serviço em nome do usuário. No entanto, o próprio endpoint de atendimento pode ter acesso a escopos de API adicionais que o agente original não está autorizado a usar.
2. Não há suporte para o consentimento do usuário final: Durante a fase beta atual, os usuários do agente não podem view ou consentir com os escopos Databricks REST API exigidos por um agente. Os usuários são responsáveis por garantir que confiam nas pessoas com permissões "Can gerenciar" no servidor endpoint para realizar ações em Databricks em seu nome.

Exemplo de ponta a ponta

O Notebook a seguir mostra como criar um agente com pesquisa vetorial usando a autenticação no local do usuário.

Em nome da autenticação do usuário com pesquisa vetorial

Open notebook in new tab

Próximas etapas

• Crie suas próprias ferramentas de agente.
• registrar um agente AI.
• Adicionar rastros a um agente AI.
• implantado um agente AI.