Logga och registrera AI-agenter

Artikel
12/21/2024

Viktigt!

Den här funktionen finns som allmänt tillgänglig förhandsversion.

Logga AI-agenter med Mosaic AI Agent Framework. Loggning av en agent är grunden för utvecklingsprocessen. Loggning registrerar en "tidpunkt" för agentens kod och konfiguration så att du kan utvärdera kvaliteten på konfigurationen.

Krav

Skapa en AI-agent innan du loggar den.

Kodbaserad loggning

Databricks rekommenderar att du använder MLflows -modeller från kodfunktioner när du loggar agenter.

I den här metoden registreras agentens kod som en Python-fil och Python-miljön registreras som en lista över paket. När agenten distribueras återställs Python-miljön och agentens kod körs för att läsa in agenten i minnet så att den kan anropas när slutpunkten anropas.

Du kan koppla den här metoden med användning av API:er för validering före distribution som mlflow.models.predict() för att säkerställa att agenten körs tillförlitligt när den distribueras för servering.

Koden som loggar agenten måste finnas i en separat anteckningsbok från agentkoden. Den här notebook-filen kallas för en drivrutinsanteckningsbok. Ett exempel på notebook-fil finns i Exempel på notebook-filer.

Utläs modellsignatur vid loggning

Under loggningen måste du definiera en MLflow-modellsignatur, som anger agentens indata- och utdataschema. Signaturen validerar indata och utdata för att säkerställa att agenten interagerar korrekt med underordnade verktyg som AI Playground och granskningsappen. Den vägleder även andra program om hur du använder agenten effektivt.

Databricks rekommenderar att du använder MLflows inferensfunktioner för modellsignatur för att automatiskt generera agentens signatur baserat på ett indataexempel som du anger. Den här metoden är enklare än att definiera signaturen manuellt.

Exemplen LangChain och PyFunc nedan använder modellsignaturinferens.

Om du hellre uttryckligen definierar en modellsignatur själv vid loggningstillfället kan du läsa MLflow-dokument – Så här loggar du modeller med signaturer.

Kodbaserad loggning med LangChain

Följande instruktioner och kodexempel visar hur du loggar en agent med LangChain.

Skapa en notebook- eller Python-fil med din kod. I det här exemplet heter anteckningsboken eller filen agent.py. Anteckningsboken eller filen måste innehålla en LangChain-agent som här kallas lc_agent.
Inkludera mlflow.models.set_model(lc_agent) i anteckningsboken eller filen.
Skapa en ny notebook-fil som ska fungera som drivrutinsanteckningsbok (kallas driver.py i det här exemplet).
I drivrutinsanteckningsboken använder du följande kod för att köra agent.py och logga resultatet till en MLflow-modell:
```
mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
Parametern resources deklarerar Databricks-hanterade resurser som behövs för att hantera agenten, till exempel ett vektorsökningsindex eller en serverslutpunkt som fungerar som grundmodell. Mer information finns i Ange resurser för automatiserad autentiseringsgenomgång.
Distribuera modellen. Se Distribuera en agent för generativ AI-program.
När serveringsmiljön agent.py läses in körs den.
När en serveringsbegäran kommer in lc_agent.invoke(...) anropas.


import mlflow

code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"

# Input example used by MLflow to infer Model Signature
input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using langchain
with mlflow.start_run():
  logged_agent_info = mlflow.langchain.log_model(
    lc_model=code_path,
    model_config=config_path, # If you specify this parameter, this configuration is used by agent code. The development_config is overwritten.
    artifact_path="agent", # This string is used as the path inside the MLflow model where artifacts are stored
    input_example=input_example, # Must be a valid input to the agent
    example_no_conversion=True, # Required
  )

print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")

# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.langchain.load_model(logged_agent_info.model_uri)
model.invoke(example)

kodbaserad loggning med PyFunc

Följande instruktioner och kodexempel visar hur du loggar en agent med PyFunc.

Skapa en notebook- eller Python-fil med din kod. I det här exemplet heter anteckningsboken eller filen agent.py. Anteckningsboken eller filen måste innehålla en PyFunc-klass med namnet PyFuncClass.
Inkludera mlflow.models.set_model(PyFuncClass) i anteckningsboken eller filen.
Skapa en ny notebook-fil som ska fungera som drivrutinsanteckningsbok (kallas driver.py i det här exemplet).
I drivrutinsanteckningsboken använder du följande kod för att köra agent.py och logga resultatet till en MLflow-modell:
```
mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
Parametern resources deklarerar Databricks-hanterade resurser som behövs för att hantera agenten, till exempel ett vektorsökningsindex eller en serverslutpunkt som fungerar som grundmodell. Mer information finns i Ange resurser för automatiserad autentiseringsgenomgång.
Distribuera modellen. Se Distribuera en agent för generativ AI-program.
När serveringsmiljön agent.py läses in körs den.
När en serveringsbegäran kommer in PyFuncClass.predict(...) anropas.

