KI-Agenten im Code erstellen

Diese Seite zeigt, wie Sie einen KI-Agent in Python mit Mosaik AI Agent Framework und beliebten Agent-Erstellungsbibliotheken wie LangGraph und OpenAI erstellen.

Anforderungen

Um Agents mithilfe des Ansatzes auf dieser Seite zu erstellen und bereitzustellen, installieren Sie Folgendes:

• databricks-agents 1.2.0 oder höher
• mlflow 3.1.3 oder höher
• Python 3.10 oder höher.
  • Verwenden Sie serverloses Compute oder Databricks Runtime 13.3 LTS oder höher, um diese Anforderung zu erfüllen.

%pip install -U -qqqq databricks-agents mlflow

Databricks empfiehlt auch die Installation von Databricks AI Bridge-Integrationspaketen für Autoren-Agents. Diese Integrationspakete bieten eine gemeinsame Ebene von APIs, die mit Databricks AI-Features wie Databricks AI/BI Genie und Vector Search interagieren, über Agent-Erstellungsframeworks und SDKs hinweg.

OpenAI

%pip install -U -qqqq databricks-openai

LangChain/LangGraph

%pip install -U -qqqq databricks-langchain

DSPy

%pip install -U -qqqq databricks-dspy

Reine Python-Agenten

%pip install -U -qqqq databricks-ai-bridge

Verwende ResponsesAgent zum Erstellen von Agenten

Databricks empfiehlt die MLflow-Schnittstelle ResponsesAgent zum Erstellen von Agenten auf Produktionsniveau. ResponsesAgent Ermöglicht Es Ihnen, Agents mit jedem Drittanbieterframework zu erstellen und sie dann mit Databricks AI-Features für robuste Protokollierung, Ablaufverfolgung, Auswertung, Bereitstellung und Überwachung zu integrieren.

Das ResponsesAgent Schema ist mit dem OpenAI-Schema Responses kompatibel. Weitere Informationen zu OpenAI Responsesfinden Sie unter OpenAI: Antworten vs. ChatCompletion.

ResponsesAgent bietet die folgenden Vorteile:

• Erweiterte Agent-Funktionen

  • Multi-Agent-Unterstützung
  • Streamingausgabe: Streamen Sie die Ausgabe in kleineren Blöcken.
  • Umfassender Nachrichtenverlauf bei Toolaufrufen: Rückgabe mehrerer Nachrichten, einschließlich zwischengeschalteter Nachrichten, für eine verbesserte Qualität und Konversationsverwaltung.
  • Unterstützung bei der Bestätigung von Toolaufrufen
  • Unterstützung von langfristigen Tools

• Vereinfachte Entwicklung, Bereitstellung und Überwachung

  • Erstellen Sie Agenten mit jedem beliebigen Framework: Verpacken Sie jeden bestehenden Agenten mithilfe der ResponsesAgent Schnittstelle, um eine sofortige Kompatibilität mit AI Playground, Agent Evaluation und Agent Monitoring zu erreichen.
  • Typisierte Erstellungsschnittstellen: Schreiben Sie Agent-Code mit typisierten Python-Klassen, und nutzen Sie dabei das AutoVervollständigen von IDE und Notebook.
  • Automatische Signaturableitung: MLflow leitet ResponsesAgent beim Protokollieren eines Agenten automatisch Signaturen ab und vereinfacht die Registrierung und Bereitstellung. Weitere Informationen finden Sie unter Ableiten der Modellsignatur während der Protokollierung.
  • Automatische Ablaufverfolgung: MLflow verfolgt Ihre predict und predict_stream Funktionen automatisch und aggregiert gestreamte Antworten, um die Auswertung und Anzeige zu erleichtern.
  • AI Gateway-erweiterte Inferencetabellen: AI-Gateway-Ableitungstabellen werden automatisch für bereitgestellte Agents aktiviert und bieten Zugriff auf detaillierte Anforderungsprotokollmetadaten.

Informationen zum Erstellen einer ResponsesAgentDatei finden Sie in den Beispielen im folgenden Abschnitt und in der MLflow-Dokumentation – ResponsesAgent für Model Serving.

