telemetry(1).py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  1. # backend/telemetry.py
  2. from fastapi import APIRouter, Request, Depends
  3. from pydantic import BaseModel
  4. import time, hashlib, hmac, os, sqlite3, json, hashlib
  5. from datetime import datetime
  6. from typing import Any, Dict, List, Optional
  7. router = APIRouter()
  8. DB_PATH = os.getenv("TPR_DB", "/data/telemetry.db")
  9. IP_SECRET = os.getenv("TPR_IP_SECRET", "change-me")
  10. def db():
  11. os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
  12. conn = sqlite3.connect(DB_PATH)
  13. try:
  14. conn.execute("PRAGMA journal_mode=WAL;")
  15. except sqlite3.OperationalError:
  16. pass # read-only FS or WAL not supported — continue anyway
  17. return conn
  18. def ip_hash(ip: str) -> str:
  19. dig = hmac.new(IP_SECRET.encode(), ip.encode(), 'sha256').hexdigest()
  20. return dig[:12]
  21. class TelemetryEvent(BaseModel):
  22. type: str
  23. ts: str
  24. sid: str
  25. ua: Optional[str] = None
  26. data: Dict[str, Any] = {}
  27. @router.on_event("startup")
  28. def init():
  29. try:
  30. with db() as conn:
  31. conn.executescript("""
  32. CREATE TABLE IF NOT EXISTS events (
  33. id INTEGER PRIMARY KEY AUTOINCREMENT,
  34. ts TEXT NOT NULL,
  35. type TEXT NOT NULL,
  36. sid TEXT,
  37. ua TEXT,
  38. ip_hash TEXT,
  39. data_json TEXT
  40. );
  41. CREATE TABLE IF NOT EXISTS ask_logs (
  42. id INTEGER PRIMARY KEY AUTOINCREMENT,
  43. ts TEXT NOT NULL,
  44. sid TEXT,
  45. ip_hash TEXT,
  46. query TEXT,
  47. normalized TEXT,
  48. allow_tps INTEGER,
  49. latency_ms INTEGER,
  50. model TEXT,
  51. ok INTEGER,
  52. topk_json TEXT,
  53. tokens_in INTEGER,
  54. tokens_out INTEGER
  55. );
  56. """)
  57. conn.commit()
  58. except Exception as e:
  59. print(f"[telemetry] DB init failed: {e} — telemetry will be disabled")
  60. # Do NOT re-raise — let the app start without telemetry
  61. @router.post("/telemetry")
  62. async def telemetry(ev: TelemetryEvent, request: Request):
  63. ip = request.client.host if request.client else "0.0.0.0"
  64. with db() as conn:
  65. conn.execute(
  66. "INSERT INTO events (ts,type,sid,ua,ip_hash,data_json) VALUES (?,?,?,?,?,json(?))",
  67. (ev.ts, ev.type, ev.sid, ev.ua, ip_hash(ip), json_dumps(ev.data))
  68. )
  69. conn.commit()
  70. return {"ok": True}
  71. # Wrap your /ask handler to also log authoritative facts:
  72. from fastapi import APIRouter
  73. import json, re
  74. ask_router = APIRouter()
  75. def normalize(q: str) -> str:
  76. return re.sub(r"\s+", " ", q.strip().lower())
  77. def json_dumps(o): return json.dumps(o, ensure_ascii=False, separators=(",",":"))
  78. @ask_router.post("/ask")
  79. async def ask(req: Dict[str, Any], request: Request):
  80. t0 = time.perf_counter()
  81. ip = request.client.host if request.client else "0.0.0.0"
  82. sid = request.headers.get("x-tpr-sid") or request.cookies.get("sid") or ""
  83. query = (req.get("query") or "").strip()
  84. allow_tps = bool(req.get("allow_tps"))
  85. # ... run retrieval/LLM as you already do ...
  86. # mock result placeholders:
  87. model = "localai/llama-3.1"
  88. topk = [{"id": "doc123", "score": 0.83}]
  89. tokens_in, tokens_out = 250, 380
  90. ok = True
  91. answer = {"text": "…", "citations": topk, "model": model, "usage": {"input_tokens": tokens_in, "output_tokens": tokens_out}}
  92. latency = int((time.perf_counter() - t0) * 1000)
  93. with db() as conn:
  94. conn.execute("""
  95. INSERT INTO ask_logs (ts,sid,ip_hash,query,normalized,allow_tps,latency_ms,model,ok,topk_json,tokens_in,tokens_out)
  96. VALUES (?,?,?,?,?,?,?,?,?,?,?,?)
  97. """, (
  98. datetime.utcnow().isoformat(),
  99. sid, ip_hash(ip), query, normalize(query), int(allow_tps),
  100. latency, model, int(ok), json_dumps(topk), tokens_in, tokens_out
  101. ))
  102. conn.commit()
  103. return answer