diff --git a/.env.example b/.env.example index dca28f0..54d0351 100644 --- a/.env.example +++ b/.env.example @@ -22,3 +22,7 @@ RERANKER_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2 LANGFUSE_PUBLIC_KEY= LANGFUSE_SECRET_KEY= LANGFUSE_HOST=https://cloud.langfuse.com + +# Async ingestion (optional — sync when unset) +# Set to enable Redis/arq queue: upload returns 202 + job_id, worker processes in background +# REDIS_URL=redis://localhost:6379 diff --git a/README.md b/README.md index 84d335d..6855e15 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,19 @@ A production-ready Retrieval-Augmented Generation system for querying PDF docume 5. Every query is **traced with Langfuse** (retrieval spans, generation spans, token usage, latency) 6. Pipeline quality is **measured with RAGAS** (faithfulness, context precision/recall, answer relevancy) +## Architecture + +![RAG Document Q&A — pipeline infographic](docs/rag-infographic.png) + +The query flow in four phases — **intake → hybrid retrieval → augmentation → generation** — plus the async indexing side-lane and the quality loop: + +- **Retrieval**: hybrid search (BM25 + vector, fused with RRF) followed by cross-encoder reranking +- **Generation**: reranked chunks are assembled into the prompt and answered by the configured LLM, with citations +- **Indexing**: uploads return `202 + job_id`; a background worker (arq + Redis) chunks, embeds and persists to ChromaDB +- **Observability**: every request is traced in Langfuse (spans, tokens, latency, user feedback); RAGAS evaluates against a versioned golden dataset + +An interactive, explorable version (dark/light themes, guided views, trace animation) lives in [`docs/rag-architecture.html`](docs/rag-architecture.html) — open it locally in a browser. The infographic source is [`docs/rag-infographic.html`](docs/rag-infographic.html). + ## Stack - **FastAPI** — REST API + web UI @@ -44,6 +57,12 @@ Every `/api/ask` call produces a trace with: No-op when Langfuse keys are not set — the app runs identically without an observability backend. +### Async ingestion (arq + Redis) +With `REDIS_URL` set, `POST /api/upload` returns `202 + job_id` instantly and a background worker (`python -m arq worker.WorkerSettings`) chunks, embeds and registers the PDF — with automatic retries (`max_tries=3`), bounded concurrency (`max_jobs=4`) and backpressure from the Redis queue. Poll `GET /api/jobs/{job_id}` for status. Without `REDIS_URL` the app processes uploads synchronously, exactly as before: zero extra infrastructure needed to run locally. + +### Feedback loop +Every `/api/ask` response includes a `trace_id`. Rate any answer with `POST /api/feedback` (score `+1`/`-1`, optional comment): feedback is stored locally (`data/feedback.jsonl`) as a tuning dataset and mirrored to Langfuse as a trace score. `GET /api/feedback/summary` returns totals and thumbs-down rate — the metric to watch after every retrieval/prompt change. + ### Evaluation (RAGAS) Answer the interview question *"how do you know your RAG works well?"* with numbers: @@ -71,10 +90,10 @@ The system is being scaled in phases, each designed to be demoable and measurabl - CI: lint + tests on every push ### 🚧 Phase 1 — Scale & reliability (in progress) -- **Async ingestion**: Redis-backed task queue (arq) for PDF processing — job status endpoint, retries, backpressure. Upload returns `202 + job_id` instead of blocking. -- **Feedback loop**: `POST /api/feedback` (👍/👎 per answer) stored in Langfuse → dataset for prompt/retrieval tuning. -- **One-command stack**: `docker compose up` brings up app + Redis + Langfuse. -- **Baseline metrics published**: RAGAS scores + p95 latency documented in this README. +- **Async ingestion**: arq/Redis queue — upload returns `202 + job_id`, worker with retries/backpressure, `GET /api/jobs/{job_id}` status. *(done)* +- **Feedback loop**: `POST /api/feedback` (👍/👎) persisted locally + Langfuse score mirroring; `/api/feedback/summary` aggregates. *(done)* +- **One-command stack**: `docker compose up` brings up app + worker + Redis. *(done)* +- **Baseline metrics published**: RAGAS scores + p95 latency documented in this README. *(pending)* ### Phase 2 — Multi-user & guardrails - **Collections / multi-tenancy**: namespaced document sets per user or project (ChromaDB collections) with per-collection queries. @@ -111,17 +130,20 @@ uvicorn main:app --reload --port 8000 docker compose up --build ``` -The API will be available at `http://localhost:8000`. +Brings up **API + arq worker + Redis** — async ingestion works out of the box. The API will be available at `http://localhost:8000`. ## API endpoints | Method | Path | Description | |--------|------|-------------| -| `POST` | `/api/upload` | Upload PDF document(s) | -| `POST` | `/api/ask` | Ask a question (returns answer + sources) | +| `POST` | `/api/upload` | Upload PDF (sync result, or `202 + job_id` when async) | +| `GET` | `/api/jobs/{job_id}` | Poll async ingestion job status | +| `POST` | `/api/ask` | Ask a question (returns answer + sources + `trace_id`) | +| `POST` | `/api/feedback` | Rate an answer 👍/👎 (score `+1`/`-1`) | +| `GET` | `/api/feedback/summary` | Feedback aggregates (thumbs-down rate) | | `GET` | `/api/documents` | List uploaded documents | | `DELETE` | `/api/documents/{id}` | Delete a document | -| `GET` | `/api/health` | Health check (includes hybrid + tracing status) | +| `GET` | `/api/health` | Health check (hybrid, tracing, ingestion mode, feedback) | ## Architecture diff --git a/docker-compose.yml b/docker-compose.yml index 75417bd..898a265 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -9,4 +9,25 @@ services: - ./data:/app/data env_file: - .env + environment: + # Inside compose, REDIS_URL defaults to the bundled redis service. + # The app falls back to sync ingestion when REDIS_URL is empty. + - REDIS_URL=${REDIS_URL:-redis://redis:6379} + depends_on: + - redis + restart: unless-stopped + + rag-worker: + build: . + command: python -m arq worker.WorkerSettings + volumes: + - ./data:/app/data + environment: + - REDIS_URL=${REDIS_URL:-redis://redis:6379} + depends_on: + - redis + restart: unless-stopped + + redis: + image: redis:7-alpine restart: unless-stopped diff --git a/docs/rag-architecture.html b/docs/rag-architecture.html new file mode 100644 index 0000000..01acee5 --- /dev/null +++ b/docs/rag-architecture.html @@ -0,0 +1,13275 @@ + + + + + + + RAG Document Q&A — How a Query Flows Diagram + + + + + + + + + + + +
+ +
+
+
+

