← AI Agents Projects
🤖 AI Agents

Build a Production AI Agent with OpenTelemetry & Braintrust Evaluation

Source: mortalapps.com

This guide provides a comprehensive, end-to-end tutorial on building a production ai agent with full observability and evaluation capabilities. You will learn to architect, implement, and deploy a robust AI agent using LangGraph, focusing on critical production engineering aspects. This project is ideal for AI engineers, MLOps professionals, and developers who need to deploy reliable, traceable, and continuously evaluated AI agents in real-world scenarios. It addresses common pain points such as lack of visibility into agent execution, difficulty in assessing performance changes, and opaque operational costs.

By following this guide, you will construct a LangGraph-based agent fully instrumented with OpenTelemetry for detailed tracing and metrics. You'll integrate Braintrust for sophisticated LLM-as-a-Judge evaluation, enabling automated regression testing and quality assurance in your CI/CD pipeline. Furthermore, the project will demonstrate per-request cost attribution, allowing you to monitor and optimize LLM expenditures effectively. The outcome is a deployable system complete with a Grafana dashboard for live monitoring of agent performance, latency, and costs.

This architecture prioritizes production readiness from the outset, leveraging battle-tested open-source and commercial tools. It provides a blueprint for building agentic systems that are not only functional but also maintainable, observable, and evaluable at scale, offering a significant advantage over ad-hoc development approaches.

What You Will Build

You will build a production-grade AI agent using the LangGraph framework, designed to respond to user queries by orchestrating an LLM and executing custom tools. The agent will be exposed via a FastAPI web service, making it accessible through a standard HTTP API. All agent interactions and internal LLM calls will be meticulously traced using OpenTelemetry, providing deep insights into execution flow, latency, and token usage. For robust state management, the agent will utilize a PostgreSQL checkpointer, ensuring conversation continuity and fault tolerance.

Automated evaluation will be a core component, with agent responses and internal reasoning logged to Braintrust. This enables LLM-as-a-Judge evaluations for continuous quality assessment and regression testing within a CI/CD pipeline. Per-request cost attribution will be implemented to precisely track LLM expenses. The entire system will be containerized using Docker and deployable to Kubernetes, with Prometheus metrics scraped and visualized through a Grafana dashboard for real-time operational monitoring.

The final architecture consists of a FastAPI application hosting the LangGraph agent. This application communicates with an LLM provider (e.g., OpenAI), a PostgreSQL database for state persistence, an OpenTelemetry Collector for trace export, and the Braintrust platform for evaluation data. Prometheus will scrape metrics from the FastAPI service, and Grafana will display these metrics alongside traces from Tempo/Jaeger.

Technology Stack

Component Technology Why This Choice
LLM Provider OpenAI (GPT-4o) Industry-leading models offering strong performance and broad API support, making it a reliable choice for agentic reasoning.
Orchestration Framework LangGraph Provides a robust state machine framework for defining complex agent workflows, enabling explicit control over multi-step reasoning and tool use.
Memory & State Management PostgreSQL with LangGraph Checkpointer Offers reliable, ACID-compliant persistence for agent conversation state, ensuring continuity and recoverability across sessions and restarts.
Tool Execution Custom Python Functions Simple, flexible, and easily extendable for defining agent capabilities, integrated directly into the LangGraph workflow.
Observability & Tracing OpenTelemetry with Jaeger/Tempo Standardized, vendor-neutral framework for collecting distributed traces, metrics, and logs, crucial for understanding and debugging agent behavior in production.
Evaluation & Testing Braintrust Specialized platform for LLM evaluation, providing LLM-as-a-Judge capabilities, dataset management, and integration with CI/CD for continuous quality assurance.
Application Server FastAPI High-performance, easy-to-use Python web framework for building APIs, well-suited for serving the agent with minimal overhead.
Deployment Docker & Kubernetes Containerization ensures consistent environments from development to production, while Kubernetes provides robust orchestration, scaling, and high availability.

Choosing the right technology stack is paramount for building a production-ready AI agent. For orchestration, LangGraph stands out due to its explicit state machine representation, which offers deterministic control over complex agentic workflows, making debugging and reasoning about agent behavior significantly easier than chain-based alternatives. While other frameworks like CrewAI offer high-level abstractions, LangGraph's granular control is invaluable for production systems where reliability and predictability are key.

OpenTelemetry was selected for observability as it provides a vendor-neutral, open standard for instrumenting applications. This allows us to collect traces, metrics, and logs without vendor lock-in, enabling flexibility to export data to various backend systems like Jaeger, Tempo, or Datadog. This is critical for understanding the black-box nature of LLM interactions and debugging multi-step agent reasoning. Alternatives like proprietary SDKs (e.g., LangSmith) offer deep integration but limit portability and control.

Braintrust is our choice for evaluation due to its robust support for LLM-as-a-Judge paradigms and seamless integration into CI/CD pipelines. This platform allows us to define evaluation datasets, run automated tests, and track performance metrics over time, which is crucial for preventing model drift and ensuring consistent agent quality. While Ragas or custom scripts can provide basic evaluation, Braintrust offers a more complete, scalable solution for continuous evaluation.

FastAPI serves as the lightweight and performant web server, enabling efficient handling of concurrent requests to the agent. Its Pydantic-based data validation simplifies API development and ensures type safety. PostgreSQL is chosen for state persistence via LangGraph's checkpointer mechanism, providing a reliable and scalable database solution for managing agent conversation history. Finally, Docker and Kubernetes offer a powerful combination for containerization and orchestration, ensuring our agent is portable, scalable, and resilient in a production environment. This holistic stack provides the necessary tools for building, deploying, and maintaining a high-quality AI agent from development to production.

Project Structure

.
├── .dockerignore
├── .env.example
├── Dockerfile
├── README.md
├── config
│   └── k8s
│       ├── agent-deployment.yaml
│       └── agent-service.yaml
├── docker-compose.yaml
├── pyproject.toml
├── src
│   ├── __init__.py
│   ├── agent.py
│   ├── config.py
│   ├── database.py
│   ├── main.py
│   ├── observability.py
│   └── tools.py
└── tests
    ├── __init__.py
    └── test_agent.py

The project structure is organized to separate concerns and facilitate development, testing, and deployment. The src/ directory contains all the core application logic. src/main.py is the FastAPI entry point, handling API requests and integrating the agent. src/agent.py defines the LangGraph agent, its state, nodes, and edges. Custom tools are housed in src/tools.py. Configuration management, including environment variables and secrets loading, is handled by src/config.py. Database interactions and the LangGraph PostgreSQL checkpointer are encapsulated in src/database.py. Crucially, src/observability.py centralizes all OpenTelemetry instrumentation and cost attribution logic.

The tests/ directory contains unit and integration tests, including the Braintrust-powered evaluation suite in test_agent.py. The config/k8s/ folder holds Kubernetes manifests for deploying the agent to a cluster. Dockerfile defines the container image, and docker-compose.yaml orchestrates local development dependencies like PostgreSQL and Jaeger. pyproject.toml manages Python dependencies and project metadata using Poetry. Finally, .env.example provides a template for environment variables, and .dockerignore optimizes Docker builds by excluding unnecessary files.

Implementation

Phase 1: Environment & Scaffold

This phase establishes the foundational project structure, sets up dependency management with Poetry, and creates a basic FastAPI application. It ensures a runnable skeleton for the agent, allowing for initial environment configuration and verification.

Yaml
pyproject.toml
[tool.poetry]
name = "production-ai-agent"
version = "0.1.0"
description = "Production-hardened AI Agent with Observability and Evaluation"
authors = ["Your Name <[email protected]>"]
readme = "README.md"

[tool.poetry.dependencies]
python = ">=3.9,<3.12"
fastapi = "^0.111.0"
uvicorn = {extras = ["standard"], version = "^0.30.1"}
langchain-openai = "^0.1.13"
langgraph = "^0.2.0"
python-dotenv = "^1.0.1"
psycopg2-binary = "^2.9.9"
sqlalchemy = "^2.0.30"
asyncpg = "^0.29.0"
braintrust = "^0.1.18"
openai = "^1.35.13"
opentelemetry-api = "^1.25.0"
opentelemetry-sdk = "^1.25.0"
opentelemetry-exporter-otlp = "^1.25.0"
opentelemetry-instrumentation-fastapi = "^0.46b0"
opentelemetry-instrumentation-logging = "^0.46b0"
langchain-core = "^0.2.19"

[tool.poetry.group.dev.dependencies]
pytest = "^8.2.2"
httpx = "^0.27.0"