ResponsesAgent Beispiele

Die folgenden Notizbücher zeigen, wie Sie Streaming und Nicht-Streaming ResponsesAgent mithilfe beliebter Bibliotheken erstellen. Informationen zum Erweitern der Funktionen dieser Agents finden Sie unter KI-Agent-Tools.

OpenAI

OpenAI Simple Chat Agent mit von Databricks gehosteten Modellen

Notebook abrufen

OpenAI-Toolaufruf-Agent unter Verwendung von Databricks-gehosteten Modellen

Notebook abrufen

OpenAI-Tool-Aufruf-Agent mit OpenAI-gehosteten Modellen

Notebook abrufen

LangGraph

Toolaufruf-Agent für LangGraph

Notebook abrufen

DSPy

DSPy Single-Turn-Tool-Calling-Agent

Notebook abrufen

Beispiel für mehrere Agents

Informationen zum Erstellen eines Multi-Agent-Systems finden Sie unter Verwendung von Genie in Multi-Agent-Systemen.

Beispiel für einen zustandsbehafteten Agenten

Informationen zum Erstellen zustandsbehafteter Agenten mit kurzfristigem und langfristigem Gedächtnis unter Verwendung von Lakebase als Erinnerungsdatenbank finden Sie unter KI-Agent-Gedächtnis.

Beispiel für einen nicht-konversationellen Agenten

Im Gegensatz zu konversationalen Agenten, die Dialoge mit mehreren Gängen verwalten, konzentrieren sich nicht-konversationale Agenten auf die effiziente Ausführung gut definierter Aufgaben. Diese optimierte Architektur ermöglicht einen höheren Durchsatz für unabhängige Anforderungen.

Informationen zum Erstellen eines nicht-konversationalen Agenten finden Sie unter Nicht-konversationale KI-Agenten mit MLflow.

Was geschieht, wenn ich bereits über einen Agenten verfüge?

Wenn Sie bereits über einen Agent verfügen, der mit LangChain, LangGraph oder einem ähnlichen Framework erstellt wurde, müssen Sie Ihren Agent nicht neu schreiben, um ihn auf Databricks zu verwenden. Wrappen Sie stattdessen einfach Ihren bestehenden Agenten mit der MLflow ResponsesAgent-Schnittstelle:

1. Schreiben Sie eine Python-Wrapperklasse, die von mlflow.pyfunc.ResponsesAgent erbt.

   Verweisen Sie innerhalb der Wrapperklasse auf den vorhandenen Agent als Attribut self.agent = your_existing_agent.

2. Die ResponsesAgent Klasse erfordert die Implementierung einer predict Methode, die eine ResponsesAgentResponse Methode zurückgibt, um Nicht-Streaming-Anforderungen zu behandeln. Im Folgenden sehen Sie ein Beispiel für das ResponsesAgentResponses Schema:

   import uuid
   # input as a dict
   {"input": [{"role": "user", "content": "What did the data scientist say when their Spark job finally completed?"}]}
   
   # output example
   ResponsesAgentResponse(
       output=[
           {
               "type": "message",
               "id": str(uuid.uuid4()),
               "content": [{"type": "output_text", "text": "Well, that really sparked joy!"}],
               "role": "assistant",
           },
       ]
   )
   

3. Konvertieren Sie in der predict Funktion die eingehenden Nachrichten von ResponsesAgentRequest in das Format, das der Agent erwartet. Nachdem der Agent eine Antwort generiert hat, konvertieren Sie die Ausgabe in ein ResponsesAgentResponse Objekt.

Sehen Sie sich die folgenden Codebeispiele an, um zu sehen, wie vorhandene Agents in ResponsesAgent konvertiert werden.

Grundlegende Konvertierung

Konvertieren Sie für Nicht-Streaming-Agents Eingaben und Ausgaben in der predict Funktion.

from uuid import uuid4

from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
)


