api: add json logging and request metadata

Author: whitebeard10

app/main.py

@@ -1,13 +1,29 @@
 import os
 import time
 import logging
-from fastapi import FastAPI, HTTPException
+import json
+import uuid
+from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
 from src.config import settings
 
-# Setup logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# Structured JSON Logging Setup
+class JsonFormatter(logging.Formatter):
+    def format(self, record):
+        log_record = {
+            "timestamp": self.formatTime(record, self.datefmt),
+            "level": record.levelname,
+            "message": record.getMessage(),
+            "module": record.module,
+            "request_id": getattr(record, "request_id", "N/A")
+        }
+        return json.dumps(log_record)
+
+logger = logging.getLogger("api_logger")
+handler = logging.StreamHandler()
+handler.setFormatter(JsonFormatter())
+logger.addHandler(handler)
+logger.setLevel(logging.INFO)
 
 app = FastAPI(title="Mistral Domain Support Assistant API")
 
@@ -22,9 +38,19 @@ class GenerateRequest(BaseModel):
 
 class GenerateResponse(BaseModel):
     response: str
-    latency: float
+    latency_seconds: float
     tokens_generated: int
     model_name: str
+    request_id: str
+    metadata: dict
+
+@app.middleware("http")
+async def add_request_id(request: Request, call_next):
+    request_id = str(uuid.uuid4())
+    request.state.request_id = request_id
+    response = await call_next(request)
+    response.headers["X-Request-ID"] = request_id
+    return response
 
 @app.on_event("startup")
 async def startup_event():
@@ -33,35 +59,61 @@ async def startup_event():
         logger.info("TESTING environment detected. Skipping real model loading.")
         return
 
-    from src.inference.engine import InferenceEngine
-    adapter_path = os.getenv("ADAPTER_PATH", None)
-    logger.info(f"Initializing Inference Engine with adapter: {adapter_path}")
-    engine = InferenceEngine(adapter_path=adapter_path)
+    try:
+        from src.inference.engine import InferenceEngine
+        adapter_path = os.getenv("ADAPTER_PATH", None)
+        logger.info(f"Initializing Inference Engine with adapter: {adapter_path}")
+        engine = InferenceEngine(adapter_path=adapter_path)
+    except Exception as e:
+        logger.error(f"Failed to load model: {str(e)}")
 
 @app.post("/generate", response_model=GenerateResponse)
-async def generate(request: GenerateRequest):
+async def generate(request: GenerateRequest, fast_request: Request):
+    request_id = fast_request.state.request_id
+    
     if engine is None:
+        logger.warning("Generation attempted but model not loaded", extra={"request_id": request_id})
         raise HTTPException(status_code=503, detail="Model not loaded")
 
+    logger.info(f"Processing generation request: {request.instruction[:50]}...", extra={"request_id": request_id})
+    
     try:
+        start_time = time.time()
         result = engine.generate(
             instruction=request.instruction,
             max_new_tokens=request.max_new_tokens
         )
+        total_latency = time.time() - start_time
+        
+        logger.info("Generation successful", extra={
+            "request_id": request_id,
+            "latency": total_latency,
+            "tokens": result['tokens_generated']
+        })
 
         return GenerateResponse(
             response=result['response'],
-            latency=result['latency'],
+            latency_seconds=round(total_latency, 3),
             tokens_generated=result['tokens_generated'],
-            model_name=settings.BASE_MODEL_NAME
+            model_name=settings.BASE_MODEL_NAME,
+            request_id=request_id,
+            metadata={
+                "temperature": request.temperature,
+                "top_p": request.top_p,
+                "timestamp": time.time()
+            }
         )
     except Exception as e:
-        logger.error(f"Generation error: {str(e)}")
+        logger.error(f"Generation error: {str(e)}", extra={"request_id": request_id})
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/health")
 async def health():
-    return {"status": "ok", "model_loaded": engine is not None}
+    return {
+        "status": "ok", 
+        "model_loaded": engine is not None,
+        "environment": "testing" if os.getenv("TESTING") else "production"
+    }
 
 if __name__ == "__main__":
     import uvicorn

tests/test_api.py

@@ -4,16 +4,13 @@
 client = TestClient(app)
 
 def test_read_health():
-    # Since the model won't be loaded in a standard test environment, 
-    # we expect the model_loaded status to be false or for the startup to fail.
-    # We can mock the engine if needed, but a simple health check is a good start.
     response = client.get("/health")
     assert response.status_code == 200
     assert "status" in response.json()
     assert response.json()["status"] == "ok"
+    assert "environment" in response.json()
 
 def test_generate_fail_no_model():
-    # Test that /generate returns 503 if engine is not initialized
     response = client.post(
         "/generate",
         json={"instruction": "How do I reset my password?"}