Pythonllm_monitoring.py
Copy
1
import time
2
from collections import defaultdict
3
from contextlib import contextmanager
4
from fastapi import FastAPI
5
from openai import OpenAI
6
 
7
app = FastAPI()
8
llm_client = OpenAI()
9
 
10
class MetricsCollector:
11
    def __init__(self):
12
        self.requests = defaultdict(int)      # model -> count
13
        self.errors = defaultdict(int)         # model -> count
14
        self.tokens_in = defaultdict(int)      # model -> total
15
        self.tokens_out = defaultdict(int)     # model -> total
16
        self.latencies: list[float] = []
17
 
18
    @contextmanager
19
    def track(self, model: str):
20
        start = time.monotonic()
21
        self.requests[model] += 1
22
        try:
23
            yield
24
        except Exception:
25
            self.errors[model] += 1
26
            raise
27
        finally:
28
            self.latencies.append(time.monotonic() - start)
29
 
30
    def record_tokens(self, model: str, input_tokens: int, output_tokens: int):
31
        self.tokens_in[model] += input_tokens
32
        self.tokens_out[model] += output_tokens
33
 
34
    def prometheus_text(self) -> str:
35
        lines = ["# HELP llm_requests_total Total LLM requests by model"]
36
        for model, count in self.requests.items():
37
            lines.append(f'llm_requests_total{{model="{model}"}} {count}')
38
        for model, count in self.errors.items():
39
            lines.append(f'llm_errors_total{{model="{model}"}} {count}')
40
        return "\n".join(lines) + "\n"
41
 
42
metrics = MetricsCollector()
43
 
44
@app.post("/chat")
45
async def chat(message: str, model: str = "gpt-4o-mini"):
46
    with metrics.track(model):
47
        resp = llm_client.chat.completions.create(
48
            model=model,
49
            messages=[{"role": "user", "content": message}],
50
        )
51
        metrics.record_tokens(model, resp.usage.prompt_tokens, resp.usage.completion_tokens)
52
        return {"content": resp.choices[0].message.content}
53
 
54
@app.get("/metrics")
55
async def get_metrics():
56
    from fastapi.responses import PlainTextResponse
57
    return PlainTextResponse(metrics.prometheus_text(), media_type="text/plain")
58