RAG Document Q&A — How a Query Flows

+
+

Hybrid retrieval (BM25 + vector + RRF) · cross-encoder reranking · traced with Langfuse · measured with RAGAS

+
+ + + + + + +
+ + RAG Document Q&A — How a Query Flows + Hybrid retrieval (BM25 + vector + RRF) · cross-encoder reranking · traced with Langfuse · measured with RAGAS + + + + + + + + + + + + + + + + + + + + + + + + + 01 / User / Client + + + 02 / External LLM + + + 03 / RAG Pipeline + + + 04 / Indexing (async) + + + 05 / Observability & Eval + + + + + 1 · Intake + + + 2 · Retrieval + + + 3 · Generate + + + + Hybrid retrieval core + + arq + Redis + + + + + + + + + + + + + + + + + User · asks a question · User / Client › 1 · Intake + + + + User + asks a question + + + + Web UI / API · FastAPI · /api/ask · User / Client › 2 · Retrieval + + + + Web UI / API + FastAPI · /api/ask + + + + Answer · grounded + cited · User / Client › 3 · Generate · trace_id + + + + Answer + grounded + cited + trace_id + + + + LLM · OpenRouter / OpenAI · External LLM › 3 · Generate + + + + LLM + OpenRouter / OpenAI + + + + Hybrid Search · BM25 + vector · RRF · RAG Pipeline › Hybrid retrieval core › 2 · Retrieval · HYBRID_SEARCH + + + + Hybrid Search + BM25 + vector · RRF + HYBRID_SEARCH + + + + Reranker · cross-encoder · RAG Pipeline › Hybrid retrieval core › 2 · Retrieval + + + + Reranker + cross-encoder + + + + Augment · prompt assembly · RAG Pipeline › 3 · Generate + + + + Augment + prompt assembly + + + + PDF Upload · /api/upload · Indexing (async) › 1 · Intake + + + + PDF Upload + /api/upload + + + + Worker · chunk + embed · Indexing (async) › arq + Redis › 2 · Retrieval + + + + Worker + chunk + embed + + + + ChromaDB · vectors + BM25 · Indexing (async) › arq + Redis › 2 · Retrieval + + + + ChromaDB + vectors + BM25 + + + + Langfuse · traces · tokens · Observability & Eval › 3 · Generate + + + + Langfuse + traces · tokens + + + + RAGAS Eval · faithfulness · recall · Observability & Eval › 2 · Retrieval · golden dataset + + + + RAGAS Eval + faithfulness · recall + golden dataset + + + + + + + query + + + + + + + + + + + trace + feedback + + + + eval + + + + + Legend + + + User UI + + + + Agent logic + + + + Policy + + + + Tool action + + + + Context / trace + + + + Cloud service + + + + External system + + + +

+ + + + + + + + +
+ + +
+
+
+
+

Why Hybrid Retrieval

