Criar um agente AI e suas ferramentas

Prévia

Esse recurso está em Pré-lançamento público.

Este artigo mostra ao senhor como criar AI agentes e ferramentas usando o Mosaic AI Agent Framework.

Saiba como usar o AI Playground para criar rapidamente protótipos de agentes de chamada de ferramentas e exportá-los para o Mosaic AI Agent Framework.

Requisitos

Entenda os conceitos de agentes e ferramentas do AI conforme descrito em O que são o sistema composto AI e os agentes AI?
A Databricks recomenda instalar a versão mais recente do cliente MLflow Python ao desenvolver agentes. Consulte Autenticação para recurso dependente para obter informações sobre os requisitos da versão mlflow.

Criar ferramentas do agente AI

AI Os agentes usam ferramentas para realizar ações além da geração de linguagem, por exemplo, para recuperar dados estruturados ou não estruturados, executar código ou conversar com serviços remotos (por exemplo, o enviar uma mensagem para email ou Slack).

Para fornecer ferramentas a um agente com o Mosaic AI Agent Framework, o senhor pode usar qualquer combinação dos seguintes métodos:

Criar ou usar funções existentes do Unity Catalog como ferramentas. Unity Catalog As funções fornecem fácil descoberta, governança e compartilhamento, e funcionam bem para aplicar transformações e agregações em grandes conjuntos de dados.
Definir ferramentas localmente como funções Python no código do agente. Essa abordagem é útil em situações em que o senhor precisa fazer chamadas para REST APIs, usar código ou biblioteca arbitrários ou executar ferramentas com latência muito baixa. Essa abordagem não tem a capacidade de descoberta e a governança integradas fornecidas pelas funções do site Unity Catalog. A Databricks recomenda que o senhor pondere essa troca ao criar seu agente para determinar qual é a melhor abordagem.

Ambas as abordagens funcionam para agentes escritos em código Python personalizado ou usando uma biblioteca de autoria de agentes como a LangGraph.

Ao definir as ferramentas, certifique-se de que a ferramenta, seus parâmetros e seu valor de retorno estejam documentados, para que o agente LLM possa entender quando e como usar a ferramenta.

Criar ferramentas de agente com funções do Unity Catalog

Esses exemplos criam ferramentas de agente AI usando funçõesUnity Catalog escritas em um ambiente Notebook ou em um editor SQL.

Execute o código a seguir em uma célula do Notebook. Ele usa a magia do %sql Notebook para criar uma função Unity Catalog chamada python_exec.

Um LLM pode usar essa ferramenta para executar o código Python fornecido a ele por um usuário.

%sql
CREATE OR REPLACE FUNCTION
main.default.python_exec (
 code STRING COMMENT 'Python code to execute. Remember to print the final result to stdout.'
)
RETURNS STRING
LANGUAGE PYTHON
DETERMINISTIC
COMMENT 'Executes Python code in the sandboxed environment and returns its stdout. The runtime is stateless and you can not read output of the previous tool executions. i.e. No such variables "rows", "observation" defined. Calling another tool inside a Python code is NOT allowed. Use standard python libraries only.'
AS $$
 import sys
 from io import StringIO
 sys_stdout = sys.stdout
 redirected_output = StringIO()
 sys.stdout = redirected_output
 exec(code)
 sys.stdout = sys_stdout
 return redirected_output.getvalue()
$$

Execute o código a seguir em um editor SQL.

Ele cria uma função do Unity Catalog chamada lookup_customer_info que um LLM poderia usar para recuperar dados estruturados de uma tabela customer_data hipotética:

CREATE OR REPLACE FUNCTION main.default.lookup_customer_info(
  customer_name STRING COMMENT 'Name of the customer whose info to look up'
)
RETURNS STRING
COMMENT 'Returns metadata about a particular customer given the customer name, including the customer email and ID. The
customer ID can be used for other queries.'
RETURN SELECT CONCAT(
    'Customer ID: ', customer_id, ', ',
    'Customer Email: ', customer_email
  )
  FROM main.default.customer_data
  WHERE customer_name = customer_name
  LIMIT 1;

Protótipo de ferramenta - chamando agentes no AI Playground

Depois de criar as funções do Unity Catalog, o senhor pode usar o AI Playground para fornecê-las a um LLM e testar o agente. O site AI Playground fornece um sandbox para protótipos de agentes de chamadas de ferramentas.

Quando estiver satisfeito com o agente AI, o senhor pode exportá-lo para desenvolvê-lo ainda mais em Python ou implantá-lo como um modelo de serviço endpoint como está.

Observação

Unity Catalog, e serverless compute Mosaic AI Agent Framework, e modelos de fundação pay-per-tokens ou modelos externos devem estar disponíveis no site workspace atual para prototipar agentes em AI Playground.

Para criar o protótipo de um endpoint de chamada de ferramenta.

No Playground, selecione um modelo com o rótulo Function calling (Chamada de função ).
Selecione Tools (Ferramentas ) e especifique os nomes das funções do site Unity Catalog em dropdown:
Chat para testar a combinação atual de LLM, ferramentas e prompt do sistema e tentar variações.

Exportação e implantado AI Playground agentes

