feat: add auth guardrails and request observability

.env.example

@@ -15,3 +15,7 @@ MAX_CONTEXT_CHARS=12000
 ENABLE_GUARDRAILS=true
 ENABLE_RAG_EVAL=false
 MAX_QUERY_LENGTH=2000
+RATE_LIMIT_PER_MINUTE=30
+AUTH_ENABLED=false
+ADMIN_API_KEY=change-me-admin
+USER_API_KEY=change-me-user

app/auth/__init__.py

@@ -0,0 +1 @@
+

app/auth/models.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel
+
+
+class UserContext(BaseModel):
+    role: str
+    user_id: str
+
+
+class LoginRequest(BaseModel):
+    api_key: str
+
+
+class LoginResponse(BaseModel):
+    role: str
+    token_type: str = "api_key"

app/auth/security.py

@@ -0,0 +1,19 @@
+from fastapi import Header, HTTPException
+
+from app.auth.models import UserContext
+from app.core.config import settings
+
+
+def resolve_user(x_api_key: str | None = Header(default=None)) -> UserContext:
+    if not settings.AUTH_ENABLED:
+        return UserContext(role="admin", user_id="local-dev")
+    if x_api_key == settings.ADMIN_API_KEY:
+        return UserContext(role="admin", user_id="admin")
+    if x_api_key == settings.USER_API_KEY:
+        return UserContext(role="user", user_id="user")
+    raise HTTPException(status_code=401, detail="Invalid or missing API key.")
+
+
+def require_admin(user: UserContext) -> None:
+    if user.role != "admin":
+        raise HTTPException(status_code=403, detail="Admin role required.")

app/core/config.py

@@ -33,6 +33,12 @@ class Settings(BaseSettings):
     ENABLE_GUARDRAILS: bool = True
     ENABLE_RAG_EVAL: bool = False
     MAX_QUERY_LENGTH: int = 2000
+    RATE_LIMIT_PER_MINUTE: int = 30
+
+    # Auth Config
+    AUTH_ENABLED: bool = False
+    ADMIN_API_KEY: str = "local-admin-key"
+    USER_API_KEY: str = "local-user-key"
 
     class Config:
         env_file = ".env"

app/engine/ingestion.py

@@ -15,6 +15,7 @@ from app.engine.document_registry import (
 )
 from app.engine.indexer import delete_document, index_documents, reset_collection
 from app.engine.loaders import SUPPORTED_EXTENSIONS, load_document
+from app.guardrails.document import filter_malicious_documents
 
 
 def ingest_documents(data_dir: str = "data/docs", force: bool = False) -> None:
@@ -40,6 +41,11 @@ def ingest_documents(data_dir: str = "data/docs", force: bool = False) -> None:
 
         doc_id = str(uuid.uuid5(uuid.NAMESPACE_URL, f"{path.name}:{content_hash}"))
         loaded = load_document(path)
+        loaded, flagged = filter_malicious_documents(loaded)
+        for flagged_source in flagged:
+            print(f"Skipping document with suspicious instructions: {flagged_source}")
+        if not loaded:
+            continue
         for document in loaded:
             document.metadata["doc_id"] = doc_id
             document.metadata["content_hash"] = content_hash

app/guardrails/document.py

@@ -0,0 +1,19 @@
+from langchain_core.documents import Document
+
+from app.guardrails.input import PROMPT_INJECTION_PATTERNS
+
+
+def document_has_malicious_instruction(document: Document) -> bool:
+    normalized = document.page_content.lower()
+    return any(pattern in normalized for pattern in PROMPT_INJECTION_PATTERNS)
+
+
+def filter_malicious_documents(documents: list[Document]) -> tuple[list[Document], list[str]]:
+    safe: list[Document] = []
+    flagged: list[str] = []
+    for document in documents:
+        if document_has_malicious_instruction(document):
+            flagged.append(document.metadata.get("source", "unknown"))
+            continue
+        safe.append(document)
+    return safe, flagged

app/guardrails/input.py

