""" sidecar/app.py — GenAI Shield V2 Sidecar Proxy (FastAPI). A language-agnostic sidecar that sits in front of any LLM API and provides: • Pre-inference guard (Prompt Guard model + regex) — runs in parallel with LLM • Sentence-level streaming — users see output word-by-word • Post-inference monitoring — each sentence checked concurrently in background • Block signal mid-stream — if output turns harmful, client is notified instantly Endpoints --------- POST /v1/chat → streaming or blocking chat with full shield GET /v1/health → liveness probe GET /v1/stats → guard model statistics GET /v1/metrics → last-request latency breakdown SSE Event Schema (stream=true) ------------------------------- { "type": "chunk", "text": "..." } { "type": "sentence", "text": "...", "sentence_id": 1 } { "type": "block_signal", "sentence_id": 3, "reason": "...", "threat_score": 85, "flags": [...] } { "type": "done", "threat_score": 5, "flags": [], "latency_ms": 420, "guard_ms": 98, "sentences": 4 } { "type": "blocked", "reason": "...", "threat_score": 100, "flags": [...], "pg_score": 0.97, "guard_ms": 102 } Configure via environment variables (see sidecar/config.py). """ import json import logging import os import sys import time from pathlib import Path from typing import AsyncGenerator, Optional import uvicorn from fastapi import FastAPI, HTTPException, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse, StreamingResponse from fastapi.staticfiles import StaticFiles from pydantic import BaseModel # ── Path fix so imports work from project root ──────────────────────────────── _ROOT = Path(__file__).parent.parent if str(_ROOT) not in sys.path: sys.path.insert(0, str(_ROOT)) from sidecar.config import ( GATE_GUARD_TIMEOUT_SEC, GEMINI_API_KEY, GEMINI_MODEL, LLM_BACKEND, LOG_LEVEL, MONITOR_BLOCK_THRESHOLD, MONITOR_WORKERS, OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_MODEL, PROMPT_GUARD_MODEL_DIR, SENTENCE_MIN_CHARS, SIDECAR_HOST, SIDECAR_PORT, SYSTEM_PROMPT, ) from sidecar.gate import BlockEvent, ShieldGate, TokenEvent from sidecar.sentence_splitter import SentenceEvent, SentenceSplitter from sidecar.stream_monitor import BlockSignal, StreamMonitor from sidecar.pipeline_events import RequestTrace, subscribe, unsubscribe # Existing shield modules from prompt_guard_engine import PromptGuardEngine from prompt_guard_text_guard import PromptGuardTextGuard from text_monitor import TextMonitor # ── Logging ─────────────────────────────────────────────────────────────────── logging.basicConfig( level = getattr(logging, LOG_LEVEL.upper(), logging.INFO), format = "[%(asctime)s] %(levelname)-8s %(name)s — %(message)s", datefmt = "%H:%M:%S", ) log = logging.getLogger("sidecar") # ── FastAPI app ─────────────────────────────────────────────────────────────── app = FastAPI( title = "GenAI Shield Sidecar", description = "Transparent LLM proxy with pre/post-inference security screening", version = "2.0.0", ) app.add_middleware( CORSMiddleware, allow_origins = ["*"], allow_methods = ["*"], allow_headers = ["*"], ) # Serve static files if the sidecar runs standalone _STATIC_DIR = _ROOT / "static" _TEMPLATES_DIR = _ROOT / "templates" if _STATIC_DIR.exists(): app.mount("/static", StaticFiles(directory=str(_STATIC_DIR)), name="static") # ── Initialise shield components ────────────────────────────────────────────── log.info("Loading Prompt Guard engine...") _PG_ENGINE = PromptGuardEngine(model_path=Path(PROMPT_GUARD_MODEL_DIR)).load() _GUARD = PromptGuardTextGuard(_PG_ENGINE) log.info("Prompt Guard ready.") # ── Initialise LLM adapter ──────────────────────────────────────────────────── if LLM_BACKEND == "openai": from openai_adapter import OpenAIAdapter _ADAPTER = OpenAIAdapter( api_key = OPENAI_API_KEY, base_url = OPENAI_BASE_URL, model = OPENAI_MODEL, system_prompt = SYSTEM_PROMPT, ) else: from gemini_adapter import GeminiAdapter _ADAPTER = GeminiAdapter( api_key = GEMINI_API_KEY, model_name = GEMINI_MODEL, system_prompt = SYSTEM_PROMPT, ) log.info("LLM adapter: %s (%s)", LLM_BACKEND, _ADAPTER.get_model_name()) # ── Shared monitor (stateful — tracks behavioural drift across requests) ────── _TEXT_MONITOR = TextMonitor(_ADAPTER, system_prompt=SYSTEM_PROMPT) _STREAM_MONITOR = StreamMonitor(_TEXT_MONITOR, block_threshold=MONITOR_BLOCK_THRESHOLD, max_workers=MONITOR_WORKERS) # ── Last-request metrics (lightweight, single-threaded access via asyncio) ──── _LAST_METRICS: dict = {} # ── Request schema ──────────────────────────────────────────────────────────── class ChatRequest(BaseModel): prompt: str stream: bool = True system_prompt: Optional[str] = None source: Optional[str] = "sidecar" # ── Routes ──────────────────────────────────────────────────────────────────── @app.get("/") async def root(): """Serve the sidecar streaming UI (standalone mode).""" ui_file = _TEMPLATES_DIR / "sidecar.html" if ui_file.exists(): return FileResponse(str(ui_file), media_type="text/html") return {"message": "GenAI Shield Sidecar", "docs": "/docs"} @app.get("/dataflow") async def dataflow_ui(): """Serve the real-time data flow visualization dashboard.""" ui_file = _TEMPLATES_DIR / "dataflow.html" if ui_file.exists(): return FileResponse(str(ui_file), media_type="text/html") return {"message": "dataflow.html not found"} @app.get("/v1/pipeline-stream") async def pipeline_stream(): """ SSE stream of structured pipeline telemetry events. The data flow dashboard subscribes here to get real-time stage data. """ async def _gen(): import asyncio q = subscribe() try: while True: try: # Poll queue with a short timeout so we can yield keepalives payload = q.get_nowait() yield f"data: {json.dumps(payload)}\n\n" except Exception: # No event — send keepalive comment yield ": keepalive\n\n" await asyncio.sleep(0.5) except asyncio.CancelledError: pass finally: unsubscribe(q) return StreamingResponse( _gen(), media_type = "text/event-stream", headers = {"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, ) @app.get("/v1/health") async def health(): return { "status": "ok", "guard_ready": _PG_ENGINE.ready, "model": _ADAPTER.get_model_name(), "backend": LLM_BACKEND, } @app.get("/v1/stats") async def stats(): return _PG_ENGINE.stats() @app.get("/v1/metrics") async def metrics(): return _LAST_METRICS or {"message": "No requests processed yet"} @app.post("/v1/chat") async def chat(req: ChatRequest): """ Main chat endpoint. - stream=true → Server-Sent Events (SSE) with sentence-level output - stream=false → Blocking JSON response (legacy-compatible) """ if not req.prompt.strip(): raise HTTPException(status_code=400, detail="Empty prompt") sys_prompt = req.system_prompt or SYSTEM_PROMPT if req.stream: return StreamingResponse( _stream_handler(req.prompt, sys_prompt, req.source or "sidecar"), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", "X-Accel-Buffering": "no", # disable nginx buffering }, ) else: return await _blocking_handler(req.prompt, sys_prompt, req.source or "sidecar") # ── Streaming handler ───────────────────────────────────────────────────────── async def _stream_handler( prompt: str, sys_prompt: str, source: str, ) -> AsyncGenerator[str, None]: """ Full streaming pipeline: Gate (guard ∥ LLM) → SentenceSplitter → StreamMonitor (background) """ t_total = time.perf_counter() gate = ShieldGate(_GUARD, _ADAPTER, guard_timeout_sec=GATE_GUARD_TIMEOUT_SEC) splitter = SentenceSplitter(min_chars=SENTENCE_MIN_CHARS) _STREAM_MONITOR.reset() # ── Telemetry trace for this request ────────────────────────────────── trace = RequestTrace() trace.on_request_in(prompt) guard_ms_ref = 0.0 all_flags: list = [] threat_score = 0 block_fired = False sentences_sent = 0 total_tokens = 0 def sse(event_dict: dict) -> str: """Format a dict as an SSE data line.""" return f"data: {json.dumps(event_dict)}\n\n" async for gate_event in gate.run(prompt, sys_prompt, trace=trace): # ── Guard blocked the prompt ─────────────────────────────────────── if isinstance(gate_event, BlockEvent): all_flags = gate_event.flags threat_score = gate_event.threat_score guard_ms_ref = gate_event.guard_ms yield sse({ "type": "blocked", "reason": gate_event.reason, "threat_score": gate_event.threat_score, "flags": gate_event.flags, "pg_score": gate_event.pg_score, "guard_ms": gate_event.guard_ms, }) block_fired = True break # ── LLM token received ───────────────────────────────────────────── if isinstance(gate_event, TokenEvent): total_tokens += 1 splitter_events = splitter.feed(gate_event.text) for ev in splitter_events: if isinstance(ev, type(ev)) and ev.type == "chunk": yield sse({"type": "chunk", "text": ev.text}) elif ev.type == "sentence": sentences_sent += 1 trace.on_sentence_ready(ev.sentence_id, ev.text) yield sse({ "type": "sentence", "text": ev.text, "sentence_id": ev.sentence_id, }) # Submit to background monitor (non-blocking) trace.on_monitor_start(ev.sentence_id) await _STREAM_MONITOR.submit(ev.sentence_id, ev.text, prompt) if block_fired: total_ms = round((time.perf_counter() - t_total) * 1000, 2) trace.on_request_done(threat_score, all_flags, blocked=True) _update_metrics(threat_score, all_flags, guard_ms_ref, 0, 0, total_ms) return # ── Stream ended — flush remaining buffer ────────────────────────────── for ev in splitter.flush(): sentences_sent += 1 trace.on_sentence_ready(ev.sentence_id, ev.text) yield sse({"type": "sentence", "text": ev.text, "sentence_id": ev.sentence_id}) trace.on_monitor_start(ev.sentence_id) await _STREAM_MONITOR.submit(ev.sentence_id, ev.text, prompt) trace.on_stream_done(total_tokens, sentences_sent) # ── Collect background monitor results ───────────────────────────────── signals = await _STREAM_MONITOR.collect(timeout=1.5) for sig in signals: threat_score = max(threat_score, sig.threat_score) all_flags.extend(sig.flags) trace.on_monitor_result(sig.sentence_id, sig.threat_score, sig.flags, blocked=True) yield sse({ "type": "block_signal", "sentence_id": sig.sentence_id, "reason": sig.reason, "threat_score": sig.threat_score, "flags": sig.flags, }) total_ms = round((time.perf_counter() - t_total) * 1000, 2) trace.on_request_done(threat_score, list(set(all_flags)), blocked=False) yield sse({ "type": "done", "threat_score": threat_score, "flags": list(set(all_flags)), "latency_ms": total_ms, "sentences": sentences_sent, }) _update_metrics(threat_score, all_flags, 0, 0, total_ms, total_ms) # ── Blocking handler (non-streaming, backward-compatible) ───────────────────── async def _blocking_handler(prompt: str, sys_prompt: str, source: str) -> dict: """ Non-streaming path — guard first, then full LLM call, then monitor. Compatible with existing /genai-chat behaviour. """ import asyncio t_start = time.perf_counter() # Guard (in thread — synchronous) loop = asyncio.get_event_loop() guard_result = await loop.run_in_executor(None, _GUARD.screen, prompt) guard_ms = round((time.perf_counter() - t_start) * 1000, 2) pg_score = guard_result.get("checks", {}).get("prompt_guard", {}).get("malicious_score", 0.0) if guard_result["blocked"]: return { "blocked": True, "response": None, "reason": guard_result["reason"], "threat_score": guard_result["threat_score"], "flags": guard_result["flags"], "pg_score": pg_score, "latency_breakdown": {"guard_ms": guard_ms, "model_ms": 0, "monitor_ms": 0}, } # LLM call (blocking adapter) t_model = time.perf_counter() response = await loop.run_in_executor(None, _ADAPTER.chat, prompt, sys_prompt) model_ms = round((time.perf_counter() - t_model) * 1000, 2) # Post-monitor t_monitor = time.perf_counter() mon_result = await loop.run_in_executor(None, _TEXT_MONITOR.analyze, prompt, response) monitor_ms = round((time.perf_counter() - t_monitor) * 1000, 2) total_ms = round(guard_ms + model_ms + monitor_ms, 2) threat_score = max(guard_result["threat_score"], mon_result["threat_score"]) all_flags = guard_result["flags"] + mon_result["flags"] _update_metrics(threat_score, all_flags, guard_ms, model_ms, monitor_ms, total_ms) return { "blocked": False, "response": response, "threat_score": threat_score, "flags": all_flags, "pg_score": pg_score, "latency_ms": total_ms, "model": _ADAPTER.get_model_name(), "latency_breakdown": { "guard_ms": guard_ms, "model_ms": model_ms, "monitor_ms": monitor_ms, }, } # ── Metrics helper ───────────────────────────────────────────────────────────── def _update_metrics(threat_score, flags, guard_ms, model_ms, monitor_ms, total_ms): global _LAST_METRICS _LAST_METRICS = { "threat_score": threat_score, "flags": flags, "guard_ms": guard_ms, "model_ms": model_ms, "monitor_ms": monitor_ms, "total_ms": total_ms, "model": _ADAPTER.get_model_name(), "timestamp": time.strftime("%H:%M:%S"), } # ── Entry point ─────────────────────────────────────────────────────────────── if __name__ == "__main__": log.info("Starting GenAI Shield Sidecar on %s:%d", SIDECAR_HOST, SIDECAR_PORT) uvicorn.run( "sidecar.app:app", host = SIDECAR_HOST, port = SIDECAR_PORT, log_level = LOG_LEVEL.lower(), reload = False, )