Skapa AI-agenter i kod
Den här artikeln visar hur du skapar en AI-agent i kod med hjälp av MLflow ChatModel. Azure Databricks använder MLflow-ChatModel för att säkerställa kompatibilitet med Databricks AI-agentfunktioner som utvärdering, spårning och distribution.
Vad är ChatModel?
ChatModel är en MLflow-klass som är utformad för att förenkla skapandet av konversations-AI-agenter. Det tillhandahåller ett standardiserat gränssnitt för att skapa modeller som är kompatibla med OpenAI:s ChatCompletion API-.
ChatModel utökar OpenAI:s ChatCompletion-schema. Med den här metoden kan du upprätthålla bred kompatibilitet med plattformar som stöder ChatCompletion-standarden, samtidigt som du lägger till dina egna anpassade funktioner.
Med hjälp av ChatModelkan utvecklare skapa agenter som är kompatibla med Databricks- och MLflow-verktyg för agentspårning, utvärdering och livscykelhantering, vilket är viktigt för distribution av produktionsklara modeller.
Se MLflow: Komma igång med ChatModel.
Krav
Databricks rekommenderar att du installerar den senaste versionen av MLflow Python-klienten när du utvecklar agenter.
Om du vill skapa och distribuera agenter med hjälp av metoden i den här artikeln måste du uppfylla följande krav:
- Installera
databricks-agentsversion 0.15.0 och senare - Installera
mlflowversion 2.20.0 och senare
%pip install -U -qqqq databricks-agents>=0.15.0 mlflow>=2.20.0
Skapa en ChatModel-agent
Du kan skapa din agent som en underklass av mlflow.pyfunc.ChatModel. Den här metoden ger följande fördelar:
- Gör att du kan skriva agentkod som är kompatibel med ChatCompletion-schemat med hjälp av inskrivna Python-klasser.
- MLflow kommer automatiskt att identifiera en signatur som är kompatibel med chattkomplettering när agenten loggas, även utan en
input_example. Detta förenklar registreringen och distributionen av agenten. Se Fastställ modellsignatur under loggning.
Följande kod körs bäst i en Databricks-notebook-fil. Notebook-filer ger en bekväm miljö för att utveckla, testa och iterera på din agent.
Klassen MyAgent utökar mlflow.pyfunc.ChatModeloch implementerar den nödvändiga predict-metoden. Detta säkerställer kompatibilitet med Mosaic AI Agent Framework.
Klassen innehåller även de valfria metoderna _create_chat_completion_chunk och predict_stream för att hantera strömmande utdata.
from dataclasses import dataclass
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 last_user_question_text.split(" "):
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, model_input=model_input)
print(response)
Medan agentklassen MyAgent definieras i en notebook-fil bör du skapa en separat drivrutinsanteckningsbok. Drivrutinsanteckningsboken loggar agenten till Model Registry och distribuerar agenten med hjälp av modellservern.
Den här separationen följer Databricks rekommenderade arbetsflöde för loggningsmodeller med MLflows modeller från kodmetoden.
Exempel: Wrap LangChain i ChatModel
Om du har en befintlig LangChain-modell och vill integrera den med andra Mosaic AI-agentfunktioner kan du omsluta den i ett MLflow-ChatModel för att säkerställa kompatibilitet.
Det här kodexemplet utför följande steg för att inkapsla en LangChain-runnable som en ChatModel:
- Omslut LangChains slutresultat med
mlflow.langchain.output_parsers.ChatCompletionOutputParserför att skapa ett slutförandemönster för chatten - Klassen
LangchainAgentutökarmlflow.pyfunc.ChatModeloch implementerar två viktiga metoder:-
predict: Hanterar synkrona förutsägelser genom att anropa kedjan och returnera ett formaterat svar. -
predict_stream: Hanterar strömningsförutsägelser genom att anropa kedjan och ge svarssegment.
-
from mlflow.langchain.output_parsers import ChatCompletionOutputParser
from mlflow.pyfunc import ChatModel
from typing import Optional, Dict, List, Generator
from mlflow.types.llm import (
ChatCompletionResponse,
ChatCompletionChunk
)
chain = (
<your chain here>
| ChatCompletionOutputParser()
)
class LangchainAgent(ChatModel):
def _prepare_messages(self, messages: List[ChatMessage]):
return {"messages": [m.to_dict() for m in messages]}
def predict(
self, context, messages: List[ChatMessage], params: ChatParams
) -> ChatCompletionResponse:
question = self._prepare_messages(messages)
response_message = self.chain.invoke(question)
return ChatCompletionResponse.from_dict(response_message)
def predict_stream(
self, context, messages: List[ChatMessage], params: ChatParams
) -> Generator[ChatCompletionChunk, None, None]:
question = self._prepare_messages(messages)
for chunk in chain.stream(question):
yield ChatCompletionChunk.from_dict(chunk)
Använd parametrar för att konfigurera agenten
I Agent Framework kan du använda parametrar för att styra hur agenter körs. På så sätt kan du snabbt iterera med olika egenskaper hos din agent utan att ändra koden. Parametrar är nyckel/värde-par som du definierar i en Python-ordlista eller en .yaml fil.
Om du vill konfigurera koden skapar du en ModelConfig, en uppsättning nyckel/värde-parametrar.
ModelConfig är antingen en Python-ordlista eller en .yaml fil. Du kan till exempel använda en ordlista under utvecklingen och sedan konvertera den till en .yaml fil för produktionsdistribution och CI/CD. Mer information om ModelConfigfinns i MLflow-dokumentationen.
Ett exempel ModelConfig visas nedan.
llm_parameters:
max_tokens: 500
temperature: 0.01
model_serving_endpoint: databricks-dbrx-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
Om du vill anropa konfigurationen från koden använder du något av följande:
# 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-dbrx-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
value = model_config.get('sample_param')
Ange retriever-schema
AI-agenter använder ofta hämtningar, en typ av agentverktyg som hittar och returnerar relevanta dokument med hjälp av ett Vector Search-index. Mer information om hämtningar finns i AI-agentverktyg för ostrukturerad hämtning.
För att säkerställa att hämtningarna spåras korrekt anropar du mlflow.models.set_retriever_schema när du definierar din agent i kod. Använd set_retriever_schema för att mappa kolumnnamnen i den returnerade tabellen till MLflows förväntade fält, till exempel primary_key, text_columnoch doc_uri.
# Define the retriever's schema by providing your column names
# These strings should be read from a config dictionary
mlflow.models.set_retriever_schema(
name="vector_search",
primary_key="chunk_id",
text_column="text_column",
doc_uri="doc_uri"
# other_columns=["column1", "column2"],
)
Not
Kolumnen doc_uri är särskilt viktig när du utvärderar hämtarens prestanda.
doc_uri är huvudidentifieraren för dokument som returneras av hämtaren, så att du kan jämföra dem med utvärderingsuppsättningar för grundsanning. Se Utvärderingsuppsättningar.
Du kan också ange ytterligare kolumner i hämtarens schema genom att ange en lista med kolumnnamn med fältet other_columns.
Om du har flera sökmotorer kan du definiera flera scheman genom att använda unika namn för varje schema.
Anpassade indata och utdata
Vissa scenarier kan kräva ytterligare agentindata, till exempel client_type och session_id, eller utdata som hämtning av källlänkar som inte ska ingå i chatthistoriken för framtida interaktioner.
I de här scenarierna stöder MLflow ChatModel internt utökade förfrågningar och svar om slutförande av OpenAI-chattar med fälten ChatParamscustom_input och custom_output.
Se följande exempel för att lära dig hur du skapar anpassade indata och utdata för PyFunc- och LangGraph-agenter.
Varning
Granskningsappen Agentutvärdering stöder för närvarande inte återgivning av spårningar för agenter med ytterligare indatafält.
Anpassade PyFunc-scheman
Följande notebook-filer visar ett anpassat schemaexempel med PyFunc.
PyFunc- anpassad schemaagentanteckningsbok
PyFunc anpassad schemadrivrutinsanteckningsbok
anpassade scheman för LangGraph
Följande notebook-filer visar ett anpassat schemaexempel med LangGraph. 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 för LangGraph
Anpassad schemadrivrutinsanteckningsbok för LangGraph
Tillhandahålla custom_inputs i AI Playground- och agentgranskningsappen
Om din agent accepterar ytterligare indata med hjälp av fältet custom_inputs kan du manuellt ange dessa indata i både AI Playground- och agentgranskningsappen.
I antingen AI Playground eller agentgranskningsappen, välj kugghjulsikonen
.Aktivera anpassade_indata.
Ange ett JSON-objekt som matchar agentens definierade indataschema.
Spridning av direktuppspelningsfel
Mosaic AI vidarebefordrar eventuella fel som uppstår vid strömning med sista tokenen under databricks_output.error. Det är upp till den anropande klienten att korrekt hantera och visa felet.
{
"delta": …,
"databricks_output": {
"trace": {...},
"error": {
"error_code": BAD_REQUEST,
"message": "TimeoutException: Tool XYZ failed to execute"
}
}
}
Exempel på notebook-filer
Dessa notebook-filer skapar en enkel "Hello, world"-kedja för att illustrera skapandet av en agent i Databricks. Det första exemplet skapar en enkel kedja, och det andra exemplet på notebook-filen visar hur du använder parametrar för att minimera kodändringar under utvecklingen.