Aller au contenu principal
Dernière mise à jour le

Traçage PydanticAI

PydanticAI est un framework Python pour le développement d'applications d'IA générative de qualité production, avec un typage fort et des APIs ergonomiques. Il s'articule autour de modèles Pydantic pour imposer une structure et une validation tout au long des workflows d'agents.

MLflow Tracing s'intègre à PydanticAI pour enregistrer les étapes de l'agent, les appels d'outils et les invocations de modèles avec leurs entrées et sorties typées. Activez-le avec mlflow.pydantic_ai.autolog.

remarque

Sur les clusters de compute Serverless, la journalisation automatique pour les cadres de traçage genAI n'est pas activée automatiquement. Vous devez activer explicitement l'autologging en appelant la fonction mlflow.<library>.autolog() appropriée pour les intégrations spécifiques que vous souhaitez suivre.

Ce que cette intégration fournit :

  • Appels d'agent avec invites, kwargs et réponses de sortie.
  • Requêtes LLM journalisant le nom du modèle, le prompt, les paramètres et la réponse
  • Exécutions d'outils capturant le nom de l'outil, les arguments et les métriques d'utilisation
  • Appels et listes de serveurs MCP pour le traçage des invocations d'outils
  • Métadonnées de span : latence, erreurs et liaison d'ID d'exécution

Prérequis​

Pour utiliser MLflow Tracing avec PydanticAI, vous devez installer MLflow et les packages PydanticAI pertinents.

Pour les environnements de développement, installez le package MLflow complet avec les modules complémentaires Databricks et PydanticAI :

Bash
pip install --upgrade "mlflow[databricks]>=3.1" pydantic-ai openai

Le package mlflow[databricks] complet inclut toutes les fonctionnalités pour le développement local et l’expérimentation sur Databricks.

remarque

MLflow 3 est recommandé pour une expérience de suivi optimale.

Avant d'exécuter les exemples, vous devrez configurer votre environnement :

Pour les utilisateurs en dehors des notebooks Databricks : Définissez vos variables d'environnement Databricks :

Bash
export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"

Pour les utilisateurs dans les Notebooks Databricks : ces identifiants sont définis automatiquement pour vous.

Exemple d'utilisation​

Python
import os
from dataclasses import dataclass
from typing import Any

from httpx import AsyncClient

from pydantic_ai import Agent, ModelRetry, RunContext


@dataclass
class Deps:
    client: AsyncClient
    weather_api_key: str | None
    geo_api_key: str | None


weather_agent = Agent(
    # Switch to your favorite LLM
    "google-gla:gemini-2.0-flash",
    # 'Be concise, reply with one sentence.' is enough for some models (like openai) to use
    # the below tools appropriately, but others like anthropic and gemini require a bit more direction.
    system_prompt=(
        "Be concise, reply with one sentence."
        "Use the `get_lat_lng` tool to get the latitude and longitude of the locations, "
        "then use the `get_weather` tool to get the weather."
    ),
    deps_type=Deps,
    retries=2,
    instrument=True,
)


@weather_agent.tool
async def get_lat_lng(
    ctx: RunContext[Deps], location_description: str
) -> dict[str, float]:
    """Get the latitude and longitude of a location.

    Args:
        ctx: The context.
        location_description: A description of a location.
    """
    if ctx.deps.geo_api_key is None:
        return {"lat": 51.1, "lng": -0.1}

    params = {
        "q": location_description,
        "api_key": ctx.deps.geo_api_key,
    }
    r = await ctx.deps.client.get("https://geocode.maps.co/search", params=params)
    r.raise_for_status()
    data = r.json()

    if data:
        return {"lat": data[0]["lat"], "lng": data[0]["lon"]}
    else:
        raise ModelRetry("Could not find the location")


