Dela via

Definiera en agents indata- och utdataschema

  • Artikel

MLflow-modellsignaturer definierar kraven för indata- och utdatascheman för AI-agenten. Modellsignaturen talar om för interna och externa komponenter hur de ska interagera med din agent. Modellsignaturen är en verifieringskontroll för att säkerställa att indata följer schemakraven.

Om du till exempel vill använda granskningsappen För agentutvärdering måste din agent följa indataschemat för agentutvärdering.

Indatascheman som stöds

Mosaic AI Agent Framework stöder följande indatascheman.

OpenAI-schema för chattens slutförande

Kommentar

Databricks rekommenderar OpenAI-chattens slutförandeschema eftersom det är allmänt använt och samverkande med många agentramverk och program. Om OpenAI-chattens slutförandeschema inte uppfyller dina behov kan du definiera ditt eget schema. Se Anpassade agentscheman.

  • (Rekommenderas) Databricks rekommenderar att du använder OpenAI-chattens slutförandeschema. OpenAI-chattens slutförandeschema bör ha en matris med objekt som en messages parameter. Det här formatet passar bäst för RAG-program.

    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 rekommenderas för chattprogram med flera svängar, särskilt när du vill hantera den aktuella frågan och historiken separat.

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

Språk för langchainuttryck

Om din agent använder LangChain kan du skriva din kedja i Språk för LangChain-uttryck. I din kedjedefinitionskod kan du använda en itemgetter för att hämta meddelanden eller query objekt history beroende på ditt indataformat.

Utdatascheman som stöds

Mosaic AI Agent Framework stöder följande utdatascheman.

ChatCompletionResponse

(Rekommenderas) ChatCompletionResponse rekommenderas för kunder med OpenAI-svarsformatskompatibilitet.

LangChain – ChatCompletionsOutputParser

Om din agent använder LangChain använder du ChatCompletionsOutputParser() från MLflow som det sista kedjesteget. Detta formaterar LangChain AI-meddelandet till ett agentkompatibelt format.


  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 – kommentera indata- och utdataklasser

Om du använder PyFunc rekommenderar Databricks att du använder typtips för att kommentera predict() funktionen med indata- och utdataklasser som är underklasser av klasser som definieras i mlflow.models.rag_signatures.

Du kan skapa ett utdataobjekt från dataklassen i predict(). Det returnerade objektet måste omvandlas till en ordlisterepresentation för att säkerställa att det kan serialiseras.


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

Explicita och härledda signaturer

MLflow kan härleda indata- och utdataschemat för din agent vid körning och skapa en signatur automatiskt. Om du använder indata- och utdatascheman som stöds är de härledda signaturerna kompatibla med Agent Framework. Information om scheman som stöds finns i Indatascheman som stöds.

Men om du använder ett anpassat agentschema måste du uttryckligen definiera din modellsignatur enligt anvisningarna i Anpassade agentscheman.

Anpassade agentscheman

Du kan anpassa en agent schema för att skicka och returnera ytterligare fält till och från agenten genom att skapa en underklass av ett indata-/utdataschema som stöds. Lägg sedan till de extra nycklarna custom_inputs och custom_outputs för att innehålla de ytterligare fälten. Se kodexempel för Pyfunc och Langchain och en UI-baserad metod för att använda anpassade indata.

Om du vill använda databricks-agents SDK, Databricks-klient-UIs som AI Playground och Granskningsappen och andra Mosaic AI Agent Framework-funktioner måste agentens schema uppfylla följande krav:

  • Agenten måste använda mlflow version 2.17.1 eller senare.
  • I agentanteckningsboken markerar du ytterligare fält som lagts till i underklassen som Optional och tilldelar standardvärden.
  • I drivrutinsanteckningsboken skapar du en ModelSignature användning infer_signature med instanser av dina underklasser.
  • Skapa ett indataexempel i drivrutinsanteckningsboken genom att anropa asdict i underklassen.

Anpassade PyFunc-scheman

Förutom kraven ovan måste PyFunc-baserade agenter också uppfylla följande krav för att interagera med Mosaic AI-agentfunktioner.

Anpassade schemakrav för PyFunc

I agentanteckningsboken måste funktionerna förutsäga och förutsäga dataström uppfylla följande krav:

  • Ha typtips för indataunderklassen.
  • Använd punkt notation för att komma åt dataklassfält (till exempel använda model_input.custom_input.id i stället model_input["custom_inputs"]för ).
  • Returnera en dictionary. Du kan anropa asdict en instans av underklassen för att formatera returen som en ordlista.

Följande notebook-filer visar ett anpassat schemaexempel med PyFunc.

PyFunc- anpassad schemaagentanteckningsbok

Hämta notebook-fil

PyFunc anpassad schemadrivrutinsanteckningsbok

Hämta notebook-fil

Anpassade scheman för Langchain

Följande notebook-filer visar ett anpassat schemaexempel med LangChain. Du kan ändra funktionen wrap_output i notebook-filerna för att parsa och extrahera information från meddelandeströmmen.

Anteckningsbok för anpassad schemaagent i Langchain

Hämta notebook-fil

Anteckningsbok för anpassad schemadrivrutin för Langchain

Hämta notebook-fil

Ange custom_inputs i appen AI Playground och agentgranskning

Om du definierar ett anpassat agentschema med ytterligare indata med hjälp av custom_inputs fältet kan du manuellt ange dessa indata i både AI Playground och agentgranskningsappen. Om inga anpassade indata anges använder agenten de standardvärden som anges i schemat.

  1. Välj kugghjulsikonen i AI Playground eller Agentgranskningsappen Kugghjulsikon.

  2. Aktivera custom_inputs.

  3. Ange ett JSON-objekt som matchar agentens definierade indataschema.

    Tillhandahåll custom_inputs på AI-lekplatsen.

JSON-objektet måste matcha agentens indataschema. Om du till exempel har en custom_inputs dataklass definierad enligt följande:

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

Sedan måste JSON-strängen som du anger i fältet custom_inputs ange värden för id och user, enligt följande exempel:

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