AI Agent System Design: Production Best Practices

import json
import uuid

from fastapi import Depends, FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

app = FastAPI()


class RunRequest(BaseModel):
    message: str = Field(min_length=1, max_length=8000)
    idempotency_key: str | None = None


def sse(event: str, data: dict) -> str:
    return f"event: {event}\ndata: {json.dumps(data)}\n\n"


@app.post("/v1/threads/{thread_id}/runs/stream")
async def stream_run(thread_id: str, req: RunRequest, user=Depends(get_current_user)):
    request_id = str(uuid.uuid4())
    config = {
        "configurable": {"thread_id": thread_id, "user_id": user.id, "groups": user.groups},
        "recursion_limit": 25,                     # hard cap on graph steps
        "metadata": {"request_id": request_id, "user_id": user.id},
        "tags": ["prod", "assistant-v3"],
    }

    async def events():
        inputs = {"messages": [("user", req.message)], "user_id": user.id}
        async for mode, chunk in graph.astream(inputs, config, stream_mode=["updates", "messages"]):
            if mode == "messages":
                token, _meta = chunk
                if token.content:
                    yield sse("token", {"text": token.content})
            elif "__interrupt__" in chunk:          # approval needed
                yield sse("approval_required", {"payload": chunk["__interrupt__"][0].value})
            else:
                yield sse("step", {"node": next(iter(chunk))})
        yield sse("done", {"request_id": request_id})

    return StreamingResponse(events(), media_type="text/event-stream",
                             headers={"X-Request-ID": request_id})

import re

INJECTION_PATTERNS = [r"ignore (all|previous|prior) instructions", r"system prompt", r"you are now"]
MAX_CHARS = 8000


def input_guard(state: AgentState) -> dict:
    text = state["messages"][-1].content
    if len(text) > MAX_CHARS:
        return {"blocked": True, "block_reason": "input_too_long"}
    if any(re.search(p, text, re.I) for p in INJECTION_PATTERNS):
        return {"blocked": True, "block_reason": "possible_injection"}
    # L2: call a classifier (e.g. Azure Prompt Shields) here; keep rules as a cheap first pass
    return {"blocked": False}


def citations_valid(answer: str, retrieved_ids: set[str]) -> bool:
    cited = set(re.findall(r"\[([\w\-#\.]+)\]", answer))
    return bool(cited) and cited.issubset(retrieved_ids)

# pip install "langgraph>=1.0" langgraph-checkpoint-postgres
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.types import Command, RetryPolicy

with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
    checkpointer.setup()                          # run once (migrations)
    builder.add_node("tools", tool_node,
                     retry_policy=RetryPolicy(max_attempts=3))   # transient failures
    graph = builder.compile(checkpointer=checkpointer)

    cfg = {"configurable": {"thread_id": "t-123"}, "recursion_limit": 25}
    graph.invoke({"messages": [("user", "Open a ticket: VPN is down")]}, cfg)

    # Graph paused on interrupt(); inspect what it wants
    snapshot = graph.get_state(cfg)
    print(snapshot.tasks[0].interrupts[0].value)

    # Resume with the human decision
    graph.invoke(Command(resume={"approved": True}), cfg)

    # Time travel: list past checkpoints to replay or fork a bad run
    history = list(graph.get_state_history(cfg))

from typing import Literal
from pydantic import BaseModel, Field


class RouteDecision(BaseModel):
    route: Literal["rag", "action", "chitchat", "out_of_scope"]
    confidence: float = Field(ge=0, le=1)
    reason: str


router_llm = small_llm.with_structured_output(RouteDecision)


def router(state: AgentState) -> dict:
    d = router_llm.invoke([
        ("system", ROUTER_PROMPT),                 # versioned, with few-shot examples
        ("user", state["messages"][-1].content),
    ])
    route = d.route if d.confidence >= 0.6 else "clarify"
    return {"route": route}

import asyncio
from typing import Literal

from langchain_core.tools import tool
from pydantic import BaseModel, Field

MAX_TOOL_CHARS = 4000


class TicketArgs(BaseModel):
    title: str = Field(max_length=120, description="Short summary of the issue")
    description: str = Field(max_length=2000)
    priority: Literal["low", "medium", "high"] = "medium"


@tool(args_schema=TicketArgs)
async def create_ticket(title: str, description: str, priority: str = "medium") -> str:
    """Create an IT support ticket. Use only when the user explicitly asks to open
    a ticket. Do NOT use for questions about existing tickets (use get_ticket)."""
    try:
        ticket_id = await asyncio.wait_for(
            ticketing.create(title, description, priority, idempotency_key=current_run_key()),
            timeout=10,
        )
    except asyncio.TimeoutError:
        return "ERROR: ticketing system timed out. Tell the user and do not retry."
    return f"Created ticket {ticket_id}"


def cap(text: str) -> str:
    return text if len(text) <= MAX_TOOL_CHARS else text[:MAX_TOOL_CHARS] + "\n...[truncated]"

from datetime import datetime
from pydantic import BaseModel


