Crie um agente AI e o implante em um modelo específico

info

Para novos casos de uso, Databricks recomenda a implantação de agentes no Databricks Apps para controle total sobre o código do agente, a configuração do servidor e o fluxo de trabalho de implantação. Veja Criar um agente AI e implantá-lo em Databricks Apps.

Esta página mostra como criar um agente AI em Python usando o Agent Framework e bibliotecas populares de criação de agentes, como LangGraph e OpenAI.

Requisitos

dica

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

Para criar e implantar agentes usando a abordagem desta página, instale o seguinte:

• databricks-agents 1.2.0 ou superior
• mlflow 3.1.3 ou acima
• Python 3.10 ou superior.
  • Para atender a esse requisito, utilize compute serverless ou Databricks Runtime 13.3 LTS ou superior.

%pip install -U -qqqq databricks-agents mlflow

Databricks também recomenda a instalação do pacote de integração Databricks AI Bridge para criar agentes. Esses pacotes de integração fornecem uma camada compartilhada de APIs que interagem com recursos AI Databricks , como Databricks AI/BI Genie e o Vector Search, em diferentes frameworks de criação de agentes e SDKs.

OpenAI

%pip install -U -qqqq databricks-openai

LangChain/LangGraph

%pip install -U -qqqq databricks-langchain

DSPy

%pip install -U -qqqq databricks-dspy

Pure Python agents

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

Use ResponsesAgent para criar agentes

A Databricks recomenda a interface MLflow ResponsesAgent para criar agentes de nível de produção. ResponsesAgent permite que você crie agentes com qualquer estrutura de terceiros e, em seguida, os integre ao recurso AI Databricks para recursos robustos de registro, rastreamento, avaliação, implantação e monitoramento.

O esquema ResponsesAgent é compatível com o esquema Responses da OpenAI. Para saber mais sobre OpenAI Responses, consulte OpenAI: Respostas vs. ChatCompletion.

nota

A interface antiga ChatAgent ainda é suportada no Databricks. No entanto, para novos agentes, a Databricks recomenda o uso da versão mais recente do MLflow e da interface ResponsesAgent .

Consulte o esquema de agente de entrada e saída legado (servindo modelo).

ResponsesAgent Oferece os seguintes benefícios:

• Capacidades avançadas do agente

  • Suporte multiagente
  • saída de transmissão : transmissão da saída em pedaços menores.
  • Histórico completo de mensagens de chamada de ferramentas : retorna várias mensagens, incluindo mensagens intermediárias de chamada de ferramentas, para melhor qualidade e gerenciamento da conversa.
  • Suporte para confirmação de chamada de ferramenta
  • Suporte de ferramentas de longa duração

• Desenvolvimento, implantação e monitoramento simplificados.

  • Crie agentes usando qualquer framework : Encapsule qualquer agente existente usando a interface ResponsesAgent para obter compatibilidade imediata com AI Playground, Avaliação de Agentes e Monitoramento de Agentes.
  • Interfaces de autoria tipadas : Escreva código de agente usando classes Python tipadas, aproveitando os recursos de autocompletar do IDE e do Notebook.
  • Inferência automática de assinatura : o MLflow infere automaticamente assinaturas ResponsesAgent ao registrar um agente, simplificando o registro e a implantação. Consulte Inferir assinatura do modelo durante o registro.
  • Rastreamento automático : MLflow rastreia automaticamente suas funções predict e predict_stream , agregando respostas de transmissão para facilitar a avaliação e a exibição.
  • Tabelas de inferência aprimoradasAI Gateway : as tabelas de inferência do AI Gateway são ativadas automaticamente para agentes implantados, fornecendo acesso a metadados detalhados log solicitações.

Para aprender como criar um ResponsesAgent, veja os exemplos na seção seguinte e a documentaçãoMLflow - ResponsesAgent para servir modelo.

ResponsesAgent exemplos

O Caderno a seguir mostra como criar transmissão e não transmissão ResponsesAgent usando biblioteca popular. Para aprender como expandir as capacidades desses agentes, consulte Ferramentas para agentesAI.

OpenAI

Agente de bate-papo simples da OpenAI usando modelos hospedados no Databricks

Open notebook in new tab

Agente de chamada de ferramenta OpenAI MCP

Open notebook in new tab

Agente de chamada de ferramentas da OpenAI usando modelos hospedados no Databricks

Open notebook in new tab

Agente de chamada de ferramentas da OpenAI usando modelos hospedados pela OpenAI

Open notebook in new tab

LangGraph

Agente de chamada de ferramentas LangGraph MCP

Open notebook in new tab

DSPy

