FastAPI Basic

This example shows a FastAPI server that enforces policies on chat messages.

Full code

import os

from fastapi import FastAPI
from pydantic import BaseModel

from njiraai import NjiraAI
from njiraai.middleware import create_middleware

app = FastAPI()

njira = NjiraAI(
    api_key=os.environ.get("NJIRA_API_KEY", "dev-key"),
    project_id=os.environ.get("NJIRA_PROJECT_ID", "dev-project"),
    mode="active",
    fail_strategy="fail_closed",
)

# Add middleware for context propagation and trace flushing
app.add_middleware(create_middleware(njira))


class ChatRequest(BaseModel):
    message: str


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


class ErrorResponse(BaseModel):
    error: str
    reasons: list[str] = []
    trace_id: str = ""


@app.post("/chat", response_model=ChatResponse | ErrorResponse)
async def chat(request: ChatRequest):
    # Pre-enforcement
    pre = await njira.enforce_pre(
        input_data=request.message,
        metadata={"endpoint": "/chat", "source": "user"},
    )

    if pre["verdict"] == "block":
        return ErrorResponse(
            error="Message blocked by policy",
            reasons=pre.get("reasons", []),
            trace_id=pre.get("traceId", ""),
        )

    effective_input = (
        pre.get("modifiedInput", request.message)
        if pre["verdict"] == "modify"
        else request.message
    )

    # Start span for LLM call
    span_id = njira.start_span(
        name="llm-call",
        span_type="llm",
        input_data={"prompt": effective_input},
    )

    try:
        # Simulate LLM response
        llm_response = f'I received: "{effective_input}"'
        njira.end_span(span_id, output=llm_response)

        # Post-enforcement
        post = await njira.enforce_post(
            output=llm_response,
            metadata={"endpoint": "/chat"},
        )

        if post["verdict"] == "block":
            return ErrorResponse(
                error="Response blocked by policy",
                reasons=post.get("reasons", []),
                trace_id=post.get("traceId", ""),
            )

        final_response = (
            post.get("modifiedOutput", llm_response)
            if post["verdict"] == "modify"
            else llm_response
        )

        return ChatResponse(
            response=final_response,
            trace_id=pre.get("traceId", ""),
        )

    except Exception as e:
        njira.span_error(span_id, e)
        return ErrorResponse(error="Internal error")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)

Run it

cd sdks/python/examples/fastapi-basic
uv sync
NJIRA_API_KEY=your-key NJIRA_PROJECT_ID=your-project uv run python main.py

Test it

curl -X POST http://localhost:8000/chat \
  -H "Content-Type: application/json" \
  -d '{"message": "Hello, world!"}'

What's happening

1. Middleware sets up request context and flushes traces after each request
2. enforce_pre validates input against policies
3. Span captures LLM call timing and I/O
4. enforce_post validates output before returning