[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"


src/config.py
import os
from dotenv import load_dotenv
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class Config:
    def __init__(self):
        load_dotenv()
        self.OPENAI_API_KEY: str = os.environ.get("OPENAI_API_KEY", "")
        if not self.OPENAI_API_KEY:
            logger.error("OPENAI_API_KEY not found in environment variables.")
            raise ValueError("OPENAI_API_KEY is required.")

        self.BRAINTRUST_API_KEY: str = os.environ.get("BRAINTRUST_API_KEY", "")
        self.BRAINTRUST_PROJECT: str = os.environ.get("BRAINTRUST_PROJECT", "production-ai-agent")

        self.DB_URI: str = os.environ.get("DB_URI", "postgresql+asyncpg://user:password@localhost:5432/agentdb")
        self.OTEL_EXPORTER_OTLP_ENDPOINT: str = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317")
        self.SERVICE_NAME: str = os.environ.get("SERVICE_NAME", "production-ai-agent")

        logger.info("Configuration loaded successfully.")


src/main.py
import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel

from src.config import Config

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

config = Config()

class AgentInput(BaseModel):
    message: str
    thread_id: str | None = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Starting up application...")
    # Any startup logic can go here
    yield
    logger.info("Shutting down application...")
    # Any shutdown logic can go here

app = FastAPI(title="Production AI Agent", version="0.1.0", lifespan=lifespan)

@app.get("/health")
async def health_check():
    return {"status": "ok", "service": config.SERVICE_NAME}

@app.post("/invoke")
async def invoke_agent(input: AgentInput, request: Request):
    logger.info(f"Received invocation for thread_id: {input.thread_id} with message: {input.message}")
    # This will be replaced by actual agent invocation in later phases
    try:
        response = f"Agent received: {input.message}"
        return {"response": response, "thread_id": input.thread_id}
    except Exception as e:
        logger.exception(f"Error invoking agent: {e}")
        raise HTTPException(status_code=500, detail=str(e))


.env.example
OPENAI_API_KEY="sk-..."
BRAINTRUST_API_KEY="your-braintrust-api-key"
DB_URI="postgresql+asyncpg://user:password@localhost:5432/agentdb"
OTEL_EXPORTER_OTLP_ENDPOINT="http://localhost:4317"
SERVICE_NAME="production-ai-agent-service"


docker-compose.yaml
version: '3.8'
services:
  postgres:
    image: postgres:16-alpine
    environment:
      POSTGRES_USER: user
      POSTGRES_PASSWORD: password
      POSTGRES_DB: agentdb
    ports:
      - "5432:5432"
    volumes:
      - postgres_data:/var/lib/postgresql/data
    healthcheck:
      test: ["CMD-SHELL", "pg_isready", "-U", "user", "-d", "agentdb"]
      interval: 5s
      timeout: 5s
      retries: 5

  jaeger:
    image: jaegertracing/all-in-one:latest
    ports:
      - "16686:16686" # Jaeger UI
      - "4317:4317"   # OTLP gRPC receiver
      - "4318:4318"   # OTLP HTTP receiver
    environment:
      COLLECTOR_OTLP_ENABLED: true

volumes:
  postgres_data:
This phase lays the groundwork for our production AI agent. We start with pyproject.toml to manage Python dependencies using Poetry, specifying essential libraries like fastapi, uvicorn, langgraph, langchain-openai, python-dotenv, psycopg2-binary, sqlalchemy, asyncpg, braintrust, openai, and opentelemetry components. Poetry ensures reproducible builds and dependency isolation. src/config.py is introduced to centralize application configuration. It uses python-dotenv to load environment variables from a .env file, making secret management and environment-specific settings straightforward. Crucially, it includes validation for OPENAI_API_KEY, raising a ValueError if missing to prevent silent failures. Default values are provided for other variables, such as DB_URI and OTEL_EXPORTER_OTLP_ENDPOINT, to enable local development out-of-the-box. src/main.py initializes a FastAPI application. The lifespan context manager handles startup and shutdown logic, providing hooks for future initialization (e.g., database connections, OpenTelemetry setup). A basic /health endpoint is added for readiness checks, and a placeholder /invoke endpoint demonstrates how the agent will receive messages. Error handling is included to catch exceptions during invocation and return appropriate HTTP responses. The .env.example file serves as a template for environment variables, guiding users on necessary configurations. Finally, docker-compose.yaml is provided to set up a local development environment. It includes a PostgreSQL database for state persistence and a Jaeger instance for tracing, configured to accept OTLP (OpenTelemetry Protocol) traces. This docker-compose setup simplifies local testing of the agent's persistence and observability features without requiring a full Kubernetes cluster.
Expected outcome You will have a Python project structure with Poetry, a `src` directory containing basic FastAPI application files, and a `docker-compose.yaml` to spin up a PostgreSQL database and a Jaeger tracing backend. The FastAPI application will be runnable locally, serving a health check and a placeholder `/invoke` endpoint.
How to test
  1. Run poetry install to set up dependencies. 2. Create a .env file based on .env.example and populate OPENAI_API_KEY. 3. Run docker-compose up -d to start PostgreSQL and Jaeger. 4. Execute poetry run uvicorn src.main:app --reload in your terminal. 5. Navigate to http://localhost:8000/health in your browser or use curl http://localhost:8000/health to verify the health check. 6. Send a POST request to http://localhost:8000/invoke with {"message": "hello", "thread_id": "test-123"} to confirm the placeholder endpoint works.

Phase 2: Core LangGraph Agent Loop

This phase implements the fundamental LangGraph agent, defining its state, nodes for LLM interaction and tool calling, and the conditional routing logic. It establishes the basic reactive loop where the agent processes input, decides on an action, and executes it.

Python
src/tools.py
import logging
import operator
from typing import Any, Dict

logger = logging.getLogger(__name__)

def search_tool(query: str) -> Dict[str, Any]:
    """A tool to perform a dummy search operation."""
    logger.info(f"Executing search_tool with query: {query}")
    try:
        # Simulate an external API call or computation
        if "error" in query.lower():
            raise ValueError("Simulated search error")
        result = {"result": f"Found information about '{query}'. This is a dummy response."}
        logger.info(f"Search tool completed for query: {query}")
        return result
    except Exception as e:
        logger.error(f"Error in search_tool for query '{query}': {e}")
        return {"error": str(e)}


src/agent.py
import logging
from typing import Annotated, Any, Dict, List, Literal, Optional, Sequence, Tuple, TypedDict, Union

from langchain_core.messages import BaseMessage, FunctionMessage, HumanMessage, ToolMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.tools import tool as langchain_tool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph

from src.config import Config
from src.tools import search_tool

logger = logging.getLogger(__name__)
config = Config()

# 1. Define the Agent State
class AgentState(TypedDict):
    input: str
    chat_history: Annotated[List[BaseMessage], operator.add]
    agent_outcome: Optional[str]
    tool_calls: List[Dict[str, Any]]
    tool_output: Optional[str]
    turn_count: int

# 2. Define Tools
@langchain_tool
def dummy_search(query: str) -> str:
    """Searches for information on the web. Input should be a search query string."""
    logger.info(f"Agent calling dummy_search tool with query: {query}")
    result = search_tool(query)
    if "error" in result:
        logger.error(f"dummy_search tool returned an error: {result['error']}")
        return f"Error during search: {result['error']}"
    return str(result["result"])

tools = [dummy_search]

# 3. Define the LLM
llm = ChatOpenAI(model="gpt-4o", temperature=0, api_key=config.OPENAI_API_KEY)
llm_with_tools = llm.bind_tools(tools)

# 4. Define Graph Nodes
class Agent: # Renamed from AgentNode to Agent for clarity and to align with LangGraph patterns
    def __init__(self, llm_with_tools):
        self.llm_with_tools = llm_with_tools
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a helpful AI assistant. You have access to tools to answer questions."),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])

    def __call__(self, state: AgentState) -> Dict[str, Any]:
        logger.info(f"Agent node invoked. Turn count: {state.get('turn_count', 0)}")
        current_history = state["chat_history"] + [HumanMessage(content=state["input"])]
        response = self.llm_with_tools.invoke({"chat_history": current_history, "input": state["input"]})
        
        new_state = {"chat_history": [response], "turn_count": state.get("turn_count", 0) + 1}

        tool_calls = []
        if hasattr(response, "tool_calls") and response.tool_calls:
            tool_calls = response.tool_calls
            logger.info(f"Agent decided to call tools: {tool_calls}")
            new_state["tool_calls"] = tool_calls
        else:
            logger.info("Agent decided to respond directly.")
            new_state["agent_outcome"] = response.content
        
        return new_state

def execute_tools(state: AgentState) -> Dict[str, Any]:
    logger.info("Execute tools node invoked.")
    tool_calls = state["tool_calls"]
    tool_output_messages = []
    for tool_call in tool_calls:
        tool_name = tool_call["name"]
        tool_args = tool_call["args"]
        try:
            if tool_name == dummy_search.name:
                output = dummy_search.invoke(tool_args)
                tool_output_messages.append(ToolMessage(content=output, tool_call_id=tool_call["id"]))
                logger.info(f"Tool '{tool_name}' executed successfully. Output: {output[:50]}...")
            else:
                logger.warning(f"Unknown tool: {tool_name}")
                tool_output_messages.append(ToolMessage(content=f"Unknown tool: {tool_name}", tool_call_id=tool_call["id"]))
        except Exception as e:
            logger.error(f"Error executing tool '{tool_name}': {e}")
            tool_output_messages.append(ToolMessage(content=f"Error executing {tool_name}: {str(e)}", tool_call_id=tool_call["id"]))

    return {"chat_history": tool_output_messages, "tool_output": str(tool_output_messages)}

# 5. Define Graph Edges and Entry Point
def should_continue(state: AgentState) -> Literal["continue", "end"]:
    if state.get("tool_calls"):
        logger.info("Decision: Continue to tool execution.")
        return "continue"
    else:
        logger.info("Decision: End, agent has a direct answer.")
        return "end"

# Build the graph
def create_agent_graph():
    workflow = StateGraph(AgentState)

    # Add nodes
    workflow.add_node("agent", Agent(llm_with_tools))
    workflow.add_node("tools", execute_tools)

    # Set entry point
    workflow.set_entry_point("agent")

    # Add edges
    workflow.add_conditional_edges(
        "agent",
        should_continue,
        {"continue": "tools", "end": END}
    )
    workflow.add_edge("tools", "agent")

    app = workflow.compile()
    logger.info("LangGraph agent compiled successfully.")
    return app


src/main.py (updated)
import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.graph import GraphCompiler

from src.config import Config
from src.agent import create_agent_graph, AgentState

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

config = Config()

class AgentInput(BaseModel):
    message: str
    thread_id: str | None = None

# Placeholder for checkpointer, will be implemented in a later phase
class NoopCheckpointSaver(BaseCheckpointSaver):
    async def aget_tuple(self, config: Dict[str, Any]) -> Optional[Tuple[Config, List[BaseMessage]]]:
        return None

    async def aput_tuple(self, config: Dict[str, Any], checkpoint_tuple: Tuple[Config, List[BaseMessage]]) -> str:
        return "noop_thread_id"

    async def alist(self, config: Dict[str, Any]) -> List[Dict[str, Any]]:
        return []

agent_app = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Starting up application...")
    global agent_app
    try:
        # For now, we'll use a NoopCheckpointSaver. This will be replaced.
        agent_app = create_agent_graph().with_config({"recursion_limit": 5, "checkpointer": NoopCheckpointSaver()})
        logger.info("LangGraph agent initialized.")
    except Exception as e:
        logger.critical(f"Failed to initialize LangGraph agent: {e}")
        raise RuntimeError(f"Agent initialization failed: {e}")
    yield
    logger.info("Shutting down application...")

app = FastAPI(title="Production AI Agent", version="0.1.0", lifespan=lifespan)

@app.get("/health")
async def health_check():
    return {"status": "ok", "service": config.SERVICE_NAME}

@app.post("/invoke")
async def invoke_agent(input: AgentInput, request: Request):
    logger.info(f"Received invocation for thread_id: {input.thread_id} with message: {input.message}")
    if agent_app is None:
        raise HTTPException(status_code=503, detail="Agent not initialized.")

    thread_id = input.thread_id if input.thread_id else "default_thread"
    try:
        # The agent_app expects a config dict with thread_id for checkpointer
        # Even with NoopCheckpointSaver, the structure is expected.
        result = await agent_app.ainvoke(
            {"input": input.message, "chat_history": [], "turn_count": 0},
            config={
                "configurable": {"thread_id": thread_id},
                "tags": ["agent_invocation"],
            }
        )
        final_response = result.get("agent_outcome", "No direct response from agent.")
        if not final_response and result.get("tool_output"):
            final_response = f"Tool executed. Output: {result['tool_output']}"
        return {"response": final_response, "thread_id": thread_id}
    except Exception as e:
        logger.exception(f"Error invoking agent: {e}")
        raise HTTPException(status_code=500, detail=f"Agent invocation failed: {str(e)}")
This phase introduces the core LangGraph agent. First, src/tools.py defines a simple search_tool function, which simulates an external search. This function includes basic logging and error handling, returning a structured dictionary to indicate success or failure. This separation of tools is crucial for modularity and reusability. src/agent.py is where the LangGraph magic happens. We define AgentState using TypedDict to explicitly type the state that flows through our graph, including input, chat_history, agent_outcome, tool_calls, tool_output, and turn_count. The Annotated type with "operator": "append" for chat_history is a LangGraph-specific instruction to append new messages to the history rather than overwriting it, which is essential for conversational agents. The dummy_search function is wrapped with LangChain's @langchain_tool decorator, making it discoverable and callable by the LLM. The llm_with_tools is created by binding these tools to the ChatOpenAI instance, enabling the LLM to choose when to use them. The Agent class acts as a node in our graph, responsible for invoking the LLM with the current state and deciding whether to call a tool or directly respond. The execute_tools function is another node, responsible for actually running the tools the LLM decided to use. The should_continue function defines the conditional logic for graph edges: if the agent decides to call a tool, the graph transitions to execute_tools; otherwise, it ends. The create_agent_graph function assembles these components into a StateGraph, defining the nodes and edges to create a complete agent workflow. The workflow.compile() step optimizes and finalizes the graph. src/main.py is updated to integrate this agent. A NoopCheckpointSaver is temporarily used to satisfy LangGraph's requirement for a checkpointer, which will be replaced with a persistent one later. The agent_app is initialized within the FastAPI lifespan function, ensuring it's ready before handling requests. The /invoke endpoint now calls agent_app.ainvoke(), passing the user message and a thread_id for state management. The config dictionary passed to ainvoke is crucial for LangGraph's internal state management and future checkpointer integration. The response logic extracts the agent_outcome or tool_output to return to the user.
Expected outcome You will have a functional LangGraph agent integrated into the FastAPI application. When you send a query that might require a tool (e.g., 'search for weather'), the agent will invoke the `dummy_search` tool. If the query is simple (e.g., 'hello'), the agent will provide a direct LLM response.
How to test
  1. Ensure docker-compose up -d is running. 2. Restart the FastAPI application: poetry run uvicorn src.main:app --reload. 3. Send a POST request to http://localhost:8000/invoke with {"message": "What is the capital of France?", "thread_id": "test-thread-1"}. The agent should respond directly. 4. Send another POST request with {"message": "Search for current news headlines.", "thread_id": "test-thread-2"}. The agent should indicate it used the dummy_search tool and return its simulated output. Observe the logs for Agent calling dummy_search tool messages.

Phase 3: OpenTelemetry & Cost Attribution

This phase focuses on instrumenting the agent with OpenTelemetry for comprehensive tracing and metrics, including LLM calls and tool execution. It also implements token counting and cost attribution to provide granular visibility into the operational expenses of the agent.

Python
src/observability.py
import os
import logging
import operator
from typing import Any, Dict, Optional

from opentelemetry import trace
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.logging import LoggingInstrumentor
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from opentelemetry.semconv.ai import SpanAttributes, LLMRequestTypeValues, LLMModelNameValues, LLMOperationTypeValues

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult

logger = logging.getLogger(__name__)

class OpenTelemetryConfig:
    def __init__(self):
        self.otel_exporter_otlp_endpoint: str = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317")
        self.service_name: str = os.environ.get("SERVICE_NAME", "production-ai-agent-service")
        self.openai_model_cost_per_token: Dict[str, Dict[str, float]] = {
            "gpt-4o": {"input": 0.000005, "output": 0.000015},
            "gpt-4o-mini": {"input": 0.0000005, "output": 0.0000015},
        }

def setup_opentelemetry(app):
    config = OpenTelemetryConfig()
    logger.info(f"Setting up OpenTelemetry for service: {config.service_name} with OTLP endpoint: {config.otel_exporter_otlp_endpoint}")

    resource = Resource.create({"service.name": config.service_name})
    provider = TracerProvider(resource=resource)
    processor = BatchSpanProcessor(OTLPSpanExporter(endpoint=config.otel_exporter_otlp_endpoint))
    provider.add_span_processor(processor)
    trace.set_tracer_provider(provider)

    FastAPIInstrumentor.instrument_app(app, tracer_provider=provider)
    LoggingInstrumentor().instrument(set_logging_format=True)
    logger.info("OpenTelemetry setup complete.")

# Custom Callback Handler for LangChain/LLM Observability and Cost Attribution
class LLMTraceCallbackHandler(BaseCallbackHandler):
    def __init__(self, tracer_name: str = "llm_tracer", model_name: str = "gpt-4o"):
        self.tracer = trace.get_tracer(tracer_name)
        self.model_name = model_name
        self.model_costs = OpenTelemetryConfig().openai_model_cost_per_token
        logger.info(f"Initialized LLMTraceCallbackHandler for model: {self.model_name}")

    def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> None:
        current_span = trace.get_current_span()
        if current_span.is_recording():
            current_span.set_attribute(SpanAttributes.LLM_VENDOR, LLMModelNameValues.OPENAI.value)
            current_span.set_attribute(SpanAttributes.LLM_MODEL_NAME, self.model_name)
            current_span.set_attribute(SpanAttributes.LLM_REQUEST_TYPE, LLMRequestTypeValues.COMPLETION.value)
            current_span.set_attribute(SpanAttributes.LLM_OPERATION_TYPE, LLMOperationTypeValues.CHAT.value)
            current_span.set_attribute("llm.prompts", str(prompts))
            logger.debug(f"LLM start for model {self.model_name}")

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        current_span = trace.get_current_span()
        if current_span.is_recording():
            token_usage = response.llm_output.get("token_usage") if response.llm_output else None
            if token_usage:
                input_tokens = token_usage.get("prompt_tokens", 0)
                output_tokens = token_usage.get("completion_tokens", 0)
                total_tokens = token_usage.get("total_tokens", 0)

                current_span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, input_tokens)
                current_span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, output_tokens)
                current_span.set_attribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, total_tokens)

                input_cost = input_tokens * self.model_costs.get(self.model_name, {}).get("input", 0)
                output_cost = output_tokens * self.model_costs.get(self.model_name, {}).get("output", 0)
                total_cost = input_cost + output_cost

                current_span.set_attribute("llm.cost.input_usd", input_cost)
                current_span.set_attribute("llm.cost.output_usd", output_cost)
                current_span.set_attribute("llm.cost.total_usd", total_cost)
                logger.info(f"LLM end for model {self.model_name}. Tokens: {total_tokens}, Cost: ${total_cost:.6f}")
            else:
                logger.warning("No token usage found in LLM response.")

    def on_llm_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None:
        current_span = trace.get_current_span()
        if current_span.is_recording():
            current_span.set_attribute("error", True)
            current_span.set_attribute("error.message", str(error))
            logger.error(f"LLM error: {error}")