class Chunk(BaseModel):
    id: str                       # f"{doc_id}#{chunk_no}" — stable, enables idempotent upsert
    doc_id: str
    content: str                  # with contextual header prepended
    title: str
    section_path: list[str]       # ["HR Policy", "Leave", "Parental leave"]
    page: int | None
    source_url: str
    doc_type: str                 # policy | manual | faq
    version: str
    effective_date: datetime | None
    allowed_groups: list[str]     # ACL — enforced at query time
    content_hash: str             # skip re-embedding unchanged chunks
    embedding_model: str          # e.g. "text-embedding-3-large@2024"
    content_vector: list[float]

from azure.search.documents import SearchClient
from azure.search.documents.models import VectorizedQuery


def search(client: SearchClient, query: str, user_groups: list[str], top: int = 8) -> list[dict]:
    groups = ",".join(user_groups)
    results = client.search(
        search_text=query,                                   # BM25 side
        vector_queries=[VectorizedQuery(vector=embed(query),
                                        k_nearest_neighbors=50,
                                        fields="content_vector")],
        filter=f"allowed_groups/any(g: search.in(g, '{groups}'))",   # security trimming
        query_type="semantic",                               # semantic reranker
        semantic_configuration_name="default",
        top=top,
        select=["id", "content", "title", "source_url", "page"],
    )
    return [dict(r) for r in results]


def rrf(rankings: list[list[str]], k: int = 60) -> list[str]:
    """Reciprocal Rank Fusion: score = sum(1 / (k + rank))."""
    scores: dict[str, float] = {}
    for ranking in rankings:
        for rank, doc_id in enumerate(ranking, start=1):
            scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (k + rank)
    return sorted(scores, key=scores.get, reverse=True)

# Option A — LangSmith: environment only, LangGraph traces automatically
# LANGSMITH_TRACING=true  LANGSMITH_API_KEY=...  LANGSMITH_PROJECT=assistant-prod

# Option B — Langfuse callback handler (v3 SDK path shown; verify for your version)
from langfuse.langchain import CallbackHandler

langfuse_handler = CallbackHandler()
config = {
    "configurable": {"thread_id": thread_id},
    "callbacks": [langfuse_handler],
    "metadata": {"request_id": request_id, "user_id": user_id,
                 "tenant_id": tenant_id, "prompt_version": "agent@12"},
}
graph.invoke(inputs, config)

import json
import pytest

GOLDEN = [json.loads(l) for l in open("evals/golden.jsonl")]
THRESHOLDS = {"recall_at_10": 0.85, "faithfulness": 0.90, "router_acc": 0.95}


def recall_at_k(retrieved: list[str], relevant: list[str], k: int = 10) -> float:
    return len(set(retrieved[:k]) & set(relevant)) / max(len(relevant), 1)


@pytest.fixture(scope="session")
def results():
    return [run_pipeline(case["question"], user=case["user"]) for case in GOLDEN]


def test_retrieval_recall(results):
    scores = [recall_at_k(r.retrieved_ids, c["relevant_ids"]) for r, c in zip(results, GOLDEN)]
    assert sum(scores) / len(scores) >= THRESHOLDS["recall_at_10"]


def test_faithfulness(results):
    scores = [judge_faithfulness(r.answer, r.context) for r in results]   # calibrated judge
    assert sum(scores) / len(scores) >= THRESHOLDS["faithfulness"]


def test_router(results):
    acc = sum(r.route == c["expected_route"] for r, c in zip(results, GOLDEN)) / len(GOLDEN)
    assert acc >= THRESHOLDS["router_acc"]

# pip install "langgraph>=1.0" langchain-openai langchain-core
from __future__ import annotations

import re
import uuid
from typing import Annotated, Literal, TypedDict

from langchain_core.messages import AIMessage, SystemMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langchain_openai import AzureChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver   # PostgresSaver in production
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langgraph.types import Command, interrupt
from pydantic import BaseModel, Field

# ---------- State ----------
class AgentState(TypedDict, total=False):
    messages: Annotated[list, add_messages]
    route: str
    blocked: bool
    retrieved_ids: list[str]
    guard_retries: int


# ---------- Models (all calls should go through your gateway) ----------
llm = AzureChatOpenAI(azure_deployment="gpt-4o", api_version="2024-10-21", temperature=0)
small_llm = AzureChatOpenAI(azure_deployment="gpt-4o-mini", api_version="2024-10-21", temperature=0)


# ---------- Stubs ----------
class FakeRetriever:
    def search(self, query: str, groups: list[str], top_k: int = 5) -> list[dict]:
        return [{"id": "it-policy#4", "text": "Laptops are refreshed every 3 years.",
                 "source": "IT-Policy.pdf p.4"}]

retriever = FakeRetriever()


# ---------- Tools ----------
@tool
def search_docs(query: str, config: RunnableConfig) -> str:
    """Search internal policy and technical documents. Use for any question about
    company policies, procedures or specs. Cite results as [id]."""
    groups = config["configurable"].get("groups", [])          # identity from API layer
    hits = retriever.search(query, groups=groups)
    return "\n".join(f"[{h['id']}] {h['text']} (source: {h['source']})" for h in hits)