+
+
    +
  • • Pure vector search misses exact terms: names, codes, numbers
  • +
  • • BM25 catches keywords; vectors catch semantics
  • +
  • • Reciprocal Rank Fusion merges both rankings
  • +
  • • Cross-encoder reranking gives final precision
  • +
+
+ +
+
+
+

Production Touches

+
+
    +
  • • Async ingestion: upload returns 202 + job_id instantly
  • +
  • • Worker with retries (×3) and bounded concurrency
  • +
  • • Every query traced: spans, token usage, latency
  • +
  • • Feedback loop (±1) stored as a tuning dataset
  • +
+
+ +
+
+
+

How Quality Is Proven

+
+
    +
  • • RAGAS: faithfulness, relevancy, precision, recall
  • +
  • • Versioned golden dataset for before/after comparisons
  • +
  • • CI runs lint + tests on every push
  • +
  • • Thumbs-down rate watched after every change
  • +
+
+
+ + + +
+ + + + diff --git a/docs/rag-architecture.png b/docs/rag-architecture.png new file mode 100644 index 0000000..a1c3b34 Binary files /dev/null and b/docs/rag-architecture.png differ diff --git a/docs/rag-architecture.workflow.json b/docs/rag-architecture.workflow.json new file mode 100644 index 0000000..363e52a --- /dev/null +++ b/docs/rag-architecture.workflow.json @@ -0,0 +1,108 @@ +{ + "schema_version": 1, + "diagram_type": "workflow", + "meta": { + "title": "RAG Document Q&A — How a Query Flows", + "subtitle": "Hybrid retrieval (BM25 + vector + RRF) · cross-encoder reranking · traced with Langfuse · measured with RAGAS", + "visual_preset": "signal-flow", + "quality_profile": "showcase", + "views": [ + { + "id": "query-path", + "label": "Query to answer", + "focus": ["user", "chatui", "hybrid", "rerank", "prompt", "llm", "answer"], + "note": "Follow a natural-language question from the UI to a grounded, cited answer." + }, + { + "id": "indexing-path", + "label": "Async indexing", + "focus": ["pdf", "worker", "chroma"], + "note": "PDFs are queued, chunked, embedded and persisted by a background worker with retries." + }, + { + "id": "quality-loop", + "label": "Observability & eval", + "focus": ["chatui", "langfuse", "chroma", "ragas"], + "note": "Every request is traced; RAGAS scores faithfulness and context quality against a golden dataset." + } + ], + "output": "docs/rag-architecture.html", + "viewBox": [720, 900] + }, + "lanes": [ + { "id": "client", "label": "User / Client" }, + { "id": "llmapi", "label": "External LLM" }, + { "id": "pipeline", "label": "RAG Pipeline" }, + { "id": "ingest", "label": "Indexing (async)" }, + { "id": "observe", "label": "Observability & Eval" } + ], + "phases": [ + { "id": "intake", "label": "1 · Intake", "fromCol": 0, "toCol": 0 }, + { "id": "retrieval", "label": "2 · Retrieval", "fromCol": 1, "toCol": 4, "variant": "emphasis" }, + { "id": "generate", "label": "3 · Generate", "fromCol": 5, "toCol": 5, "variant": "dashed" } + ], + "groups": [ + { "id": "query_core", "label": "Hybrid retrieval core", "lane": "pipeline", "fromCol": 1, "toCol": 4, "variant": "emphasis" }, + { "id": "indexing_group", "label": "arq + Redis", "lane": "ingest", "fromCol": 2, "toCol": 3, "variant": "dashed" } + ], + "mainPath": ["user", "chatui", "hybrid", "rerank", "prompt", "llm", "answer"], + "nodes": [ + { "id": "user", "lane": "client", "col": 0, "type": "external", "label": "User", "sublabel": "asks a question" }, + { "id": "chatui", "lane": "client", "col": 1, "type": "frontend", "label": "Web UI / API", "sublabel": "FastAPI · /api/ask" }, + { "id": "answer", "lane": "client", "col": 5, "type": "backend", "label": "Answer", "sublabel": "grounded + cited", "tag": "trace_id" }, + { "id": "llm", "lane": "llmapi", "col": 5, "type": "cloud", "label": "LLM", "sublabel": "OpenRouter / OpenAI" }, + { "id": "hybrid", "lane": "pipeline", "col": 1, "type": "backend", "label": "Hybrid Search", "sublabel": "BM25 + vector · RRF", "tag": "HYBRID_SEARCH" }, + { "id": "rerank", "lane": "pipeline", "col": 4, "type": "backend", "label": "Reranker", "sublabel": "cross-encoder" }, + { "id": "prompt", "lane": "pipeline", "col": 5, "type": "backend", "label": "Augment", "sublabel": "prompt assembly" }, + { "id": "pdf", "lane": "ingest", "col": 0, "type": "external", "label": "PDF Upload", "sublabel": "/api/upload" }, + { "id": "worker", "lane": "ingest", "col": 2, "type": "messagebus", "label": "Worker", "sublabel": "chunk + embed" }, + { "id": "chroma", "lane": "ingest", "col": 3, "type": "database", "label": "ChromaDB", "sublabel": "vectors + BM25" }, + { "id": "langfuse", "lane": "observe", "col": 5, "type": "database", "label": "Langfuse", "sublabel": "traces · tokens" }, + { "id": "ragas", "lane": "observe", "col": 3, "type": "security", "label": "RAGAS Eval", "sublabel": "faithfulness · recall", "tag": "golden dataset" } + ], + "edges": [ + { "id": "user-query", "from": "user", "to": "chatui", "variant": "default" }, + { "id": "query-retrieve", "from": "chatui", "to": "hybrid", "label": "query", "variant": "emphasis" }, + { "id": "hybrid-rerank", "from": "hybrid", "to": "rerank", "variant": "emphasis", "route": "straight" }, + { "id": "rerank-prompt", "from": "rerank", "to": "prompt", "variant": "emphasis" }, + { "id": "prompt-llm", "from": "prompt", "to": "llm", "variant": "emphasis" }, + { "id": "llm-answer", "from": "llm", "to": "answer", "variant": "emphasis", "role": "return" }, + { "id": "upload-worker", "from": "pdf", "to": "worker", "variant": "default" }, + { "id": "worker-store", "from": "worker", "to": "chroma", "variant": "default", "route": "straight" }, + { "id": "store-search", "from": "chroma", "to": "hybrid", "variant": "default", "role": "branch" }, + { "id": "prompt-trace", "from": "prompt", "to": "langfuse", "label": "trace + feedback", "variant": "dashed", "fromSide": "bottom", "toSide": "top", "route": "drop" }, + { "id": "store-eval", "from": "chroma", "to": "ragas", "label": "eval", "variant": "dashed", "role": "branch" } + ], + "cards": [ + { + "dot": "cyan", + "title": "Why Hybrid Retrieval", + "items": [ + "Pure vector search misses exact terms: names, codes, numbers", + "BM25 catches keywords; vectors catch semantics", + "Reciprocal Rank Fusion merges both rankings", + "Cross-encoder reranking gives final precision" + ] + }, + { + "dot": "violet", + "title": "Production Touches", + "items": [ + "Async ingestion: upload returns 202 + job_id instantly", + "Worker with retries (×3) and bounded concurrency", + "Every query traced: spans, token usage, latency", + "Feedback loop (±1) stored as a tuning dataset" + ] + }, + { + "dot": "emerald", + "title": "How Quality Is Proven", + "items": [ + "RAGAS: faithfulness, relevancy, precision, recall", + "Versioned golden dataset for before/after comparisons", + "CI runs lint + tests on every push", + "Thumbs-down rate watched after every change" + ] + } + ] +} diff --git a/docs/rag-infographic.html b/docs/rag-infographic.html new file mode 100644 index 0000000..5571f1c --- /dev/null +++ b/docs/rag-infographic.html @@ -0,0 +1,359 @@ + + + + + + + + +

