Esquema de agente de entrada e saída legado

nota

Databricks descontinuou os esquemas de agente AI ChatModel, SplitChatMessageRequest e StringResponse. A Databricks recomenda a migração para o esquema ChatAgent recomendado para agentes de autor. Veja Author AI agents in code.

AI Os agentes devem aderir a requisitos específicos de esquema de entrada e saída para serem compatíveis com outros recursos em Databricks. Este artigo explica como usar as assinaturas e interfaces de criação de agentes legados: interface ChatModel, o esquema de entrada SplitChatMessageRequest e o esquema de saída StringResponse.

Crie um agente legado do ChatModel

important

Databricks recomenda a interface ChatAgent para criar agentes ou gerar aplicativos AI. Para migrar do ChatModel para o ChatAgent, consulte a documentação do MLflow - Migrar do ChatModel para o ChatAgent.

ChatModel é uma interface de criação de agentes legados no MLflow que estende o esquema ChatCompletion da OpenAI, permitindo que o senhor mantenha a compatibilidade com plataformas que suportam o padrão ChatCompletion e, ao mesmo tempo, adicione funcionalidades personalizadas. MLflowConsulte: Começando com o ChatModel para obter mais detalhes.

Criar seu agente como uma subclasse de mlflow.pyfunc.ChatModel oferece os seguintes benefícios:

• Ativa a saída do agente de transmissão ao invocar um agente atendido (ignorando {stream: true} no corpo da solicitação).
• Ativa automaticamente as tabelas de inferência do gateway AI quando seu agente é atendido, fornecendo acesso a metadados aprimorados da solicitação log, como o nome do solicitante.
• Permite que o senhor escreva código de agente compatível com o esquema ChatCompletion usando classes Python tipadas.
• O MLflow infere 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.

Python

import re
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 re.findall(r"\S+\s*", last_user_question_text):
            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.

Esquema de entrada SplitChatMessageRequest (obsoleto)

SplitChatMessagesRequest permite que o senhor passe a consulta atual e o histórico separadamente como entrada do agente.

Python

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

Esquema de saída StringResponse (obsoleto)

StringResponse permite que o senhor retorne a resposta do agente como um objeto com um único campo de strings content:

{"content": "This is an example string response"}