FastAPI Integration

from contextlib import asynccontextmanager
from fastapi import FastAPI
from rdk import init, shutdown

@asynccontextmanager
async def lifespan(app: FastAPI):
    init(batch_size=10, flush_interval=5.0)
    yield
    shutdown()

app = FastAPI(title="My LLM API", lifespan=lifespan)

from contextlib import asynccontextmanager
from typing import Any

from fastapi import FastAPI, HTTPException
from langchain_anthropic import ChatAnthropic
from langchain_core.messages import HumanMessage, ToolMessage
from langchain_core.tools import tool
from pydantic import BaseModel

from rdk import flush, init, observe, shutdown


@asynccontextmanager
async def lifespan(app: FastAPI):
    init(batch_size=10, flush_interval=5.0)
    yield
    shutdown()


app = FastAPI(title="LLM API with Tracing", lifespan=lifespan)


# Request/Response models
class ChatRequest(BaseModel):
    message: str
    model: str = "claude-sonnet-4-6"


class ChatResponse(BaseModel):
    response: str
    model: str


# Chat endpoint with tracing
@app.post("/chat", response_model=ChatResponse)
@observe(name="api-chat")
async def chat_endpoint(request: ChatRequest) -> ChatResponse:
    try:
        llm = ChatAnthropic(model=request.model)
        response = await llm.ainvoke(request.message)

        return ChatResponse(
            response=response.content,
            model=request.model,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# Health check
@app.get("/health")
async def health_check():
    return {"status": "healthy"}


# Manual flush endpoint (useful for debugging)
@app.post("/flush")
async def flush_traces():
    flush()
    return {"status": "flushed"}

class AgentRequest(BaseModel):
    question: str


class AgentResponse(BaseModel):
    answer: str
    tools_used: list[str]


@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Search results for: {query}"


@tool
def calculate(expr: str) -> str:
    """Calculate math expression."""
    return str(eval(expr))


tools = [search, calculate]
tool_map = {t.name: t for t in tools}


@app.post("/agent", response_model=AgentResponse)
@observe(name="api-agent")
async def agent_endpoint(request: AgentRequest) -> AgentResponse:
    llm = ChatAnthropic(model="claude-sonnet-4-6")
    llm_with_tools = llm.bind_tools(tools)

    messages: list[Any] = [HumanMessage(content=request.question)]
    tools_used = []

    for _ in range(5):  # Max iterations
        response = await llm_with_tools.ainvoke(messages)
        messages.append(response)

        if not response.tool_calls:
            return AgentResponse(
                answer=response.content,
                tools_used=tools_used,
            )

        for tc in response.tool_calls:
            result = tool_map[tc["name"]].invoke(tc["args"])
            tools_used.append(tc["name"])
            messages.append(ToolMessage(
                content=result,
                tool_call_id=tc["id"]
            ))

    return AgentResponse(
        answer="Max iterations reached",
        tools_used=tools_used,
    )

import uuid
from fastapi import Request
from rdk import observe, get_current_trace_id

@app.post("/chat")
@observe(name="api-chat")
async def chat_with_request_id(
    request: Request,
    body: ChatRequest
) -> ChatResponse:
    request_id = request.headers.get("X-Request-ID", str(uuid.uuid4()))
    trace_id = get_current_trace_id()
    # Correlate request_id with trace_id in logs
    ...

from fastapi import HTTPException
from rdk import observe

@app.post("/chat")
@observe(name="api-chat", tags=["chat"])
async def chat_endpoint(request: ChatRequest) -> ChatResponse:
    try:
        # Your logic
        ...
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail="Internal error")

export RDK_API_KEY="your-api-key"
export ANTHROPIC_API_KEY="your-anthropic-key"

uvicorn app:app --reload

import pytest
from rdk import init, shutdown

@pytest.fixture(autouse=True)
def rdk_test():
    init(mode="test")
    yield
    shutdown()