class MyWrappedAgent(ResponsesAgent):
    def __init__(self, agent):
        # Reference your existing agent
        self.agent = agent

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        # Convert incoming messages to your agent's format
        # prep_msgs_for_llm is a function you write to convert the incoming messages
        messages = self.prep_msgs_for_llm([i.model_dump() for i in request.input])

        # Call your existing agent (non-streaming)
        agent_response = self.agent.invoke(messages)

        # Convert your agent's output to ResponsesAgent format, assuming agent_response is a str
        output_item = (self.create_text_output_item(text=agent_response, id=str(uuid4())),)

        # Return the response
        return ResponsesAgentResponse(output=[output_item])

Streaming mit Wiederverwenden von Code

Für Streaming-Agents können Sie geschickt Logik wiederverwenden, um das Duplizieren des Codes zu vermeiden, der Nachrichten konvertiert.

from typing import Generator
from uuid import uuid4

from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
    ResponsesAgentStreamEvent,
)


class MyWrappedStreamingAgent(ResponsesAgent):
    def __init__(self, agent):
        # Reference your existing agent
        self.agent = agent

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        """Non-streaming predict: collects all streaming chunks into a single response."""
        # Reuse the streaming logic and collect all output items
        output_items = []
        for stream_event in self.predict_stream(request):
            if stream_event.type == "response.output_item.done":
                output_items.append(stream_event.item)

        # Return all collected items as a single response
        return ResponsesAgentResponse(output=output_items)

    def predict_stream(
        self, request: ResponsesAgentRequest
    ) -> Generator[ResponsesAgentStreamEvent, None, None]:
        """Streaming predict: the core logic that both methods use."""
        # Convert incoming messages to your agent's format
        # prep_msgs_for_llm is a function you write to convert the incoming messages, included in full examples linked below
        messages = self.prep_msgs_for_llm([i.model_dump() for i in request.input])

        # Stream from your existing agent
        item_id = str(uuid4())
        aggregated_stream = ""
        for chunk in self.agent.stream(messages):
            # Convert each chunk to ResponsesAgent format
            yield self.create_text_delta(delta=chunk, item_id=item_id)
            aggregated_stream += chunk

        # Emit an aggregated output_item for all the text deltas with id=item_id
        yield ResponsesAgentStreamEvent(
            type="response.output_item.done",
            item=self.create_text_output_item(text=aggregated_stream, id=item_id),
        )

Migrieren von ChatCompletions

Wenn Ihr vorhandener Agent die OpenAI ChatCompletions-API verwendet, können Sie ihn migrieren, ResponsesAgent ohne seine Kernlogik neu zu schreiben. Fügen Sie einen Wrapper hinzu, der:

1. Konvertiert eingehende ResponsesAgentRequest Nachrichten in das ChatCompletions von Ihrem Agent erwartete Format.
2. ChatCompletions Übersetzt Ausgaben in das ResponsesAgentResponse Schema.
3. Unterstützt das Streaming optional, indem inkrementelle Deltas von ChatCompletions objekten zugeordnet werden ResponsesAgentStreamEvent .

from typing import Generator
from uuid import uuid4

from databricks.sdk import WorkspaceClient
from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
    ResponsesAgentStreamEvent,
)


# Legacy agent that outputs ChatCompletions objects
class LegacyAgent:
    def __init__(self):
        self.w = WorkspaceClient()
        self.OpenAI = self.w.serving_endpoints.get_open_ai_client()

    def stream(self, messages):
        for chunk in self.OpenAI.chat.completions.create(
            model="databricks-claude-sonnet-4-5",
            messages=messages,
            stream=True,
        ):
            yield chunk.to_dict()


