|
41 | 41 | import logging |
42 | 42 | import os |
43 | 43 | import re |
| 44 | +import threading |
44 | 45 | import time |
45 | 46 | import uuid |
46 | 47 | from datetime import datetime, timezone |
| 48 | +from pathlib import Path |
47 | 49 |
|
48 | 50 | import numpy as np |
49 | 51 | from fastapi import FastAPI, HTTPException, Request |
@@ -89,6 +91,34 @@ def filter(self, record): |
89 | 91 | for handler in logging.getLogger().handlers: |
90 | 92 | handler.addFilter(_RequestIDFilter()) |
91 | 93 |
|
| 94 | +_query_log_lock = threading.Lock() |
| 95 | +_query_log_path: Path | None = None |
| 96 | + |
| 97 | + |
| 98 | +def _init_query_log(): |
| 99 | + global _query_log_path |
| 100 | + log_dir = Path(os.environ.get("PIXELRAG_QUERY_LOG_DIR", "logs")) |
| 101 | + log_dir.mkdir(parents=True, exist_ok=True) |
| 102 | + _query_log_path = log_dir / "queries.jsonl" |
| 103 | + |
| 104 | + |
| 105 | +def _log_query(req: "SearchRequest", request_id: str): |
| 106 | + if _query_log_path is None: |
| 107 | + return |
| 108 | + record = { |
| 109 | + "ts": datetime.now(timezone.utc).isoformat(), |
| 110 | + "request_id": request_id, |
| 111 | + "queries": [q.text for q in req.queries], |
| 112 | + "has_image": [q.image is not None for q in req.queries], |
| 113 | + "n_docs": req.n_docs, |
| 114 | + "department": req.department, |
| 115 | + } |
| 116 | + line = json.dumps(record, ensure_ascii=False) + "\n" |
| 117 | + with _query_log_lock: |
| 118 | + with open(_query_log_path, "a") as f: |
| 119 | + f.write(line) |
| 120 | + |
| 121 | + |
92 | 122 | app = FastAPI(title="PixelRAG Search API") |
93 | 123 |
|
94 | 124 | app.add_middleware( |
@@ -540,6 +570,8 @@ async def search(req: SearchRequest): |
540 | 570 | time.time() - t0, |
541 | 571 | ) |
542 | 572 |
|
| 573 | + _log_query(req, _request_id_ctx.get()) |
| 574 | + |
543 | 575 | return SearchResponse(results=results) |
544 | 576 |
|
545 | 577 |
|
@@ -616,6 +648,7 @@ async def tile_by_id(article_id: int, tile_index: int, chunk_index: int): |
616 | 648 |
|
617 | 649 | def load(args): |
618 | 650 | """Load index, metadata, model, and articles.json.""" |
| 651 | + _init_query_log() |
619 | 652 | import torch |
620 | 653 |
|
621 | 654 | device = args.device |
|
0 commit comments