class TicketArgs(BaseModel):
    title: str = Field(max_length=120)
    description: str = Field(max_length=2000)
    priority: Literal["low", "medium", "high"] = "medium"


@tool(args_schema=TicketArgs)
def create_ticket(title: str, description: str, priority: str = "medium") -> str:
    """Create an IT support ticket. Only when the user explicitly asks to open one."""
    decision = interrupt({"action": "create_ticket",
                          "args": {"title": title, "description": description, "priority": priority}})
    # Everything below runs only after resume. Side effects go AFTER interrupt().
    if not decision.get("approved"):
        return "The user rejected creating this ticket."
    ticket_id = f"INC-{uuid.uuid4().hex[:6].upper()}"         # ticketing.create(..., idempotency_key=...)
    return f"Created ticket {ticket_id}"


READ_TOOLS = [search_docs]
ALL_TOOLS = [search_docs, create_ticket]
TOOLS_BY_ROUTE = {"rag": READ_TOOLS, "action": ALL_TOOLS}


# ---------- Nodes ----------
INJECTION = re.compile(r"ignore (all|previous|prior) instructions|reveal .*system prompt", re.I)

def input_guard(state: AgentState) -> dict:
    text = state["messages"][-1].content
    return {"blocked": len(text) > 8000 or bool(INJECTION.search(text)), "guard_retries": 0}


class RouteDecision(BaseModel):
    route: Literal["rag", "action", "out_of_scope"]
    confidence: float = Field(ge=0, le=1)


def router(state: AgentState) -> dict:
    d = small_llm.with_structured_output(RouteDecision).invoke([
        SystemMessage("Classify the request: rag = question about company docs, "
                      "action = user wants a ticket created, out_of_scope = anything else."),
        state["messages"][-1],
    ])
    return {"route": d.route if d.confidence >= 0.6 else "out_of_scope"}


SYSTEM = SystemMessage("You are the internal IT/policy assistant. Answer only from tool results. "
                       "Cite every factual claim as [id]. If nothing relevant is found, say so.")

def agent(state: AgentState) -> dict:
    tools = TOOLS_BY_ROUTE[state["route"]]                      # route scopes the tools
    response = llm.bind_tools(tools).invoke([SYSTEM, *state["messages"]])
    return {"messages": [response]}


def output_guard(state: AgentState) -> dict:
    answer = state["messages"][-1].content
    if state["route"] == "rag" and not re.search(r"\[[\w\-#\.]+\]", answer):
        return {"guard_retries": state.get("guard_retries", 0) + 1,
                "messages": [SystemMessage("Your answer had no citations. Answer again with [id] citations, "
                                           "or say the documents don't cover it.")]}
    return {"guard_retries": -1}                                # -1 = passed


def refuse(state: AgentState) -> dict:
    return {"messages": [AIMessage("I can help with company policies and IT tickets. "
                                   "I can't help with that request.")]}


# ---------- Edges ----------
def after_input_guard(state: AgentState) -> str:
    return "refuse" if state["blocked"] else "router"

def after_router(state: AgentState) -> str:
    return "refuse" if state["route"] == "out_of_scope" else "agent"

def after_agent(state: AgentState) -> str:
    last = state["messages"][-1]
    return "tools" if getattr(last, "tool_calls", None) else "output_guard"

def after_output_guard(state: AgentState) -> str:
    r = state["guard_retries"]
    return END if r == -1 or r >= 2 else "agent"


# ---------- Graph ----------
builder = StateGraph(AgentState)
builder.add_node("input_guard", input_guard)
builder.add_node("router", router)
builder.add_node("agent", agent)
builder.add_node("tools", ToolNode(ALL_TOOLS))
builder.add_node("output_guard", output_guard)
builder.add_node("refuse", refuse)

builder.add_edge(START, "input_guard")
builder.add_conditional_edges("input_guard", after_input_guard, ["router", "refuse"])
builder.add_conditional_edges("router", after_router, ["agent", "refuse"])
builder.add_conditional_edges("agent", after_agent, ["tools", "output_guard"])
builder.add_edge("tools", "agent")
builder.add_conditional_edges("output_guard", after_output_guard, ["agent", END])
builder.add_edge("refuse", END)

graph = builder.compile(checkpointer=InMemorySaver())


# ---------- Run (tracing: set LANGSMITH_TRACING=true, or pass a Langfuse handler in callbacks) ----------
if __name__ == "__main__":
    cfg = {"configurable": {"thread_id": "demo-1", "user_id": "u1", "groups": ["it", "all"]},
           "recursion_limit": 25,
           "metadata": {"request_id": str(uuid.uuid4()), "prompt_version": "agent@1"}}

    out = graph.invoke({"messages": [("user", "Please open a ticket: my VPN keeps dropping")]}, cfg)
    if "__interrupt__" in out:
        print("Approval needed:", out["__interrupt__"][0].value)
        out = graph.invoke(Command(resume={"approved": True}), cfg)
    print(out["messages"][-1].content)