def create_custom_span(name: str, attributes: Optional[Dict[str, Any]] = None):
    tracer = trace.get_tracer("custom_tracer")
    span = tracer.start_as_current_span(name, attributes=attributes)
    return span


src/main.py (updated)
import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.graph import GraphCompiler
from opentelemetry import trace

from src.config import Config
from src.agent import create_agent_graph, AgentState
from src.observability import setup_opentelemetry, LLMTraceCallbackHandler, create_custom_span

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

config = Config()

class AgentInput(BaseModel):
    message: str
    thread_id: str | None = None

class NoopCheckpointSaver(BaseCheckpointSaver):
    async def aget_tuple(self, config: Dict[str, Any]) -> Optional[Tuple[Config, List[BaseMessage]]]:
        return None

    async def aput_tuple(self, config: Dict[str, Any], checkpoint_tuple: Tuple[Config, List[BaseMessage]]) -> str:
        return "noop_thread_id"

    async def alist(self, config: Dict[str, Any]) -> List[Dict[str, Any]]:
        return []

agent_app = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Starting up application...")
    global agent_app
    try:
        setup_opentelemetry(app) # Initialize OpenTelemetry
        # Pass the custom callback handler to the LLM within create_agent_graph
        agent_app = create_agent_graph().with_config({"recursion_limit": 5, "checkpointer": NoopCheckpointSaver()})
        logger.info("LangGraph agent initialized.")
    except Exception as e:
        logger.critical(f"Failed to initialize LangGraph agent: {e}")
        raise RuntimeError(f"Agent initialization failed: {e}")
    yield
    logger.info("Shutting down application...")