@@ -0,0 +1,51 @@
+import re
+import time
+from collections import defaultdict, deque
+
+from fastapi import HTTPException
+
+from app.core.config import settings
+
+
+_requests_by_client: dict[str, deque[float]] = defaultdict(deque)
+
+
+PROMPT_INJECTION_PATTERNS = [
+    "ignore previous instructions",
+    "ignore the instructions above",
+    "forget all previous",
+    "system prompt",
+    "developer message",
+    "bypass system",
+    "disregard instructions",
+    "reveal hidden",
+]
+
+
+def validate_query(query: str) -> None:
+    if not query.strip():
+        raise HTTPException(status_code=400, detail="Query cannot be empty.")
+    if len(query) > settings.MAX_QUERY_LENGTH:
+        raise HTTPException(status_code=400, detail="Query is too long.")
+    normalized = query.lower()
+    for pattern in PROMPT_INJECTION_PATTERNS:
+        if pattern in normalized:
+            raise HTTPException(status_code=400, detail="Prompt injection attempt detected.")
+
+
+def enforce_rate_limit(client_id: str) -> None:
+    if settings.RATE_LIMIT_PER_MINUTE <= 0:
+        return
+    now = time.time()
+    bucket = _requests_by_client[client_id]
+    while bucket and now - bucket[0] > 60:
+        bucket.popleft()
+    if len(bucket) >= settings.RATE_LIMIT_PER_MINUTE:
+        raise HTTPException(status_code=429, detail="Rate limit exceeded.")
+    bucket.append(now)
+
+
+def contains_pii(text: str) -> bool:
+    email = r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}"
+    phone = r"\b(?:\+?\d{1,3}[-.\s]?)?(?:\d{10}|\d{3}[-.\s]\d{3}[-.\s]\d{4})\b"
+    return bool(re.search(email, text) or re.search(phone, text))

app/guardrails/output.py

@@ -0,0 +1,12 @@
+import re
+
+
+def redact_sensitive_data(text: str) -> str:
+    text = re.sub(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}", "[REDACTED_EMAIL]", text)
+    text = re.sub(r"\b(?:\+?\d{1,3}[-.\s]?)?(?:\d{10}|\d{3}[-.\s]\d{3}[-.\s]\d{4})\b", "[REDACTED_PHONE]", text)
+    return text
+
+
+def refuses_system_prompt_leak(text: str) -> bool:
+    normalized = text.lower()
+    return "system prompt" in normalized or "developer message" in normalized

app/main.py

@@ -1,18 +1,24 @@
 import asyncio
-import time
-from fastapi import FastAPI, HTTPException
+from fastapi import Depends, FastAPI, HTTPException, Request
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from guardrails import Guard
 from langchain_core.prompts import PromptTemplate
 from langchain_ollama import ChatOllama
 
+from app.auth.models import LoginRequest, LoginResponse, UserContext
+from app.auth.security import require_admin, resolve_user
 from app.core.config import settings
 from app.core.dependencies import check_ollama, check_qdrant
 from app.core.logging import configure_logging, logger
+from app.engine.document_registry import load_registry
+from app.engine.ingestion import delete_indexed_document, ingest_documents, reset_index
 from app.engine.context_builder import build_context, format_sources
 from app.graph.workflow import compile_workflow
+from app.guardrails.input import enforce_rate_limit, validate_query
+from app.guardrails.output import redact_sensitive_data
 from app.guardrails.validators import DetectPromptInjection
+from app.observability.metrics import RequestMetrics, log_request_metrics, timed_stage
 
 configure_logging()
 app = FastAPI(title=settings.PROJECT_NAME)
@@ -34,6 +40,10 @@ class ChatResponse(BaseModel):
     answer: str
     sources: list[SourceCitation] = []
 
+class IngestionRequest(BaseModel):
+    data_dir: str = "data/docs"
+    force: bool = False
+
 @app.get("/health")
 async def health_endpoint():
     return {"status": "ok", "project": settings.PROJECT_NAME}
@@ -48,11 +58,41 @@ async def ready_endpoint():
         "qdrant": qdrant,
     }
 
+@app.post("/auth/login", response_model=LoginResponse)
+async def login_endpoint(request: LoginRequest):
+    if request.api_key == settings.ADMIN_API_KEY:
+        return LoginResponse(role="admin")
+    if request.api_key == settings.USER_API_KEY:
+        return LoginResponse(role="user")
+    raise HTTPException(status_code=401, detail="Invalid API key.")
+
+@app.get("/documents")
+async def documents_endpoint(user: UserContext = Depends(resolve_user)):
+    return {"documents": list(load_registry().values()), "role": user.role}
+
+@app.post("/admin/ingest")
+async def admin_ingest_endpoint(request: IngestionRequest, user: UserContext = Depends(resolve_user)):
+    require_admin(user)
+    ingest_documents(data_dir=request.data_dir, force=request.force)
+    return {"status": "ok", "message": "Ingestion completed."}
+
+@app.delete("/admin/documents/{doc_id}")
+async def admin_delete_document_endpoint(doc_id: str, user: UserContext = Depends(resolve_user)):
+    require_admin(user)
+    delete_indexed_document(doc_id)
+    return {"status": "ok", "doc_id": doc_id}
+
+@app.post("/admin/reset")
+async def admin_reset_endpoint(user: UserContext = Depends(resolve_user)):
+    require_admin(user)
+    reset_index()
+    return {"status": "ok", "message": "Index reset."}
+
 @app.post("/chat", response_model=ChatResponse)