@weather_agent.tool
async def get_weather(ctx: RunContext[Deps], lat: float, lng: float) -> dict[str, Any]:
    """Get the weather at a location.

    Args:
        ctx: The context.
        lat: Latitude of the location.
        lng: Longitude of the location.
    """

    if ctx.deps.weather_api_key is None:
        return {"temperature": "21 °C", "description": "Sunny"}

    params = {
        "apikey": ctx.deps.weather_api_key,
        "location": f"{lat},{lng}",
        "units": "metric",
    }
    r = await ctx.deps.client.get(
        "https://api.tomorrow.io/v4/weather/realtime", params=params
    )
    r.raise_for_status()
    data = r.json()

    values = data["data"]["values"]
    # https://docs.tomorrow.io/reference/data-layers-weather-codes
    code_lookup = {
        1000: "Clear, Sunny",
        1100: "Mostly Clear",
        1101: "Partly Cloudy",
        1102: "Mostly Cloudy",
        1001: "Cloudy",
        2000: "Fog",
        2100: "Light Fog",
        4000: "Drizzle",
        4001: "Rain",
        4200: "Light Rain",
        4201: "Heavy Rain",
        5000: "Snow",
        5001: "Flurries",
        5100: "Light Snow",
        5101: "Heavy Snow",
        6000: "Freezing Drizzle",
        6001: "Freezing Rain",
        6200: "Light Freezing Rain",
        6201: "Heavy Freezing Rain",
        7000: "Ice Pellets",
        7101: "Heavy Ice Pellets",
        7102: "Light Ice Pellets",
        8000: "Thunderstorm",
    }
    return {
        "temperature": f'{values["temperatureApparent"]:0.0f}°C',
        "description": code_lookup.get(values["weatherCode"], "Unknown"),
    }


async def main():
    async with AsyncClient() as client:
        weather_api_key = os.getenv("WEATHER_API_KEY")
        geo_api_key = os.getenv("GEO_API_KEY")
        deps = Deps(
            client=client, weather_api_key=weather_api_key, geo_api_key=geo_api_key
        )
        result = await weather_agent.run(
            "What is the weather like in London and in Wiltshire?", deps=deps
        )
        print("Response:", result.output)


# If you are running this on a notebook
await main()

# Uncomment this is you are using an IDE or Python script.
# asyncio.run(main())

Utilisez PydanticAI comme d'habitude (agents, outils et orchestrations). Les traces apparaîtront dans l'expérimentation associée.

Connectez-vous à un serveur MCP​

L'exemple suivant ci-dessous démontre comment exécuter un serveur MCP utilisant PydanticAI avec le traçage MLflow activé. Toutes les opérations d'invocation d'outils et de listing sont automatiquement capturées comme des spans de trace dans l'interface utilisateur.

Python
import mlflow
import asyncio

mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("MCP Server")
mlflow.pydantic_ai.autolog()

from pydantic_ai import Agent
from pydantic_ai.mcp import MCPServerStdio

server = MCPServerStdio(
    "deno",
    args=[
        "run",
        "-N",
        "-R=node_modules",
        "-W=node_modules",
        "--node-modules-dir=auto",
        "jsr:@pydantic/mcp-run-python",
        "stdio",
    ],
)

agent = Agent("openai:gpt-4o", mcp_servers=[server], instrument=True)


async def main():
    async with agent.run_mcp_servers():
        result = await agent.run("How many days between 2000-01-01 and 2025-03-18?")
    print(result.output)
    # > There are 9,208 days between January 1, 2000, and March 18, 2025.


# If you are running this on a notebook
await main()

# Uncomment this is you are using an IDE or Python script.
# asyncio.run(main())

Suivi de l'utilisation des jetons​

MLflow 3.2.0+ Enregistre les totaux d'utilisation des jetons dans les informations de trace et par appel dans les attributs d'étendue.

Python
import mlflow

last_trace_id = mlflow.get_last_active_trace_id()
trace = mlflow.get_trace(trace_id=last_trace_id)

print(trace.info.token_usage)
for span in trace.data.spans:
    usage = span.get_attribute("mlflow.chat.tokenUsage")
    if usage:
        print(span.name, usage)

Désactiver le suivi automatique​

Désactivez avec mlflow.pydantic_ai.autolog(disable=True) ou globalement avec mlflow.autolog(disable=True).

Sur cette page