app = FastAPI(title="Production AI Agent", version="0.1.0", lifespan=lifespan)

@app.get("/health")
async def health_check():
    return {"status": "ok", "service": config.SERVICE_NAME}

@app.post("/invoke")
async def invoke_agent(input: AgentInput, request: Request):
    logger.info(f"Received invocation for thread_id: {input.thread_id} with message: {input.message}")
    if agent_app is None:
        raise HTTPException(status_code=503, detail="Agent not initialized.")

    thread_id = input.thread_id if input.thread_id else "default_thread"
    
    with create_custom_span("agent.invoke", {"thread_id": thread_id, "user_input": input.message}) as span:
        try:
            # LangGraph's ainvoke takes callbacks directly in the config
            result = await agent_app.ainvoke(
                {"input": input.message, "chat_history": [], "turn_count": 0},
                config={
                    "configurable": {"thread_id": thread_id},
                    "tags": ["agent_invocation"],
                    "callbacks": [LLMTraceCallbackHandler(model_name=config.OPENAI_MODEL_NAME)] # Pass callback here
                }
            )
            final_response = result.get("agent_outcome", "No direct response from agent.")
            if not final_response and result.get("tool_output"):
                final_response = f"Tool executed. Output: {result['tool_output']}"
            
            span.set_attribute("agent.response", final_response)
            return {"response": final_response, "thread_id": thread_id}
        except Exception as e:
            span.set_attribute("error", True)
            span.set_attribute("error.message", str(e))
            logger.exception(f"Error invoking agent: {e}")
            raise HTTPException(status_code=500, detail=f"Agent invocation failed: {str(e)}")


# --- src/agent.py (updated LLM initialization and tool execution) ---
import logging
from typing import Annotated, Any, Dict, List, Literal, Optional, Sequence, Tuple, TypedDict, Union

from langchain_core.messages import BaseMessage, FunctionMessage, HumanMessage, ToolMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.tools import tool as langchain_tool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph
from opentelemetry import trace

from src.config import Config
from src.tools import search_tool
from src.observability import LLMTraceCallbackHandler, create_custom_span # Import create_custom_span

logger = logging.getLogger(__name__)
config = Config()

# 1. Define the Agent State
class AgentState(TypedDict):
    input: str
    chat_history: Annotated[List[BaseMessage], operator.add]
    agent_outcome: Optional[str]
    tool_calls: List[Dict[str, Any]]
    tool_output: Optional[str]
    turn_count: int

# 2. Define Tools
@langchain_tool
def dummy_search(query: str) -> str:
    """Searches for information on the web. Input should be a search query string."""
    with create_custom_span("tool.dummy_search", {"query": query}) as span: # Custom span for tool
        logger.info(f"Agent calling dummy_search tool with query: {query}")
        result = search_tool(query)
        if "error" in result:
            logger.error(f"dummy_search tool returned an error: {result['error']}")
            span.set_attribute("error", True)
            span.set_attribute("error.message", result['error'])
            return f"Error during search: {result['error']}"
        span.set_attribute("tool.output", result["result"])
        return str(result["result"])

tools = [dummy_search]

# 3. Define the LLM
# Model name moved to config for easier management
llm = ChatOpenAI(model=config.OPENAI_MODEL_NAME, temperature=0, api_key=config.OPENAI_API_KEY)
llm_with_tools = llm.bind_tools(tools).with_config({"callbacks": [LLMTraceCallbackHandler(model_name=config.OPENAI_MODEL_NAME)]}) # Pass callback here

# 4. Define Graph Nodes
class Agent:
    def __init__(self, llm_with_tools):
        self.llm_with_tools = llm_with_tools
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a helpful AI assistant. You have access to tools to answer questions."),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])

    def __call__(self, state: AgentState) -> Dict[str, Any]:
        with create_custom_span("agent.decision_making", {"turn_count": state.get('turn_count', 0)}) as span: # Custom span for agent node
            logger.info(f"Agent node invoked. Turn count: {state.get('turn_count', 0)}")
            current_history = state["chat_history"] + [HumanMessage(content=state["input"])]
            response = self.llm_with_tools.invoke({"chat_history": current_history, "input": state["input"]})
            
            new_state = {"chat_history": [response], "turn_count": state.get("turn_count", 0) + 1}

            tool_calls = []
            if hasattr(response, "tool_calls") and response.tool_calls:
                tool_calls = response.tool_calls
                logger.info(f"Agent decided to call tools: {tool_calls}")
                span.set_attribute("agent.action", "call_tool")
                new_state["tool_calls"] = tool_calls
            else:
                logger.info("Agent decided to respond directly.")
                span.set_attribute("agent.action", "respond_directly")
                new_state["agent_outcome"] = response.content
            
            return new_state

def execute_tools(state: AgentState) -> Dict[str, Any]:
    with create_custom_span("agent.tool_execution") as span: # Custom span for tool execution node
        logger.info("Execute tools node invoked.")
        tool_calls = state["tool_calls"]
        tool_output_messages = []
        for tool_call in tool_calls:
            tool_name = tool_call["name"]
            tool_args = tool_call["args"]
            try:
                if tool_name == dummy_search.name:
                    output = dummy_search.invoke(tool_args)
                    tool_output_messages.append(ToolMessage(content=output, tool_call_id=tool_call["id"]))
                    logger.info(f"Tool '{tool_name}' executed successfully. Output: {output[:50]}...")
                else:
                    logger.warning(f"Unknown tool: {tool_name}")
                    tool_output_messages.append(ToolMessage(content=f"Unknown tool: {tool_name}", tool_call_id=tool_call["id"]))
            except Exception as e:
                logger.error(f"Error executing tool '{tool_name}': {e}")
                span.set_attribute("error", True)
                span.set_attribute("error.message", str(e))
                tool_output_messages.append(ToolMessage(content=f"Error executing {tool_name}: {str(e)}", tool_call_id=tool_call["id"]))

        return {"chat_history": tool_output_messages, "tool_output": str(tool_output_messages)}

