Freigeben über

Definieren des Eingabe- und Ausgabeschemas eines Agents

  • Artikel

MLflow-Modellsignaturen definieren die Eingabe- und Ausgabeschemaanforderungen für Ihren KI-Agent. Die Modellsignatur teilt internen und externen Komponenten mit, wie sie mit Ihrem Agent interagieren. Die Modellsignatur fungiert als Überprüfungsprüfung, um sicherzustellen, dass Eingaben den Schemaanforderungen entsprechen.

Um beispielsweise die Agent Evaluation Review-App zu verwenden, muss Ihr Agent das Eingabeschema für die Agentauswertung einhalten.

Unterstützte Eingabeschemas

Mosaik AI Agent Framework unterstützt die folgenden Eingabeschemas.

OpenAI-Chatabschlussschema

Hinweis

Databricks empfiehlt das OpenAI-Chatabschlussschema, da es weit verbreitet ist und mit vielen Agent-Frameworks und -Anwendungen interoperabel ist. Wenn das OpenAI-Chatabschlussschema Ihre Anforderungen nicht erfüllt, können Sie Ihr eigenes Schema definieren. Siehe benutzerdefinierte Agentschemas.

  • (Empfohlen) Databricks empfiehlt die Verwendung des OpenAI-Chatabschlussschemas. Das OpenAI-Chatabschlussschema sollte ein Array von Objekten als messages Parameter aufweisen. Dieses Format eignet sich am besten für RAG-Anwendungen.

    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 wird für Multi-Turn-Chatanwendungen empfohlen, insbesondere, wenn Sie die aktuelle Abfrage und den aktuellen Verlauf separat verwalten möchten.

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

Langbundausdruckssprache

Wenn Ihr Agent LangChain verwendet, können Sie Ihre Kette in der LangChain-Ausdruckssprache schreiben. In Ihrem Kettendefinitionscode können Sie je nach Eingabeformat die itemgetter Nachrichten oder query history Objekte abrufen.

Unterstützte Ausgabeschemas

Mosaik AI Agent Framework unterstützt die folgenden Ausgabeschemas.

ChatCompletionResponse

(Empfohlen) ChatCompletionResponse wird für Kunden mit OpenAI-Antwortformatinteroperabilität empfohlen.

LangChain - ChatCompletionsOutputParser

Wenn Ihr Agent LangChain verwendet, verwenden Sie ChatCompletionsOutputParser() MLflow als letzten Kettenschritt. Dadurch wird die LangChain AI-Nachricht in ein agentkompatibles Format formatiert.


  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 – Kommentieren von Eingabe- und Ausgabeklassen

Wenn Sie PyFunc verwenden, empfiehlt Databricks die Verwendung von Typhinweisen, um die predict() Funktion mit Eingabe- und Ausgabedatenklassen zu kommentieren, die Unterklassen von Klassen sind, die in mlflow.models.rag_signatures definiert sind.

Sie können ein Ausgabeobjekt aus der Datenklasse innerhalb von predict() erstellen, um sicherzustellen, dass das Format befolgt wird. Das zurückgegebene Objekt muss in eine Wörterbuchdarstellung umgewandelt werden, um sicherzustellen, dass es serialisiert werden kann.


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

Explizite und abgeleitete Signaturen

MLflow kann das Eingabe- und Ausgabeschema Ihres Agents zur Laufzeit ableiten und automatisch eine Signatur erstellen. Wenn Sie unterstützte Eingabe- und Ausgabeschemas verwenden, sind die abgeleiteten Signaturen mit dem Agent Framework kompatibel. Informationen zu unterstützten Schemas finden Sie unter "Unterstützte Eingabeschemas".

Wenn Sie jedoch ein benutzerdefiniertes Agentschema verwenden, müssen Sie Ihre Modellsignatur explizit gemäß den Anweisungen in benutzerdefinierten Agent-Schemas definieren.

Benutzerdefinierte Agentschemas

Sie können das Schema eines Agents anpassen, um zusätzliche Felder an den Agent zu übergeben und vom Agent zurückzugeben, indem Sie eine Unterklasse eines unterstützten Eingabe-/Ausgabeschemas erstellen. Fügen Sie dann die zusätzlichen Schlüssel custom_inputs hinzu, und custom_outputs fügen Sie die zusätzlichen Felder hinzu. Codebeispiele für Pyfunc und Langchain finden Sie unten.

Um das databricks-agents SDK, Databricks-Client-UIs wie den AI Playground und die Review-App und andere Mosaik AI Agent Framework-Features zu verwenden, muss das Schema Ihres Agents die folgenden Anforderungen erfüllen:

  • Der Agent muss Version 2.17.1 oder höher verwenden mlflow .
  • Markieren Sie im Agent-Notizbuch zusätzliche Felder, die in Ihrer Unterklasse hinzugefügt wurden, und Optional weisen Sie Standardwerte zu.
  • Erstellen Sie im Treibernotizbuch eine ModelSignature Verwendung infer_signature mit Instanzen Ihrer Unterklassen.
  • Erstellen Sie im Treibernotizbuch ein Eingabebeispiel, indem Sie Ihre Unterklasse aufrufen asdict .

Benutzerdefinierte PyFunc-Schemas

Zusätzlich zu den oben genannten Anforderungen müssen PyFunc-basierte Agents auch die folgenden Anforderungen erfüllen, um mit Mosaik AI-Agent-Features zu interagieren.

Benutzerdefinierte Schemaanforderungen für PyFunc

Im Agent-Notizbuch müssen die Funktionen für die Vorhersage und Vorhersage des Datenstroms die folgenden Anforderungen erfüllen:

  • Geben Sie Eingabehinweise für Ihre Eingabeunterklasse ein.
  • Verwenden Sie die Punktnotation, um auf Datenklassenfelder zuzugreifen (z. B. anstelle model_input.custom_input.id von model_input["custom_inputs"]).
  • Eine dictionary zurückgeben. Sie können eine Instanz Ihrer Unterklasse aufrufen asdict , um die Rückgabe als Wörterbuch zu formatieren.

Die folgenden Notizbücher zeigen ein benutzerdefiniertes Schemabeispiel mit PyFunc.

Benutzerdefiniertes PyFunc-Schema-Agent-Notizbuch

Notebook abrufen

Benutzerdefiniertes PyFunc-Schematreibernotizbuch

Notebook abrufen

Benutzerdefinierte Langchain-Schemas

Die folgenden Notizbücher zeigen ein benutzerdefiniertes Schemabeispiel mit LangChain. Sie können die wrap_output-Funktion in den Notizbüchern ändern, um Informationen aus dem Nachrichtendatenstrom zu analysieren und zu extrahieren.

Benutzerdefiniertes Schema-Agent-Notizbuch für die Langchain

Notebook abrufen

Benutzerdefiniertes Schematreibernotizbuch der Langchain

Notebook abrufen