# Wrapper that converts the legacy agent to a ResponsesAgent
class MyWrappedStreamingAgent(ResponsesAgent):
    def __init__(self, agent):
        # `agent` is your existing ChatCompletions agent
        self.agent = agent

    def prep_msgs_for_llm(self, messages):
        # dummy example of prep_msgs_for_llm
        # real example of prep_msgs_for_llm included in full examples linked below
        return [{"role": "user", "content": "Hello, how are you?"}]

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        """Non-streaming predict: collects all streaming chunks into a single response."""
        # Reuse the streaming logic and collect all output items
        output_items = []
        for stream_event in self.predict_stream(request):
            if stream_event.type == "response.output_item.done":
                output_items.append(stream_event.item)

        # Return all collected items as a single response
        return ResponsesAgentResponse(output=output_items)

    def predict_stream(
        self, request: ResponsesAgentRequest
    ) -> Generator[ResponsesAgentStreamEvent, None, None]:
        """Streaming predict: the core logic that both methods use."""
        # Convert incoming messages to your agent's format
        messages = self.prep_msgs_for_llm([i.model_dump() for i in request.input])

        # process the ChatCompletion output stream
        agent_content = ""
        tool_calls = []
        msg_id = None
        for chunk in self.agent.stream(messages):  # call the underlying agent's stream method
            delta = chunk["choices"][0]["delta"]
            msg_id = chunk.get("id", None)
            content = delta.get("content", None)
            if tc := delta.get("tool_calls"):
                if not tool_calls:  # only accommodate for single tool call right now
                    tool_calls = tc
                else:
                    tool_calls[0]["function"]["arguments"] += tc[0]["function"]["arguments"]
            elif content is not None:
                agent_content += content
                yield ResponsesAgentStreamEvent(**self.create_text_delta(content, item_id=msg_id))

        # aggregate the streamed text content
        yield ResponsesAgentStreamEvent(
            type="response.output_item.done",
            item=self.create_text_output_item(agent_content, msg_id),
        )

        for tool_call in tool_calls:
            yield ResponsesAgentStreamEvent(
                type="response.output_item.done",
                item=self.create_function_call_item(
                    str(uuid4()),
                    tool_call["id"],
                    tool_call["function"]["name"],
                    tool_call["function"]["arguments"],
                ),
            )


agent = MyWrappedStreamingAgent(LegacyAgent())

for chunk in agent.predict_stream(
    ResponsesAgentRequest(input=[{"role": "user", "content": "Hello, how are you?"}])
):
    print(chunk)

Vollständige Beispiele finden Sie unter ResponsesAgent Beispielen.

Streaming-Antworten

Streaming ermöglicht Es Agents, Antworten in Echtzeitblöcken zu senden, anstatt auf die vollständige Antwort zu warten. Um Streaming mit ResponsesAgentzu implementieren, geben Sie eine Reihe von Delta-Ereignissen aus, gefolgt von einem endgültigen Abschlussereignis:

1. Übertragen von Delta-Ereignissen: Senden Sie mehrere output_text.delta Ereignisse mit demselben item_id, um Textblöcke in Echtzeit zu streamen.
2. Fertig stellen: Senden Sie ein endgültiges response.output_item.done Ereignis mit demselben item_id wie die Delta-Ereignisse, die den vollständigen endgültigen Ausgabetext enthalten.

Jedes Delta-Ereignis streamt einen Textabschnitt an den Client. Das letzte fertige Ereignis enthält den vollständigen Antworttext und signalisiert Databricks folgendes:

• Verfolgen Sie die Ergebnisse Ihres Agents mit der MLflow-Ablaufverfolgung
• Aggregierte gestreamte Antworten in KI-Gateway-Ableitungstabellen
• Zeige die vollständige Ausgabe in der AI Playground-Oberfläche an

Streamingfehlerausbreitung

Mosaic AI propagiert alle Fehler, die beim Streaming mit dem letzten Token unter databricks_output.error auftreten. Es liegt bei dem aufrufenden Client, diesen Fehler ordnungsgemäß zu behandeln und anzuzeigen.

{
  "delta": …,
  "databricks_output": {
    "trace": {...},
    "error": {
      "error_code": BAD_REQUEST,
      "message": "TimeoutException: Tool XYZ failed to execute."
    }
  }
}

Erweiterte Funktionen

Benutzerdefinierte Eingaben und Ausgaben

Einige Szenarien erfordern möglicherweise zusätzliche Eingaben von Agenten, wie z. B. client_type und session_id, oder Ausgaben wie Abrufquellenlinks, die nicht in den Chatverlauf für zukünftige Interaktionen einbezogen werden sollten.

Für diese Szenarien unterstützt MLflow ResponsesAgent die Felder custom_inputs und custom_outputsnativ. Sie können die benutzerdefinierten Eingaben über request.custom_inputs in allen oben in "ResponsesAgent Examples" verknüpften Beispielen abrufen.

