Creare agenti di intelligenza artificiale nel codice

Questo articolo illustra come creare un agente di intelligenza artificiale nel codice usando MLflow ChatModel. Azure Databricks sfrutta mlflow ChatModel per garantire la compatibilità con le funzionalità dell'agente di intelligenza artificiale di Databricks, ad esempio valutazione, traccia e distribuzione.

Che cos'è ChatModel?

ChatModel è una classe MLflow progettata per semplificare la creazione di agenti di intelligenza artificiale conversazionale. Fornisce un'interfaccia standardizzata per la creazione di modelli compatibili con 'API ChatCompletion di OpenAI.

ChatModel estende lo schema ChatCompletion di OpenAI. Questo approccio consente di mantenere un'ampia compatibilità con le piattaforme che supportano lo standard ChatCompletion, aggiungendo al tempo stesso funzionalità personalizzate.

Usando ChatModel, gli sviluppatori possono creare agenti compatibili con gli strumenti Databricks e MLflow per la gestione del ciclo di vita, monitoraggio, e valutazione degli agenti, essenziali per la distribuzione di modelli pronti per la produzione.

Consultare MLflow: Introduzione a ChatModel.

Requisiti

Databricks consiglia di installare la versione più recente del client Python MLflow durante lo sviluppo di agenti.

Per creare e distribuire agenti usando l'approccio in questo articolo, è necessario soddisfare i requisiti seguenti:

• Installare databricks-agents versione 0.15.0 e successive
• Installare mlflow versione 2.20.0 e successive

%pip install -U -qqqq databricks-agents>=0.15.0 mlflow>=2.20.0

Creare un agente ChatModel

È possibile creare l'agente come sottoclasse di mlflow.pyfunc.ChatModel. Questo metodo offre i vantaggi seguenti:

• Consente di scrivere codice agente compatibile con lo schema ChatCompletion usando classi Python tipate.
• MLflow dedurrà automaticamente una firma compatibile con il completamento della chat durante la registrazione dell'agente, anche senza un input_example. Questo semplifica il processo di registrazione e distribuzione dell'agente. Consultare Inferire la firma del modello durante il processo di logging.

Il codice seguente viene eseguito in modo ottimale in un notebook di Databricks. I notebook offrono un ambiente pratico per lo sviluppo, il test e l'iterazione del tuo agente.

La classe MyAgent estende mlflow.pyfunc.ChatModel, implementando il metodo predict necessario. Ciò garantisce la compatibilità con Mosaic AI Agent Framework.

La classe include anche i metodi facoltativi _create_chat_completion_chunk e predict_stream per gestire gli output di streaming.

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)

Mentre la classe agente MyAgent è definita in un unico notebook, è necessario creare un notebook del driver separato. Il notebook del driver registra l'agente nel Registro modelli e distribuisce l'agente usando Model Serving.

Questa separazione segue il flusso di lavoro raccomandato da Databricks per la registrazione dei modelli mediante i modelli di tracciatura di MLflow secondo la metodologia "Models from Code".

Esempio: Eseguire il wrapping di LangChain in ChatModel

Se si dispone di un modello LangChain esistente e si vuole integrarlo con altre funzionalità dell'agente di Mosaic AI, è possibile eseguirne il wrapping in un ChatModel MLflow per garantire la compatibilità.

Questo esempio di codice esegue i passaggi seguenti per eseguire il wrapping di un elemento LangChain eseguibile come ChatModel:

1. Eseguire il wrapping dell'output finale di LangChain con mlflow.langchain.output_parsers.ChatCompletionOutputParser per produrre una firma di output di completamento della chat
2. La classe LangchainAgent estende mlflow.pyfunc.ChatModel e implementa due metodi chiave:
   • predict: gestisce le stime sincrone richiamando la catena e restituendo una risposta formattata.
   • predict_stream: gestisce le predizioni di streaming richiamando la catena e generando segmenti di risposte.

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)

Usare i parametri per configurare l'agente

In Agent Framework è possibile usare i parametri per controllare la modalità di esecuzione degli agenti. In questo modo è possibile eseguire rapidamente l'iterazione in base a diverse caratteristiche dell'agente senza modificare il codice. I parametri sono coppie chiave-valore definite in un dizionario Python o in un file di .yaml.