import mlflow
from mlflow.models.resources import (
    DatabricksServingEndpoint,
    DatabricksVectorSearchIndex,
)

code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"

# Input example used by MLflow to infer Model Signature
input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model=agent_notebook_path,
    artifact_path="agent",
    input_example=input_example,
    resources=resources_path,
    example_no_conversion=True,
    resources=[
      DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
    ]
  )

print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")

# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.pyfunc.load_model(logged_agent_info.model_uri)
model.invoke(example)

Ange resurser för automatiskt godkännande vid autentisering

AI-agenter behöver ofta autentisera till andra resurser för att slutföra uppgifter. En agent kan till exempel behöva komma åt ett Vector Search-index för att köra frågor mot ostrukturerade data.

Enligt beskrivningen i Autentisering för beroende resurserstöder modellservern autentisering till både Databricks-hanterade och externa resurser när du distribuerar agenten.

För de vanligaste Databricks-resurstyperna har Databricks stöd för och rekommenderar att du deklarerar resursberoenden för agenten i förväg under loggningen. Detta möjliggör automatisk autentiseringsöverföring när du distribuerar agenten – Databricks hanterar automatiskt etablering, rotation och hantering av kortlivade autentiseringsuppgifter för säker åtkomst till dessa resursberoenden från agentens slutpunkt.

Om du vill aktivera automatisk autentiseringsgenomströmning anger du beroende resurser med hjälp av resources-parametern för API:et log_model() enligt följande kod.

import mlflow
from mlflow.models.resources import (
    DatabricksVectorSearchIndex,
    DatabricksServingEndpoint,
    DatabricksSQLWarehouse,
    DatabricksFunction,
    DatabricksGenieSpace,
    DatabricksTable,
)

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model=agent_notebook_path,
    artifact_path="agent",
    input_example=input_example,
    example_no_conversion=True,
    # Specify resources for automatic authentication passthrough
    resources=[
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
      DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
      DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
      DatabricksSQLWarehouse(warehouse_id="your_warehouse_id"),
      DatabricksFunction(function_name="ml.tools.python_exec"),
      DatabricksGenieSpace(genie_space_id="your_genie_space_id"),
      DatabricksTable(table_name="your_table_name"),
    ]
  )

Databricks rekommenderar att du specificerar resources för alla agent-smaker manuellt.

Kommentar

Om du inte anger resurser när du loggar LangChain-agenter med mlflow.langchain.log_model(...)utför MLflow bästa möjliga automatiska slutsatsdragning av resurser. Detta kanske inte samlar in alla beroenden, vilket kan leda till auktoriseringsfel när agenten används för servering eller förfrågningar.

I följande tabell visas de Databricks-resurser som stöder automatisk autentiseringsgenomströmning och den lägsta mlflow version som krävs för att logga resursen.

Resurstyp	Lägsta `mlflow` version som krävs för att logga resursen
Index för vektorsökning	Kräver `mlflow` 2.13.1 eller senare
Modell som betjänar slutpunkt	Kräver `mlflow` 2.13.1 eller senare
SQL-lager	Kräver `mlflow` 2.16.1 eller senare
Unity Catalog-funktion	Kräver `mlflow` 2.16.1 eller senare
Genie-utrymme	Kräver `mlflow` 2.17.1 eller senare
Unity Catalog-tabell	Kräver `mlflow` 2.18.0 eller senare

Registrera agenten i Unity Catalog

Innan du distribuerar agenten måste du registrera agenten i Unity Catalog. Genom att registrera agenten paketeras den som en modell i Unity Catalog. Därför kan du använda Unity Catalog-behörigheter för auktorisering för resurser i agenten.

import mlflow

mlflow.set_registry_uri("databricks-uc")

catalog_name = "test_catalog"
schema_name = "schema"
model_name = "agent_name"

model_name = catalog_name + "." + schema_name + "." + model_name
uc_model_info = mlflow.register_model(model_uri=logged_agent_info.model_uri, name=model_name)

Dela via