Älteres Eingabe- und Ausgabe-Agent-Schema

Hinweis

Databricks hat die veralteten AI-Agent-Schemen ChatModel, SplitChatMessageRequest und StringResponse abgekündigt. Databricks empfiehlt die Migration zum empfohlenen ChatAgent Schema zur Erstellung von Agents. Siehe Erstellen von KI-Agents im Code.

KI-Agents müssen bestimmte Eingabe- und Ausgabeschemaanforderungen einhalten, um mit anderen Features in Azure Databricks kompatibel zu sein. In diesem Artikel wird erläutert, wie Sie die Legacy-Agent-Erstellungssignaturen und -schnittstellen verwenden: ChatModel Schnittstelle, SplitChatMessageRequest Eingabeschema und StringResponse Ausgabeschema.

Erstellen eines älteren ChatModel-Agents

Von Bedeutung

Databricks empfiehlt die ChatAgent Schnittstelle zum Erstellen von Agenten oder KI-Apps. Informationen zum Migrieren von ChatModel zu ChatAgent finden Sie in der MLflow-Dokumentation – Migrieren von ChatModel zu ChatAgent.

ChatModel ist eine ältere Agent-Erstellungsschnittstelle in MLflow, die das ChatCompletion-Schema von OpenAI erweitert, sodass Sie die Kompatibilität mit Plattformen beibehalten können, die den ChatCompletion-Standard unterstützen, während sie benutzerdefinierte Funktionen hinzufügen. Weitere Details finden Sie unter MLflow: Erste Schritte mit ChatModel .

Die Erstellung Ihres Agents als Unterklasse von mlflow.pyfunc.ChatModel bietet die folgenden Vorteile:

• Aktiviert die Ausgabe des Streaming-Agents beim Aufrufen eines bereitgestellten Agents (Umgehung {stream: true} im Anforderungstext).
• Aktiviert automatisch KI-Gateway-Inference-Tabellen, wenn Ihr Agent bereitsteht, und ermöglicht den Zugriff auf erweiterte Anforderungsprotokoll-Metadaten, z. B. den Namen des Anforderers.
• Ermöglicht Ihnen das Schreiben von Agentcode, der mit dem ChatCompletion-Schema mit typierten Python-Klassen kompatibel ist.
• MLflow leitet automatisch eine chat-kompatible Signatur ab, wenn der Agent protokolliert wird, auch ohne input_example. Dies vereinfacht das Registrieren und Bereitstellen des Agents. Weitere Informationen finden Sie unter Ableiten der Modellsignatur während der Protokollierung.

Der folgende Code wird am besten in einem Databricks-Notizbuch ausgeführt. Notebooks bieten eine praktische Umgebung zum Entwickeln, Testen und Iterieren Ihres Agents.

Die MyAgent Klasse erweitert mlflow.pyfunc.ChatModel, wobei die erforderliche predict-Methode implementiert wird. Dadurch wird die Kompatibilität mit Dem Mosaik AI Agent Framework sichergestellt.

Die Klasse enthält auch die optionalen Methoden _create_chat_completion_chunk und predict_stream zum Behandeln von Streamingausgaben.

# Install the latest version of mlflow
%pip install -U mlflow
dbutils.library.restartPython()

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, messages=model_input.messages, params=None)
print(response)

Während die Agentklasse MyAgent in einem Notizbuch definiert ist, sollten Sie ein separates Treibernotizbuch erstellen. Das Treibernotizbuch protokolliert den Agenten im Modellregister und stellt ihn mithilfe von Model Serving bereit.

Diese Trennung folgt dem von Databricks empfohlenen Workflow für die Protokollierung von Modellen mithilfe der Methodik 'Models from Code' von MLflow.

SplitChatMessageRequest-Eingabeschema (veraltet)

mit SplitChatMessagesRequest können Sie die aktuelle Abfrage und den aktuellen Verlauf separat als Agenteingabe übergeben.

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

StringResponse-Ausgabeschema (veraltet)

StringResponse ermöglicht es Ihnen, die Antwort des Agents als Objekt mit einem einzelnen Zeichenfolgenfeld content zurückzugeben:

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