sidecar/app.py

"""
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,
    )