Skip to content

README.md

Latest commit

 

History

History
312 lines (258 loc) · 16.1 KB

README.md

File metadata and controls

312 lines (258 loc) · 16.1 KB
Raw

ATLAS — Industrial Agentic RAG for Aurora Station

Watch the walkthrough: Agentic RAG — Architectural Patterns for Knowledge-Driven AI Systems (video) — the patterns this platform implements, presented by the author.

ATLAS is a self-correcting research analyst that answers questions about Aurora Station, a fictional six-person Mars research outpost. It combines three classic Agentic RAG patterns in one LangGraph loop:

  1. Routing + query decomposition — a structured router picks the right knowledge collection(s); multi-collection questions are decomposed into standalone sub-questions.
  2. Corrective RAG (CRAG, arXiv:2401.15884) — every retrieved chunk is graded for relevance; weak evidence triggers an anchored query rewrite and escalation to an archive index (plus optional Tavily web search).
  3. Self-RAG groundedness reflection (arXiv:2310.11511) — the draft answer is checked against the context; ungrounded drafts send the agent back for another retrieval round, under a hard retry budget enforced in code.

This repository is the industrial re-implementation of the academic notebook in ../atlas-agentic-rag/: same agent behavior, production engineering — config-driven settings, persistent vector stores, a REST + SSE API, containerization, structured logging, an eval harness, and a deterministic fake-LLM mode that runs the whole system (including the test suite) with zero API keys.

Architecture

The two design frames this implementation is built on (from Agentic RAG — Architectural Patterns for Knowledge-Driven AI Systems):

Inside a knowledge retrieval agent

Anatomy of the knowledge agent — query understanding, retrieval, preprocessing and ranking, reasoning and generation, with provenance collected continuously across every stage. In ATLAS that rail is the bracketed source citations plus the run trace.

Retrieval as a control loop

Retrieval as a control loop — perceive, plan, act (retrieve), evaluate, adapt, synthesize. ATLAS implements the loop concretely: routing and decomposition, CRAG grading, anchored rewriting with archive escalation, and groundedness reflection under a hard retry budget.

                        ┌──────────┐   route=multiple   ┌────────────┐
        question ──────▶│  router  │───────────────────▶│ decompose  │──┐
                        └──────────┘                    └────────────┘  │
                             │ route=ops|science|crew                   ▼
                             │                                    ┌──────────┐
                             └───────────────────────────────────▶│ retrieve │
                                                                  └──────────┘
                                                                        │
                                                                        ▼
                    ┌────────────────────────────────────────────  ┌────────────┐
                    │ enough relevant chunks                       │   grade    │
                    ▼                                              └────────────┘
              ┌───────────┐      grounded=yes                            │ weak,
              │  generate │──▶ ┌────────────────┐        END             │ budget left
              └───────────┘    │ groundedness   │──────▶ (answer)        ▼
                               └────────────────┘                 ┌──────────────┐
                                       │ no, budget left          │ rewrite_query│
                                       └─────────────────────────▶│ (anchored)   │
                                                                  └──────────────┘
                                                                         │
                                                                         ▼
                                                              ┌───────────────────┐
                                                              │ retrieve_archive  │──▶ grade
                                                              │ (+ optional web)  │    (loop)
                                                              └───────────────────┘

Rendered directly from the compiled graph (atlas graph):

graph TD;
	__start__ --> router;
	router -.-> decompose;
	router -.-> retrieve;
	decompose --> retrieve;
	retrieve --> grade_documents;
	grade_documents -.-> generate;
	grade_documents -.-> rewrite_query;
	rewrite_query --> retrieve_archive;
	retrieve_archive --> grade_documents;
	generate --> check_groundedness;
	check_groundedness -.-> rewrite_query;
	check_groundedness -.-> __end__;
Loading

Academic → industrial mapping