# 5. Define Graph Edges and Entry Point
def should_continue(state: AgentState) -> Literal["continue", "end"]:
    if state.get("tool_calls"):
        logger.info("Decision: Continue to tool execution.")
        return "continue"
    else:
        logger.info("Decision: End, agent has a direct answer.")
        return "end"

# Build the graph
def create_agent_graph():
    workflow = StateGraph(AgentState)

    # Add nodes
    workflow.add_node("agent", Agent(llm_with_tools))
    workflow.add_node("tools", execute_tools)

    # Set entry point
    workflow.set_entry_point("agent")

    # Add edges
    workflow.add_conditional_edges(
        "agent",
        should_continue,
        {"continue": "tools", "end": END}
    )
    workflow.add_edge("tools", "agent")

    app = workflow.compile()
    logger.info("LangGraph agent compiled successfully.")
    return app


src/config.py (updated)
import os
from dotenv import load_dotenv
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class Config:
    def __init__(self):
        load_dotenv()
        self.OPENAI_API_KEY: str = os.environ.get("OPENAI_API_KEY", "")
        if not self.OPENAI_API_KEY:
            logger.error("OPENAI_API_KEY not found in environment variables.")
            raise ValueError("OPENAI_API_KEY is required.")

        self.OPENAI_MODEL_NAME: str = os.environ.get("OPENAI_MODEL_NAME", "gpt-4o")

        self.BRAINTRUST_API_KEY: str = os.environ.get("BRAINTRUST_API_KEY", "")
        self.BRAINTRUST_PROJECT: str = os.environ.get("BRAINTRUST_PROJECT", "production-ai-agent")

        self.DB_URI: str = os.environ.get("DB_URI", "postgresql+asyncpg://user:password@localhost:5432/agentdb")
        self.OTEL_EXPORTER_OTLP_ENDPOINT: str = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317")
        self.SERVICE_NAME: str = os.environ.get("SERVICE_NAME", "production-ai-agent-service")

        logger.info("Configuration loaded successfully.")


# .env.example (updated)
# OPENAI_API_KEY="sk-..."
# OPENAI_MODEL_NAME="gpt-4o"
# BRAINTRUST_API_KEY="your-braintrust-api-key"
# DB_URI="postgresql+asyncpg://user:password@localhost:5432/agentdb"
# OTEL_EXPORTER_OTLP_ENDPOINT="http://localhost:4317"
# SERVICE_NAME="production-ai-agent-service"
This phase significantly enhances our agent with production-grade observability and cost attribution. src/observability.py is introduced to encapsulate all OpenTelemetry setup and custom tracing logic. The setup_opentelemetry function initializes the TracerProvider, configures a BatchSpanProcessor to send traces to an OTLP endpoint (e.g., Jaeger), and instruments FastAPI and Python's logging module. This automatically generates spans for incoming HTTP requests and logs. The LLMTraceCallbackHandler is a custom BaseCallbackHandler from LangChain that integrates with OpenTelemetry. It intercepts LLM calls, creating spans with relevant SpanAttributes (vendor, model name, request type). Crucially, in on_llm_end, it extracts token usage from the LLM response and calculates the estimated cost based on predefined model pricing. These token counts and costs are added as attributes to the current OpenTelemetry span, providing granular cost attribution per LLM call. Error handling is also included in on_llm_error to mark spans as erroneous. create_custom_span is a utility function to easily create custom spans for specific operations within our agent, like tool execution or agent decision-making. This allows us to trace internal logic beyond what automatic instrumentation provides. src/agent.py is updated to utilize this. The dummy_search tool, the Agent node's __call__ method, and the execute_tools function are wrapped with create_custom_span to generate detailed traces for these critical steps. The LLMTraceCallbackHandler is now directly passed to the ChatOpenAI instance via .with_config({"callbacks": [...]}) to ensure all LLM calls are traced. src/main.py is updated to call setup_opentelemetry(app) during the FastAPI lifespan startup. The invoke_agent endpoint also wraps the entire agent invocation in a custom span, providing an overarching trace for each user request. The LLMTraceCallbackHandler is passed to the agent_app.ainvoke configuration, ensuring that LLM calls within the LangGraph execution are also traced. The src/config.py and .env.example are updated to include OPENAI_MODEL_NAME, making the LLM model configurable.
Expected outcome When you invoke the agent, OpenTelemetry traces will be generated for the FastAPI request, the LangGraph agent's decision-making, tool execution, and individual LLM calls. These traces will include attributes detailing LLM model usage (prompt/completion tokens) and estimated costs. You will be able to visualize these traces and their attributes in Jaeger (running via `docker-compose`).
How to test
  1. Ensure docker-compose up -d is running. 2. Restart the FastAPI application: poetry run uvicorn src.main:app --reload. 3. Send a POST request to http://localhost:8000/invoke with {"message": "What is the capital of France?", "thread_id": "trace-test-1"}. 4. Send another POST request with {"message": "Search for current weather in London.", "thread_id": "trace-test-2"}. 5. Open Jaeger UI at http://localhost:16686. Select production-ai-agent-service from the service dropdown and click 'Find Traces'. You should see traces for each invocation, with spans for FastAPI, agent decision, LLM calls, and tool execution. Inspect the LLM spans for llm.usage.total_tokens and llm.cost.total_usd attributes.

Phase 4: Braintrust Evaluation & Automated Testing

This phase integrates Braintrust for LLM-as-a-Judge evaluation, enabling automated testing and continuous quality assessment of the agent. It demonstrates how to log agent runs to Braintrust and execute evaluation tests within a Python test suite.

Python
tests/test_agent.py
import pytest
import httpx
import os
import braintrust as bt
from dotenv import load_dotenv

# Load environment variables for Braintrust API key
load_dotenv()

BRAINTRUST_API_KEY = os.environ.get("BRAINTRUST_API_KEY")
BRAINTRUST_PROJECT = os.environ.get("BRAINTRUST_PROJECT", "production-ai-agent")
AGENT_SERVICE_URL = os.environ.get("AGENT_SERVICE_URL", "http://localhost:8000")

if not BRAINTRUST_API_KEY:
    pytest.skip("BRAINTRUST_API_KEY not set, skipping Braintrust tests", allow_module_level=True)

# Initialize Braintrust
bt.init(api_key=BRAINTRUST_API_KEY, project=BRAINTRUST_PROJECT)

@pytest.fixture(scope="module")
def client():
    return httpx.Client(base_url=AGENT_SERVICE_URL)

@bt.test(project=BRAINTRUST_PROJECT, experiment="agent-functional-test")
@pytest.mark.asyncio
async def test_agent_direct_response(client: httpx.Client, bt_record: bt.Span):
    input_message = "What is the capital of France?"
    thread_id = "test-direct-response-1"
    
    bt_record.log(input=input_message, expected="Paris")

    try:
        response = client.post(
            "/invoke", 
            json={
                "message": input_message,
                "thread_id": thread_id
            }
        )
        response.raise_for_status()
        output = response.json()["response"]
        
        bt_record.log(output=output)
        bt_record.span.log(metrics=bt.Metrics(accuracy=(1 if "paris" in output.lower() else 0)))
        assert "paris" in output.lower()
    except httpx.HTTPStatusError as e:
        bt_record.log(error=f"HTTP Error: {e.response.status_code} - {e.response.text}")
        pytest.fail(f"HTTP error: {e}")
    except Exception as e:
        bt_record.log(error=str(e))
        pytest.fail(f"Agent invocation failed: {e}")

@bt.test(project=BRAINTRUST_PROJECT, experiment="agent-tool-use-test")
@pytest.mark.asyncio
async def test_agent_tool_use(client: httpx.Client, bt_record: bt.Span):
    input_message = "Search for the latest news on AI."
    thread_id = "test-tool-use-1"
    
    bt_record.log(input=input_message, expected="Found information about 'latest news on AI'.")

    try:
        response = client.post(
            "/invoke", 
            json={
                "message": input_message,
                "thread_id": thread_id
            }
        )
        response.raise_for_status()
        output = response.json()["response"]
        
        bt_record.log(output=output)
        bt_record.span.log(metrics=bt.Metrics(tool_used=(1 if "dummy response" in output.lower() else 0)))
        assert "dummy response" in output.lower()
    except httpx.HTTPStatusError as e:
        bt_record.log(error=f"HTTP Error: {e.response.status_code} - {e.response.text}")
        pytest.fail(f"HTTP error: {e}")
    except Exception as e:
        bt_record.log(error=str(e))
        pytest.fail(f"Agent invocation failed: {e}")

