Skip to content

CLAUDE.md

Latest commit

 

History

History
600 lines (465 loc) · 42.4 KB

CLAUDE.md

File metadata and controls

600 lines (465 loc) · 42.4 KB
Raw

CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Project Overview

OpenRag is a modular Retrieval-Augmented Generation (RAG) framework built with FastAPI, Ray for distributed computing, and Milvus as the vector database. It provides document ingestion, chunking, embedding, and retrieval capabilities with an OpenAI-compatible API.

Project Layout

openrag/        # Python package (application code only): core/ services/ api/ di/ prompts/
conf/           # YAML configuration
infra/          # All deployment infrastructure
  docker/       #   api.Dockerfile, ray.Dockerfile, ui.Dockerfile (build from repo root)
  compose/      #   docker-compose.yaml + service configs (grafana, prometheus, milvus, .env.example)
  scripts/      #   entrypoint.sh and other deployment scripts
  ansible/      #   Ansible playbooks
  charts/       #   Helm charts (openrag-stack)
  cluster.yaml  #   Ray cluster config
scripts/        # Developer/operational CLI tools (check_layer_imports.py, data_indexer.py, postgres-init/)
tests/          # Integration tests (api_tests/, integration/)
docs/           # Documentation (Astro site + refactoring docs)
ui/             # Admin frontend (React SPA), served by the admin-ui container (infra/docker/ui.Dockerfile)
extern/         # Git submodules + compose service includes