Explaining RAG Document Q&A
to your Team

+ +
+
RAG Pipeline
+ +
+ + +
+

Data
injection

+
+
💬Chat UIType your query
+
+
👤User
+
📄DocumentsPDF · POST /api/upload
+
+
+ +
+ + +
+
+

Phase 1

+

Query intake & orchestration

+ Intake +
+
+
⚡FastAPIPOST /api/ask
+
+
✅Validationschema + limits
+
+
Hybrid
search
ON?
+
+ Async indexing +
⚙️Workerarq + Redis · 202 + job_id
+
🧩Chunk + EmbedLangChain · sentence-transformers
+
+
+
+ +
+ + +
+
+

Phase 2

+

Hybrid retrieval

+ Retrieve +
+
+
🔍BM25 Searchkeyword match · exact terms
+
🧬Vector Searchsemantic similarity · ChromaDB
+
🔀RRF Fusionmerge both rankings
+
🎯Rerankcross-encoder MiniLM
+
+
+ +
+ + +
+
+

Phase 3

+

Context augmentation

+ Augment +
+
+
+
Retrieved Data Augmentation
+
Reranked chunks from ChromaDB
+
Add system prompt
+
Augment user query
+
+
+
+ +
+ + +
+
+

Phase 4

+

Grounded generation

+ Generate +
+
+
🤖LLM GenerationOpenRouter / OpenAI
+
+
📄Answer + Citationsgrounded in sources
+
+
👍Feedback ±1tuning dataset
+
+ Quality loop +
📊Langfusetraces · tokens · latency
+
📈RAGASfaithfulness · golden dataset
+
+
+
+ +
+ + +
+

