Compartir a través de

Definición del esquema de entrada y salida de un agente

  • Artículo

Las firmas de modelo de MLflow definen los requisitos de esquema de entrada y salida para el agente de IA. La firma del modelo indica a los componentes internos y externos cómo interactuar con el agente. La firma del modelo es una comprobación de validación para asegurarse de que las entradas cumplen los requisitos de esquema.

Por ejemplo, para usar la aplicación de revisión evaluación del agente, el agente debe cumplir el esquema de entrada de evaluación del agente.

Esquemas de entrada admitidos

Mosaic AI Agent Framework admite los siguientes esquemas de entrada.

Esquema de finalización de chat de OpenAI

Nota:

Databricks recomienda el esquema de finalización del chat de OpenAI porque se usa ampliamente e interoperable con muchos marcos y aplicaciones de agente. Si el esquema de finalización del chat de OpenAI no satisface sus necesidades, puede definir su propio esquema. Consulte Esquemas de agente personalizados.

  • (Recomendado) Databricks recomienda usar el esquema de finalización del chat de OpenAI. El esquema de finalización del chat de OpenAI debe tener una matriz de objetos como parámetro messages . Este formato es el mejor para las aplicaciones 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",
        },
    ]
}

SplitChatMessageRequest

SplitChatMessagesRequest se recomienda para las aplicaciones de chat multiturno, especialmente cuando desea administrar la consulta y el historial actuales por separado.

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

Lenguaje de expresiones langchain

Si el agente usa LangChain, puede escribir la cadena en lenguaje de expresiones LangChain. En el código de definición de cadena, puede usar para itemgetter obtener los mensajes o query history objetos en función del formato de entrada.

Esquemas de salida admitidos

Mosaic AI Agent Framework admite los siguientes esquemas de salida.

ChatCompletionResponse

(Recomendado) Se recomienda ChatCompletionResponse para los clientes con interoperabilidad de formato de respuesta openAI.

LangChain: ChatCompletionsOutputParser

Si el agente usa LangChain, use ChatCompletionsOutputParser() desde MLflow como paso de cadena final. Esto da formato al mensaje LangChain AI en un formato compatible con el agente.


  from mlflow.langchain.output_parsers import ChatCompletionsOutputParser

  chain = (
      {
          "user_query": itemgetter("messages")
          | RunnableLambda(extract_user_query_string),
          "chat_history": itemgetter("messages") | RunnableLambda(extract_chat_history),
      }
      | RunnableLambda(DatabricksChat)
      | ChatCompletionsOutputParser()
  )

PyFunc: anota las clases de entrada y salida

Si usa PyFunc, Databricks recomienda utilizar sugerencias de tipo para anotar la función predict() con clases de datos de entrada y salida que sean subclases de las clases definidas en mlflow.models.rag_signatures.

Puede construir un objeto de salida a partir de la clase de datos dentro predict()de . El objeto que se devuelve debe transformarse en una representación de diccionario para garantizar que se pueda serializar.


  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))]
        ))

Firmas explícitas e inferidas

MLflow puede deducir el esquema de entrada y salida del agente en tiempo de ejecución y crear automáticamente una firma. Si usa esquemas de entrada y salida admitidos, las firmas deducidas son compatibles con Agent Framework. Para obtener información sobre los esquemas admitidos, consulte Esquemas de entrada admitidos.

Sin embargo, si usa un esquema de agente personalizado, debe definir explícitamente la firma del modelo según las instrucciones de Esquemas de agente personalizados.

Esquemas de agente personalizados

Puede personalizar el esquema de un agente para pasar y devolver campos adicionales hacia y desde el agente mediante la creación de una subclase de un esquema de entrada y salida admitidos. A continuación, agregue las claves custom_inputs adicionales y custom_outputs para contener los campos adicionales. Consulte ejemplos de código para Pyfunc y Langchain y un método basado en la interfaz de usuario para usar entradas personalizadas.

Para usar el SDK, las databricks-agents interfaces de usuario de cliente de Databricks, como AI Playground y la aplicación de revisión, y otras características de Mosaic AI Agent Framework, el esquema del agente debe cumplir los siguientes requisitos:

  • El agente debe usar mlflow la versión 2.17.1 o superior.
  • En el cuaderno del agente, marque campos adicionales agregados en la subclase como Optional y asigne valores predeterminados.
  • En el cuaderno del controlador, cree un ModelSignature uso infer_signature con instancias de las subclases.
  • En el cuaderno del controlador, cree un ejemplo de entrada llamando a asdict en la subclase.

Esquemas personalizados de PyFunc

Además de los requisitos anteriores, los agentes basados en PyFunc también deben cumplir los siguientes requisitos para interactuar con las características del agente de Mosaic AI.

Requisitos de esquema personalizado de PyFunc

En el cuaderno del agente, las funciones de predicción y predicción de flujos deben cumplir los siguientes requisitos:

  • Tenga sugerencias de tipo para la subclase de entrada.
  • Use la notación de puntos para acceder a los campos dataclass (por ejemplo, use model_input.custom_input.id en lugar de model_input["custom_inputs"]).
  • Devolver dictionary. Puede llamar a asdict en una instancia de la subclase para dar formato a la devolución como diccionario.

Los cuadernos siguientes muestran un ejemplo de esquema personalizado mediante PyFunc.

Cuaderno del agente de esquema personalizado de PyFunc

Obtener el cuaderno

Cuaderno del controlador de esquema personalizado de PyFunc

Obtener el cuaderno

Esquemas personalizados de Langchain

Los cuadernos siguientes muestran un ejemplo de esquema personalizado mediante LangChain. Puede modificar la función wrap_output en los cuadernos para analizar y extraer información de la secuencia de mensajes.

Cuaderno del agente de esquema personalizado langchain

Obtener el cuaderno

Cuaderno del controlador de esquema personalizado langchain

Obtener el cuaderno

Proporcionar custom_inputs en la aplicación ai Playground y la aplicación de revisión del agente

Si define un esquema de agente personalizado con entradas adicionales mediante el custom_inputs campo , puede proporcionar manualmente estas entradas tanto en ai Playground como en la aplicación de revisión del agente. Si no se proporciona ninguna entrada personalizada, el agente usa los valores predeterminados especificados en el esquema.

  1. En ai Playground o la aplicación de revisión del agente, seleccione el icono icono de engranajede engranaje .

  2. Habilite custom_inputs.

  3. Proporcione un objeto JSON que coincida con el esquema de entrada definido del agente.

    Proporcione custom_inputs en el área de juegos de IA.

El objeto JSON debe coincidir con el esquema de entrada del agente. Por ejemplo, si tiene una custom_inputs clase de datos definida como se indica a continuación:

@dataclass
class CustomInputs():
  id: int = 0
  user: str = "default"

A continuación, la cadena JSON que escriba en el campo custom_inputs debe proporcionar valores para id y user, como se muestra en el ejemplo siguiente:

{
  "id": 123
  "user": "dev_test",
}