Prompt templates ship inside the package at openrag/prompts/templates/*.txt and are loaded into DEFAULT_SEEDS by openrag/prompts/__init__.py.

Common Commands

Development

# Install dependencies
uv sync

# Run the application locally (requires Docker services).
# The compose stack and its service configs live under infra/compose/.
cd infra/compose
docker compose up -d                 # GPU deployment
docker compose --profile cpu up -d   # CPU deployment

# Run with rebuild for development
docker compose up --build -d

Testing

# Run all unit tests (fast, no infra needed)
uv run pytest tests/unit/

# Run a single test file
uv run pytest tests/unit/core/models/test_chunk.py

# Run tests matching a pattern
uv run pytest -k "test_chunk"

# Integration tests (need running services) / load tests
uv run pytest tests/integration/
uv run pytest tests/load/

Linting

uv run ruff check openrag/ tests/
uv run ruff format openrag/ tests/

Documentation Site

npm i
npm run dev  # Start dev server at http://localhost:4321/openrag

Architecture

Core Components

The main application entry point is openrag/api/main.py which creates a FastAPI app with Ray initialization.

Ray Actors (distributed components):

  • Indexer (openrag/services/workers/indexer_pool.py) - Handles document ingestion, chunking, and insertion into vector DB
  • TaskStateManager (openrag/services/workers/task_state.py) - Tracks async task states: QUEUED → SERIALIZING → COMPLETED (or FAILED or CANCELLED)
  • Vectordb / MilvusDB (openrag/services/storage/milvus_store.py) - Vector database operations with hybrid search (dense + BM25 sparse)
  • DocSerializer (openrag/services/workers/parsers/doc_serializer.py) - Serializes files to Document objects using appropriate loaders
  • MarkerPool / MarkerWorker (openrag/services/workers/parsers/marker_workers.py) - Pool of workers for PDF processing with Marker

Pipeline Classes:

  • RagPipeline (openrag/services/orchestrators/query_service.py) - Orchestrates retrieval and LLM generation
  • RetrieverPipeline (openrag/core/retrieval/pipeline.py) - Handles document retrieval and reranking
  • RAGMapReduce (openrag/services/orchestrators/query_service.py) - Map-reduce for processing large document sets

Document Processing Flow

  1. Files uploaded via POST /indexer/partition/{partition}/file/{file_id} (multipart with file=@…)
  2. Indexer.add_file() serializes file to Document using appropriate loader
  3. Chunker splits document into chunks with contextual metadata
  4. Embedder generates vectors via VLLM (OpenAI-compatible API)
  5. Chunks inserted into Milvus with partition-based organization

File Loaders (openrag/services/workers/parsers/legacy_loaders/)

Each file type has a dedicated loader that converts to markdown:

  • MarkerLoader (default for PDF, in pdf_loaders/marker.py) - Supports OCR, complex layouts, tables
  • DocxLoader, PPTXLoader, DocLoader - Office formats (uses MarkItDown library)
  • ImageLoader - VLM-powered image captioning
  • VideoAudioLoader - Audio transcription via Whisper
  • MarkdownLoader, TextLoader (txt_loader.py) - Markdown and plain text files

Loader base class: All loaders inherit from BaseLoader (base.py) which provides:

  • self.image_captioning - whether image captioning is enabled (use this, not self.config.loader["image_captioning"])
  • self.config - Hydra config access
  • get_image_description(image_data) - Low-level VLM captioning (accepts PIL Image, HTTP URL, or data URI)
  • caption_images(images, desc) - Caption a list of PIL images concurrently with progress bar
  • replace_markdown_images_with_captions(content, ...) - Find and replace markdown image references with captions
  • Class regex patterns: HTTP_IMAGE_PATTERN, DATA_URI_IMAGE_PATTERN

Loader image captioning pattern: Loaders that process images must check self.image_captioning before captioning. Use the shared methods above rather than duplicating captioning logic. Access additional loader config via self.config.loader.get("option_name", default).

Image handling approaches:

  • PDF/DOCX/PPTX: Extract binary image data from file, pass to VLM directly
  • Markdown: Parse image URLs from text; HTTP URLs require IMAGE_CAPTIONING_URL=true

Source Citation Filtering

The RAG pipeline filters out false-positive sources by having the LLM self-report which sources it actually used:

  1. format_context() (openrag/core/prompts/chat_prompt_builder.py) numbers each source ([Source 1], [Source 2], ...) in the context and returns (formatted_text, included_indices) — the indices track which docs fit within the token budget
  2. Prompt templates (openrag/prompts/templates/*.txt) instruct the LLM to append [Sources: 1, 3, 5] at the end of its response
  3. extract_and_strip_sources_block() (openrag/core/utils/source_filtering.py) strips this tag from the response before sending to the client
  4. filter_sources_by_citations() (openrag/core/utils/source_filtering.py) filters the source metadata to only include cited sources; if no [Sources: ...] tag is found at all, every presented source is kept instead (a missing tag means the model didn't report citations, not that it used none)
  5. For streaming, the OpenAI router buffers the last 100 chars to catch the sources tag before it reaches the client

The extra field in API responses is a JSON object with these keys. It was a JSON-encoded string up to and including v2.2.2 — a breaking change for readers written against the old shape, which must stop calling json.loads on it:

  • sources — legacy field, kept as-is for existing clients (e.g. Twake): cited sources, or every presented source as a fallback when the model reported no citations — neither a [Sources: ...] tag nor an inline [Source N] marker — or the request skips citation reporting (structured output, direct LLM, web search with no results).
  • presented_sources — every source actually shown to the LLM (after format_context()/format_web_context() truncation), regardless of citation. Always present; a client can fall back to this ("sources consulted") when nothing was cited.
  • cited_sources — strictly what the model cited, in the tag or inline [Source N] markers; unlike sources, this never falls back to "everything" — it's [] whenever nothing was cited or citations weren't reported. Chainlit uses this field directly so its source panel never presents uncited retrieval candidates.
  • citations_reported (bool) — true when the model reported citations: a [Sources: ...] tag with sources or none, or inline [Source N] markers (an empty [Sources: ] counts as no tag); false otherwise, which is exactly when sources falls back to keeping everything. Lets a client tell "the model cited every source" apart from "the model didn't report citations at all".
  • all_retrieved_sources — the complete retrieval set, captured before the context-token-budget truncation, so it also includes documents/web results that didn't fit in the prompt (and, on the map-reduce path, the original retrieved docs rather than the LLM-generated summaries). Only included when the request sets metadata.include_all_retrieved_sources: true — it's debug/eval telemetry, gated off by default since retrieval is uncapped up to retriever.top_k while the context budget only fits a handful of documents.

Each document source entry (build_document_source_link) is shaped:

{
  "source_type": "document",
  "chunk": { "...the chunk's metadata, copied verbatim..." },
  "rerank_score": 0.646,
  "chunk_url": "https://host/extract/<chunk id>",
  "file_url":  "https://host/static/<chunk id>"
}

chunk holds what the chunk carries; the siblings are what the server computed about it. Two consequences worth knowing: source_type, chunk_url and file_url are authoritative and are scrubbed from chunk as well as overridden at the top level, so metadata can't spoof them in either place (the guard is structural now, not a manual overwrite); and file_url is still omitted entirely when the chunk has no source.

Web entries (source_type: "web") are unchanged and flat — url, title, snippet, no chunk. Clients switch on source_type, as before.

rerank_score — the raw score the reranker gave that chunk — sits beside chunk, not inside it: it describes how this query ranked the chunk, and the same chunk retrieved by a different query scores differently.

It gets there via ScoredChunk (openrag/core/models/retrieval_result.py), a Chunk subclass holding vector_score / rerank_score / combined_score as typed fields. _rerank_chunks (openrag/core/retrieval/pipeline.py) returns ScoredChunk.from_chunk(...) instead of the bare chunk, and ScoredChunk.to_langchain() folds the non-null scores into metadata at the boundary the API response is built from, and build_document_source_link lifts them back out to sit beside chunk. Because it subclasses Chunk, every list[Chunk] signature through retrieval, expansion and RRF stays valid. Only rerank_score is populated today — the vector score is still dropped in vector_store_searcher._dict_to_chunk (score is in its skip set), and nothing computes a combined score.

The three key names are shared as RETRIEVAL_SCORE_KEYS (core/utils/consts.py) because promoting a metadata key to an authoritative sibling is only safe if nothing else can put it there. Milvus collections have a dynamic field, so upload metadata is persisted verbatim — {"rerank_score": 0.99} would come back on every read. _dict_to_chunk therefore drops those keys along with score, and ScoredChunk.to_langchain() clears them from inherited metadata before stamping its typed fields. A score in an API response was set by that retrieval, never by whoever uploaded the file.

Three caveats:

  • The key is absent, not null, when no reranker ran (reranker disabled, or a web source — web results are built separately and never reranked). Such chunks stay plain Chunks and carry no score field at all.
  • Its scale is provider-dependent: Infinity/vLLM return a relevance_score, TEI a score with raw_scores: false (0–1). Compare within one response, never across deployments, and don't threshold on an absolute value.
  • On the multi-query path it does not explain the ordering. get_relevant_docs reranks each sub-query's list separately and then fuses them with RRF, so the final order is the RRF rank; a chunk retrieved by several sub-queries keeps the score from whichever list RRF saw first.

Prometheus Metrics

GET /metrics (openrag/api/routers/admin/monitoring.py) serves the default prometheus_client registry: HTTP counters/histogram recorded by api/middleware/instrumentation.py, plus the API-process side of the Tier-1 metrics (openrag_ingest_tasks, inference, tokens, circuit breakers). Anything recorded inside a Ray actor goes through ray.util.metrics instead and is exported by Ray's metrics agent, not /metrics; specs live in core/observability/metric_specs.py, and tests/unit/core/observability/test_metric_emission.py fails the build on a spec nothing writes. Full list and query rules: docs/content/docs/documentation/metrics_reference.md. The path is in DEFAULT_BYPASS_PATHS (no user token needed) and the route enforces its own METRICS_TOKEN (server.metrics_token, blank = unset) via require_metrics_token; admin tokens are deliberately not accepted there — one mechanism, no fallback. It fails closed: token unset and METRICS_ALLOW_UNAUTHENTICATED (server.metrics_allow_unauthenticated) false → 403 on every scrape, with a startup warning from describe_metrics_access. The opt-in exists because the API port is exactly what the Ingress / admin-ui proxy forwards (review on PR #914), so "no token" must never silently mean "open"; a configured token always wins over the opt-in. The admin UI's System > Metrics tab reads GET /monitoring/metrics (admin_router, require_admin, an API prefix) instead — same exposition, separate audience, so the scrape path never touches the Postgres token lookup and an admin never holds the scrape secret. The config is read through load_config() rather than the request container so a scrape keeps working while the container is degraded. Compose: the monitoring overlay writes METRICS_TOKEN into the Prometheus container via a configs.content entry (Compose ≥ 2.23.1) and fails fast without it; the admin-ui nginx returns 404 on /metrics. Helm: openrag.metrics.* (pod annotations + optional ServiceMonitor with bearerTokenFromSecret), env.secrets.METRICS_TOKEN, env.config.METRICS_ALLOW_UNAUTHENTICATED. Docs: docs/content/docs/documentation/prometheus_metrics.md.

Grafana dashboards live once, in infra/charts/openrag-stack/dashboards/ (Helm's .Files.Glob reads only inside the chart): the Compose overlay mounts that directory, and the chart renders one sidecar ConfigMap per file (monitoring.dashboards.*, templates/grafana-dashboards.yaml). monitoring.bundled (default off) installs kube-prometheus-stack as a sub-chart aliased kubePrometheusStack — vllm-stack carries its own, and two same-named sub-charts make Helm coalesce one's defaults into the other — and implies the dashboards, the alert rules (OpenRagTargetDown then leaves the stack's own openrag-monitoring-*/openrag-grafana jobs out) plus the API ServiceMonitor with the bearer, so it refuses to render without METRICS_TOKEN. It also opens the admission-webhook port in the default-deny NetworkPolicy (networkPolicy.webhookFrom, empty = any source): the caller is the API server, which no selector can name. That port is on the webhook's own pod (kubePrometheusStack.prometheusOperator.admissionWebhooks.deployment.enabled: true), not the operator's, whose listener also serves /debug/pprof/. The bundled Grafana's sidecar label/labelValue/folderAnnotation can't be derived from monitoring.dashboards.*, so grafana-dashboards.yaml fails the render when they differ, and when monitoring.dashboards.labels restates the watched label (it would render the key twice, last value wins). Every bundled workload, the certificate hook Jobs included, has requests and a memory limit, no CPU limit (upstream sets none: BestEffort, and a ResourceQuota requiring requests refuses the hook Jobs, failing the install); the numbers come from a kind run, and the Prometheus comment in values.yaml carries the series-count model behind its 1Gi/3Gi. The control-plane targets (kubeControllerManager, kubeScheduler, kubeEtcd, kubeProxy) are off: on a managed control plane their absent(up) alerts fire forever. Its cluster-scoped objects have fixed names, so it is one bundled release per cluster. Two tests render the pinned kube-prometheus-stack archive itself (skipped where it isn't fetched, as in unit CI), so an upstream key renamed by a version bump fails them instead of silently dropping a setting. Under Argo CD the webhook's CA is written by a PostSync hook, which waits for the whole Application to be Healthy, so validation is skipped (failurePolicy: Ignore) until then; admissionWebhooks.certManager.enabled avoids the hooks. Tests: tests/unit/infra/test_monitoring_delivery.py; docs: kubernetes.md, Monitoring.

API Routers (openrag/api/routers/)

  • user/chat.py - OpenAI-compatible /v1/chat/completions endpoint
  • admin/indexing.py - Document ingestion endpoints
  • user/search.py - Semantic search endpoints
  • admin/partitions.py - Partition management (multi-tenant document collections)
  • admin/users.py - User and membership management
  • admin/jobs.py - Task queue monitoring
  • admin/workspaces.py - Workspace CRUD and file management
  • admin/tools.py - Tools like extractText at /v1/tools/execute (tool param requires JSON: {"name": "extractText"})

User Management & Authentication

The system uses token-based authentication with role-based access control (RBAC) for multi-tenant partition access.

Database Schema (PostgreSQL with SQLAlchemy, in openrag/services/persistence/schema.py):

  • users - User accounts with id, external_user_id, display_name, token (SHA-256 hashed), is_admin, file_quota, file_count, managed_by_config (provisioned from auth.seed_users)
  • files - File records with file_id, partition_name, file_metadata, created_by (FK to users), relationship_id, parent_id
  • partition_memberships - Join table linking users to partitions with roles (owner, editor, viewer)
  • partitions - Document collections with cascade delete to files and memberships
  • workspaces - Named file subsets within a partition for scoped search/chat. workspace_id is unique per partition only ((partition_name, workspace_id)), so every lookup is keyed on both; a workspace-scoped multi-partition search whose id matches several searchable partitions fails with AmbiguousWorkspaceError (422) instead of picking one
  • workspace_files - Join table linking workspaces to files by their integer PKs (workspaces.id, files.id), never by the per-partition string ids

Authentication Flow (AuthMiddleware from openrag/api/middleware/auth.py, registered in openrag/api/main.py):

  1. Token extracted from Authorization: Bearer <token> header (or ?token= query param for /static routes)
  2. Token hashed with SHA-256, looked up in database
  3. User info and accessible partitions set on request.state.user and request.state.user_partitions
  4. Bypassed for: /docs, /openapi.json, /redoc, /health_check, /version, /chainlit/* — except in OIDC mode the three docs paths (/docs, /redoc, /openapi.json) are login-gated instead of public (see Middleware Behavior below)
  5. If AUTH_TOKEN env var is not set, defaults to admin user (id=1) for all requests

Role Hierarchy (openrag/services/orchestrators/auth_service.py):

ROLE_HIERARCHY = {"viewer": 1, "editor": 2, "owner": 3}

Permission Dependencies (openrag/api/dependencies/auth.py):

  • require_admin - User must have is_admin=True
  • require_partition_viewer / require_partition_editor / require_partition_owner - Check partition membership role
  • SUPER_ADMIN_MODE=true env var allows admin users (is_admin=True) to bypass partition checks; regular users remain restricted to their partition memberships

User API Endpoints (/users/):

Endpoint Method Auth Description
/users/ GET Admin List all users
/users/info GET Any Get current user info
/users/ POST Admin Create user (returns token once)
/users/{user_id} DELETE Admin Delete user (cannot delete id=1)
/users/{user_id}/regenerate_token POST Admin/self Regenerate API token
/users/{user_id}/quota PATCH Admin Update user file quota

Partition Membership Endpoints (/partition/{partition}/users):

Endpoint Method Auth Description
/partition/{partition}/users GET Owner List partition members
/partition/{partition}/users POST Owner Add user with role
/partition/{partition}/users/{user_id} DELETE Owner Remove user
/partition/{partition}/users/{user_id} PATCH Owner Update user role

Core Implementation (PartitionFileManager in openrag/services/persistence/partition_repo.py):

# User operations (called via MilvusDB Ray actor)
await vectordb.create_user.remote(display_name="Name", is_admin=False)
await vectordb.get_user_by_token.remote(token)
await vectordb.regenerate_user_token.remote(user_id)

# Membership operations
await vectordb.add_partition_member.remote(partition, user_id, role="editor")
await vectordb.update_partition_member_role.remote(partition, user_id, "owner")
await vectordb.list_partition_members.remote(partition)

Token Format: "or-" + secrets.token_hex(16) (34-char string, shown only once on creation/regeneration)

Bootstrap: On startup, ensures admin user (id=1) exists using AUTH_TOKEN env var or generates a random token.

Seed users (#1153): auth.seed_users (or the SEED_USERS env var, rendered by the chart from openrag.seedUsers) declares operator-managed accounts, provisioned by PgUserRepository.ensure_seed_users at API startup only, with each token read from the env var the entry names. Seeding only upserts rows it created (users.managed_by_config) and listed memberships; it never deletes, and a removed entry just has its token revoked. Full semantics: docs/content/docs/documentation/user_auth.md.

Multi-Partition Search: Users can search across all their accessible partitions:

  • Search endpoint: GET /search?partitions=all&text=query
  • Chat completions: POST /v1/chat/completions with "model": "openrag-all"
  • For regular users, all resolves to their partition memberships only
  • For admins with SUPER_ADMIN_MODE=true, all resolves to all system partitions
  • Model prefix is openrag- (legacy: ragondin-)

Web Search Integration

Optional web search augmentation via the Staan API, allowing the LLM to combine RAG document context with live web results.

Configuration (conf/config.yaml → websearch: block, env vars):

  • WEBSEARCH_API_TOKEN — provider API token; if unset, web search is silently disabled
  • WEBSEARCH_BASE_URL — provider endpoint (default: Staan API)
  • WEBSEARCH_TOP_K — number of web results (default: 5)
  • WEBSEARCH_LANG — search language/market (default: fr-FR)

How it works:

  • Client sends metadata: {"websearch": true} in the chat completion request
  • Combined mode (partition + websearch): RAG retrieval and web search run concurrently via asyncio.gather(); web results are appended after document sources with continuous [Source N] numbering
  • Web-only mode (no partition + websearch): skips RAG retrieval entirely, uses web results as sole context; if no results (token unset / search fails), falls back to plain direct LLM mode
  • Source entries include source_type: "document" or source_type: "web" in the extra.sources response

Key files:

  • openrag/services/websearch/ — WebSearchService (service.py), BaseWebSearchProvider (base.py), StaanProvider (providers/staan.py)
  • openrag/core/prompts/chat_prompt_builder.py — format_web_context() formats web results as numbered source blocks
  • openrag/services/orchestrators/query_service.py — _prepare_for_web_only(), web search logic in _prepare_for_chat_completion()
  • openrag/api/routers/user/chat.py — __prepare_sources() merges document and web sources

Indexing Status Callbacks

Indexing is asynchronous, so a client can either poll the task-status URL or hand OpenRag a URL to notify once the task settles.

Request fields (multipart form, on POST and PUT /indexer/partition/{partition}/file/{file_id}):

  • callback_url — POSTed once the task reaches a terminal state
  • callback_token — sent as Authorization: Bearer <token> on that POST

Body: {"partition", "file_id", "status": "success"|"error", "metadata"}. metadata is the upload metadata echoed back minus UPLOAD_METADATA_SERVER_KEYS (core/utils/consts.py: source, filename, original_filename, file_size, file_id, content_sha256) — an exclusion, not a fixed field list, so any caller-supplied field (cozy-stack's revision marker is doc_rev) travels through unrecomputed and under whatever name the caller gave it; a key the caller never sent is simply absent, not echoed back as null. The exclusion exists because the target is a caller-supplied URL and _build_metadata merges those server-computed keys into the same dict — source (the server's on-disk path) is the load-bearing one. test_build_metadata_only_adds_keys_in_upload_metadata_server_keys (tests/unit/services/orchestrators/test_indexing_service.py) fails the build if _build_metadata ever injects a key the constant doesn't cover.

Guarantees: one attempt, no retries, a 5 s deadline on the request, never raises — a failed callback is logged and cannot change the indexing outcome. No callback_url → strict no-op. A user-cancelled task sends nothing (only a real failure notifies "error"), which is why the sender keys off the return value of set_failed_if_not_cancelled and the pool's pre-flight handler catches Exception, not BaseException.

Not guaranteed: delivery. The send is awaited on the worker's slot, so a blackholing target costs up to 5 s of indexing throughput per file, and an actor lost before the task settles notifies nothing at all. Clients keep a timeout and fall back to polling the task-status URL.

callback_url is checked against the SSRF guard (is_safe_url) in the router and again in the sender. The literal check covers the scheme, loopback/private/link-local and reserved addresses, alternate numeric IPv4 spellings, IDNA-equivalent separators, and a single DNS root dot. DNS hostnames are not resolved by this callback check; use network egress controls when hostname destinations must be restricted. HTTP and HTTPS are both allowed. INDEXING_CALLBACK_ALLOW_PRIVATE_URLS=true / indexing_callback.allow_private_urls lifts the address half of the guard for dev stacks whose target is a local instance; the scheme check always applies. Keep it off in production: with it on, any user allowed to upload can make the server POST to an internal address.

Rolling deploys: the worker actors are named, detached and get_if_exists, so changing process_file's remote contract requires bumping _INDEXER_ACTOR_PROTOCOL_VERSION in indexer_pool.py (v3 → v4 for callback_url/callback_token; v5 added worker-ref-registration wait and TSM set_state fencing; v7 folds in a second, independent v6 lineage — STT-preset-aware registry hydration plus the _active_indexation_config contextvar — that landed on develop under the same version string while this branch's own v6 was in flight; v8 (develop) added max_restarts on the dispatcher and workers, since Ray only applies actor options when it creates the actor and get_if_exists=True would otherwise silently keep the previous release's restart policy (#846); v8 (this branch, independently) covers TaskStateManager bounding its in-memory retention and being replaced during bootstrap when an older actor lacks that support — that replacement changes the TaskStateManager actor id, which strands the dispatcher's and workers' cached handles to it unless the whole generation rolls together; v9 folds in both independent v8 lineages). Without the bump, new replicas attach to the previous release's actors and every submit raises TypeError. Old generations are retired with services/workers/retire_indexer_generation.py.

Key files:

  • openrag/services/workers/indexing_callback.py — send_indexing_callback() (was webhook.py; the target is a normal authenticated route now, not an unauthenticated webhook trigger)
  • openrag/core/utils/url_safety.py — is_safe_url(url, *, allow_private_hosts=False), shared with the MCP index_url tool (which keeps the strict default). The web-search content fetcher (services/websearch/content_fetcher.py) does not import this — it has its own older, un-synced _is_safe_url, so hardening this module does not automatically harden that one.
  • openrag/core/config/indexation.py — IndexingCallbackConfig

Both fields travel the same chain as the rest of an indexing job: api/routers/admin/indexing.py → services/orchestrators/indexing_service.py → core/indexing/dispatcher.py (port) → services/workers/dispatcher.py → indexer_pool.py → indexer_actor.py → indexing_callback.py.

File Quota System

Per-user file quota enforcement tracked via the file_count and file_quota columns on users, and created_by on files.

How it works:

  • files.created_by records which user uploaded each file (nullable for pre-migration files)
  • users.file_count is incremented/decremented in application code (in PartitionFileManager) — no SQL triggers
  • Decrements use func.greatest(file_count - N, 0) to prevent negative values from race conditions
  • delete_partition queries per-uploader counts before cascade delete, then bulk decrements
  • Quota check (check_user_file_quota in openrag/api/dependencies/auth.py) runs on upload, considering both indexed files and pending tasks

Quota logic (file_quota column):

  • None → use global default (DEFAULT_FILE_QUOTA env var, default -1)
  • < 0 → unlimited
  • >= 0 → specific limit
  • Admins always bypass quota checks

Key design decisions:

  • Counts are tracked per uploader (whoever calls the upload API), not per partition owner
  • created_by uses ondelete="SET NULL" so deleting a user doesn't cascade-delete their files
  • Indexer.delete_file and MilvusDB.delete_file/delete_partition don't need a user_id parameter — the uploader is looked up from files.created_by

Migration: openrag/services/persistence/migrations/alembic/versions/c224d4befe71_add_file_count_and_file_quota.py

Alembic Migration Idempotency

Base.metadata.create_all() runs at app startup (PartitionFileManager.__init__ in openrag/services/persistence/partition_repo.py), so a freshly bootstrapped database already contains the full current-model schema before alembic ever touches it. Migrations must therefore be idempotent — re-applying an ADD COLUMN / CREATE TABLE / CREATE INDEX against an already-existing object would raise DuplicateColumn / DuplicateTable.

Guard every schema-mutating op with an inspector-based existence check (table_exists, column_exists, index_exists, fk_exists), in both upgrade() and downgrade(). For migrations that convert a column type, also short-circuit if the column is already the target type.

Configuration

Configuration is a single YAML file validated with Pydantic models:

  • Main config: conf/config.yaml
  • Loaded by openrag/core/config/loader.py (load_config() exposed from openrag/core/config/__init__.py)
  • Pydantic config classes live in openrag/core/config/ (root.py, auth.py, chunking.py, retrieval.py, indexation.py, endpoints.py, mcp.py, infrastructure.py, base.py)

Environment variables override config values (see infra/compose/.env.example).

Testing Structure

All tests live in a separate tests/ tree (zero test files inside the openrag/ package):

  • Unit tests: tests/unit/**/test_*.py (pytest, mirrors the package structure; no external services needed)
  • Integration tests: tests/integration/api/*.py (HTTP endpoint tests, requires running server) and tests/integration/repos/*.py (repo/store tests)
  • Robot Framework tests: tests/integration/robot/api/*.robot
  • Load/benchmark tests: tests/load/
  • Shared fixtures: tests/unit/conftest.py (mock ports), tests/unit/api/conftest.py (ASGI client), plus per-suite conftests
  • Test config lives in pyproject.toml ([tool.pytest.ini_options]): testpaths = ["tests"], pythonpath = ["./openrag"], and the env block sets PROMPTS_DIR=./openrag/prompts/templates

Running integration tests locally with act:

# Run API tests using GitHub Actions locally
act -j api-tests -W .github/workflows/api_tests.yml --bind

Mock VLLM for CI: tests/api_tests/api_run/mock_vllm.py provides fake embeddings and completions endpoints (streaming and non-streaming) for testing without a real LLM. Pydantic request models use ConfigDict(extra="allow") to accept vendor-specific fields like extra_body.

Key Patterns

Ray Actor Access

# Get actor references
vectordb = ray.get_actor("Vectordb", namespace="openrag")
indexer = ray.get_actor("Indexer", namespace="openrag")
task_state_manager = ray.get_actor("TaskStateManager", namespace="openrag")

# Call remote methods
await vectordb.async_search.remote(query=query, partition=partition)

Ray Actor Timeout and Cancellation

Use the centralized utility for calling Ray actors with proper timeout and cancellation handling:

from services.workers.ray_utils import call_ray_actor_with_timeout

result = await call_ray_actor_with_timeout(
    future=actor.method.remote(args),
    timeout=TIMEOUT_SECONDS,
    task_description="Description for error messages",
)

This handles:

  • Timeout with ray.wait() and ray.cancel()
  • asyncio.CancelledError propagation
  • RayTaskError and TaskCancelledError handling

Custom Exceptions

All custom exceptions inherit from OpenRAGError (openrag/core/utils/exceptions.py):

  • VDBError subclasses for vector database errors
  • EmbeddingError for embedding failures

Logging

Uses Loguru with structured logging:

from core.utils.logging import get_logger
logger = get_logger()
logger.bind(file_id=file_id, partition=partition).info("Message")

stderr is the only sink (core/utils/logging.py); there is no log file. LOG_FORMAT=text (default) is the colorized terminal format, LOG_FORMAT=json writes one flat JSON object per line (json_record/json_sink: ts, level, logger, function, line, msg, exception, then every bound extra at the top level, collisions prefixed extra_) and routes stdlib logging (uvicorn, Ray) through loguru via InterceptHandler (idempotent intercept_stdlib_logging(level): one interceptor on the root, foreign handlers detached not closed, root at loguru's numeric level so libraries' isEnabledFor(DEBUG) guards hold, chatty loggers capped at WARNING only while still at NOTSET). RequestIdMiddleware binds request_id into the loguru context for the whole request; the unhandled-500 handler binds it explicitly because it runs after that scope unwinds. Ray relays worker output onto the API stream with a (Actor pid=N) prefix that the collector strips; RAY_DEDUP_LOGS=0 and RAY_COLOR_PREFIX=0 are required for that relay to carry valid JSON (overlay and chart set them), and with ray.enabled=true the chart sets RAY_LOG_TO_STDERR=1 so workers write to their own pod's stderr (the relay only carries the current driver job's actors). Shipping: compose infra/compose/logging.docker-compose.yaml (Alloy, infra/compose/alloy/config.alloy), Helm sets LOG_FORMAT=json and relies on the platform DaemonSet. Docs: docs/content/docs/documentation/loki_logs.md.

Import Conventions

Use absolute imports from the openrag/ directory (which is the Python path root):

# Correct - absolute imports
from services.workers.ray_utils import call_ray_actor_with_timeout
from core.utils.logging import get_logger
from core.config import load_config

# Avoid relative imports across packages
# from .ray_utils import ...  # Only within same package

OIDC Authentication (OpenID Connect)

OpenRag supports two authentication modes, controlled by the AUTH_MODE environment variable:

Token Mode (AUTH_MODE=token, default):

  • Bearer token authentication via Authorization: Bearer <AUTH_TOKEN> header
  • Existing behavior unchanged
  • Suitable for programmatic access, CI/CD, and testing
  • Admin user (id=1) created with AUTH_TOKEN env var or random token on bootstrap

OIDC Mode (AUTH_MODE=oidc):

  • OpenID Connect Authorization Code + PKCE flow
  • Users authenticate via an external IdP (Keycloak, LemonLDAP::NG, etc.)
  • Browser UI (Chainlit, Indexer) redirects to IdP login
  • Opaque session tokens stored in openrag_session httpOnly cookie
  • Bearer users.token still accepted for programmatic access

Env Variables (required when AUTH_MODE=oidc):

Variable Purpose Example
OIDC_ENDPOINT Issuer URL for auto-discovery https://idp.example.com/realms/openrag
OIDC_CLIENT_ID Client registered at IdP openrag
OIDC_CLIENT_SECRET Client secret (provided by IdP)
OIDC_REDIRECT_URI Callback URL — the front door that serves the UI and reaches the backend /auth/callback (the admin-ui / proxy port, not necessarily APP_PORT); must match IdP config https://openrag.example.com/auth/callback
OIDC_TOKEN_ENCRYPTION_KEY Fernet key for token encryption (generate via: python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())")

Optional Env Variables:

Variable Default Purpose
OIDC_CLAIM_SOURCE id_token Where to read claims for claim mapping: id_token (verified JWT) or userinfo (/userinfo endpoint)
OIDC_CLAIM_MAPPING (none) CSV of db_field:claim pairs to sync IdP claims into the users row on every login (whitelist: display_name, email). Unset = no post-login update.
OIDC_SCOPES openid email profile offline_access Space-separated scope list (include offline_access for refresh tokens)
OIDC_POST_LOGOUT_REDIRECT_URI — URL the IdP sends the user to after RP-initiated logout. No default (an OpenRag URL would re-trigger OIDC login)
OIDC_AUTO_PROVISION_LOGIN false When true, an unknown sub triggers on-the-fly creation of a non-admin user from the ID-token claims (name/preferred_username → display_name, email → email). Default keeps the strict admin-pre-provisioning policy below.

User Matching & Provisioning:

When a user logs in via OIDC, matching is exclusively by users.external_user_id == sub (the stable OIDC claim). There is no email fallback. If the sub is unknown, the callback either:

  • returns 403 "User not registered" (default — admins must pre-create every user), or
  • creates a non-admin user from the ID-token claims when OIDC_AUTO_PROVISION_LOGIN=true. Auto-provisioned users inherit the default file quota; is_admin is always false (operators can promote afterwards via /users/{id} or /users/).

Optionally, if OIDC_CLAIM_MAPPING is set, after a successful match the callback reads the configured claims (from the ID token or /userinfo, per OIDC_CLAIM_SOURCE) and updates the user row. The writable whitelist is strict — only display_name and email are allowed; is_admin, external_user_id, file_quota, token are never writable via claim mapping.

Admin Pre-provisioning: Admins create users with the external_user_id matching the IdP's sub claim for that user. Example:

curl -X POST http://localhost:8080/users/ \
  -H "Authorization: Bearer <AUTH_TOKEN>" \
  -H "Content-Type: application/json" \
  -d '{"display_name": "Alice", "external_user_id": "kc-alice-uuid", "is_admin": false}'

Database Schema:

Columns on users table relevant to OIDC:

  • external_user_id (String, unique, nullable): Must equal the IdP's sub for OIDC matching
  • email (String, unique, nullable): Pure metadata; populated manually or via claim mapping. Not used for matching.

New table oidc_sessions:

  • session_token_hash (unique): SHA-256 of the opaque session token
  • user_id (FK): User this session belongs to
  • sid (nullable): OIDC session identifier (used for back-channel logout)
  • sub (required): OIDC sub claim (stable user identifier)
  • id_token_encrypted, access_token_encrypted, refresh_token_encrypted: Fernet-encrypted IdP tokens
  • access_token_expires_at, session_expires_at: Token expiry times
  • revoked_at (nullable): Set on back-channel logout or manual revocation

Auth Endpoints (all bypass the normal middleware):

Method Path Purpose
GET /auth/login Start Authorization Code + PKCE flow; redirects to IdP
GET /auth/callback IdP callback; creates session, sets cookie; redirects to next_url
POST /auth/backchannel-logout IdP-driven logout (OIDC spec); revokes sessions by sid
GET /auth/logout RP-initiated logout; invalidates session + redirects to IdP
GET /auth/me (debug) Returns current user and session expiry

Session Management:

  • Session token: URL-safe opaque token (secrets.token_urlsafe(32) — ~43 chars from 32 bytes of randomness), hashed (SHA-256) before storage
  • Cookie: openrag_session (httpOnly, Secure if HTTPS, SameSite=Lax, Path=/, no Domain=)
  • TTL: Aligned with access_token_expires_at; auto-refresh if refresh_token available (<60s before expiry)
  • Revocation: Via back-channel logout or manual invalidation

Middleware Behavior:

  • UI paths (/, /chainlit, /static) in OIDC mode without auth → 302 redirect to /auth/login?next=...
  • Interactive docs (/docs, /redoc, /openapi.json): public in token mode (bypassed), but login-gated in OIDC mode — an unauthenticated browser is 302-redirected to /auth/login; a valid session renders them. This stops the full API surface + schema from being served anonymously in production. The set is AuthBypassConfig.oidc_gated_paths (default ("/docs", "/redoc", "/openapi.json")); override it to () to keep docs public under OIDC.
  • API paths (/v1, /indexer, /search, etc.) without auth:
    • Token mode → 403 {"detail": "Missing token"} (no bearer) or 403 {"detail": "Invalid token"} (unknown bearer). The 403 status is a legacy contract the robot suite asserts (tests/api/).
    • OIDC mode → 401 {"detail": "Unauthenticated"} (no usable session/bearer and the path isn't a UI redirect target).
  • Programmatic access: Bearer users.token accepted in both modes

See Also: Full configuration and troubleshooting guide at docs/content/docs/documentation/oidc.md (quick start: docs/content/docs/documentation/sso-quickstart.md).