Per configurare il codice, creare un ModelConfig, un set di parametri chiave-valore. ModelConfig è un dizionario Python o un file .yaml. Ad esempio, è possibile usare un dizionario durante lo sviluppo e convertirlo in un file .yaml per la distribuzione di produzione e CI/CD. Per informazioni dettagliate su ModelConfig, vedere la documentazione MLflow.

Di seguito è riportato un esempio di ModelConfig.

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

Per chiamare la configurazione dal codice, usare una delle opzioni seguenti:

# 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')

Impostare lo schema del retriever

Gli agenti di intelligenza artificiale usano spesso i retriever, un tipo di strumento agente che trova e restituisce documenti pertinenti usando un indice di ricerca vettoriale. Per ulteriori informazioni sui sistemi di recupero, vedere Strumenti non strutturati per agenti di intelligenza artificiale.

Per assicurarsi che i retriever vengano tracciati correttamente, chiamare mlflow.models.set_retriever_schema quando si definisce l'agente nel codice. Usare set_retriever_schema per eseguire il mapping dei nomi delle colonne nella tabella restituita ai campi previsti di MLflow, ad esempio primary_key, text_columne 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"],
)

Nota

La colonna doc_uri è particolarmente importante quando si valutano le prestazioni del retriever. doc_uri è l'identificatore principale per i documenti restituiti dal retriever, consentendo di confrontarli con i set di valutazione della verità di base. Vedere Set di valutazione.

È anche possibile specificare colonne aggiuntive nello schema del retriever fornendo un elenco di nomi di colonna con il campo other_columns.

Se si dispone di più retriever, è possibile definire più schemi usando nomi univoci per ogni schema di recupero.

input e output personalizzati

Alcuni scenari possono richiedere input aggiuntivi da parte degli agenti, come client_type e session_id, oppure output come i link delle fonti di recupero che non devono essere inclusi nella cronologia della chat per le interazioni future.

Per questi scenari, MLflow ChatModel supportano in modo nativo l'aumento delle richieste e delle risposte di completamento delle chat OpenAI con i campi ChatParamscustom_input e custom_output.

Vedere gli esempi seguenti per informazioni su come creare input e output personalizzati per gli agenti PyFunc e LangGraph.

Avvertimento

L'app Valutazione agente attualmente non supporta il rendering delle tracce per gli agenti con campi di input aggiuntivi.

Schemi personalizzati PyFunc

I notebook seguenti illustrano un esempio di schema personalizzato usando PyFunc.

Notebook dell'agente dello schema personalizzato PyFunc

Prendi il notebook

Notebook del driver dello schema personalizzato PyFunc

Prendi il notebook

schemi personalizzati LangGraph

I notebook seguenti illustrano un esempio di schema personalizzato usando LangGraph. È possibile modificare la funzione wrap_output nei notebook per analizzare ed estrarre informazioni dal flusso di messaggi.

Notebook dell'agente dello schema personalizzato LangGraph

Prendi il notebook

Notebook del driver di schema personalizzato di LangGraph

Prendi il notebook

Fornire custom_inputs nel Playground per l'intelligenza artificiale e nell'app di revisione degli agenti

Se l'agente accetta input aggiuntivi usando il campo custom_inputs, è possibile specificare manualmente questi input sia nel Playground AI che nell'app di revisione dell'agente .

1. In AI Playground o nell'app di revisione agente, selezionare l'icona a forma di ingranaggio .

2. Abilitare custom_inputs.

3. Fornire un oggetto JSON che corrisponda allo schema di input definito dell'agente.

   

Propagazione degli errori di streaming

La propagazione di Mosaic AI di qualsiasi errore riscontrato durante lo streaming con l'ultimo token sotto databricks_output.error. Spetta al client chiamante gestire e individuare correttamente questo errore.

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

Notebook di esempio

Questi notebook creano una semplice catena "Hello, world" per illustrare la creazione di un agente in Databricks. Il primo esempio crea una catena semplice e il secondo notebook di esempio illustra come usare i parametri per ridurre al minimo le modifiche al codice durante lo sviluppo.

Notebook con catena semplice

Prendi il notebook

Quaderno sui cambi a catena semplice

Prendi il notebook

Notebook con catena con parametri

Prendi il notebook

Notebook per driver parametrizzato di concatenamento

Prendi il notebook

Passaggi successivi

• Creare strumenti dell'agente personalizzati.
• Registrare un agente di intelligenza artificiale.
• Aggiungere tracce a un agente di intelligenza artificiale.
• Distribuire un agente di intelligenza artificiale.