Autor AI agentes em código
Este artigo mostra como criar um agente AI em código, usando MLflow ChatModel
. Databricks aproveita MLflow ChatModel
para garantir a compatibilidade com Databricks AI recurso de agente como avaliação, rastreamento e implantação.
O que é ChatModel
?
ChatModel
é uma classe MLflow projetada para simplificar a criação de agentes de conversação AI. Ele fornece uma interface padronizada para a criação de modelos compatíveis com a API ChatCompletion da OpenAI.
ChatModel
estende o esquema ChatCompletion do OpenAI. Essa abordagem permite que você mantenha ampla compatibilidade com plataformas compatíveis com o padrão ChatCompletion, além de adicionar sua própria funcionalidade personalizada.
Ao usar ChatModel
, os desenvolvedores podem criar agentes compatíveis com as ferramentas Databricks e MLflow para acompanhamento, avaliação e gerenciamento do ciclo de vida dos agentes, que são essenciais para a implantação de modelos prontos para produção.
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 atender aos seguintes requisitos:
Instale
databricks-agents
versão 0.15.0 e acimaInstalar
mlflow
versão 2.20.0 e acima
%pip install -U -qqqq databricks-agents>=0.15.0 mlflow>=2.20.0
Crie um agente do ChatModel
Você pode criar seu agente como uma subclasse de mlflow.pyfunc.ChatModel. Esse método oferece os seguintes benefícios:
Permite que o senhor escreva código de agente compatível com o esquema ChatCompletion usando classes Python tipadas.
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 do modelo durante o registro.
O código a seguir é melhor executado em um notebook Databricks. O Notebook oferece um ambiente conveniente para o desenvolvimento, teste e iteração de seu agente.
A classe MyAgent
estende mlflow.pyfunc.ChatModel
, implementando o método predict
necessário. 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 as saídas de transmissão.
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
seja definida em um Notebook, o senhor deve criar um Notebook de driver separado. O driver Notebook logs o agente para Model Registry e implantado o agente usando o modelo servindo.
Essa separação segue o fluxo de trabalho recomendado pela Databricks para modelos de registro usando a metodologia Models from Code do MLflow.
Exemplo: Envolver LangChain em ChatModel
Se o senhor tiver um modelo LangChain existente e quiser integrá-lo a outro recurso de agente Mosaic AI, poderá envolvê-lo em um MLflow ChatModel
para garantir a compatibilidade.
Este exemplo de código executa as seguintes etapas para envolver um runnable LangChain como ChatModel
:
Envolva a saída final da LangChain com
mlflow.langchain.output_parsers.ChatCompletionOutputParser
para produzir uma assinatura de saída de conclusão de bate-papoA classe
LangchainAgent
estendemlflow.pyfunc.ChatModel
e implementa dois métodos key:predict
: processa previsões síncronas invocando a cadeia e retornando uma resposta formatada.predict_stream
: Trata as previsões de transmissão invocando a cadeia e produzindo pedaços 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)
Use parâmetros para configurar o agente
Na Agent Framework, você pode utilizar parâmetros para controlar como os agentes são executados. Isso possibilita que você faça uma iteração rápida variando as características do agente sem alterar o código. 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
, um conjunto de parâmetros key-value. ModelConfig
é um dicionário Python ou um arquivo .yaml
. 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.
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 do recuperador
AI Os agentes costumam usar recuperadores, um tipo de ferramenta de agente que localiza e retorna documentos relevantes usando um índice de pesquisa vetorial. Para obter mais informações sobre recuperadores, consulte Recuperação não estruturada AI agent tools.
Para garantir que os recuperadores sejam rastreados corretamente, chame mlflow.models.set_retriever_schema quando você define seu agente em 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_column
e 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 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.
Você também pode especificar colunas adicionais no esquema do seu retriever fornecendo uma lista de nomes de colunas com o campo other_columns
.
Se você tiver vários recuperadores, poderá definir vários esquemas usando nomes exclusivos para cada esquema de recuperação.
Entradas e saídas personalizadas
Alguns cenários podem exigir entradas adicionais dos agentes, 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 ChatModel
suporta nativamente o aumento das solicitações e respostas de conclusão de bate-papo da OpenAI com os campos ChatParams custom_input
e custom_output
.
Veja os exemplos a seguir para saber como criar entradas e saídas personalizadas para agentes PyFunc e LangGraph.
Aviso
Atualmente, o aplicativo Agent Evaluation Review não oferece suporte à renderização de rastreamentos para agentes com campos de entrada adicionais.
Esquemas personalizados do PyFunc
O Notebook a seguir mostra um exemplo de esquema personalizado usando PyFunc.
Esquemas personalizados do LangGraph
O Notebook a seguir mostra um exemplo de esquema personalizado usando o LangGraph. O senhor pode modificar a função wrap_output
no Notebook para analisar e extrair informações da transmissão da mensagem.
Fornecer 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
, o senhor poderá fornecer manualmente essas entradas no AI Playground e no aplicativo de avaliaçã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 seu agente.
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.
{
"delta": …,
"databricks_output": {
"trace": {...},
"error": {
"error_code": BAD_REQUEST,
"message": "TimeoutException: Tool XYZ failed to execute"
}
}
}
Notebooks de Exemplo
Esses Notebooks criam uma cadeia simples de "Hello, world" para ilustrar a criação de um agente em Databricks. O primeiro exemplo cria uma cadeia simples, e o segundo exemplo do Notebook ilustra como usar parâmetros para minimizar as alterações de código durante o desenvolvimento.