Agente de chamada de ferramenta de turno único DSPy

Open notebook in new tab

Exemplo multiagente

Para aprender como criar um sistema multiagente, consulte Usar Genie em sistemas multiagentes (servindo modelo).

E se eu já tiver um agente?

Se você já possui um agente desenvolvido com LangChain, LangGraph ou um framework similar, não precisa reescrevê-lo para usá-lo no Databricks. Em vez disso, basta encapsular seu agente existente com a interface MLflow ResponsesAgent :

1. Escreva uma classe wrapper em Python que herde de mlflow.pyfunc.ResponsesAgent.

   Dentro da classe wrapper, faça referência ao agente existente como um atributo self.agent = your_existing_agent.

2. A classe ResponsesAgent requer a implementação de um método predict que retorna um ResponsesAgentResponse para lidar com solicitações de não transmissão. Segue abaixo um exemplo do esquema ResponsesAgentResponses :

   Python

   import uuid
   # input as a dict
   {"input": [{"role": "user", "content": "What did the data scientist say when their Spark job finally completed?"}]}
   
   # output example
   ResponsesAgentResponse(
       output=[
           {
               "type": "message",
               "id": str(uuid.uuid4()),
               "content": [{"type": "output_text", "text": "Well, that really sparked joy!"}],
               "role": "assistant",
           },
       ]
   )

3. Na função predict , converta as mensagens recebidas de ResponsesAgentRequest para o formato que o agente espera. Após o agente gerar uma resposta, converta sua saída em um objeto ResponsesAgentResponse .

Consulte os exemplos de código a seguir para ver como converter agentes existentes para ResponsesAgent:

Basic conversion

Para agentes não transmissivos, converta entradas e saídas na função predict .

Python

from uuid import uuid4

from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
)


class MyWrappedAgent(ResponsesAgent):
    def __init__(self, agent):
        # Reference your existing agent
        self.agent = agent

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        # Convert incoming messages to your agent's format
        # prep_msgs_for_llm is a function you write to convert the incoming messages
        messages = self.prep_msgs_for_llm([i.model_dump() for i in request.input])

        # Call your existing agent (non-streaming)
        agent_response = self.agent.invoke(messages)

        # Convert your agent's output to ResponsesAgent format, assuming agent_response is a str
        output_item = (self.create_text_output_item(text=agent_response, id=str(uuid4())),)

        # Return the response
        return ResponsesAgentResponse(output=[output_item])

Streaming with code re-use

Para agentes de segurança, você pode ser inteligente e reutilizar a lógica para evitar duplicar o código que converte mensagens:

Python

from typing import Generator
from uuid import uuid4

from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
    ResponsesAgentStreamEvent,
)


class MyWrappedStreamingAgent(ResponsesAgent):
    def __init__(self, agent):
        # Reference your existing agent
        self.agent = agent

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        """Non-streaming predict: collects all streaming chunks into a single response."""
        # Reuse the streaming logic and collect all output items
        output_items = []
        for stream_event in self.predict_stream(request):
            if stream_event.type == "response.output_item.done":
                output_items.append(stream_event.item)

        # Return all collected items as a single response
        return ResponsesAgentResponse(output=output_items)

    def predict_stream(
        self, request: ResponsesAgentRequest
    ) -> Generator[ResponsesAgentStreamEvent, None, None]:
        """Streaming predict: the core logic that both methods use."""
        # Convert incoming messages to your agent's format
        # prep_msgs_for_llm is a function you write to convert the incoming messages, included in full examples linked below
        messages = self.prep_msgs_for_llm([i.model_dump() for i in request.input])

        # Stream from your existing agent
        item_id = str(uuid4())
        aggregated_stream = ""
        for chunk in self.agent.stream(messages):
            # Convert each chunk to ResponsesAgent format
            yield self.create_text_delta(delta=chunk, item_id=item_id)
            aggregated_stream += chunk

        # Emit an aggregated output_item for all the text deltas with id=item_id
        yield ResponsesAgentStreamEvent(
            type="response.output_item.done",
            item=self.create_text_output_item(text=aggregated_stream, id=item_id),
        )

Migrate from ChatCompletions

Se o seu agente atual usa a API OpenAI ChatCompletions, você pode migrá-lo para ResponsesAgent sem reescrever sua lógica principal. Adicione um elemento encapsulador que:

1. Converte mensagens ResponsesAgentRequest recebidas para o formato ChatCompletions que seu agente espera.
2. Traduz as saídas ChatCompletions para o esquema ResponsesAgentResponse .
3. Opcionalmente, suporta transmissões mapeando deltas incrementais de ChatCompletions em objetos ResponsesAgentStreamEvent .