@bt.test(project=BRAINTRUST_PROJECT, experiment="agent-error-handling-test")
@pytest.mark.asyncio
async def test_agent_tool_error_handling(client: httpx.Client, bt_record: bt.Span):
    input_message = "Search for something that will cause an error."
    thread_id = "test-tool-error-1"

    bt_record.log(input=input_message, expected="Error during search")

    try:
        response = client.post(
            "/invoke", 
            json={
                "message": input_message,
                "thread_id": thread_id
            }
        )
        response.raise_for_status()
        output = response.json()["response"]

        bt_record.log(output=output)
        bt_record.span.log(metrics=bt.Metrics(handled_error=(1 if "error during search" in output.lower() else 0)))
        assert "error during search" in output.lower()
    except httpx.HTTPStatusError as e:
        bt_record.log(error=f"HTTP Error: {e.response.status_code} - {e.response.text}")
        pytest.fail(f"HTTP error: {e}")
    except Exception as e:
        bt_record.log(error=str(e))
        pytest.fail(f"Agent invocation failed: {e}")


# .env.example (updated)
# OPENAI_API_KEY="sk-..."
# OPENAI_MODEL_NAME="gpt-4o"
# BRAINTRUST_API_KEY="your-braintrust-api-key"
# BRAINTRUST_PROJECT="production-ai-agent"
# DB_URI="postgresql+asyncpg://user:password@localhost:5432/agentdb"
# OTEL_EXPORTER_OTLP_ENDPOINT="http://localhost:4317"
# SERVICE_NAME="production-ai-agent-service"
# AGENT_SERVICE_URL="http://localhost:8000"
This phase introduces automated evaluation and testing using Braintrust, a crucial component for maintaining agent quality in production. The tests/test_agent.py file contains our test suite, leveraging pytest for test execution and httpx for making asynchronous HTTP requests to our FastAPI agent. Key to this phase is the integration with Braintrust. We load the BRAINTRUST_API_KEY and BRAINTRUST_PROJECT from environment variables, ensuring that sensitive information is not hardcoded. The bt.init() call initializes the Braintrust SDK, connecting our local tests to the Braintrust platform. The @bt.test decorator from Braintrust is used to wrap our pytest functions. This decorator automatically logs the test run's metadata, input, and output to Braintrust, allowing for centralized tracking and analysis. Each test function (test_agent_direct_response, test_agent_tool_use, test_agent_tool_error_handling) defines a specific scenario: a direct LLM response, an agent using a tool, and an agent handling a tool error. Inside each test, bt_record.log(input=..., expected=...) records the test case's input and expected outcome. After invoking the agent via an HTTP request, the actual output is logged. Finally, bt_record.span.log(metrics=bt.Metrics(...)) is used to log custom metrics for each test run (e.g., accuracy for direct responses, tool_used for tool invocations, handled_error for error scenarios). These metrics are vital for LLM-as-a-Judge evaluations and tracking performance trends in Braintrust. Error handling within the tests ensures that network issues or agent failures are properly captured and reported to Braintrust, providing a comprehensive view of test outcomes. The AGENT_SERVICE_URL is added to .env.example for configurable test endpoint, and a pytest.fixture provides an httpx.Client for clean API calls. This setup allows us to run these tests locally or integrate them into a CI/CD pipeline, continuously evaluating the agent's performance against predefined criteria.
Expected outcome You will be able to run `pytest` tests that invoke your agent and automatically log results, inputs, outputs, and custom metrics to your Braintrust project. You can then view these test runs and their evaluations in the Braintrust UI, gaining insights into agent performance and detecting regressions.
How to test
  1. Ensure docker-compose up -d is running and the FastAPI application is running (poetry run uvicorn src.main:app --reload). 2. Populate BRAINTRUST_API_KEY and BRAINTRUST_PROJECT in your .env file. 3. Run the tests: poetry run pytest tests/test_agent.py. 4. Observe the console output for test results. 5. Log in to your Braintrust account and navigate to the specified project (production-ai-agent). You should see new test runs under the 'Experiments' section, corresponding to the bt.test decorators, with logged inputs, outputs, and metrics.

Phase 5: PostgreSQL Checkpointer & Production Configuration

This phase implements the LangGraph PostgreSQL checkpointer for persistent agent state, ensuring conversation continuity across restarts. It also refines the production configuration, preparing the agent for robust deployment and operation.

Python
src/database.py
import os
import logging
from typing import AsyncGenerator

from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import declarative_base
from sqlalchemy import text

logger = logging.getLogger(__name__)

Base = declarative_base()

class DBManager:
    def __init__(self, db_uri: str):
        self.db_uri = db_uri
        self.engine = create_async_engine(db_uri, echo=False, future=True)
        self.async_session_maker = async_sessionmaker(self.engine, expire_on_commit=False, class_=AsyncSession)
        logger.info(f"Database engine initialized for URI: {db_uri.split('@')[-1]}")

    async def create_all_tables(self):
        async with self.engine.begin() as conn:
            # LangGraph's PostgresSaver creates its own tables, so we just ensure connection
            # We can run a dummy query to ensure the connection is active
            await conn.execute(text("SELECT 1"))
            logger.info("Database connection verified.")

    async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
        async with self.async_session_maker() as session:
            try:
                yield session
            except Exception as e:
                logger.error(f"Database session error: {e}")
                await session.rollback()
                raise
            finally:
                await session.close()

db_manager: DBManager | None = None

async def initialize_db_manager(db_uri: str):
    global db_manager
    if db_manager is None:
        db_manager = DBManager(db_uri)
        await db_manager.create_all_tables()
        logger.info("DBManager initialized and tables checked.")
    return db_manager


src/main.py (updated)
import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver # For local dev, consider switching to PostgresSaver
from langgraph.checkpoint.postgres import PostgresSaver # Import PostgresSaver
from langgraph.graph import GraphCompiler
from opentelemetry import trace

from src.config import Config
from src.agent import create_agent_graph, AgentState
from src.observability import setup_opentelemetry, LLMTraceCallbackHandler, create_custom_span
from src.database import initialize_db_manager # Import DB initialization

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

config = Config()

class AgentInput(BaseModel):
    message: str
    thread_id: str | None = None

# Global variable for the checkpointer
checkpointer: PostgresSaver | None = None
agent_app = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Starting up application...")
    global agent_app, checkpointer
    try:
        setup_opentelemetry(app)
        
        # Initialize DB manager and PostgresSaver
        db_manager = await initialize_db_manager(config.DB_URI)
        checkpointer = PostgresSaver.from_async_engine(db_manager.engine)
        logger.info("PostgreSQL checkpointer initialized.")

        agent_app = create_agent_graph().with_config({
            "recursion_limit": 5,
            "checkpointer": checkpointer
        })
        logger.info("LangGraph agent initialized with PostgreSQL checkpointer.")
    except Exception as e:
        logger.critical(f"Failed to initialize application: {e}", exc_info=True)
        raise RuntimeError(f"Application initialization failed: {e}")
    yield
    logger.info("Shutting down application...")

app = FastAPI(title="Production AI Agent", version="0.1.0", lifespan=lifespan)

@app.get("/health")
async def health_check():
    return {"status": "ok", "service": config.SERVICE_NAME}

@app.post("/invoke")
async def invoke_agent(input: AgentInput, request: Request):
    logger.info(f"Received invocation for thread_id: {input.thread_id} with message: {input.message}")
    if agent_app is None:
        raise HTTPException(status_code=503, detail="Agent not initialized.")
    if checkpointer is None:
        raise HTTPException(status_code=503, detail="Checkpointer not initialized.")

    thread_id = input.thread_id if input.thread_id else "default_thread"
    
    with create_custom_span("agent.invoke", {"thread_id": thread_id, "user_input": input.message}) as span:
        try:
            # LangGraph's ainvoke takes callbacks directly in the config
            result = await agent_app.ainvoke(
                {"input": input.message, "chat_history": [], "turn_count": 0},
                config={
                    "configurable": {"thread_id": thread_id},
                    "tags": ["agent_invocation"],
                    "callbacks": [LLMTraceCallbackHandler(model_name=config.OPENAI_MODEL_NAME)]
                }
            )
            final_response = result.get("agent_outcome", "No direct response from agent.")
            if not final_response and result.get("tool_output"):
                final_response = f"Tool executed. Output: {result['tool_output']}"
            
            span.set_attribute("agent.response", final_response)
            return {"response": final_response, "thread_id": thread_id}
        except Exception as e:
            span.set_attribute("error", True)
            span.set_attribute("error.message", str(e))
            logger.exception(f"Error invoking agent: {e}")
            raise HTTPException(status_code=500, detail=f"Agent invocation failed: {str(e)}")
