Counting OpenAI Tokens per Request in FastAPI — Without Touching My Use Cases

I'm rebuilding evorag, a Python retrieval service: it chunks documents, embeds them with OpenAI's text-embedding-3-small, stores vectors in PostgreSQL + pgvector, and returns the most relevant chunks for a query. Answer generation is out of scope. The plan is to connect it to my NestJS rag-service, with evorag handling retrieval and rag-service handling the API and generation, but today the two run independently.

This week's goal was observability. For every request I wanted to know three things:

1. Which request a log line belongs to (a request ID that also works across both services)
2. How long it took — and how much of that was OpenAI vs. my own code
3. How many embedding tokens it consumed, because tokens are money

The catch: evorag follows Clean Architecture, and I didn't want any of this to leak into the domain or use-case layers.

The constraint: the dependency rule

The project is laid out like this:

src/evorag/
  domain/          Chunk entity (plain frozen dataclass)
  ingestion/       IngestDocument use case + ports
  retrieval/       SearchChunks use case + ports
  shared/          EmbeddingProvider port
  observability/   JSON formatter, per-request context
  adapters/        FastAPI, OpenAI, pgvector
  main.py          composition root

Source code dependencies point inward only. A use case knows an EmbeddingProvider protocol, not OpenAI:

class EmbeddingProvider(Protocol):
    async def embed(self, texts: list[str]) -> list[list[float]]: ...

Here's the problem. The token count is only known inside the OpenAI adapter (response.usage.total_tokens). The per-request log line is written in the HTTP middleware. Between them sits the use case.

The naive fix is to make embed() return (vectors, tokens) and thread the number up through IngestDocument.execute() back to the route. That would change the port, the use case, every fake in my tests, and the use case would suddenly "know" about billing metrics. That's exactly what the dependency rule is supposed to prevent.

So I needed a side channel: something request-scoped that the adapter can write to and the middleware can read from, with the use case in the middle never noticing.

Step 1: one JSON line per event

First, logs. I stuck to the standard library and wrote a small formatter that emits JSON Lines:

class JsonFormatter(logging.Formatter):
    def format(self, record: logging.LogRecord) -> str:
        payload = {
            "ts": datetime.fromtimestamp(record.created, UTC).isoformat(),
            "level": record.levelname,
            "logger": record.name,
            "event": record.getMessage(),
            "request_id": get_request_id(),
        }
        payload.update(getattr(record, "fields", {}))
        return json.dumps(payload)

The message is an event name (request.completed, embedding.created), and structured data goes in via extra={"fields": {...}}. Any log tool can then filter by event or aggregate latency_ms without regex.

Step 2: request ID with contextvars

contextvars is the async-safe way to hold "current request" state in Python. Each request gets its own value, even when hundreds are interleaved on one event loop.

_request_id: ContextVar[str | None] = ContextVar("request_id", default=None)

def set_request_id(request_id: str) -> None:
    _request_id.set(request_id)

def get_request_id() -> str | None:
    return _request_id.get()

The middleware reuses an incoming X-Request-ID header if there is one (so once rag-service calls evorag, the same ID can follow a request across both services), otherwise it generates one, and it echoes the ID back in the response:

async def log_requests(request: Request, call_next) -> Response:
    request_id = request.headers.get("x-request-id") or uuid.uuid4().hex
    set_request_id(request_id)
    ...
    response = await call_next(request)
    response.headers["x-request-id"] = request_id

Because the formatter calls get_request_id() itself, every log line from anywhere in the app gets the ID automatically. No logger needs to be passed around.

Step 3: token counting — and why it's not a ContextVar[int]

The obvious approach for tokens would be the same pattern as the request ID: a ContextVar[int], reset to 0 in the middleware, incremented in the adapter with .set(current + tokens), read back in the middleware after call_next.

That looks right, but it doesn't work. A minimal reproduction:

counter: ContextVar[int] = ContextVar("c", default=0)
holder: ContextVar[dict | None] = ContextVar("h", default=None)

@app.middleware("http")
async def mw(request, call_next):
    counter.set(0)
    h = {"n": 0}
    holder.set(h)
    response = await call_next(request)
    print(counter.get(), h["n"])
    return response

@app.get("/x")
async def x():
    counter.set(counter.get() + 42)
    holder.get()["n"] += 42
    return {}

Output:

0 42

The int version reads back 0. The object version reads back 42.

