deploy CodeSecAudit RAG service

Dockerfile

@@ -0,0 +1,23 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY rag_service/ rag_service/
+
+RUN pip install --no-cache-dir \
+    fastapi>=0.100.0 \
+    uvicorn[standard]>=0.20.0 \
+    sentence-transformers>=2.2.0 \
+    numpy>=1.24.0 \
+    requests>=2.31.0 \
+    pydantic>=2.0.0
+
+ENV RAG_BUILD_ON_START=true
+ENV RAG_DATASET_REPO=OMCHOKSI108/CodeSecAudit-RAG
+ENV RAG_CORPUS_FILE=rag/rag_corpus.jsonl.gz
+ENV RAG_EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2
+ENV RAG_INDEX_PATH=/tmp/codesec_rag_index
+
+EXPOSE 7860
+
+CMD ["uvicorn", "rag_service.main:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -0,0 +1,52 @@
+---
+title: CodeSecAudit RAG Service
+emoji: 🛡️
+colorFrom: blue
+colorTo: purple
+sdk: docker
+app_port: 7860
+pinned: false
+---
+
+# CodeSecAudit RAG Service
+
+Remote RAG microservice for OWASP cheat sheet retrieval. Deployed as a Hugging Face Space with Docker SDK.
+
+## Deploy
+
+1. Create a new Space at https://huggingface.co/new-space
+2. Choose **Docker** (not Streamlit/Gradio)
+3. Set Space SDK to **Docker**
+4. Copy these files:
+   - `Dockerfile`
+   - `start.sh`
+   - The entire `rag_service/` directory
+
+   Or use the automated script: `python scripts/deploy_hf_rag_space.py`
+
+5. Add Secrets in Space Settings:
+
+   | Secret | Required | Description |
+   |---|---|---|
+   | `RAG_API_KEY` | **Production** | Shared API key for request auth. If empty, the service is **public** — anyone can search. |
+   | `RAG_DATASET_REPO` | No | HF dataset repo (default: `OMCHOKSI108/CodeSecAudit-RAG`) |
+   | `RAG_EMBEDDING_MODEL` | No | Sentence-transformer model (default: `sentence-transformers/all-MiniLM-L6-v2`) |
+
+   > **Production**: Always set `RAG_API_KEY`. Without it, the service is public and anyone with the URL can query your RAG index.
+
+6. Space will build and start on port 7860.
+
+## Health check
+
+```bash
+curl https://your-space.hf.space/health
+```
+
+## Search
+
+```bash
+curl -X POST https://your-space.hf.space/rag/search \
+  -H "Content-Type: application/json" \
+  -H "X-CodeSec-RAG-Key: your-key" \
+  -d '{"query":"sql injection prepared statements","top_k":3}'
+```

rag_service/README.md

@@ -0,0 +1,26 @@
+# CodeSecAudit RAG Service
+
+Remote RAG microservice for OWASP cheat sheet retrieval.
+
+## Run locally
+
+```bash
+pip install fastapi uvicorn sentence-transformers numpy requests pydantic
+uvicorn rag_service.main:app --port 7860
+```
+
+## API
+
+| Endpoint | Method | Description |
+|---|---|---|
+| `/` | GET | Service info |
+| `/health` | GET | Health check |
+| `/rag/search` | POST | Search RAG corpus |
+
+### POST /rag/search
+
+```json
+{"query": "sql injection prepared statements", "top_k": 3}
+```
+
+Set `RAG_API_KEY` env var and pass `X-CodeSec-RAG-Key` header to protect the endpoint.

rag_service/index.py

@@ -0,0 +1,98 @@
+import gzip
+import json
+import logging
+import os
+import time
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import requests
+from sentence_transformers import SentenceTransformer
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_DATASET_REPO = "OMCHOKSI108/CodeSecAudit-RAG"
+DEFAULT_CORPUS_FILE = "rag/rag_corpus.jsonl.gz"
+DEFAULT_EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+
+
+class RagIndex:
+    def __init__(self):
+        self.dataset_repo: str = os.getenv("RAG_DATASET_REPO", DEFAULT_DATASET_REPO)
+        self.corpus_file: str = os.getenv("RAG_CORPUS_FILE", DEFAULT_CORPUS_FILE)
+        self.embed_model_name: str = os.getenv("RAG_EMBEDDING_MODEL", DEFAULT_EMBED_MODEL)
+        self.max_results: int = int(os.getenv("RAG_MAX_RESULTS", "8"))
+
+        self._model: Optional[SentenceTransformer] = None
+        self._chunks: list[dict] = []
+        self._embeddings: Optional[np.ndarray] = None
+
+    @property
+    def model(self) -> SentenceTransformer:
+        if self._model is None:
+            logger.info("Loading embedding model: %s", self.embed_model_name)
+            self._model = SentenceTransformer(self.embed_model_name)
+        return self._model
+
+    def load_corpus(self) -> int:
+        url = f"https://huggingface.co/datasets/{self.dataset_repo}/resolve/main/{self.corpus_file}"
+        logger.info("Downloading corpus from %s", url)
+        resp = requests.get(url, timeout=120)
+        resp.raise_for_status()
+        raw = gzip.decompress(resp.content)
+        chunks = []
+        for line in raw.decode("utf-8").splitlines():
+            line = line.strip()
+            if line:
+                chunks.append(json.loads(line))
+        self._chunks = chunks
+        logger.info("Loaded %d chunks", len(chunks))
+        return len(chunks)
+
+    def build_index(self) -> None:
+        logger.info("Embedding %d chunks with %s ...", len(self._chunks), self.embed_model_name)
+        texts = [c.get("content", "") or c.get("text", "") for c in self._chunks]
+        self._embeddings = self.model.encode(texts, show_progress_bar=True)
+        logger.info("Index built: shape=%s", self._embeddings.shape)
+
+    def search(self, query: str, top_k: int = 3) -> list[dict]:
+        if self._embeddings is None or not self._chunks:
+            return []
+
+        k = min(top_k, self.max_results, len(self._chunks))
+        query_vec = self.model.encode([query])[0]
+        scores = np.dot(self._embeddings, query_vec) / (
+            np.linalg.norm(self._embeddings, axis=1) * np.linalg.norm(query_vec) + 1e-12
+        )
+        top_indices = np.argsort(scores)[::-1][:k]
+
+        results = []
+        for rank, idx in enumerate(top_indices, 1):
+            c = self._chunks[idx]
+            results.append({
+                "rank": rank,
+                "score": float(scores[idx]),
+                "title": c.get("title", ""),
+                "section_title": c.get("section_title", "") or c.get("section", ""),
+                "cwe_id": c.get("cwe_id", ""),
+                "content": c.get("content", "") or c.get("text", ""),
+                "source_file": c.get("source_file", "") or c.get("source", ""),
+            })
+        return results
+
+    @property
+    def is_loaded(self) -> bool:
+        return self._embeddings is not None
+
+    @property
+    def embedding_count(self) -> int:
+        return len(self._embeddings) if self._embeddings is not None else 0
+
+    @property
+    def document_count(self) -> int:
+        return len(self._chunks)
+
+    @property
+    def total_chunks(self) -> int:
+        return len(self._chunks)

rag_service/main.py

@@ -0,0 +1,107 @@
+import logging
+import os
+import time
+
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+
+from rag_service.index import DEFAULT_DATASET_REPO, DEFAULT_EMBED_MODEL, RagIndex
+from rag_service.schemas import (
+    SearchMetadata,
+    SearchRequest,
+    SearchResponse,
+    SearchResultItem,
+)
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+START_TIME = time.time()
+
+app = FastAPI(
+    title="CodeSecAudit RAG Service",
+    description="Remote RAG retrieval service for OWASP cheat sheets",
+    version="0.6.0",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+index = RagIndex()
+api_key = os.getenv("RAG_API_KEY", "")
+
+
+def _verify_key(request: Request) -> None:
+    if not api_key:
+        return
+    key = request.headers.get("X-CodeSec-RAG-Key", "")
+    if key != api_key:
+        raise HTTPException(status_code=401, detail="Unauthorized")
+
+
+@app.on_event("startup")
+def startup():
+    build_on_start = os.getenv("RAG_BUILD_ON_START", "true").lower() in ("1", "true", "yes")
+    if build_on_start:
+        logger.info("RAG_BUILD_ON_START=true — loading corpus and building index")
+        try:
+            n = index.load_corpus()
+            logger.info("Loaded %d chunks from dataset", n)
+            index.build_index()
+            logger.info("Index built successfully (%d chunks, %d embeddings)",
+                        index.total_chunks, index.embedding_count)
+        except Exception as e:
+            logger.error("Failed to build index on startup: %s", e)
+    else:
+        logger.info("RAG_BUILD_ON_START=false — index not loaded")
+
+
+@app.get("/")
+def root():
+    return {
+        "service": "CodeSecAudit RAG Service",
+        "version": "0.6.0",
+        "status": "ok",
+        "embedding_model": os.getenv("RAG_EMBEDDING_MODEL", DEFAULT_EMBED_MODEL),
+        "dataset_repo": os.getenv("RAG_DATASET_REPO", DEFAULT_DATASET_REPO),
+        "total_chunks": index.total_chunks,
+        "api_key_protected": bool(api_key),
+        "uptime_seconds": int(time.time() - START_TIME),
+    }
+
+
+@app.get("/health")
+def health():
+    return {
+        "status": "ok",
+        "index_loaded": index.is_loaded,
+        "total_chunks": index.total_chunks,
+        "embedding_count": index.embedding_count,
+        "embedding_model": os.getenv("RAG_EMBEDDING_MODEL", DEFAULT_EMBED_MODEL),
+        "uptime_seconds": int(time.time() - START_TIME),
+    }
+
+
+@app.post("/rag/search", response_model=SearchResponse)
+def search(req: SearchRequest, request: Request):
+    _verify_key(request)
+    t0 = time.time()
+    results = index.search(req.query, top_k=req.top_k)
+    elapsed = (time.time() - t0) * 1000
+
+    return SearchResponse(
+        query=req.query,
+        top_k=req.top_k,
+        results=[SearchResultItem(**r) for r in results],
+        metadata=SearchMetadata(
+            embedding_model=os.getenv("RAG_EMBEDDING_MODEL", DEFAULT_EMBED_MODEL),
+            dataset_repo=os.getenv("RAG_DATASET_REPO", DEFAULT_DATASET_REPO),
+            total_chunks=index.total_chunks,
+            query_time_ms=round(elapsed, 2),
+        ),
+    )

rag_service/schemas.py

@@ -0,0 +1,30 @@
+from pydantic import BaseModel, Field
+
+
+class SearchRequest(BaseModel):
+    query: str = Field(min_length=1, description="Search query")
+    top_k: int = Field(default=3, ge=1, le=20, description="Number of results")
+
+
+class SearchResultItem(BaseModel):
+    rank: int
+    score: float
+    title: str = ""
+    section_title: str = ""
+    cwe_id: str = ""
+    content: str = ""
+    source_file: str = ""
+
+
+class SearchMetadata(BaseModel):
+    embedding_model: str
+    dataset_repo: str
+    total_chunks: int = 0
+    query_time_ms: float = 0.0
+
+
+class SearchResponse(BaseModel):
+    query: str
+    top_k: int
+    results: list[SearchResultItem] = Field(default_factory=list)
+    metadata: SearchMetadata

start.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+echo "=== CodeSecAudit RAG Service ==="
+echo "Dataset : ${RAG_DATASET_REPO:-OMCHOKSI108/CodeSecAudit-RAG}"
+echo "Model   : ${RAG_EMBEDDING_MODEL:-sentence-transformers/all-MiniLM-L6-v2}"
+echo "Port    : 7860"
+echo ""
+
+uvicorn rag_service.main:app --host 0.0.0.0 --port 7860