""" main.py — SentinelAI-RAG FastAPI Production Server. Endpoints: - POST /v1/rag/query: Guardrails -> Cache -> SIMD Top-K RAG -> LLM Router -> Output Guard -> Cache Store - POST /v1/rag/stream: Server-Sent Events (SSE) streaming query response - GET /v1/health: Readiness, liveness, SIMD engine status, and SLA compliance - GET /v1/metrics: Real-time Prometheus/JSON metrics (cost reduction, cache hit rate, latencies) """ import time import json import logging from contextlib import asynccontextmanager from typing import Optional, List, Dict, Any from pathlib import Path from fastapi import FastAPI, HTTPException, Request, Depends, status from fastapi.responses import JSONResponse, StreamingResponse, FileResponse from fastapi.staticfiles import StaticFiles from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field from app.config import ( PDF_PATH, CACHE_SIMILARITY_THRESHOLD, RAG_TOP_K, PII_REDACTION_ENABLED, PROMPT_INJECTION_ENABLED, OUTPUT_GUARD_ENABLED, LLM_PRICING, AVG_INPUT_TOKENS, AVG_OUTPUT_TOKENS, ) from app.guardrails.pii_redactor import PIIRedactor from app.guardrails.prompt_injection import PromptInjectionGuard from app.guardrails.output_guard import OutputGuard from app.guardrails.rate_limiter import RateLimiter from app.embeddings.onnx_embedder import ONNXEmbedder from app.rag.pdf_chunker import PDFChunker from app.rag.retriever import SIMDRetriever, SIMD_AVAILABLE from app.rag.query_rewriter import QueryRewriter from app.rag.quality_evaluator import RetrievalQualityEvaluator from app.gateway.semantic_cache import SemanticCache, is_negative_or_refusal_response from app.gateway.circuit_breaker import CircuitBreaker from app.gateway.llm_router import LLMRouter logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s") logger = logging.getLogger("sentinelai.server") # Global subsystem state embedder: Optional[ONNXEmbedder] = None retriever: Optional[SIMDRetriever] = None query_rewriter: Optional[QueryRewriter] = None quality_evaluator: Optional[RetrievalQualityEvaluator] = None cache: Optional[SemanticCache] = None llm_router: Optional[LLMRouter] = None pii_redactor: Optional[PIIRedactor] = None injection_guard: Optional[PromptInjectionGuard] = None output_guard: Optional[OutputGuard] = None rate_limiter: Optional[RateLimiter] = None # Metrics aggregation SERVER_START_TIME = time.time() query_latencies: List[float] = [] total_queries_served = 0 total_cost_saved_usd = 0.0 total_baseline_cost_usd = 0.0 def initialize_pipeline(): """Initialize embedder, PDF chunker, retriever, cache, and guardrails.""" global embedder, retriever, cache, llm_router global pii_redactor, injection_guard, output_guard, rate_limiter global query_rewriter, quality_evaluator logger.info("Initializing SentinelAI-RAG subsystems...") # Guardrails pii_redactor = PIIRedactor() injection_guard = PromptInjectionGuard() output_guard = OutputGuard() rate_limiter = RateLimiter(max_requests_per_minute=120, burst=20) # Advanced RAG: Query Rewriter & Quality Evaluator query_rewriter = QueryRewriter() quality_evaluator = RetrievalQualityEvaluator() # Embedder embedder = ONNXEmbedder() # Retriever & Cache retriever = SIMDRetriever(top_k=RAG_TOP_K) cache = SemanticCache(similarity_threshold=CACHE_SIMILARITY_THRESHOLD) cache.clear() # Ensure pristine state on startup without stale or negative entries # Router llm_router = LLMRouter() # Ingest PDF if found if PDF_PATH.exists(): logger.info(f"Ingesting PDF manual from {PDF_PATH}...") chunker = PDFChunker() parsed_doc = chunker.parse(str(PDF_PATH)) texts = [chunk.text for chunk in parsed_doc.chunks] if texts: logger.info(f"Embedding {len(texts)} chunks from PDF...") chunk_embeddings = embedder.embed(texts) retriever.ingest(parsed_doc.chunks, chunk_embeddings) logger.info(f"Successfully ingested {len(texts)} PDF chunks into SIMD retriever.") else: logger.warning(f"PDF corpus not found at {PDF_PATH}. Starting with empty retriever.") @asynccontextmanager async def lifespan(app: FastAPI): initialize_pipeline() yield logger.info("Shutting down SentinelAI-RAG service.") app = FastAPI( title="SentinelAI-RAG Enterprise Gateway", description="High-Throughput AI Gateway, Guardrail Engine & C++17 AVX2 Semantic Cache Proxy", version="1.0.0", lifespan=lifespan, ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) STATIC_DIR = Path(__file__).resolve().parent / "static" INDEX_HTML = STATIC_DIR / "index.html" # Serve ChatGPT-Style conversational interface at GET / @app.get("/", response_class=FileResponse) async def serve_chat_ui(): """Serve the ChatGPT-style conversational web interface.""" if INDEX_HTML.exists(): return FileResponse(str(INDEX_HTML), media_type="text/html") raise HTTPException(status_code=404, detail="Web UI index.html not found") if STATIC_DIR.exists(): app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static") # --- Request & Response Models --- class QueryRequest(BaseModel): query: str = Field(..., description="The user's enterprise query string") client_id: Optional[str] = Field("default_client", description="Client identifier for rate limiting") top_k: Optional[int] = Field(None, description="Override default Top-K retrieval count") bypass_cache: Optional[bool] = Field(False, description="Force bypass of semantic cache") user_provider: Optional[str] = Field(None, description="Dynamic BYOK provider override (groq, openai, simulator)") user_model: Optional[str] = Field(None, description="Dynamic model override") user_api_key: Optional[str] = Field(None, description="Ephemeral BYOK API key") class RetrievedContextItem(BaseModel): chunk_id: int page_number: int section_title: str similarity_score: float snippet: str class GuardrailStats(BaseModel): pii_detected_count: int pii_findings: List[str] prompt_injection_blocked: bool injection_risk_score: float output_guard_passed: bool output_issues: List[str] class LatencyBreakdown(BaseModel): guardrail_latency_ms: float embedding_latency_ms: float cache_lookup_latency_ms: float retrieval_latency_ms: float llm_latency_ms: float output_guard_latency_ms: float total_pipeline_latency_ms: float class QueryResponse(BaseModel): model_config = {"protected_namespaces": ()} query: str sanitized_query: str response: str cache_hit: bool cache_similarity: float retrieved_contexts: List[RetrievedContextItem] guardrails: GuardrailStats timing: LatencyBreakdown estimated_cost_usd: float cost_saved_usd: float model_used: str retrieval_quality_score: Optional[float] = Field(None, description="Pre-generation retrieval quality score (0.0 to 1.0)") rewritten_query: Optional[str] = Field(None, description="Expanded/rewritten query used for hybrid retrieval") # --- API Endpoints --- @app.post("/v1/rag/query", response_model=QueryResponse) async def query_endpoint(req: QueryRequest, request: Request): """ Execute full Enterprise RAG lifecycle: Rate Limit -> PII Redact -> Prompt Injection Check -> Embedding -> AVX2 Semantic Cache Lookup -> SIMD Top-K RAG -> LLM Router -> Output Guard -> Cache Store """ global total_queries_served, total_cost_saved_usd, total_baseline_cost_usd t_start = time.perf_counter() # 1. Rate Limiting rl_res = rate_limiter.check(client_id=req.client_id) if not rl_res.allowed: raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=f"Rate limit exceeded. Retry after {rl_res.retry_after_seconds:.1f}s", ) # 2. Input Guardrails: PII Redaction t_g_start = time.perf_counter() pii_res = pii_redactor.redact(req.query) if PII_REDACTION_ENABLED else None active_query = pii_res.redacted_text if pii_res else req.query pii_count = pii_res.redaction_count if pii_res else 0 pii_findings = [f"{t}: {val}" for t, val in (pii_res.findings if pii_res else [])] # 3. Input Guardrails: Prompt Injection Detection inj_res = injection_guard.analyze(active_query) if PROMPT_INJECTION_ENABLED else None t_guardrail = (time.perf_counter() - t_g_start) * 1000.0 if inj_res and inj_res.is_blocked: t_total = (time.perf_counter() - t_start) * 1000.0 return QueryResponse( query=req.query, sanitized_query=active_query, response="[BLOCKED] Your query violated enterprise security policies (Prompt Injection / Jailbreak attempt).", cache_hit=False, cache_similarity=0.0, retrieved_contexts=[], guardrails=GuardrailStats( pii_detected_count=pii_count, pii_findings=pii_findings, prompt_injection_blocked=True, injection_risk_score=inj_res.risk_score, output_guard_passed=True, output_issues=["PROMPT_INJECTION_INTERCEPTED"], ), timing=LatencyBreakdown( guardrail_latency_ms=t_guardrail, embedding_latency_ms=0.0, cache_lookup_latency_ms=0.0, retrieval_latency_ms=0.0, llm_latency_ms=0.0, output_guard_latency_ms=0.0, total_pipeline_latency_ms=t_total, ), estimated_cost_usd=0.0, cost_saved_usd=0.0, model_used="guardrail-block", ) # 4. Generate Single Query Embedding (Zero-double-embedding principle) t_emb_start = time.perf_counter() query_vec = embedder.embed_single(active_query) t_embedding = (time.perf_counter() - t_emb_start) * 1000.0 # 5. Semantic Cache Lookup (AVX2 SIMD) t_cache_start = time.perf_counter() cache_hit = False cache_sim = 0.0 cached_response = None if not req.bypass_cache: cache_res = cache.lookup(query_vec) if cache_res.hit: cache_hit = True cache_sim = cache_res.similarity_score cached_response = cache_res.response t_cache = (time.perf_counter() - t_cache_start) * 1000.0 # Baseline cost estimation for non-cached call baseline_cost = ((AVG_INPUT_TOKENS / 1_000_000) * LLM_PRICING["gpt-4o-mini"]["input"] + (AVG_OUTPUT_TOKENS / 1_000_000) * LLM_PRICING["gpt-4o-mini"]["output"]) if cache_hit and cached_response: t_total = (time.perf_counter() - t_start) * 1000.0 query_latencies.append(t_total) total_queries_served += 1 total_cost_saved_usd += baseline_cost total_baseline_cost_usd += baseline_cost return QueryResponse( query=req.query, sanitized_query=active_query, response=cached_response, cache_hit=True, cache_similarity=cache_sim, retrieved_contexts=[], guardrails=GuardrailStats( pii_detected_count=pii_count, pii_findings=pii_findings, prompt_injection_blocked=False, injection_risk_score=inj_res.risk_score if inj_res else 0.0, output_guard_passed=True, output_issues=[], ), timing=LatencyBreakdown( guardrail_latency_ms=t_guardrail, embedding_latency_ms=t_embedding, cache_lookup_latency_ms=t_cache, retrieval_latency_ms=0.0, llm_latency_ms=0.0, output_guard_latency_ms=0.0, total_pipeline_latency_ms=t_total, ), estimated_cost_usd=0.0, cost_saved_usd=baseline_cost, model_used="semantic-cache-avx2", ) # 6. Pre-Retrieval Query Rewriting & Expansion (Level 3 Production RAG) rewritten = query_rewriter.rewrite(active_query) if query_rewriter else None expanded_search_query = rewritten.expanded_query if rewritten else active_query # 7. Production Hybrid Search (BM25 + AVX2 SIMD + RRF + Cross-Attention Reranking) t_ret_start = time.perf_counter() rag_res = retriever.hybrid_search(query_vec, query_text=expanded_search_query, top_k=req.top_k) t_retrieval = (time.perf_counter() - t_ret_start) * 1000.0 # 8. Pre-Generation Retrieval Quality Gating quality_report = quality_evaluator.evaluate(active_query, rag_res.chunks, rag_res.scores) if quality_evaluator else None quality_score = quality_report.quality_score if quality_report else 0.85 context_str = "\n\n".join([f"[{c.section_title} | Page {c.page_number}]: {c.text}" for c in rag_res.chunks]) context_items = [ RetrievedContextItem( chunk_id=c.chunk_id, page_number=c.page_number, section_title=c.section_title, similarity_score=score, snippet=c.text[:200] + "..." if len(c.text) > 200 else c.text, ) for c, score in zip(rag_res.chunks, rag_res.scores) ] # 9. LLM Generation via Router (with dynamic user BYOK credentials) user_provider = request.headers.get("X-User-Provider") or req.user_provider user_model = request.headers.get("X-User-Model") or req.user_model user_api_key = request.headers.get("X-User-Api-Key") or req.user_api_key t_llm_start = time.perf_counter() llm_resp = await llm_router.generate( prompt=active_query, context=context_str, query=active_query, user_provider=user_provider, user_model=user_model, user_api_key=user_api_key, ) t_llm = (time.perf_counter() - t_llm_start) * 1000.0 # 10. Output Guardrails t_og_start = time.perf_counter() og_res = output_guard.validate(llm_resp.text) if OUTPUT_GUARD_ENABLED else None t_og = (time.perf_counter() - t_og_start) * 1000.0 final_text = og_res.sanitized_text if og_res else llm_resp.text # 11. Store In Semantic Cache for Future Hits (Guarded against negative/refusal caching) if ( og_res and og_res.is_safe and len(rag_res.chunks) > 0 and not is_negative_or_refusal_response(final_text) ): cache.store( query_text=active_query, response_text=final_text, query_embedding=query_vec, model_used=llm_resp.model, metadata={"chunks_count": len(rag_res.chunks)}, ) t_total = (time.perf_counter() - t_start) * 1000.0 query_latencies.append(t_total) total_queries_served += 1 total_baseline_cost_usd += baseline_cost return QueryResponse( query=req.query, sanitized_query=active_query, response=final_text, cache_hit=False, cache_similarity=cache_sim, retrieved_contexts=context_items, guardrails=GuardrailStats( pii_detected_count=pii_count, pii_findings=pii_findings, prompt_injection_blocked=False, injection_risk_score=inj_res.risk_score if inj_res else 0.0, output_guard_passed=og_res.is_safe if og_res else True, output_issues=og_res.issues if og_res else [], ), timing=LatencyBreakdown( guardrail_latency_ms=t_guardrail, embedding_latency_ms=t_embedding, cache_lookup_latency_ms=t_cache, retrieval_latency_ms=t_retrieval, llm_latency_ms=t_llm, output_guard_latency_ms=t_og, total_pipeline_latency_ms=t_total, ), estimated_cost_usd=llm_resp.estimated_cost_usd, cost_saved_usd=0.0, model_used=llm_resp.model, retrieval_quality_score=quality_score, rewritten_query=expanded_search_query, ) @app.post("/v1/rag/stream") async def stream_endpoint(req: QueryRequest, request: Request): """ Server-Sent Events (SSE) streaming query response with guardrail checks. """ user_provider = request.headers.get("X-User-Provider") or req.user_provider user_model = request.headers.get("X-User-Model") or req.user_model user_api_key = request.headers.get("X-User-Api-Key") or req.user_api_key # Rate Limiting rl_res = rate_limiter.check(client_id=req.client_id) if not rl_res.allowed: raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=f"Rate limit exceeded. Retry after {rl_res.retry_after_seconds:.1f}s", ) # Input Guardrails pii_res = pii_redactor.redact(req.query) if PII_REDACTION_ENABLED else None active_query = pii_res.redacted_text if pii_res else req.query inj_res = injection_guard.analyze(active_query) if PROMPT_INJECTION_ENABLED else None if inj_res and inj_res.is_blocked: async def block_gen(): yield "data: " + json.dumps({"token": "[BLOCKED: Prompt injection detected]"}) + "\n\n" yield "data: [DONE]\n\n" return StreamingResponse(block_gen(), media_type="text/event-stream") query_vec = embedder.embed_single(active_query) # Check cache first if not req.bypass_cache: cache_res = cache.lookup(query_vec) if cache_res.hit: async def cache_gen(): for word in cache_res.response.split(): yield "data: " + json.dumps({"token": word + " "}) + "\n\n" yield "data: [DONE]\n\n" # Hybrid Retrieval (BM25 + AVX2 SIMD + RRF + Cross-Attention Reranking) rewritten = query_rewriter.rewrite(active_query) if query_rewriter else None expanded_search_query = rewritten.expanded_query if rewritten else active_query rag_res = retriever.hybrid_search(query_vec, query_text=expanded_search_query, top_k=req.top_k) context_str = "\n\n".join([f"[{c.section_title}]: {c.text}" for c in rag_res.chunks]) async def token_generator(): collected_tokens = [] async for token in llm_router.generate_stream( prompt=active_query, context=context_str, query=active_query, user_provider=user_provider, user_model=user_model, user_api_key=user_api_key, ): collected_tokens.append(token) yield f"data: {json.dumps({'token': token})}\n\n" full_resp = "".join(collected_tokens) cache.store( query_text=active_query, response_text=full_resp, query_embedding=query_vec, model_used=user_model or "streaming-model", ) yield "data: [DONE]\n\n" return StreamingResponse(token_generator(), media_type="text/event-stream") @app.get("/v1/health") async def health_check(): """ Comprehensive health status of all subsystems. """ uptime = time.time() - SERVER_START_TIME simd_active = SIMD_AVAILABLE if simd_active: try: import simd_engine simd_active = simd_engine.is_simd_available() except Exception: simd_active = False return { "status": "healthy", "uptime_seconds": round(uptime, 2), "subsystems": { "simd_avx2_engine": { "compiled_and_loaded": SIMD_AVAILABLE, "hardware_avx2_fma": simd_active, }, "retriever": { "indexed_chunks": retriever.chunk_count if retriever else 0, "ready": retriever is not None and retriever.chunk_count > 0, }, "semantic_cache": { "current_entries": cache.size if cache else 0, "hit_rate": cache.stats["hit_rate"] if cache else 0.0, }, "llm_router": { "providers_active": llm_router.stats["providers_configured"] if llm_router else 0, "local_simulator_ready": True, }, "guardrails": { "pii_redactor": "active", "prompt_injection_guard": "active", "output_guard": "active", "rate_limiter": "active", }, }, } @app.get("/v1/metrics") async def metrics_endpoint(): """ Export full operational metrics for monitoring and SLAs. """ lats = sorted(query_latencies) if query_latencies else [0.0] n = len(lats) p50 = lats[int(n * 0.50)] p95 = lats[int(n * 0.95)] if n > 1 else lats[0] p99 = lats[int(n * 0.99)] if n > 1 else lats[0] cost_decrease_pct = 0.0 if total_baseline_cost_usd > 0: cost_decrease_pct = (total_cost_saved_usd / total_baseline_cost_usd) * 100.0 return { "total_queries_served": total_queries_served, "latency_ms": { "p50": round(p50, 2), "p95": round(p95, 2), "p99": round(p99, 2), "mean": round(sum(lats) / max(n, 1), 2), }, "cache_metrics": cache.stats if cache else {}, "cost_metrics": { "total_baseline_cost_usd": round(total_baseline_cost_usd, 6), "total_cost_saved_usd": round(total_cost_saved_usd, 6), "cost_decrease_percentage": round(cost_decrease_pct, 2), }, "guardrail_metrics": { "pii": pii_redactor.stats if pii_redactor else {}, "prompt_injection": injection_guard.stats if injection_guard else {}, "output_guard": output_guard.stats if output_guard else {}, "rate_limiter": rate_limiter.stats if rate_limiter else {}, }, "retriever_metrics": retriever.stats if retriever else {}, } @app.post("/v1/cache/clear") async def clear_cache_endpoint(): """Clear all entries in the semantic cache.""" global cache if cache: cache.clear() logger.info("Semantic cache cleared via /v1/cache/clear") return {"status": "ok", "message": "Semantic cache cleared successfully", "current_size": 0} if __name__ == "__main__": import uvicorn uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=False)