Depois de adicionar ferramentas e testar o agente, exporte o agente do Playground para Python Notebook:

Clique em Export agent code para gerar o Python Notebook Notebook que ajuda o senhor a desenvolver e implantar o agente AI.

Depois de exportar o código do agente, o senhor verá três arquivos salvos no site workspace:
- agent Notebook: Contém o código Python que define seu agente usando LangChain.
- driver Notebook: Contém o código Python para log, rastrear, registrar e implantar o agente AI usando o Mosaic AI Agent Framework.
- config.yml: Contém informações de configuração sobre seu agente, incluindo definições de ferramentas.
Abra o agent Notebook para ver o código LangChain que define seu agente. Use esse Notebook para testar e iterar o agente programaticamente, como definir mais ferramentas ou ajustar os parâmetros do agente.

Observação

O código exportado pode ter um comportamento diferente da sua sessão de AI playground. Databricks recomenda que o senhor execute o Notebook exportado para iterar e depurar mais, avaliar a qualidade do agente e, em seguida, implantar o agente para compartilhar com outras pessoas.
Quando estiver satisfeito com os resultados do agente, o senhor pode executar o driver Notebook para log e implantar seu agente em um modelo de serviço endpoint.

Definir um agente no código

Além de gerar código de agente a partir do AI Playground, o senhor também pode definir um agente em código, usando estruturas como LangChain ou código Python. Para implantar um agente usando o Agent Framework, sua entrada deve estar em conformidade com um dos formatos de entrada e saída suportados.

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 depois 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')

Formatos de entrada suportados

A seguir estão os formatos de entrada compatíveis com seu agente.

(Recomendado) Consultas utilizando o esquema de conclusão de bate-papo OpenAI. Deve ter uma matriz de objetos como parâmetro messages . Este formato é melhor para aplicações RAG.

question = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-Augmented Generation?",
        },
        {
            "role": "assistant",
            "content": "RAG, or Retrieval Augmented Generation, is a generative AI design pattern that combines a large language model (LLM) with external knowledge retrieval. This approach allows for real-time data connection to generative AI applications, improving their accuracy and quality by providing context from your data to the LLM during inference. Databricks offers integrated tools that support various RAG scenarios, such as unstructured data, structured data, tools & function calling, and agents.",
        },
        {
            "role": "user",
            "content": "How to build RAG for unstructured data",
        },
    ]
}

SplitChatMessagesRequest. Recomendado para aplicações de bate-papo com vários turnos, especialmente quando o usuário deseja gerenciar a consulta atual e o histórico separadamente.

question = {
    "query": "What is MLflow",
    "history": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?"
        },
        {
            "role": "assistant",
            "content": "RAG is"
        }
    ]
}

Para a LangChain, a Databricks recomenda escrever sua cadeia em LangChain Expression Language. Em seu código de definição de cadeia, você pode utilizar um itemgetter para receber as mensagens ou os objetos query ou history, dependendo do formato de entrada que estiver utilizando.

Formatos de saída suportados

Seu agente deve ter um dos seguintes formatos de saída compatíveis:

(Recomendado) ChatCompletionResponse. Esse formato é recomendado para clientes com interoperabilidade de formato de resposta OpenAI.
StringObjectResponse. Esse formato é o mais fácil e simples de interpretar.

Para LangChain, use StringResponseOutputParser() ou ChatCompletionsOutputParser() de MLflow como sua cadeia final de passos. Isso formata a mensagem LangChain AI em um formato compatível com o agente.

  from mlflow.langchain.output_parsers import StringResponseOutputParser, ChatCompletionsOutputParser

  chain = (
      {
          "user_query": itemgetter("messages")
          | RunnableLambda(extract_user_query_string),
          "chat_history": itemgetter("messages") | RunnableLambda(extract_chat_history),
      }
      | RunnableLambda(fake_model)
      | StringResponseOutputParser() # use this for StringObjectResponse
      # ChatCompletionsOutputParser() # or use this for ChatCompletionResponse
  )

Se você estiver utilizando PyFunc, a Databricks recomenda utilizar dicas de tipo para anotar a função predict() com classes de dados de entrada e saída que são subclasses de classes definidas em mlflow.models.rag_signatures.

Você pode construir um objeto de saída a partir da classe de dados dentro de predict() para garantir que o formato seja seguido. O objeto retornado deve ser transformado em uma representação de dicionário para garantir que possa ser serializado.

  from mlflow.models.rag_signatures import ChatCompletionRequest, ChatCompletionResponse, ChainCompletionChoice, Message

  class RAGModel(PythonModel):
    ...
      def predict(self, context, model_input: ChatCompletionRequest) -> ChatCompletionResponse:
        ...
        return asdict(ChatCompletionResponse(
            choices=[ChainCompletionChoice(message=Message(content=text))]
        ))

Notebooks de Exemplo

Esses notebooks criam uma cadeia simples "Hello, world" para ilustrar como criar uma aplicação de cadeia no Databricks. O primeiro exemplo cria uma cadeia simples. O segundo notebook de exemplo ilustra como utilizar parâmetros para minimizar as alterações de código durante o desenvolvimento.