Crie agentes de IA através de código
Este artigo mostra como criar um agente de IA em código, usando MLflow ChatModel. O Azure Databricks aproveita o MLflow ChatModel para garantir a compatibilidade com os recursos do agente de IA do Databricks, como avaliação, rastreamento e implantação.
O que é ChatModel?
ChatModel é uma classe MLflow projetada para simplificar a criação de agentes de IA conversacionais. Ele fornece uma interface padronizada para a construção de modelos compatíveis com API ChatCompletion da OpenAI.
ChatModel estende o esquema ChatCompletion da OpenAI. Essa abordagem permite que você mantenha uma ampla compatibilidade com plataformas que suportam o padrão ChatComplete, ao mesmo tempo em que adiciona sua própria funcionalidade personalizada
Usando ChatModel, os desenvolvedores podem criar agentes compatíveis com ferramentas Databricks e MLflow para rastreamento, avaliação e gerenciamento do ciclo de vida dos agentes, que são essenciais para a implantação de modelos prontos para produção.
Consulte MLflow: Introdução ao ChatModel.
Requerimentos
O Databricks recomenda a instalação da versão mais recente do cliente Python MLflow ao desenvolver agentes.
Para criar e implantar agentes usando a abordagem neste artigo, você deve atender aos seguintes requisitos:
- Instale o
databricks-agentsversão 0.15.0 e superior - Instale o
mlflowversão 2.20.0 e superior
%pip install -U -qqqq databricks-agents>=0.15.0 mlflow>=2.20.0
Criar um agente ChatModel
Você pode criar seu agente como uma subclasse de mlflow.pyfunc.ChatModel. Este método oferece os seguintes benefícios:
- Permite que você escreva código de agente compatível com o esquema ChatCompletion usando classes Python digitadas.
- O MLflow inferirá automaticamente uma assinatura compatível com a conclusão do chat ao registrar o agente, mesmo sem um
input_example. Isso simplifica o processo de registro e implantação do agente. Consulte Inferir assinatura de modelo durante o registro em log.
O código a seguir é melhor executado em um bloco de anotações Databricks. Os notebooks fornecem um ambiente conveniente para desenvolver, testar e iterar em seu agente.
A classe MyAgent estende mlflow.pyfunc.ChatModel, implementando o método necessário predict. Isso garante a compatibilidade com o Mosaic AI Agent Framework.
A classe também inclui os métodos opcionais _create_chat_completion_chunk e predict_stream para lidar com saídas de streaming.
from dataclasses import dataclass
from typing import Optional, Dict, List, Generator
from mlflow.pyfunc import ChatModel
from mlflow.types.llm import (
# Non-streaming helper classes
ChatCompletionRequest,
ChatCompletionResponse,
ChatCompletionChunk,
ChatMessage,
ChatChoice,
ChatParams,
# Helper classes for streaming agent output
ChatChoiceDelta,
ChatChunkChoice,
)
class MyAgent(ChatModel):
"""
Defines a custom agent that processes ChatCompletionRequests
and returns ChatCompletionResponses.
"""
def predict(self, context, messages: list[ChatMessage], params: ChatParams) -> ChatCompletionResponse:
last_user_question_text = messages[-1].content
response_message = ChatMessage(
role="assistant",
content=(
f"I will always echo back your last question. Your last question was: {last_user_question_text}. "
)
)
return ChatCompletionResponse(
choices=[ChatChoice(message=response_message)]
)
def _create_chat_completion_chunk(self, content) -> ChatCompletionChunk:
"""Helper for constructing a ChatCompletionChunk instance for wrapping streaming agent output"""
return ChatCompletionChunk(
choices=[ChatChunkChoice(
delta=ChatChoiceDelta(
role="assistant",
content=content
)
)]
)
def predict_stream(
self, context, messages: List[ChatMessage], params: ChatParams
) -> Generator[ChatCompletionChunk, None, None]:
last_user_question_text = messages[-1].content
yield self._create_chat_completion_chunk(f"Echoing back your last question, word by word.")
for word in last_user_question_text.split(" "):
yield self._create_chat_completion_chunk(word)
agent = MyAgent()
model_input = ChatCompletionRequest(
messages=[ChatMessage(role="user", content="What is Databricks?")]
)
response = agent.predict(context=None, model_input=model_input)
print(response)
Embora a classe de agente MyAgent esteja definida em um bloco de anotações, você deve criar um bloco de anotações de driver separado. O notebook do driver regista o agente no Registo de Modelos e implementa o agente utilizando o Model Serving.
Essa separação segue o fluxo de trabalho recomendado pela Databricks para registrar modelos usando a metodologia Models from Code do MLflow.
Exemplo: Wrap LangChain no ChatModel
Se você tiver um modelo LangChain existente e quiser integrá-lo com outros recursos do agente Mosaic AI, poderá envolvê-lo em um ChatModel MLflow para garantir a compatibilidade.
Este exemplo de código executa as seguintes etapas para encapsular um LangChain executável como um ChatModel:
- Envolva o resultado final do LangChain com
mlflow.langchain.output_parsers.ChatCompletionOutputParserpara obter uma assinatura de conclusão de chat. - A classe
LangchainAgentestendemlflow.pyfunc.ChatModele implementa dois métodos principais:-
predict: Manipula previsões síncronas invocando a cadeia e retornando uma resposta formatada. -
predict_stream: Gere previsões de streaming invocando a cadeia e gerando partes de respostas.
-
from mlflow.langchain.output_parsers import ChatCompletionOutputParser
from mlflow.pyfunc import ChatModel
from typing import Optional, Dict, List, Generator
from mlflow.types.llm import (
ChatCompletionResponse,
ChatCompletionChunk
)
chain = (
<your chain here>
| ChatCompletionOutputParser()
)
class LangchainAgent(ChatModel):
def _prepare_messages(self, messages: List[ChatMessage]):
return {"messages": [m.to_dict() for m in messages]}
def predict(
self, context, messages: List[ChatMessage], params: ChatParams
) -> ChatCompletionResponse:
question = self._prepare_messages(messages)
response_message = self.chain.invoke(question)
return ChatCompletionResponse.from_dict(response_message)
def predict_stream(
self, context, messages: List[ChatMessage], params: ChatParams
) -> Generator[ChatCompletionChunk, None, None]:
question = self._prepare_messages(messages)
for chunk in chain.stream(question):
yield ChatCompletionChunk.from_dict(chunk)
Usar parâmetros para configurar o agente
No Agent Framework, você pode usar parâmetros para controlar como os agentes são executados. Isso permite que você itere rapidamente variando as características do seu agente sem alterar o código. Os parâmetros são pares chave-valor que você define em um dicionário Python ou em um arquivo .yaml.
Para configurar o código, crie um ModelConfig, um conjunto de parâmetros chave-valor.
ModelConfig é um dicionário Python ou um arquivo .yaml. Por exemplo, você pode usar um dicionário durante o desenvolvimento e, em seguida, convertê-lo em um arquivo de .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.
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
Para chamar a configuração do seu código, use uma das seguintes opções:
# 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
value = model_config.get('sample_param')
Definir esquema de recuperador
Os agentes de IA geralmente usam retrievers, um tipo de ferramenta de agente que encontra e retorna documentos relevantes usando um índice de Pesquisa Vetorial. Para obter mais informações sobre retrievers, consulte Unstructured retrieval AI agent tools.
Para garantir que os retrievers sejam rastreados corretamente, chame mlflow.models.set_retriever_schema quando definir seu agente no código. Use set_retriever_schema para mapear os nomes das colunas na tabela retornada para os campos esperados do MLflow, como primary_key, text_columne doc_uri.
# Define the retriever's schema by providing your column names
# These strings should be read from a config dictionary
mlflow.models.set_retriever_schema(
name="vector_search",
primary_key="chunk_id",
text_column="text_column",
doc_uri="doc_uri"
# other_columns=["column1", "column2"],
)
Observação
A coluna doc_uri é especialmente importante ao avaliar o desempenho do recuperador.
doc_uri é o identificador principal dos documentos devolvidos pelo sistema de recuperação de informação, permitindo a comparação com conjuntos de avaliação de referência. Consulte Conjuntos de avaliação
Você também pode especificar colunas adicionais no esquema do retriever fornecendo uma lista de nomes de colunas com o campo other_columns.
Se você tiver vários retrievers, poderá definir vários esquemas usando nomes exclusivos para cada esquema retriever.
Entradas e saídas personalizadas
Alguns cenários podem exigir entradas adicionais de agentes, como client_type e session_id, ou saídas, como links para a fonte de recuperação que não devem ser incluídos no histórico de bate-papo para interações futuras.
Para estes cenários, o MLflow ChatModel oferece suporte nativo para aprimorar as solicitações e respostas de conclusão de bate-papo do OpenAI com os campos ChatParamscustom_input e custom_output.
Veja os exemplos a seguir para saber como criar entradas e saídas personalizadas para agentes PyFunc e LangGraph.
Advertência
Atualmente, o aplicativo de revisão de Avaliação de Agentes não oferece suporte à renderização de rastros para agentes com campos de entrada adicionais.
esquemas personalizados do PyFunc
Os blocos de anotações a seguir mostram um exemplo de esquema personalizado usando PyFunc.
Notebook do agente de esquema personalizado PyFunc
Notebook do driver de esquema personalizado PyFunc
esquemas personalizados LangGraph
Os blocos de anotações a seguir mostram um exemplo de esquema personalizado usando LangGraph. Você pode modificar a função wrap_output nos blocos de anotações para analisar e extrair informações do fluxo de mensagens.
Notebook do agente de esquema personalizado LangGraph
Obter de bloco de notas
Notebook do driver de esquema personalizado LangGraph
Forneça custom_inputs no AI Playground e no aplicativo de revisão de agentes
Se o seu agente aceitar entradas adicionais usando o campo custom_inputs, poderá fornecer manualmente essas entradas tanto no AI Playground quanto no aplicativo de revisão do agente .
No AI Playground ou no aplicativo Agent Review, selecione o ícone de engrenagem
.Habilite custom_inputs.
Forneça um objeto JSON que corresponda ao esquema de entrada definido pelo agente.
Propagação de erros de streaming
Mosaic AI propaga quaisquer erros encontrados durante o streaming com o último token sob databricks_output.error. Cabe ao cliente chamador a responsabilidade de lidar corretamente e tornar visível este erro.
{
"delta": …,
"databricks_output": {
"trace": {...},
"error": {
"error_code": BAD_REQUEST,
"message": "TimeoutException: Tool XYZ failed to execute"
}
}
}
<a i
Exemplos de blocos de notas
Esses blocos de anotações criam uma cadeia simples "Olá, mundo" para ilustrar a criação de um agente no Databricks. O primeiro exemplo cria uma cadeia simples e o segundo bloco de anotações de exemplo ilustra como usar parâmetros para minimizar alterações de código durante o desenvolvimento.