In den folgenden Notizbüchern erfahren Sie, wie Sie benutzerdefinierte Eingaben und Ausgaben festlegen.

Bereitstellen custom_inputs in der AI-Playground- und Überprüfungs-App

Wenn Ihr Agent zusätzliche Eingaben mithilfe des custom_inputs Felds akzeptiert, können Sie diese Eingaben sowohl im AI-Playground als auch in der Rezensions-App manuell bereitstellen.

1. Wählen Sie in der AI Playground- oder in der Agent Review-App das Zahnradsymbol aus.

2. Aktivieren Sie custom_inputs.

3. Stellen Sie ein JSON-Objekt bereit, das dem definierten Eingabeschema Ihres Agents entspricht.

   

Benutzerdefinierte Retrieverschemas angeben

KI-Agents verwenden häufig Retriever, um unstrukturierte Daten aus Vektorsuchindizes zu finden und abzufragen. Beispiele für Retriever-Tools finden Sie unter Retriever-Tools für unstrukturierte Daten erstellen und tracken .

Verfolgen Sie diese Retriever in Ihrem Agent mit MLflow RETRIEVER, um Databricks-Produktfeatures zu aktivieren, einschließlich:

• Automatisches Anzeigen von Links zu abgerufenen Quelldokumenten in der AI Playground-Benutzeroberfläche
• Automatisierte Ausführung der Abrufübereinstimmung und der Relevanzbewertung in der Agentenbewertung

Wenn Ihr Agent Retriever-Spanne mit einem benutzerdefinierten Schema umfasst, rufen Sie auf mlflow.models.set_retriever_schema , wenn Sie Ihren Agent im Code definieren. Dadurch werden die Ausgabespalten des Retrievers den erwarteten Feldern von MLflow (primary_key, text_column, doc_uri) zugeordnet.

import mlflow
# Define the retriever's schema by providing your column names
# For example, the following call specifies the schema of a retriever that returns a list of objects like
# [
#     {
#         'document_id': '9a8292da3a9d4005a988bf0bfdd0024c',
#         'chunk_text': 'MLflow is an open-source platform, purpose-built to assist machine learning practitioners...',
#         'doc_uri': 'https://mlflow.org/docs/latest/index.html',
#         'title': 'MLflow: A Tool for Managing the Machine Learning Lifecycle'
#     },
#     {
#         'document_id': '7537fe93c97f4fdb9867412e9c1f9e5b',
#         'chunk_text': 'A great way to get started with MLflow is to use the autologging feature. Autologging automatically logs your model...',
#         'doc_uri': 'https://mlflow.org/docs/latest/getting-started/',
#         'title': 'Getting Started with MLflow'
#     },
# ...
# ]
mlflow.models.set_retriever_schema(
    # Specify the name of your retriever span
    name="mlflow_docs_vector_search",
    # Specify the output column name to treat as the primary key (ID) of each retrieved document
    primary_key="document_id",
    # Specify the output column name to treat as the text content (page content) of each retrieved document
    text_column="chunk_text",
    # Specify the output column name to treat as the document URI of each retrieved document
    doc_uri="doc_uri",
    # Specify any other columns returned by the retriever
    other_columns=["title"],
)

Bereitstellungsüberlegungen

Vorbereiten auf das Databricks-Modellserving

Databricks setzt ResponsesAgent in einer verteilten Umgebung für das Bereitstellen von Modellen in Databricks Model Serving ein. Dies bedeutet, dass bei einem mehrstufigen Gespräch möglicherweise nicht dieselbe Replik alle Anfragen verarbeitet. Beachten Sie die folgenden Auswirkungen auf die Verwaltung des Agent-Status:

• Vermeiden Sie die lokale Zwischenspeicherung: Bei der Bereitstellung eines ResponsesAgent sollten Sie nicht davon ausgehen, dass dasselbe Replikat alle Anforderungen in einer Multi-Turn-Unterhaltung verarbeitet. Rekonstruieren Sie den internen Zustand mithilfe eines Schemas für ResponsesAgentRequest-Wörterbücher für jeden Teil.