Academic notebook cell Industrial module
Inline Document lists data/corpus/<collection>/*.md + atlas.corpus loader
InMemoryVectorStore globals atlas.stores: protocol + Chroma (persistent, default), JSON file, in-memory adapters
getpass key loading atlas.config.Settings (pydantic-settings, ATLAS_ env prefix, .env)
MAX_RETRIES / MIN_RELEVANT constants ATLAS_MAX_RETRIES / ATLAS_MIN_RELEVANT settings, enforced in graph edges
Pydantic schemas + prompts (verbatim) atlas.schemas, atlas.prompts
Chain globals (router_llm, grader, ...) atlas.chains.build_chains()
Notebook node functions atlas.nodes.NodeFactory (injected chains/stores/settings, per-node latency)
builder.compile() atlas.graph.build_graph()
ask() print helper atlas.service.AtlasService.query() / query_stream() with run_id + QueryResult
Gradio cell atlas.ui rebuilt on the service (same streaming UX)
Report-card cell (pandas) evals/report_card.py → console table + machine-readable evals/results.json
— (not in notebook) FastAPI app (/v1/query, SSE stream, health/readiness), Dockerfile, CI
— (not in notebook) Deterministic fake LLM/embeddings for tests and CI

Layout

atlas-industrial/
├── data/corpus/{ops,science,crew,archive}/*.md   # 17 seed documents, verbatim
├── src/atlas/                  # the package (src layout, `pip install -e .`)
│   ├── config.py  logging.py  schemas.py  prompts.py  llm.py  chains.py
│   ├── state.py   nodes.py    graph.py    corpus.py   ingest.py  service.py
│   ├── stores/{base,memory,jsonstore,chroma}.py
│   ├── api/{models,app}.py
│   ├── cli.py                  # atlas ask|serve|ingest|eval|seed|graph
│   └── ui.py                   # Gradio chat UI on AtlasService
├── evals/{questions.json,report_card.py}
└── tests/                      # pytest, all runnable with zero API keys

Quickstart

Local venv

make setup                       # python3 -m venv .venv && pip install -e ".[dev]"
export OPENAI_API_KEY=sk-...     # or copy .env.example to .env

# Zero-key smoke run first (deterministic fake provider):
ATLAS_LLM_PROVIDER=fake .venv/bin/atlas seed
ATLAS_LLM_PROVIDER=fake .venv/bin/atlas ask \
  "Who leads the greenhouse experiment, and what did its most recent harvest produce?"

# Real run:
.venv/bin/atlas seed
.venv/bin/atlas ask "What actually happened on Sol 188?"
.venv/bin/atlas serve            # FastAPI on :8000
.venv/bin/atlas ui  (python -m atlas.ui)   # Gradio chat UI

Docker

docker compose up --build        # API on http://localhost:8000, chroma data on a volume

API

Endpoint Description
POST /v1/query Full agent run → structured QueryResult (answer, sources, trace, retries, run_id, latency)
GET /v1/query/stream?question=... Same run as Server-Sent Events: one trace event per node, then a final result event
GET /health Liveness probe
GET /ready Readiness probe (503 until every collection is non-empty)
GET /v1/collections Collection names + document counts
GET /v1/metrics Process counters: requests, errors, avg latency, tokens, est. cost
GET /graph Agent graph as Mermaid 
curl -s localhost:8000/v1/query -H 'content-type: application/json' \
  -d '{"question": "How long does a message to Earth take?"}' | jq .result.answer

curl -N "localhost:8000/v1/query/stream?question=What%20happened%20on%20Sol%20188%3F"
# event: trace
# data: {"run_id":"...","node":"router","lines":["[router]          route = ops -- ..."],...}
# ...
# event: result
# data: {"answer":"...","sources":["arc-02"],"retries":1,...}

Errors are structured: {"error": {"code": "...", "message": "...", "run_id": "..."}}.

Configuration

All settings are environment variables with the ATLAS_ prefix (see .env.example):

Variable Default Meaning
ATLAS_LLM_PROVIDER openai openai or fake (deterministic, no keys)
ATLAS_OPENAI_MODEL gpt-4o-mini Chat model
ATLAS_EMBEDDING_MODEL text-embedding-3-small Embedding model
ATLAS_TEMPERATURE 0.0 LLM temperature
ATLAS_MAX_RETRIES 2 Hard corrective-loop budget (enforced in code)
ATLAS_MIN_RELEVANT 1 Chunks that must pass grading before generation
ATLAS_RETRIEVE_K / ATLAS_ARCHIVE_K 2 / 3 Retrieval fan-out per index / archive
ATLAS_STORE_BACKEND chroma chroma (persistent), json (file), memory
ATLAS_PERSIST_DIR data/vectorstore Where chroma/json stores persist
ATLAS_CORPUS_DIR data/corpus Seed corpus location
ATLAS_CHUNK_SIZE / ATLAS_CHUNK_OVERLAP 1000 / 100 Ingestion chunking
ATLAS_USE_WEB_SEARCH false Tavily escalation in retrieve_archive (needs TAVILY_API_KEY)
ATLAS_LOG_FORMAT / ATLAS_LOG_LEVEL text / INFO json for structured logs
ATLAS_API_KEYS (empty) Comma-separated API keys; empty = auth disabled
ATLAS_RATE_LIMIT_RPM 0 Requests/min per key or IP; 0 = disabled
ATLAS_INPUT_PRICE_PER_1M / ATLAS_OUTPUT_PRICE_PER_1M 0.15 / 0.60 USD per 1M tokens (gpt-4o-mini defaults)
ATLAS_TRACE_DIR (unset) When set, appends run records to runs.jsonl
ATLAS_LANGSMITH_PROJECT (unset) Exported as LANGSMITH_PROJECT
ATLAS_STORE_DIR (unset) Alias overriding ATLAS_PERSIST_DIR
ATLAS_API_HOST / ATLAS_API_PORT 0.0.0.0 / 8000 Uvicorn bind
ATLAS_CORS_ORIGINS * Comma-separated CORS allowlist

Security

Both features are off by default — the zero-config dev flow stays open.

API-key authentication (src/atlas/api/auth.py): set ATLAS_API_KEYS to a comma-separated list and every /v1/* endpoint requires Authorization: Bearer <key> or X-API-Key: <key>. /health, /ready and /graph stay public. Comparison is constant-time; failures return a structured 401 and are logged without ever logging the key.

ATLAS_API_KEYS=alice-key,bob-key atlas serve
curl -s -o /dev/null -w '%{http_code}\n' localhost:8000/v1/metrics            # 401
curl -s localhost:8000/v1/metrics -H 'X-API-Key: alice-key'                   # 200

Rate limiting (src/atlas/api/ratelimit.py): set ATLAS_RATE_LIMIT_RPM to a positive integer to apply a sliding-window limit per API key (or per client IP when auth is off). Over-limit requests get a structured 429 with a Retry-After header. In-memory, stdlib-only, thread- and async-safe; state is per-process.

Observability

  • Structured logs: ATLAS_LOG_FORMAT=json emits one JSON object per record with run_id, node, and latency fields.
  • Per-run usage: every QueryResult carries a usage block (prompt / completion / total tokens, estimated_cost_usd) built from OpenAI usage_metadata by a LangChain callback handler. Prices are configurable via ATLAS_INPUT_PRICE_PER_1M / ATLAS_OUTPUT_PRICE_PER_1M. In fake mode (or when a model reports no usage) usage is null.
  • Process metrics: GET /v1/metrics returns cumulative counters — requests total, errors, average latency, tokens, estimated cost.
  • LangSmith (zero code): set LANGSMITH_TRACING=true and LANGSMITH_API_KEY and every LLM call is traced by LangChain natively. ATLAS_LANGSMITH_PROJECT is exported as LANGSMITH_PROJECT when set.
  • Built-in trace file (no external service): set ATLAS_TRACE_DIR and every run appends one JSON line to <trace_dir>/runs.jsonl — run_id, question (truncated), route, retries, per-node latencies, usage, timestamp.

Evals (RAGAS)

In addition to the built-in ISUSE report card (make eval), a RAGAS harness is available as an optional extra (kept out of the core dependencies because it is heavy):

pip install 'atlas-rag[ragas]'
export OPENAI_API_KEY=sk-...
make ragas          # python evals/ragas_eval.py

It runs the report-card questions through the service and scores faithfulness and answer relevancy, writing evals/ragas_results.json. Without the extra or a key it prints a clear message and exits 2 (CI-safe skip).

Ingesting your own documents

atlas ingest --dir /path/to/docs --collection ops

.md/.txt files are chunked with RecursiveCharacterTextSplitter and upserted with deterministic ids (<source>:<chunk>), so re-ingesting is idempotent. PDFs work with the optional extra: pip install ".[pdf]".

Evals

.venv/bin/python evals/report_card.py --fake   # deterministic smoke, no keys
.venv/bin/python evals/report_card.py          # real ISUSE judging (needs OPENAI_API_KEY)

Runs the six report-card questions from evals/questions.json through the full agent, scores each answer with the Self-RAG usefulness judge (ISUSE, 1–5), prints a console table and writes evals/results.json (mean usefulness, grounded count, per-question route/loops/abstention vs. expectation) — machine-readable for CI.

How fake mode works

ATLAS_LLM_PROVIDER=fake swaps in FakeChatModel + FakeEmbeddings (atlas/llm.py). The fake model implements with_structured_output and answers each chain by keyword rules on the rendered prompt: routing scores topic word-sets, the grader uses stemmed token overlap plus the incident-specificity rule from the prompt ("Sol 188" questions reject chunks that never mention 188), the generator quotes the best-scoring surviving chunk (its bracketed source id is the citation) and abstains verbatim when the context is empty, and the groundedness/usefulness judges apply matching heuristics. An optional script (AtlasService(settings, llm_script={"grounded": ["no", "yes"]})) forces specific verdicts in order, which is how tests exercise the reflection loop and the retry budget deterministically. FakeEmbeddings is an MD5-hashed bag-of-words embedder, so retrieval has real lexical behavior and is stable across runs and processes. One known approximation: the report-card question "What movies does the crew watch on their rest sols?" lexically matches the movie-nights policy chunk, so fake mode quotes it instead of abstaining — the abstention path is covered by the Wi-Fi question and the retry-budget tests.

Tests and CI

make test    # pytest — 36 tests, all fake-mode, zero API keys
make lint    # ruff check src tests evals

GitHub Actions (.github/workflows/ci.yml) runs the same ruff + pytest on every push. The Dockerfile builds a non-root python:3.12-slim image whose entrypoint is uvicorn atlas.api.app:create_app --factory.