Why: @app.middleware("http") is Starlette's BaseHTTPMiddleware, and call_next runs the rest of the app in a separate task. A new task gets a copy of the current context. The endpoint can read everything the middleware set (that's why the request ID works), but a .set() inside the endpoint only changes the copy. The middleware never sees it.

What does cross the boundary is a mutable object. The copied context still points to the same object, so mutating it in place is visible on both sides.

So the fix is to store a small holder object in the ContextVar and mutate it, never re-set it:

@dataclass
class RequestMetrics:
    embedding_tokens: int = 0

_metrics: ContextVar[RequestMetrics | None] = ContextVar("metrics", default=None)

def start_metrics() -> RequestMetrics:
    metrics = RequestMetrics()
    _metrics.set(metrics)
    return metrics

def add_embedding_tokens(tokens: int) -> None:
    metrics = _metrics.get()
    if metrics is None:   # e.g. called outside a request (scripts, tests)
        return
    metrics.embedding_tokens += tokens

The middleware creates the holder before call_next and keeps a reference to it:

    metrics = start_metrics()
    started = time.perf_counter()
    response = await call_next(request)
    latency_ms = round((time.perf_counter() - started) * 1000, 2)

    logger.info("request.completed", extra={"fields": {
        "method": request.method,
        "path": request.url.path,
        "status": response.status_code,
        "latency_ms": latency_ms,
        "embedding_tokens": metrics.embedding_tokens,
    }})

And the OpenAI adapter, the only place that knows about tokens, reports them:

        started = time.perf_counter()
        response = await self._client.embeddings.create(...)
        latency_ms = round((time.perf_counter() - started) * 1000, 2)

        tokens = response.usage.total_tokens
        add_embedding_tokens(tokens)
        logger.info("embedding.created", extra={"fields": {
            "model": self._model,
            "inputs": len(texts),
            "tokens": tokens,
            "latency_ms": latency_ms,
        }})

The EmbeddingProvider port didn't change. IngestDocument and SearchChunks didn't change. The test fakes didn't change. Observability lives entirely in the outer ring: the HTTP adapter and the OpenAI adapter, both talking to a small observability module.

What the logs showed

One POST /documents now produces two lines (some fields trimmed):

{"event": "embedding.created", "request_id": "2a2c1c38...", "model": "text-embedding-3-small", "inputs": 1, "tokens": 18, "latency_ms": 390.19}
{"event": "request.completed", "request_id": "2a2c1c38...", "method": "POST", "path": "/documents", "status": 201, "latency_ms": 439.52, "embedding_tokens": 18}

Measuring the OpenAI call separately from the whole request turned out to be the most useful part. In my early measurements the embedding call took 89–97% of total request time, while evorag's own work (chunking, DB write, serialization) stayed around 50 ms.

That changes what I'll work on next. Optimizing my own code is pointless right now; the real levers are around the external call: timeouts, retries, and moving ingestion off the request path (a RabbitMQ consumer is on the roadmap).

Testing it

The observability tests use a tiny FastAPI app with the real middleware and a fake endpoint, so no database or API key is needed:

def make_app():
    app = FastAPI()
    app.middleware("http")(log_requests)

    @app.get("/fake")
    async def fake():
        add_embedding_tokens(30)
        add_embedding_tokens(12)
        return {"ok": "yes"}

    return app

def test_middleware_sums_embedding_tokens(caplog):
    caplog.set_level(logging.INFO, logger="evorag")
    TestClient(make_app()).get("/fake")

    [record] = [r for r in caplog.records if r.getMessage() == "request.completed"]
    assert record.fields["embedding_tokens"] == 42

This test is also a regression guard for the trap above: if someone "simplifies" RequestMetrics back into a plain ContextVar[int], it fails with 0 != 42.

Other tests cover generating a request ID, reusing an incoming X-Request-ID, and the formatter with and without extra fields.

Takeaways

• Clean Architecture makes cross-cutting concerns a design question, not an afterthought. "Where does this number get produced, and where is it consumed?" decided the whole solution.
• contextvars is the right tool for request-scoped data in async Python, but values set in a child task don't flow back to the parent. If you need to collect data, store a mutable holder.
• Measure external calls separately. One extra latency_ms field told me more about where to spend effort than any amount of profiling my own code would have.
• The standard library is enough to start. No logging framework, no tracing stack yet; JSON Lines on stdout can be picked up by any tool later.

Next up for evorag: timeouts and retries around the OpenAI calls, a RabbitMQ consumer for asynchronous ingestion, and a small evaluation set to calibrate chunk size and the similarity threshold properly.

The code is on GitHub: YusufOztoprak/evorag

Originally published on my blog.