-async def chat_endpoint(request: ChatRequest):
-    started_at = time.perf_counter()
-    if len(request.query) > settings.MAX_QUERY_LENGTH:
-        raise HTTPException(status_code=400, detail="Query is too long.")
+async def chat_endpoint(request: ChatRequest, http_request: Request, user: UserContext = Depends(resolve_user)):
+    metrics = RequestMetrics()
+    enforce_rate_limit(http_request.client.host if http_request.client else user.user_id)
+    validate_query(request.query)
     if settings.ENABLE_GUARDRAILS:
         try:
             input_guard.validate(request.query)
@@ -61,19 +101,20 @@ async def chat_endpoint(request: ChatRequest):
 
     initial_state = {"question": request.query, "run_count": 0}
     try:
-        final_state = rag_agent.invoke(initial_state)
-        answer = final_state.get("generation", "Unable to compile answer.")
+        with timed_stage(metrics, "rag_workflow"):
+            final_state = rag_agent.invoke(initial_state)
+        answer = redact_sensitive_data(final_state.get("generation", "Unable to compile answer."))
         sources = final_state.get("sources", [])
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-    logger.info("chat completed latency_ms=%s", round((time.perf_counter() - started_at) * 1000, 2))
+    log_request_metrics(metrics, route="/chat", sources=len(sources), model=settings.OLLAMA_MODEL)
     return ChatResponse(query=request.query, answer=answer, sources=sources)
 
 @app.post("/chat/stream")
-async def chat_stream_endpoint(request: ChatRequest):
-    if len(request.query) > settings.MAX_QUERY_LENGTH:
-        raise HTTPException(status_code=400, detail="Query is too long.")
+async def chat_stream_endpoint(request: ChatRequest, http_request: Request, user: UserContext = Depends(resolve_user)):
+    enforce_rate_limit(http_request.client.host if http_request.client else user.user_id)
+    validate_query(request.query)
     if settings.ENABLE_GUARDRAILS:
         try:
             input_guard.validate(request.query)
@@ -81,9 +122,10 @@ async def chat_stream_endpoint(request: ChatRequest):
             raise HTTPException(status_code=400, detail=str(getattr(e, "message", e)))
 
     async def token_generator():
-        started_at = time.perf_counter()
+        metrics = RequestMetrics()
         initial_state = {"question": request.query, "run_count": 0}
-        final_state = rag_agent.invoke(initial_state)
+        with timed_stage(metrics, "rag_workflow"):
+            final_state = rag_agent.invoke(initial_state)
         documents = final_state.get("documents", [])
         
         if not documents:
@@ -102,9 +144,9 @@ async def chat_stream_endpoint(request: ChatRequest):
 
         async for chunk in rag_chain.astream({"context": context, "question": request.query}):
             if chunk.content:
-                yield chunk.content
+                yield redact_sensitive_data(chunk.content)
                 await asyncio.sleep(0.01)
         yield format_sources(documents)
-        logger.info("stream completed latency_ms=%s", round((time.perf_counter() - started_at) * 1000, 2))
+        log_request_metrics(metrics, route="/chat/stream", sources=len(documents), model=settings.OLLAMA_MODEL)
 
     return StreamingResponse(token_generator(), media_type="text/event-stream")

app/observability/__init__.py

@@ -0,0 +1 @@
+

app/observability/metrics.py

@@ -0,0 +1,36 @@
+import logging
+import time
+import uuid
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+
+logger = logging.getLogger("support_docs_copilot.metrics")
+
+
+@dataclass
+class RequestMetrics:
+    request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    started_at: float = field(default_factory=time.perf_counter)
+    stages: dict[str, float] = field(default_factory=dict)
+
+    def total_ms(self) -> float:
+        return round((time.perf_counter() - self.started_at) * 1000, 2)
+
+
+@contextmanager
+def timed_stage(metrics: RequestMetrics, stage: str):
+    started = time.perf_counter()
+    try:
+        yield
+    finally:
+        metrics.stages[stage] = round((time.perf_counter() - started) * 1000, 2)
+
+
+def log_request_metrics(metrics: RequestMetrics, **extra) -> None:
+    logger.info(
+        "request_id=%s total_ms=%s stages=%s extra=%s",
+        metrics.request_id,
+        metrics.total_ms(),
+        metrics.stages,
+        extra,
+    )