Python

from typing import Generator
from uuid import uuid4

from databricks.sdk import WorkspaceClient
from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
    ResponsesAgentStreamEvent,
)


# Legacy agent that outputs ChatCompletions objects
class LegacyAgent:
    def __init__(self):
        self.w = WorkspaceClient()
        self.OpenAI = self.w.serving_endpoints.get_open_ai_client()

    def stream(self, messages):
        for chunk in self.OpenAI.chat.completions.create(
            model="databricks-claude-sonnet-4-5",
            messages=messages,
            stream=True,
        ):
            yield chunk.to_dict()


# Wrapper that converts the legacy agent to a ResponsesAgent
class MyWrappedStreamingAgent(ResponsesAgent):
    def __init__(self, agent):
        # `agent` is your existing ChatCompletions agent
        self.agent = agent

    def prep_msgs_for_llm(self, messages):
        # dummy example of prep_msgs_for_llm
        # real example of prep_msgs_for_llm included in full examples linked below
        return [{"role": "user", "content": "Hello, how are you?"}]

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        """Non-streaming predict: collects all streaming chunks into a single response."""
        # Reuse the streaming logic and collect all output items
        output_items = []
        for stream_event in self.predict_stream(request):
            if stream_event.type == "response.output_item.done":
                output_items.append(stream_event.item)

        # Return all collected items as a single response
        return ResponsesAgentResponse(output=output_items)

    def predict_stream(
        self, request: ResponsesAgentRequest
    ) -> Generator[ResponsesAgentStreamEvent, None, None]:
        """Streaming predict: the core logic that both methods use."""
        # Convert incoming messages to your agent's format
        messages = self.prep_msgs_for_llm([i.model_dump() for i in request.input])

        # process the ChatCompletion output stream
        agent_content = ""
        tool_calls = []
        msg_id = None
        for chunk in self.agent.stream(messages):  # call the underlying agent's stream method
            delta = chunk["choices"][0]["delta"]
            msg_id = chunk.get("id", None)
            content = delta.get("content", None)
            if tc := delta.get("tool_calls"):
                if not tool_calls:  # only accommodate for single tool call right now
                    tool_calls = tc
                else:
                    tool_calls[0]["function"]["arguments"] += tc[0]["function"]["arguments"]
            elif content is not None:
                agent_content += content
                yield ResponsesAgentStreamEvent(**self.create_text_delta(content, item_id=msg_id))

        # aggregate the streamed text content
        yield ResponsesAgentStreamEvent(
            type="response.output_item.done",
            item=self.create_text_output_item(agent_content, msg_id),
        )

        for tool_call in tool_calls:
            yield ResponsesAgentStreamEvent(
                type="response.output_item.done",
                item=self.create_function_call_item(
                    str(uuid4()),
                    tool_call["id"],
                    tool_call["function"]["name"],
                    tool_call["function"]["arguments"],
                ),
            )


agent = MyWrappedStreamingAgent(LegacyAgent())

for chunk in agent.predict_stream(
    ResponsesAgentRequest(input=[{"role": "user", "content": "Hello, how are you?"}])
):
    print(chunk)

Para exemplos completos, veja ResponsesAgent exemplos.

respostas

A transmissão permite que os agentes enviem respostas em blocos reais em vez de esperar pela resposta completa. Para implementar transmissões com ResponsesAgent, emita uma série de eventos delta seguidos por um evento de conclusão final:

1. Emitir eventos delta : Enviar múltiplos eventos output_text.delta com o mesmo item_id para transmitir blocos de texto em tempo real.
2. Finalizar com evento concluído : Envie um evento final response.output_item.done com o mesmo item_id que os eventos delta contendo o texto de saída final completo.

Cada evento delta transmite um trecho de texto para o cliente. O evento final "done" contém o texto completo da resposta e sinaliza ao Databricks para fazer o seguinte:

• Rastreie a saída do seu agente com o rastreamento do MLflow.
• Agregar respostas de transmissão em tabelas de inferência AI Gateway
• Exiba a saída completa na interface do usuário do AI Playground.

propagação de erro de transmissão

Mosaic AI propaga quaisquer erros encontrados durante a transmissão com os últimos tokens sob databricks_output.error. Cabe ao cliente que fez a chamada tratar e apresentar esse erro adequadamente.

Bash

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

Recurso avançado

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 origem de recuperação que não devem ser incluídos no histórico do chat para interações futuras.

Para esses cenários, MLflow ResponsesAgent suporta nativamente os campos custom_inputs e custom_outputs. Você pode acessar as entradas personalizadas por meio de request.custom_inputs em todos os exemplos vinculados acima em Exemplos do ResponsesAgent.