Output

+
+
Generated answer according to query (with source citations)
+
📊Trace + scoresLangfuse · RAGAS report
+
+
+ +
+
+ + + + + diff --git a/docs/rag-infographic.png b/docs/rag-infographic.png new file mode 100644 index 0000000..dee634b Binary files /dev/null and b/docs/rag-infographic.png differ diff --git a/main.py b/main.py index 5f07a14..c229dcf 100644 --- a/main.py +++ b/main.py @@ -11,10 +11,11 @@ from fastapi import FastAPI, File, Form, HTTPException, UploadFile from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import HTMLResponse -from fastapi.staticfiles import StaticFiles +from fastapi.responses import HTMLResponse, JSONResponse +from rag import jobs from rag.chunker import PDFChunker +from rag.feedback import FeedbackStore from rag.hybrid import HybridRetriever from rag.llm import LLMGenerator from rag.observability import Tracer, elapsed_ms, timed @@ -58,6 +59,7 @@ model=os.getenv("OPENROUTER_MODEL", "meta-llama/llama-3.1-8b-instruct:free"), ) tracer = Tracer() +feedback = FeedbackStore(str(BASE_DIR / "data" / "feedback.jsonl"), tracer=tracer) # --- Document registry (simple JSON-based) --- REGISTRY_PATH = BASE_DIR / "data" / "registry.json" @@ -92,12 +94,19 @@ async def health(): "llm_enabled": llm.is_enabled(), "hybrid_search": USE_HYBRID, "tracing": tracer.enabled, + "ingestion_mode": jobs.describe_mode(), + "feedback": feedback.summary(), } @app.post("/api/upload") async def upload_document(file: UploadFile = File(...)): - """Upload a PDF, chunk it, embed it, and register it.""" + """Upload a PDF for ingestion. + + Async mode (REDIS_URL set): enqueues the job, returns ``202 + job_id`` + immediately — poll ``GET /api/jobs/{job_id}`` for the result. + Sync mode (default): processes inline and returns the result directly. + """ if not file.filename or not file.filename.lower().endswith(".pdf"): raise HTTPException(400, "Only PDF files are supported") @@ -108,6 +117,23 @@ async def upload_document(file: UploadFile = File(...)): logger.info("Processing %s (%d bytes)...", file.filename, len(content)) + if jobs.async_enabled(): + try: + job_id = await jobs.enqueue_ingestion( + doc_id=doc_id, + filepath=str(filepath), + filename=file.filename, + size_bytes=len(content), + ) + except Exception as exc: + filepath.unlink(missing_ok=True) + logger.error("Failed to enqueue ingestion: %s", exc) + raise HTTPException(503, "Ingestion queue unavailable") from exc + return JSONResponse( + status_code=202, + content={"job_id": job_id, "doc_id": doc_id, "status": "queued"}, + ) + try: chunks = chunker.chunk_pdf(str(filepath)) except (FileNotFoundError, ValueError) as exc: @@ -137,6 +163,17 @@ async def upload_document(file: UploadFile = File(...)): return {"doc_id": doc_id, "filename": file.filename, "chunks": len(chunks)} +@app.get("/api/jobs/{job_id}") +async def get_job_status(job_id: str): + """Poll the status of an async ingestion job.""" + if not jobs.async_enabled(): + raise HTTPException(400, "Async ingestion is not enabled (REDIS_URL unset)") + result = await jobs.job_status(job_id) + if result["status"] == "not_found": + raise HTTPException(404, "Job not found") + return result + + @app.post("/api/ask") async def ask_question(question: str = Form(...), top_k: int = Form(5)): """Answer a question using retrieved document chunks.""" @@ -167,7 +204,13 @@ async def ask_question(question: str = Form(...), top_k: int = Form(5)): usage=llm.last_usage, latency_ms=gen_ms, ) trace.log_answer(answer, llm_used=True) - return {"question": question, "answer": answer, "sources": sources, "llm": True} + return { + "question": question, + "answer": answer, + "sources": sources, + "llm": True, + "trace_id": trace.trace_id, + } # Fallback: return retrieved chunks directly answer = "LLM not configured. Here are the most relevant chunks:\n\n" + "\n\n---\n\n".join(context_chunks) @@ -177,9 +220,43 @@ async def ask_question(question: str = Form(...), top_k: int = Form(5)): "answer": answer, "sources": sources, "llm": False, + "trace_id": trace.trace_id, } +@app.post("/api/feedback") +async def submit_feedback( + question: str = Form(...), + answer: str = Form(...), + score: int = Form(...), + comment: str = Form(""), + trace_id: Optional[str] = Form(None), +): + """Record user feedback (👍 = +1 / 👎 = -1) for an answer. + + Stored locally (data/feedback.jsonl) and mirrored as a Langfuse score + when a ``trace_id`` from ``/api/ask`` is provided. The local store is + the tuning dataset: thumbs-down answers drive prompt/retrieval fixes. + """ + try: + entry = feedback.record( + question=question, + answer=answer, + score=score, + comment=comment, + trace_id=trace_id, + ) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc + return {"recorded": True, "feedback_id": entry["feedback_id"]} + + +@app.get("/api/feedback/summary") +async def feedback_summary(): + """Aggregate feedback stats: totals and thumbs-down rate.""" + return feedback.summary() + + @app.get("/api/documents") async def list_documents(): """List all uploaded documents.""" diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..a635c5c --- /dev/null +++ b/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +pythonpath = . diff --git a/rag/feedback.py b/rag/feedback.py new file mode 100644 index 0000000..9de9763 --- /dev/null +++ b/rag/feedback.py @@ -0,0 +1,106 @@ +"""Feedback loop — capture 👍/👎 per answer and push scores to Langfuse. + +Feedback is always persisted locally (append-only JSONL) so it works with +zero external dependencies, and mirrored as a Langfuse score when a +``trace_id`` is provided and tracing is enabled. The local store doubles as +a tuning dataset (thumbs-down answers → prompt/retrieval improvements). +""" +import json +import logging +import threading +import time +import uuid +from pathlib import Path +from typing import Dict, List, Optional + +logger = logging.getLogger(__name__) + + +class FeedbackStore: + """Append-only local feedback store with Langfuse score mirroring.""" + + def __init__(self, path: str, tracer=None): + self._path = Path(path) + self._path.parent.mkdir(parents=True, exist_ok=True) + self._tracer = tracer + self._lock = threading.Lock() + + def record( + self, + question: str, + answer: str, + score: int, + comment: str = "", + trace_id: Optional[str] = None, + sources: Optional[List[Dict]] = None, + ) -> Dict: + """Persist one feedback entry and mirror it to Langfuse if possible. + + Args: + question: The original user question. + answer: The answer that was rated. + score: +1 (helpful) or -1 (not helpful). + comment: Optional free-text comment. + trace_id: Langfuse trace id to attach the score to (optional). + sources: Optional source chunks shown with the answer. + + Returns: + The stored feedback entry (includes generated ``feedback_id``). + """ + if score not in (1, -1): + raise ValueError("score must be +1 (helpful) or -1 (not helpful)") + + entry = { + "feedback_id": uuid.uuid4().hex[:12], + "ts": time.time(), + "question": question, + "answer": answer, + "score": score, + "comment": comment, + "trace_id": trace_id, + "sources": sources or [], + } + + with self._lock: + with self._path.open("a", encoding="utf-8") as fh: + fh.write(json.dumps(entry, ensure_ascii=False) + "\n") + + if trace_id and self._tracer is not None: + self._tracer.score_feedback( + trace_id=trace_id, + score=score, + comment=comment, + ) + + logger.info( + "Feedback recorded: %s score=%+d trace=%s", + entry["feedback_id"], score, trace_id or "-", + ) + return entry + + def summary(self) -> Dict: + """Aggregate stats: total, positive, negative, thumbs-down rate.""" + entries = self._read_all() + pos = sum(1 for e in entries if e["score"] == 1) + neg = sum(1 for e in entries if e["score"] == -1) + total = len(entries) + return { + "total": total, + "positive": pos, + "negative": neg, + "thumbs_down_rate": round(neg / total, 4) if total else None, + } + + def _read_all(self) -> List[Dict]: + """Read all feedback entries (small dataset — fine for a demo).""" + if not self._path.exists(): + return [] + out = [] + for line in self._path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if line: + try: + out.append(json.loads(line)) + except json.JSONDecodeError: + logger.warning("Skipping malformed feedback line") + return out diff --git a/rag/hybrid.py b/rag/hybrid.py index 8e6bfdb..9e6950c 100644 --- a/rag/hybrid.py +++ b/rag/hybrid.py @@ -9,7 +9,7 @@ import logging import re import threading -from typing import Dict, List, Optional +from typing import Dict, List from .fusion import fuse_results diff --git a/rag/jobs.py b/rag/jobs.py new file mode 100644 index 0000000..f52471b --- /dev/null +++ b/rag/jobs.py @@ -0,0 +1,79 @@ +"""Async ingestion — arq/Redis job queue for PDF processing. + +When ``REDIS_URL`` is set, ``POST /api/upload`` enqueues a job and returns +``202 + job_id`` immediately; a separate worker process (``python -m +arq worker.WorkerSettings``) chunks, embeds and registers the document with +retries and backpressure handled by arq/Redis. + +When ``REDIS_URL`` is unset the app falls back to fully synchronous +processing — same behaviour as before, zero extra infrastructure needed. +This mirrors the Langfuse no-op pattern: the feature activates via env var. +""" +import logging +import os +from typing import Dict + +logger = logging.getLogger(__name__) + + +def async_enabled() -> bool: + """Async ingestion is active only when REDIS_URL is configured.""" + return bool(os.getenv("REDIS_URL", "").strip()) + + +_pool = None + + +async def get_pool(): + """Lazily create the arq connection pool (shared across requests).""" + global _pool + if _pool is None: + from arq import create_pool + from arq.connections import RedisSettings + + _pool = await create_pool(RedisSettings.from_dsn(os.environ["REDIS_URL"])) + logger.info("arq pool connected to %s", os.environ["REDIS_URL"]) + return _pool + + +async def enqueue_ingestion(doc_id: str, filepath: str, filename: str, size_bytes: int) -> str: + """Enqueue a PDF ingestion job. Returns the arq job id.""" + pool = await get_pool() + job = await pool.enqueue_job( + "ingest_pdf", + doc_id=doc_id, + filepath=filepath, + filename=filename, + size_bytes=size_bytes, + ) + logger.info("Enqueued ingestion job %s for doc %s (%s)", job.job_id, doc_id, filename) + return job.job_id + + +async def job_status(job_id: str) -> Dict: + """Return status for an ingestion job: queued | in_progress | complete | failed | not_found.""" + from arq.jobs import Job, JobStatus + + pool = await get_pool() + job = Job(job_id, pool) + status = await job.status() + info = await job.info() + + if status == JobStatus.not_found or info is None: + return {"job_id": job_id, "status": "not_found"} + + result: Dict = {"job_id": job_id, "status": status.value} + if status == JobStatus.complete: + try: + result["result"] = await job.result() + except Exception as exc: # job raised — arq marks it complete but result re-raises + result["status"] = "failed" + result["error"] = str(exc) + elif info.enqueue_time: + result["enqueued_at"] = info.enqueue_time.isoformat() + return result + + +def describe_mode() -> str: + """Human-readable ingestion mode for the health endpoint.""" + return "async" if async_enabled() else "sync" diff --git a/rag/observability.py b/rag/observability.py index 2502d3a..0f5b730 100644 --- a/rag/observability.py +++ b/rag/observability.py @@ -61,6 +61,25 @@ def trace_question(self, question: str, top_k: int): finally: self._client.flush() + def score_feedback(self, trace_id: str, score: int, comment: str = "") -> None: + """Attach a user feedback score to an existing Langfuse trace. + + No-op when tracing is disabled or the client call fails — feedback + is always persisted locally by ``FeedbackStore`` regardless. + """ + if not self.enabled: + return + try: + self._client.score( + trace_id=trace_id, + name="user-feedback", + value=score, + comment=comment or None, + ) + self._client.flush() + except Exception as exc: + logger.warning("Langfuse score failed (%s) — local copy kept", exc) + def shutdown(self) -> None: """Flush pending events (call on app shutdown).""" if self.enabled: @@ -127,6 +146,11 @@ def log_answer(self, answer: str, llm_used: bool) -> None: return self._trace.update(output={"answer": answer, "llm": llm_used}) + @property + def trace_id(self) -> Optional[str]: + """Langfuse trace id for this request (for feedback scoring).""" + return getattr(self._trace, "id", None) if self.enabled else None + def _seconds_ago(ms: float): """Timestamp ``ms`` milliseconds in the past (for span start times).""" diff --git a/rag/retriever.py b/rag/retriever.py index dda405d..8c8d74f 100644 --- a/rag/retriever.py +++ b/rag/retriever.py @@ -3,7 +3,6 @@ from typing import Dict, List, Optional import chromadb -from chromadb.config import Settings from chromadb.utils import embedding_functions logger = logging.getLogger(__name__) diff --git a/requirements.txt b/requirements.txt index 4560c88..c249007 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,3 +10,4 @@ pypdf==5.1.0 httpx==0.27.0 rank-bm25==0.2.2 langfuse>=2.60,<3 +arq==0.26.3 diff --git a/tests/test_observability.py b/tests/test_observability.py index dce69b9..0022c70 100644 --- a/tests/test_observability.py +++ b/tests/test_observability.py @@ -1,5 +1,4 @@ """Tests for observability tracer — no-op behaviour without Langfuse config.""" -import os from rag.observability import TraceHandle, Tracer diff --git a/tests/test_phase1.py b/tests/test_phase1.py new file mode 100644 index 0000000..2effd48 --- /dev/null +++ b/tests/test_phase1.py @@ -0,0 +1,55 @@ +"""Tests for the feedback store and async-ingestion feature flag.""" +import json + +from rag.feedback import FeedbackStore +from rag.jobs import async_enabled, describe_mode + + +def test_feedback_record_and_summary(tmp_path): + store = FeedbackStore(str(tmp_path / "feedback.jsonl")) + entry = store.record(question="q1", answer="a1", score=1, trace_id="t1") + assert entry["feedback_id"] + store.record(question="q2", answer="a2", score=-1, comment="wrong source") + + summary = store.summary() + assert summary == { + "total": 2, + "positive": 1, + "negative": 1, + "thumbs_down_rate": 0.5, + } + + # Verify the JSONL on disk is parseable and complete + lines = (tmp_path / "feedback.jsonl").read_text().splitlines() + assert len(lines) == 2 + parsed = json.loads(lines[1]) + assert parsed["score"] == -1 + assert parsed["comment"] == "wrong source" + + +def test_feedback_rejects_invalid_score(tmp_path): + store = FeedbackStore(str(tmp_path / "feedback.jsonl")) + try: + store.record(question="q", answer="a", score=5) + assert False, "should have raised" + except ValueError: + pass + + +def test_feedback_empty_summary(tmp_path): + store = FeedbackStore(str(tmp_path / "feedback.jsonl")) + summary = store.summary() + assert summary["total"] == 0 + assert summary["thumbs_down_rate"] is None + + +def test_async_ingestion_disabled_by_default(monkeypatch): + monkeypatch.delenv("REDIS_URL", raising=False) + assert async_enabled() is False + assert describe_mode() == "sync" + + +def test_async_ingestion_enabled_with_redis_url(monkeypatch): + monkeypatch.setenv("REDIS_URL", "redis://localhost:6379") + assert async_enabled() is True + assert describe_mode() == "async" diff --git a/worker.py b/worker.py new file mode 100644 index 0000000..3401050 --- /dev/null +++ b/worker.py @@ -0,0 +1,95 @@ +"""arq worker — processes PDF ingestion jobs from the Redis queue. + +Run with:: + + REDIS_URL=redis://localhost:6379 python -m arq worker.WorkerSettings + +The worker builds its own component instances (chunker, retriever, hybrid +index) so it is fully independent from the API process. +""" +import json +import logging +import os +from pathlib import Path + +from rag.chunker import PDFChunker +from rag.hybrid import HybridRetriever +from rag.retriever import Retriever + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(message)s") +logger = logging.getLogger(__name__) + +BASE_DIR = Path(__file__).parent +REGISTRY_PATH = BASE_DIR / "data" / "registry.json" + + +def _load_registry() -> dict: + if REGISTRY_PATH.exists(): + return json.loads(REGISTRY_PATH.read_text()) + return {} + + +def _save_registry(reg: dict) -> None: + REGISTRY_PATH.write_text(json.dumps(reg, indent=2)) + + +async def ingest_pdf(ctx, doc_id: str, filepath: str, filename: str, size_bytes: int) -> dict: + """Chunk, embed and register one PDF. Runs inside the worker process. + + Raises on failure so arq retries with backoff (see ``max_tries`` below). + """ + chunker: PDFChunker = ctx["chunker"] + retriever: Retriever = ctx["retriever"] + hybrid: HybridRetriever = ctx["hybrid"] + + logger.info("Ingesting doc %s (%s, %d bytes)", doc_id, filename, size_bytes) + + chunks = chunker.chunk_pdf(filepath) + if not chunks: + Path(filepath).unlink(missing_ok=True) + raise ValueError(f"No extractable text found in {filename}") + + retriever.add_documents(chunks, {"doc_id": doc_id, "filename": filename}) + hybrid.sync_index() + + reg = _load_registry() + reg[doc_id] = {"filename": filename, "chunks": len(chunks), "size_bytes": size_bytes} + _save_registry(reg) + + logger.info("Ingested doc %s: %d chunks", doc_id, len(chunks)) + return {"doc_id": doc_id, "filename": filename, "chunks": len(chunks)} + + +async def startup(ctx): + """Initialise heavy components once per worker process.""" + ctx["chunker"] = PDFChunker(chunk_size=1000, chunk_overlap=200) + retriever = Retriever( + persist_dir=str(BASE_DIR / "data" / "chroma"), + collection_name="documents", + ) + ctx["retriever"] = retriever + ctx["hybrid"] = HybridRetriever(retriever) + logger.info("Worker ready") + + +def get_worker_settings(): + """Build WorkerSettings lazily so importing this module does not require arq.""" + from arq.connections import RedisSettings + + class WorkerSettings: + """arq settings: functions, retry policy, connection.""" + + functions = [ingest_pdf] + on_startup = startup + max_jobs = 4 # concurrency — protects CPU-bound embedding work + max_tries = 3 # automatic retries with backoff on failure + job_timeout = 600 # 10 min per PDF, generous for large documents + redis_settings = RedisSettings.from_dsn( + os.getenv("REDIS_URL", "redis://localhost:6379") + ) + + return WorkerSettings + + +# arq resolves the settings class by name from the module namespace. +WorkerSettings = get_worker_settings()