• Threadsicherer Zustand: Entwurfs-Agent-Zustand, der threadsicher sein soll und Konflikte in Multithread-Umgebungen verhindert.

• Zustand in der predict-Funktion initialisieren: Initialisieren Sie den Zustand jedes Mal, wenn die predict-Funktion aufgerufen wird, nicht während der Initialisierung von ResponsesAgent. Das Speichern des Zustands auf der ResponsesAgent Ebene könnte Informationen zwischen Unterhaltungen leaken und Konflikte verursachen, da ein einzelnes ResponsesAgent-Replikat Anforderungen aus mehreren Unterhaltungen verarbeiten kann.

Parametrisieren von Code für die Bereitstellung in allen Umgebungen

Parametrisieren Sie Agentcode, um denselben Agentcode in verschiedenen Umgebungen wiederzuverwenden.

Parameter sind Schlüsselwertpaare, die Sie in einem Python-Wörterbuch oder einer .yaml Datei definieren.

Um den Code zu konfigurieren, erstellen Sie ein ModelConfig Python-Wörterbuch oder eine .yaml Datei. ModelConfig ist eine Reihe von Schlüsselwertparametern, die eine flexible Konfigurationsverwaltung ermöglichen. Sie können z. B. während der Entwicklung ein Wörterbuch verwenden und dann in eine .yaml Datei für die Produktionsbereitstellung und CI/CD konvertieren.

Ein Beispiel ModelConfig wird unten gezeigt:

llm_parameters:
  max_tokens: 500
  temperature: 0.01
model_serving_endpoint: databricks-meta-llama-3-3-70b-instruct
vector_search_index: ml.docs.databricks_docs_index
prompt_template: 'You are a hello world bot. Respond with a reply to the user''s
  question that indicates your prompt template came from a YAML file. Your response
  must use the word "YAML" somewhere. User''s question: {question}'
prompt_template_input_vars:
  - question

Sie können im Agent-Code über die Datei .yaml oder ein Wörterbuch auf eine Standardkonfiguration (für die Entwicklung) verweisen:

import mlflow
# Example for loading from a .yml file
config_file = "configs/hello_world_config.yml"
model_config = mlflow.models.ModelConfig(development_config=config_file)

# Example of using a dictionary
config_dict = {
    "prompt_template": "You are a hello world bot. Respond with a reply to the user's question that is fun and interesting to the user. User's question: {question}",
    "prompt_template_input_vars": ["question"],
    "model_serving_endpoint": "databricks-meta-llama-3-3-70b-instruct",
    "llm_parameters": {"temperature": 0.01, "max_tokens": 500},
}

model_config = mlflow.models.ModelConfig(development_config=config_dict)

# Use model_config.get() to retrieve a parameter value
# You can also use model_config.to_dict() to convert the loaded config object
# into a dictionary
value = model_config.get('sample_param')

Geben Sie dann beim Protokollieren des Agents den model_config Parameter als log_model an, um einen benutzerdefinierten Satz von Parametern anzugeben, die beim Laden des protokollierten Agents verwendet werden sollen. Weitere Informationen finden Sie in der MLflow-Dokumentation – ModelConfig.

Verwenden von synchronem Code oder Callback-Mustern

Um Stabilität und Kompatibilität zu gewährleisten, verwenden Sie synchronen Code oder rückrufbasierte Muster in Ihrer Agentimplementierung.

Azure Databricks verwaltet automatisch die asynchrone Kommunikation, um eine optimale Parallelität und Leistung zu gewährleisten, wenn Sie einen Agent bereitstellen. Die Einführung benutzerdefinierter Ereignisschleifen oder asynchroner Frameworks kann zu Fehlern wie RuntimeError: This event loop is already running and caused unpredictable behavior führen.

Azure Databricks empfiehlt, beim Entwickeln von Agents asynchrone Programmierung zu vermeiden, z. B. die Verwendung von asyncio oder das Erstellen von benutzerdefinierten Ereignisschleifen.

Nächste Schritte

• Erstellen Sie Ihre eigenen Agent-Tools.
• Protokollieren eines KI-Agents.
• Hinzufügen von Ablaufverfolgungen zu einem KI-Agent
• Bereitstellen eines KI-Agents.