atenção

O aplicativo de avaliação de agentes não oferece suporte à renderização de rastreamentos para agentes com campos de entrada adicionais.

Consulte o seguinte Notebook para aprender como configurar entradas e saídas personalizadas.

Forneça custom_inputs no AI Playground e revise o aplicativo.

Se o seu agente aceitar entradas adicionais usando o campo custom_inputs , você poderá fornecer essas entradas manualmente tanto no AI Playground quanto no aplicativo de revisão.

1. Tanto no AI Playground quanto no aplicativo Agent Review, selecione o ícone de engrenagem..

2. Ativar entradas personalizadas .

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

   

Especifique esquemas de recuperação personalizados

Agentes AI geralmente usam mecanismos de recuperação para encontrar e consultar dados não estruturados em índices de busca vetorial. Para exemplos de ferramentas de recuperação, consulte Conectar agentes a dados não estruturados.

Rastreie esses recuperadores dentro do seu agente com spans RETRIEVERMLflow para habilitar o recurso de produto Databricks , incluindo:

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

nota

Databricks recomenda o uso de ferramentas de recuperação fornecidas pelo pacote Databricks AI Bridge, como databricks_langchain.VectorSearchRetrieverTool e databricks_openai.VectorSearchRetrieverTool porque elas já estão em conformidade com o esquema de recuperação MLflow . Consulte Desenvolver localmente ferramentas de recuperação de pesquisa vetorial com o AI Bridge.

Se o seu agente incluir spans de recuperação com um esquema personalizado, chame mlflow.models.set_retriever_schema ao definir seu agente no código. Isso mapeia as colunas de saída do seu recuperador 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 recuperador. doc_uri é o identificador principal para documentos retornados pelo recuperador, permitindo compará-los com conjuntos de avaliação de verdade fundamental. Consulte Conjuntos de avaliação (MLflow 2).

Considerações sobre a implantação

Prepare-se para o modelo Databricks instalado

Databricks implanta ResponsesAgents em um ambiente distribuído no Databricks servindo modelo. Isso significa que, durante uma conversa com várias interações, a mesma réplica servidora pode não processar todas as solicitações. Preste atenção às seguintes implicações para o gerenciamento do estado do agente:

• Evite o cache local : Ao implantar um ResponsesAgent, não assuma que a mesma réplica lida com todas as solicitações em uma conversa de várias etapas. Reconstrua o estado interno usando um esquema de dicionário ResponsesAgentRequest para cada turno.

• Estado thread-safe : Projete o estado do agente para ser thread-safe, evitando conflitos em ambientes multithread.

• Inicializar estado na função predict : Inicializar estado cada vez que a função predict for chamada, não durante a inicialização de ResponsesAgent . Armazenar o estado no nível ResponsesAgent pode vazar informações entre conversas e causar conflitos porque uma única réplica ResponsesAgent pode lidar com solicitações de várias conversas.

Parametrizar o código para implantação em diferentes ambientes.

Parametrizar o código do agente para reutilizar o mesmo código 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-valor que permite um gerenciamento de configuração flexível. Por exemplo, você pode usar um dicionário durante o desenvolvimento e depois convertê-lo em um arquivo .yaml para implantação de produção e CI/CD.

Um exemplo ModelConfig é mostrado abaixo:

YAML

llm_parameters:
  max_tokens: 500
  temperature: 0.01
model_serving_endpoint: databricks-meta-llama-3-3-70b-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 seu agente, você pode referenciar uma configuração default (de 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-meta-llama-3-3-70b-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 seu agente, especifique o parâmetro model_config para log_model para Especifique um conjunto personalizado de parâmetros a serem usados ao carregar o agente de logs. Ver Documentação do MLflow - ModelConfig.

Use código síncrono ou padrões de retorno de chamada.

Para garantir estabilidade e compatibilidade, utilize código síncrono ou padrões baseados em callbacks na implementação do seu agente.

Databricks gerencia automaticamente a comunicação assíncrona para fornecer concorrência e desempenho ideais quando você implanta um agente. A introdução de loops de eventos personalizados ou frameworks assíncronos pode levar a erros como RuntimeError: This event loop is already running and caused unpredictable behavior.

A Databricks recomenda evitar a programação assíncrona, como o uso de asyncio ou a criação de loops de eventos personalizados, ao desenvolver agentes.

Próximos passos

• Crie suas próprias ferramentas de agente.
• Registrar um agente AI.
• Adicionar rastros a um agente AI.
• implantou um agente AI.