Update app.py

app.py

@@ -1,9 +1,13 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
+from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+from contextlib import asynccontextmanager
 import uvicorn
-import importlib.util
 import os
-import sys
+import json
+import asyncio
+import logging
 from pathlib import Path
 
 import requests as _requests
@@ -11,11 +15,37 @@ import requests as _requests
 from core.pipeline_1.logic import PipelineLLMOnly
 from core.pipeline_2.logic import PipelineRAG
 from core.pipeline_3.logic import PipelineGraphRAG
+from services.metrics_service import MetricsService
+
+# --- SILENCE NOISY LOGGERS ---
+logging.getLogger("httpx").setLevel(logging.WARNING)
+logging.getLogger("httpcore").setLevel(logging.WARNING)
+logging.getLogger("transformers").setLevel(logging.ERROR)
+logging.getLogger("sentence_transformers").setLevel(logging.WARNING)
+logging.getLogger("pyTigerGraph").setLevel(logging.WARNING)
+logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Modern lifespan handler (replaces @on_event)"""
+    loop = asyncio.get_event_loop()
+    await loop.run_in_executor(None, shared_metrics.warmup)
+    yield
 
 app = FastAPI(
     title="SEC Dataset API",
     description="A simple FastAPI setup to fetch and interact with the PleIAs/SEC dataset.",
-    version="1.0.0"
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+# Enable CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
 # Request model
@@ -23,16 +53,23 @@ class QueryRequest(BaseModel):
     query: str
     ground_truth: str = None
 
-# Initialize Pipeline (Scaling up for full context)
+# Initialize Shared Services
+shared_metrics = MetricsService()
+
+# Initialize Pipelines
 pipeline_baseline = PipelineLLMOnly(top_n=20, max_full_text=20)
 pipeline_rag = PipelineRAG(retrieval_top_k=50, rerank_top_n=10, max_full_text=3)
 pipeline_graph = PipelineGraphRAG(rerank_top_n=10, max_full_text=3, retriever=pipeline_rag.retriever)
 
+# Inject shared metrics
+pipeline_baseline.metrics = shared_metrics
+pipeline_rag.metrics = shared_metrics
+pipeline_graph.metrics = shared_metrics
+
 @app.get("/")
 async def root():
     return {"message": "Welcome to the SEC Dataset API", "status": "running"}
 
-
 @app.get("/health")
 async def health():
     tg_host = os.environ.get("TG_HOST", "")
@@ -76,40 +113,34 @@ async def health():
         "qdrant": qdrant_status,
     }
 
+async def stream_pipeline(pipeline, query, ground_truth):
+    try:
+        for event in pipeline.run_stream(query, ground_truth):
+            yield f"data: {json.dumps(event)}\n\n"
+            await asyncio.sleep(0.01)
+    except Exception as e:
+        yield f"data: {json.dumps({'type': 'error', 'data': str(e)})}\n\n"
+
 @app.post("/query/baseline")
 async def query_baseline(request: QueryRequest):
-    """
-    Pipeline 1: LLM-Only Baseline.
-    Performs metadata filtering and LLM synthesis.
-    """
-    if not request.query:
-        raise HTTPException(status_code=400, detail="Query cannot be empty")
-
-    result = pipeline_baseline.run(request.query, ground_truth=request.ground_truth)
-    return result
-
+    return StreamingResponse(
+        stream_pipeline(pipeline_baseline, request.query, request.ground_truth),
+        media_type="text/event-stream"
+    )
 
 @app.post("/query/rag")
 async def query_rag(request: QueryRequest):
-    """
-    Pipeline 2: Hybrid RAG.
-    Hybrid vector search + cross-encoder rerank + Groq LLM synthesis.
-    """
-    if not request.query:
-        raise HTTPException(status_code=400, detail="Query cannot be empty")
-
-    result = pipeline_rag.run(request.query, ground_truth=request.ground_truth)
-    return result
-
+    return StreamingResponse(
+        stream_pipeline(pipeline_rag, request.query, request.ground_truth),
+        media_type="text/event-stream"
+    )
 
 @app.post("/query/graph")
 async def query_graph(request: QueryRequest):
-    """
-    Pipeline 3: GraphRAG.
-    TigerGraph vector search + topic linking + citation expansion + Groq LLM synthesis.
-    """
-    if not request.query:
-        raise HTTPException(status_code=400, detail="Query cannot be empty")
-
-    result = pipeline_graph.run(request.query, ground_truth=request.ground_truth)
-    return result
+    return StreamingResponse(
+        stream_pipeline(pipeline_graph, request.query, request.ground_truth),
+        media_type="text/event-stream"
+    )
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="warning")