This phase solidifies the agent's production readiness by introducing state persistence and robust configuration. src/database.py is added to manage PostgreSQL connections. It defines a DBManager class that uses SQLAlchemy's create_async_engine and async_sessionmaker for asynchronous database operations. The create_all_tables method (which in the context of LangGraph's PostgresSaver primarily ensures connectivity) and get_session for managing database sessions are provided. A global db_manager is initialized within an initialize_db_manager function to ensure a single, managed database connection pool. src/main.py is updated to integrate the PostgreSQL checkpointer. During the FastAPI lifespan startup, initialize_db_manager is called to establish the database connection. Then, PostgresSaver.from_async_engine(db_manager.engine) is used to create a LangGraph-compatible checkpointer. This checkpointer is then passed to the agent_app's configuration using .with_config({"checkpointer": checkpointer}). This ensures that all agent states, associated with a thread_id, are automatically saved to and loaded from the PostgreSQL database. This is a critical feature for long-running conversations, agent recovery from crashes, and handling concurrent users. With the PostgresSaver in place, the agent's conversation history is no longer lost between invocations or application restarts. When agent_app.ainvoke is called with a thread_id, LangGraph will automatically retrieve the existing state for that thread from PostgreSQL, append the new message, process it through the graph, and then save the updated state. This makes the agent truly stateful and resilient. The DB_URI environment variable (defined in .env.example) now points to our running PostgreSQL instance.
Expected outcome The agent will now persist conversation state in PostgreSQL. If you interact with the agent using a specific `thread_id`, restart the FastAPI application, and then send another message with the *same* `thread_id`, the agent will recall the previous conversation context. You can inspect the PostgreSQL database to see the `checkpoints` table created and populated.
How to test
  1. Ensure PostgreSQL is running via docker-compose up -d. 2. Restart the FastAPI application: poetry run uvicorn src.main:app --reload. 3. Send a POST request to http://localhost:8000/invoke with {"message": "Hello, what's your name?", "thread_id": "my-persistent-chat-1"}. 4. Send another POST request to http://localhost:8000/invoke with {"message": "Can you tell me more about it?", "thread_id": "my-persistent-chat-1"}. The agent should respond in context. 5. Crucially, stop the FastAPI application (Ctrl+C) and restart it. 6. Send a third request with {"message": "And what about your purpose?", "thread_id": "my-persistent-chat-1"}. The agent should continue the conversation, demonstrating state persistence. 7. Optionally, connect to the PostgreSQL database (e.g., docker exec -it <postgres-container-id> psql -U user -d agentdb) and run \dt to see the checkpoints table, then SELECT * FROM checkpoints; to view stored states.

Testing & Evaluation

Robust testing and evaluation are non-negotiable for production AI agents, given their non-deterministic nature. This section outlines a comprehensive approach to ensure our agent's quality and reliability.

Functional Testing

Functional testing verifies that the agent performs as expected across various scenarios. This includes unit tests for individual tools and LangGraph nodes, and integration tests for the entire agent workflow. Our tests/test_agent.py demonstrates integration testing using httpx to hit the FastAPI endpoint. We cover:

  • Direct Response: Agent correctly answers simple queries without tool use.
  • Tool Use: Agent correctly identifies when to use a tool and processes its output.
  • Error Handling: Agent gracefully handles errors during tool execution or LLM interaction.

These tests are executed via pytest and can be run in any CI/CD pipeline.

LLM-as-a-Judge Evaluation

LLM-as-a-Judge (LLM-A-J) is critical for evaluating the subjective quality of agent responses, especially when no single 'correct' answer exists. Braintrust provides a powerful platform for this. Our tests/test_agent.py integrates directly with Braintrust, logging inputs, outputs, and custom metrics for each test run. Braintrust can then run an LLM-A-J against these logged runs to score agent performance based on criteria like helpfulness, accuracy, conciseness, and safety. This allows for qualitative assessment at scale.

Automated Evaluation Code Example (from tests/test_agent.py):

import pytest
import httpx
import os
import braintrust as bt
from dotenv import load_dotenv

load_dotenv()

BRAINTRUST_API_KEY = os.environ.get("BRAINTRUST_API_KEY")
BRAINTRUST_PROJECT = os.environ.get("BRAINTRUST_PROJECT", "production-ai-agent")
AGENT_SERVICE_URL = os.environ.get("AGENT_SERVICE_URL", "http://localhost:8000")

if not BRAINTRUST_API_KEY:
    pytest.skip("BRAINTRUST_API_KEY not set, skipping Braintrust tests", allow_module_level=True)

bt.init(api_key=BRAINTRUST_API_KEY, project=BRAINTRUST_PROJECT)

@pytest.fixture(scope="module")
def client():
    return httpx.Client(base_url=AGENT_SERVICE_URL)

@bt.test(project=BRAINTRUST_PROJECT, experiment="agent-functional-test")
@pytest.mark.asyncio
async def test_agent_direct_response(client: httpx.Client, bt_record: bt.Span):
    input_message = "What is the capital of France?"
    thread_id = "test-direct-response-1"
    
    bt_record.log(input=input_message, expected="Paris")

    try:
        response = client.post(
            "/invoke", 
            json={
                "message": input_message,
                "thread_id": thread_id
            }
        )
        response.raise_for_status()
        output = response.json()["response"]
        
        bt_record.log(output=output)
        # Example: LLM-as-a-Judge can evaluate this 'accuracy' metric in Braintrust
        bt_record.span.log(metrics=bt.Metrics(accuracy=(1 if "paris" in output.lower() else 0)))
        assert "paris" in output.lower()
    except Exception as e:
        bt_record.log(error=str(e))
        pytest.fail(f"Agent invocation failed: {e}")

This snippet shows how bt.test captures the input and output of the agent, and how custom metrics like accuracy can be logged. Braintrust then uses these logs to run configurable LLM-A-J prompts, providing a score and rationale for each response.

Failure Scenario Testing

We must anticipate failures. This includes:

  • LLM Provider Outages/Timeouts: Simulate unavailable or slow LLM responses and ensure the agent's error handling (e.g., retries, fallback LLMs) is robust.
  • Tool Failures: Our dummy_search tool demonstrates returning errors. Comprehensive tests should cover all known failure modes of integrated tools.
  • Database Connectivity: Test scenarios where the PostgreSQL checkpointer is unreachable.

Edge Case Coverage

Agents can behave unpredictably with unusual inputs:

  • Empty or Malformed Input: What happens if the user sends an empty string or non-textual input?
  • Ambiguous Queries: How does the agent handle questions with multiple interpretations?
  • Prompt Injections: While not fully covered here, advanced testing should include red-teaming for prompt injection attempts.

Performance Validation

Beyond correctness, performance is key. Observability (OpenTelemetry) helps here:

  • Latency: Monitor end-to-end response times and individual span durations for bottlenecks. Use P99 latency as a key metric.
  • Throughput: Measure requests per second under load to understand scaling limits.
  • Cost: Track llm.cost.total_usd attributes to identify expensive queries or inefficient prompt engineering.

Regression Testing

Automated tests, especially those leveraging Braintrust, are crucial for regression testing. Every code change should trigger a test suite run against a golden dataset. If the LLM-A-J scores drop below a threshold, or if specific functional tests fail, the change is flagged as a regression, preventing degraded performance from reaching production. This continuous feedback loop is vital for iterative agent development.

Deployment

Deploying a production AI agent requires careful consideration of containerization, orchestration, and operational best practices. This guide details how to package our LangGraph agent using Docker and deploy it to a Kubernetes cluster.

Dockerfile

Our Dockerfile creates a lightweight, production-ready image for the FastAPI application. It uses a multi-stage build to minimize the final image size and improve security.

# Stage 1: Build dependencies
FROM python:3.10-slim-buster AS builder

WORKDIR /app

# Install poetry
RUN pip install poetry==1.8.2

# Copy project files
COPY pyproject.toml poetry.lock ./ 

# Install dependencies
RUN poetry install --no-root --no-dev

# Stage 2: Final image
FROM python:3.10-slim-buster

WORKDIR /app

# Copy installed dependencies from builder stage
COPY --from=builder /usr/local/bin/poetry /usr/local/bin/poetry
COPY --from=builder /usr/local/lib/python3.10/site-packages /usr/local/lib/python3.10/site-packages

# Copy application code
COPY src ./src
COPY .env.example ./.env.example # Copy example for reference, secrets will be mounted

# Expose port for FastAPI
EXPOSE 8000

# Set environment variables for Gunicorn and FastAPI
ENV PYTHONUNBUFFERED=1
ENV PORT=8000

# Install Gunicorn for production server
RUN pip install gunicorn==22.0.0

# Command to run the application with Gunicorn
CMD ["gunicorn", "src.main:app", "--workers", "4", "--worker-class", "uvicorn.workers.UvicornWorker", "--bind", "0.0.0.0:8000"]

This Dockerfile first installs dependencies with Poetry in a builder stage, then copies only the necessary runtime files into a slim final image. It uses Gunicorn with Uvicorn workers for production-grade concurrency and performance.

Kubernetes Manifest

Deploying to Kubernetes involves several resources: a Deployment for the agent pods, a Service to expose it, and potentially ConfigMap and Secret for configuration.

config/k8s/agent-deployment.yaml:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: production-ai-agent-deployment
  labels:
    app: production-ai-agent
spec:
  replicas: 2
  selector:
    matchLabels:
      app: production-ai-agent
  template:
    metadata:
      labels:
        app: production-ai-agent
      annotations:
        prometheus.io/scrape: "true"
        prometheus.io/port: "8000"
        prometheus.io/path: "/metrics" # If you add Prometheus metrics endpoint
    spec:
      containers:
      - name: agent
        image: your-docker-repo/production-ai-agent:latest # Replace with your image
        ports:
        - containerPort: 8000
        envFrom:
        - configMapRef:
            name: agent-config
        - secretRef:
            name: agent-secrets
        resources:
          requests:
            memory: "256Mi"
            cpu: "200m"
          limits:
            memory: "512Mi"
            cpu: "500m"
        readinessProbe:
          httpGet:
            path: /health
            port: 8000
          initialDelaySeconds: 5
          periodSeconds: 10
          timeoutSeconds: 3
          failureThreshold: 3
        livenessProbe:
          httpGet:
            path: /health
            port: 8000
          initialDelaySeconds: 15
          periodSeconds: 20
          timeoutSeconds: 5
          failureThreshold: 3

config/k8s/agent-service.yaml:

apiVersion: v1
kind: Service
metadata:
  name: production-ai-agent-service
  labels:
    app: production-ai-agent
spec:
  selector:
    app: production-ai-agent
  ports:
    - protocol: TCP
      port: 80
      targetPort: 8000
  type: ClusterIP # Use LoadBalancer for external access in cloud environments

Secrets Management

Never hardcode sensitive values. In Kubernetes, Secrets are used to manage credentials like OPENAI_API_KEY and BRAINTRUST_API_KEY. ConfigMaps store non-sensitive configuration.

# config/k8s/agent-config.yaml
apiVersion: v1
kind: ConfigMap
metadata:
  name: agent-config
data:
  DB_URI: "postgresql+asyncpg://user:password@postgres-service:5432/agentdb" # Update for your K8s Postgres service
  OTEL_EXPORTER_OTLP_ENDPOINT: "http://tempo-distributor.observability.svc.cluster.local:4317" # Update for your K8s OTLP collector
  SERVICE_NAME: "production-ai-agent-service"
  OPENAI_MODEL_NAME: "gpt-4o"
  BRAINTRUST_PROJECT: "production-ai-agent"
---
# config/k8s/agent-secrets.yaml (Base64 encoded values)
apiVersion: v1
kind: Secret
metadata:
  name: agent-secrets
type: Opaque
data:
  OPENAI_API_KEY: "<base64-encoded-openai-key>"
  BRAINTRUST_API_KEY: "<base64-encoded-braintrust-key>"

Create these with kubectl create -f config/k8s/agent-config.yaml and echo -n 'sk-...' | base64 for secrets.

Health Check and Readiness Probe

Kubernetes readinessProbe and livenessProbe are configured to use the /health endpoint. The readiness probe prevents traffic from reaching unhealthy pods, while the liveness probe restarts pods that become unresponsive, enhancing reliability.

Scaling Considerations

The Deployment specifies replicas: 2 for high availability. For dynamic scaling, Kubernetes HorizontalPodAutoscalers (HPAs) can be configured to adjust replica counts based on CPU utilization or custom metrics (e.g., number of concurrent requests, LLM token usage). For LLM-intensive workloads, consider scaling based on GPU availability if using self-hosted models, or managing rate limits with external LLM providers.

Rollback Strategy

Kubernetes deployments support automatic rollbacks. If a new deployment introduces issues, you can revert to a previous stable version using kubectl rollout undo deployment/production-ai-agent-deployment. Ensure your PostgresSaver schema changes are backward-compatible or managed with database migrations (e.g., Alembic) to prevent data corruption during rollbacks. Automated CI/CD pipelines should include integration tests and evaluation to catch regressions before deployment, making rollbacks less frequent but still available as a safety net.

Production Hardening

Hardening a production AI agent involves a multi-faceted approach, addressing security, observability, reliability, and cost. This checklist outlines key considerations for our LangGraph agent.

Security

  • Input/Output Sanitization: Implement rigorous validation and sanitization for all user inputs (message) to prevent prompt injection, data leakage, and other OWASP Top 10 for Agentic AI risks. Similarly, sanitize agent outputs before displaying them to users to prevent cross-site scripting (XSS) or other vulnerabilities.
  • Least Privilege: Ensure the Kubernetes service account and database user for the agent have only the minimum necessary permissions. For example, the PostgreSQL user should only have access to the checkpoints table and no other sensitive data.
  • Dependency Scanning: Integrate Snyk, Trivy, or similar tools into your CI/CD pipeline to scan Docker images and Python dependencies for known vulnerabilities, patching them promptly.
  • Secrets Management: Reinforce the use of Kubernetes Secrets (or equivalent cloud secrets managers like AWS Secrets Manager, Azure Key Vault) for all sensitive credentials. Never bake secrets into Docker images or commit them to source control.

Observability

  • Comprehensive Logging: Configure structured logging (e.g., JSON format) with appropriate log levels (DEBUG, INFO, WARNING, ERROR) for all agent components. Centralize logs with a system like Grafana Loki or Elastic Stack for easy searching and analysis. Our current logging setup is a good start, but consider a dedicated logging library like structlog.
  • Custom Metrics: Beyond OpenTelemetry traces, expose Prometheus metrics for key agent performance indicators, such as total invocations, average latency, LLM calls per request, tool execution success rates, and token counts. Use a library like prometheus_client in FastAPI.
  • Alerting: Set up alerts in Grafana (or your monitoring system) for critical conditions: high error rates, increased P99 latency, unexpected cost spikes, or unresponsive agent pods. Proactive alerting is key to minimizing downtime.
  • Trace Context Propagation: Ensure OpenTelemetry trace context is propagated across all internal services and, if applicable, to external LLM providers or tools to enable end-to-end tracing of complex workflows.

Reliability

  • Retry Mechanisms: Implement exponential backoff and retry logic for external API calls (LLM providers, tools, database connections) to handle transient failures. LangChain's LLM clients often have built-in retry mechanisms, but custom tools need explicit handling.
  • Circuit Breakers: For critical external dependencies, implement circuit breakers to prevent cascading failures. If a dependency is consistently failing, the circuit breaker can temporarily stop calls to it, allowing it to recover and preventing the agent from getting stuck.
  • Fallback LLMs: Configure the agent to use a cheaper or less performant fallback LLM provider or model if the primary one experiences high latency or outages. This can be integrated into the LangGraph state machine with conditional routing.
  • Graceful Shutdown: Ensure the FastAPI application handles SIGTERM signals for graceful shutdown, allowing ongoing requests to complete before the pod terminates in Kubernetes.

Rate Limiting

  • API Gateway Rate Limits: Implement rate limiting at the ingress/API gateway level (e.g., Nginx, Envoy, cloud load balancers) to protect the agent from excessive traffic and abuse.
  • Internal Rate Limits: Apply internal rate limits when interacting with external LLM providers or tools to respect their API quotas and avoid being throttled.
  • Concurrency Limits: Configure Gunicorn workers and Uvicorn threads to match the available CPU resources and avoid overloading the agent with too many concurrent requests.

Authentication

  • API Key Authentication: For direct client access, use API key authentication for the FastAPI endpoint. Implement a secure token generation and validation mechanism.
  • OAuth/RBAC for Tools: If tools access sensitive internal systems, ensure they use appropriate authentication and authorization (e.g., OAuth 2.1, Role-Based Access Control) to prevent unauthorized data access.

Governance

  • Prompt Versioning: Manage system prompts and tool descriptions under version control. Track changes and their impact on agent performance (via Braintrust evaluations). This helps in prompt drift detection.
  • Audit Trails: Log all significant agent decisions, tool calls, and LLM interactions, including who initiated the request, for auditing and compliance purposes.
  • Human-in-the-Loop (HITL): For high-stakes decisions or uncertain outputs, integrate HITL escalation points where the agent can pause and request human review or approval.

Cost Optimisation

  • Token Limits: Implement token limits for LLM requests to prevent runaway costs from excessively long prompts or responses. Our cost attribution already provides visibility.
  • Model Selection: Dynamically select the most cost-effective LLM based on the complexity of the query. For simple tasks, use cheaper models (e.g., gpt-4o-mini) and reserve powerful models (e.g., GPT-4o) for complex reasoning.
  • Caching: Cache LLM responses for identical or highly similar prompts to reduce redundant LLM calls and costs. This can be implemented at the application level or via an external cache like Redis.

Compliance

  • Data Residency: Ensure that sensitive user data and conversation history (stored in PostgreSQL) comply with data residency requirements by deploying databases in appropriate geographical regions.
  • PII Handling: Implement mechanisms to detect, redact, or anonymize Personally Identifiable Information (PII) from prompts and responses before they reach the LLM or persistent storage.