diff --git a/.env.ollama b/.env.ollama new file mode 100644 index 000000000..feaf6504e --- /dev/null +++ b/.env.ollama @@ -0,0 +1,49 @@ +# LLM +BASE_URL=http://host.docker.internal:11434/v1 +API_KEY=ollama +MODEL=qwen2.5:0.5b +LLM_SEMAPHORE=1 +TIMEOUT=300 + +# VLM (Vision support) +VLM_BASE_URL=http://host.docker.internal:11434/v1 +VLM_API_KEY=ollama +VLM_MODEL=qwen3-vl:2b +VLM_SEMAPHORE=1 + +# EMBEDDER +EMBEDDER_MODEL_NAME=dengcao/Qwen3-Embedding-0.6B:Q8_0 +EMBEDDER_BASE_URL=http://host.docker.internal:11434/v1 +EMBEDDER_API_KEY=ollama +MAX_MODEL_LEN=2048 + +# RAG CONFIG +RAG_MODE=SimpleRag +CONTEXTUAL_RETRIEVAL=false +RERANKER_ENABLED=false +MAX_OUTPUT_TOKENS=2048 + +# App Settings +# Docker Compose expands PWD from the shell; customize if your shell does not set it. +SHARED_ENV=${PWD}/.env.ollama +APP_PORT=8002 +CHAINLIT_PORT=8090 +DEFAULT_LANGUAGE=en-US +INDEXERUI_PORT=8060 + +# Chainlit conversation history (SQLAlchemy -> existing rdb PostgreSQL) +CHAINLIT_DATABASE_URL=postgresql+asyncpg://root:root_password@rdb:5432/chainlit +DATABASE_URL=postgresql://root:root_password@rdb:5432/chainlit +CHAINLIT_AUTH_SECRET=openrag_local_dev_secret_2026 + +# RAY & System +RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 +RAY_NUM_GPUS=0 +RAY_DASHBOARD_PORT=8265 +RAY_memory_usage_threshold=0.99 +RAY_memory_monitor_refresh_ms=0 +SUPER_ADMIN_MODE=true +AUTH_TOKEN=sk-1234 +SAVE_UPLOADED_FILES=true +PDFLoader=PyMuPDFLoader +WHISPER_N_WORKERS=0 diff --git a/.github/workflows/integration_tests.yml b/.github/workflows/integration_tests.yml new file mode 100644 index 000000000..e71a7c7dc --- /dev/null +++ b/.github/workflows/integration_tests.yml @@ -0,0 +1,84 @@ +name: Integration tests + +on: + push: + branches: + - main + - dev + - "refactor/hexagonal" + - "refactor/phase-**" + pull_request: + +jobs: + milvus-integration: + runs-on: ubuntu-latest + timeout-minutes: 20 + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Docker Compose + uses: docker/setup-compose-action@v1 + with: + version: v2.34.0 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install uv + run: | + curl -LsSf https://astral.sh/uv/install.sh | sh + echo "$HOME/.local/bin" >> $GITHUB_PATH + + - name: Cache uv + uses: actions/cache@v4 + with: + path: | + ~/.cache/uv + key: uv-${{ runner.os }}-3.12-${{ hashFiles('**/pyproject.toml', '**/uv.lock') }} + restore-keys: | + uv-${{ runner.os }}-3.12- + + - name: Install dependencies + run: uv sync + + - name: Start Milvus stack + working-directory: tests/integration + run: docker compose up -d --wait --wait-timeout 180 + + - name: Verify Milvus reachable from host + run: | + for i in {1..30}; do + if curl -sf http://localhost:9091/healthz 2>/dev/null; then + echo "Milvus healthy after $i attempts" + break + fi + echo "Attempt $i/30 - waiting..." + sleep 2 + done + curl -sf http://localhost:9091/healthz + + - name: Run integration tests + env: + OPENRAG_TEST_VDB_HOST: localhost + OPENRAG_TEST_VDB_PORT: "19530" + run: uv run pytest tests/integration/ -m integration -v --tb=short + + - name: Show logs on failure + if: failure() + working-directory: tests/integration + run: | + echo "=== Milvus Logs ===" + docker compose logs milvus --tail=200 + echo "=== etcd Logs ===" + docker compose logs etcd --tail=50 + echo "=== minio Logs ===" + docker compose logs minio --tail=50 + + - name: Cleanup + if: always() + working-directory: tests/integration + run: docker compose down -v diff --git a/.github/workflows/layer_guard.yml b/.github/workflows/layer_guard.yml new file mode 100644 index 000000000..4c3041bf5 --- /dev/null +++ b/.github/workflows/layer_guard.yml @@ -0,0 +1,21 @@ +name: Layer guard + +on: + push: + branches: + - "refactor/hexagonal" + - "refactor/phase-**" + pull_request: + branches: + - "refactor/hexagonal" + +jobs: + layer-import-guard: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Run layer guard + run: python scripts/check_layer_imports.py diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 6b75e654a..9d6a81297 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -2,9 +2,16 @@ name: Linting on: push: - branches: [ main, dev ] + branches: + - main + - dev + - "refactor/hexagonal" + - "refactor/phase-**" pull_request: - branches: [ main, dev ] + branches: + - main + - dev + - "refactor/hexagonal" jobs: lint: diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index 7007a29ed..51e86c64a 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -2,7 +2,11 @@ name: Unit tests on: push: - branches: [ main, dev ] + branches: + - main + - dev + - "refactor/hexagonal" + - "refactor/phase-**" pull_request: jobs: diff --git a/.gitignore b/.gitignore index 5bb3231d3..508de51de 100644 --- a/.gitignore +++ b/.gitignore @@ -59,9 +59,9 @@ volumes/* !volumes/.gitkeep # Keep the placeholder # services -services/ -services/* -!services/.gitkeep # Keep the placeholder +/services/ +/services/* +!/services/.gitkeep # Keep the placeholder *.csv *.pkl diff --git a/FORWARD_PORT_LOG.md b/FORWARD_PORT_LOG.md new file mode 100644 index 000000000..e30c56c74 --- /dev/null +++ b/FORWARD_PORT_LOG.md @@ -0,0 +1,27 @@ +# Forward Port Log + +Tracks `dev` changes during MODE 2 isolation (Phases 5-9). +Each entry: what changed on `dev`, whether it was forward-ported or deferred +to the cutover re-implementation queue. + +> Created retroactively at the start of Phase 5 (2026-04-29). The Phase 0-4 +> Mode 1 work merged from `dev` cleanly so this log starts empty. + +--- + +## Forward-ported (critical) + +_Security fixes, data-loss bugs, production outages re-implemented against +the new architecture. Each entry pairs the dev commit with the refactor +commit so a reviewer can audit equivalence._ + +(none yet) + +## Deferred to cutover (features) + +_Non-critical changes that landed on `dev` during MODE 2. These will be +re-implemented directly in the new architecture during MODE 3 (Phases +10-12) or post-cutover. List the dev PR number / commit and the target +location in the new layout._ + +(none yet) diff --git a/REFACTORING_DECISION_LOG.md b/REFACTORING_DECISION_LOG.md new file mode 100644 index 000000000..e3071c2af --- /dev/null +++ b/REFACTORING_DECISION_LOG.md @@ -0,0 +1,1410 @@ +# Refactoring Decision Log + +Records **why** decisions were made that deviate from or extend the refactoring +docs. When a decision changes the plan, update the strategy/workflow docs to +reflect the new reality — then log the reasoning here so future readers know +why the docs changed. + +Source abbreviations: +- STRATEGY = `docs/refactoring/REFACTORING_STRATEGY_v1.md` +- WORKFLOW = `docs/refactoring/REFACTORING_DEV_WORKFLOW.md` + +--- + +## Phase 0 — Scaffold + import guard + CI wiring (2026-04-21) + +**1. The guard ignores files outside the four new layer roots.** +Files under `openrag/components/`, `openrag/routers/`, `openrag/models/`, +`openrag/config/`, `openrag/utils/` are skipped. +- Why: Phase 0's verification requires existing tests to keep passing. If the + guard ran against legacy code, every old import that doesn't fit the new + rules would trip the check and block the phase. Legacy code gets migrated in + Phases 5–12 and the guard picks those files up as they move into the new + layer roots. +- Alternative considered: whitelist-only enforcement on new code (same idea, + different framing). What we chose is "enforce wherever the file lives in one + of the four roots", which is simpler. + +**2. Split CI into `layer_guard.yml` + extending existing `lint.yml` and +`unit_tests.yml`, instead of one new `refactor-ci.yml`.** +WORKFLOW's CI example is a single file with three jobs (`unit-tests`, +`layer-guard`, `docker-build`). We took a different shape. +- Why: We already have a well-set-up `unit_tests.yml` and `lint.yml`. Creating + a parallel `refactor-ci.yml` with its own unit-tests job would duplicate the + uv setup and caching. Extending the existing files adds a few lines of + config and reuses everything. +- Alternative considered: follow the WORKFLOW example literally. Rejected for + the duplication reason above. Trade-off is that refactor-specific CI isn't + all in one file. + +**3. `docker-build` CI check NOT wired in Phase 0.** +WORKFLOW lists it as a required check. +- Why: Existing `build.yml` and `build_dev.yml` workflows push images to ghcr, + which isn't what we want on every refactor push. A lightweight "docker build + only, don't push" check needs a new job. Deferred to keep Phase 0 scope + tight. Docker build was verified locally on the phase-0 tree. +- Alternative considered: add the job in this phase. Rejected for scope. + Follow-up: add a `docker-build` job in a separate PR, modelled on the + WORKFLOW CI example. + +**4. Decision log policy: log reasoning, update docs.** +When a decision deviates from the strategy/workflow docs, update the docs to +match reality, then record the reasoning here. +- Why: The docs should always reflect the current plan. The log captures + why the plan changed, not what the plan is. + +--- + +## Phase 5 — Core domain logic for retrieval, chunking, prompts (2026-04-29) + +**1. New `RetrievalSearcher` port in `core/retrieval/searcher.py`, separate from +the narrow `VectorStore` ABC.** +The retriever needs four operations (search by query string, multi-query +search, related-chunk lookup, ancestor lookup) that the Phase-4 +`VectorStore` ABC does not cover — that ABC is intentionally narrow +(`search(embedding, top_k)`). We added a transitional ABC the retriever +depends on, implemented by `services/storage/milvus_ray_shim.py` over the +legacy Ray actor. +- Why: STRATEGY §5A says the retriever should "call `VectorStore.search()` + (port method), not `vectordb.async_search.remote()`". But the legacy Ray + actor's `async_search` takes a query *string* and embeds internally; the + narrow `VectorStore.search(embedding, ...)` ABC doesn't fit. Pre-embedding + in the shim before calling Ray is impractical because the actor also owns + BM25 and surrounding-chunks semantics. A retrieval-facing port keeps the + retriever clean of Ray today and survives Phase 7 — when the Vectordb + god object is decomposed, these methods either move onto a richer + `VectorStore` or split between `VectorStore` and `ChunkRepository`. +- Alternative considered: extend `VectorStore` with the four legacy methods. + Rejected — bloats the ABC with operations that should not exist past + Phase 7. Also considered: skip the new core port and have the retriever + call the Ray actor through the shim with the legacy method names — + rejected because that leaks legacy method names into core/ and makes the + retriever harder to test. + +**2. Skipped: bringing up integration tests for the new code.** +Phase 5 ships pure-domain unit tests only (50 new tests in `core/`, no Ray / +Milvus / real LLM). The new pipeline is dormant until Phase 8 wires it. +- Why: Mode 2 forbids touching the legacy wiring; the new pipeline has + nowhere to be plugged in yet. Integration coverage will land with Phase 8 + orchestrators (or Phase 7 storage if it goes first). +- Alternative considered: stand up a fake searcher in an integration + fixture and run a full retriever-pipeline-RRF round trip. Defers the + same coverage to Phase 8 with less code; not worth the extra fixtures. + +**3. `Query`, `SearchQueries`, `TemporalPredicate` lifted into +`core/models/query.py`.** +The legacy `components/pipeline.py` defined these inline. The new +`RetrieverPipeline` consumes them — they're domain types, not pipeline +internals. +- Why: STRATEGY §2 calls these out as `pipeline.py SearchQueries → core/models/query.py`. +- Alternative considered: keep them in `core/retrieval/`. Rejected — they + describe a query in the abstract; the orchestrator (Phase 8) and the API + layer will both use them, not just retrieval. + +**4. Phase 5.15 (re-export shims) deferred to follow-up — then completed.** +The first Phase 5 commits (5A/5B/5C, 2026-04-29) created the new core/ +modules but left the legacy `components/` files intact. STRATEGY §4.1 +mandates a three-step move — create new file, update old file to +re-export from new, update consumers — and Phase 5 step 5.15 says +"Update old files to re-export from core/". We skipped that. +The follow-up sweep (2026-05-05) replaced six legacy files with shims: + +| Legacy file | Shim strategy | +|---|---| +| `components/indexer/chunker/utils.py` | Plain re-export from `core.chunking.markdown_utils` | +| `components/prompts/prompts.py` | `load_prompt(key)` adapter calling `core.prompts.template_loader.load_template_by_key` | +| `components/utils.py:format_context` + `format_web_context` | Adapters into `core.prompts.chat_prompt_builder` (rest of utils stays — Phase 6+ scope) | +| `components/indexer/chunker/chunker.py` | `BaseChunker` / `RecursiveSplitter` delegate to `core.chunking.RecursiveSplitter` via Document↔ProcessedDocument↔Chunk conversion. `ChunkContextualizer` + `ChunkerFactory` retained (5D + Phase 8). | +| `components/retriever.py` | `Single`/`MultiQuery`/`HyDe` retrievers wrap `core.retrieval.retriever` strategies. Ray actor → `MilvusRayShim`; `ChatOpenAI` → `_LangChainLLMAdapter`. `RetrieverFactory` retained. | +| `components/pipeline.py` | `Query`/`SearchQueries`/`TemporalPredicate` re-exported from `core.models.query`. `RetrieverPipeline` delegates to `core.retrieval.pipeline.RetrieverPipeline` via a `_LegacyRerankerAdapter` bridging the legacy reranker (Document-in / Document-out) to the core ABC (str-in / `(idx, score)`-out). `RagPipeline` + `RAGMODE` retained (Phase 8). | + +- Why: STRATEGY §4.1 is explicit ("Update old file to re-export from new + location"); leaving the duplication in place would let the codepaths + drift. Three CodeRabbit fixes from PR #352 (image_caption ChunkType, + page-marker semantics, chunk_table header-only flush) had to be applied + twice or only fixed in core — exactly the failure mode 5.15 prevents. + The shim pattern matches the prior-art shims for config (1329cc18) and + exceptions (a0f3d9f2). +- Alternative considered: leave the copies until Phase 8 cutover. Rejected + — the doc explicitly puts 5.15 *inside* Phase 5, and one round of + drift already happened. +- Side effect: a noqa side-effect import in `chunker.py` keeps the legacy + `components.utils` ↔ `components.indexer.utils.files` circular-import + resolving in the right order. Removed once `components.utils` is split + in Phase 6+. +- Behavioral note: image elements in the legacy chunker now stamp + `chunk_type=image_caption` (matching `core.models.chunk.ChunkType`) + instead of the previous raw `image`. No legacy reader filters on this + value, so the change is invisible to consumers. + +**5. `core/chunking/recursive.py` keeps `langchain.text_splitter.RecursiveCharacterTextSplitter` as a dependency.** +STRATEGY §3 lists core as "stdlib + pydantic + pure libs" and §4.7 limits +LangChain in core to boundary converters (`from_langchain` / `to_langchain` +on domain models). The chunker imports `RecursiveCharacterTextSplitter` +directly, which is neither stdlib nor a boundary converter. +- Why: `RecursiveCharacterTextSplitter` is a self-contained recursive + separator-based string splitter — no IO, no LLM client, no Document + semantics. Reimplementing it in core/ would be a meaningful chunk of + pure code with no behavior change, and the legacy chunker has been + using it for two years with stable output. Keeping it in for Phase 5 + preserves byte-for-byte chunk equivalence, which the strangler-fig + shim relies on for behavior parity. The import is also deferred + (inside the constructor / `split_text`) so importing the module + without LangChain installed doesn't fail. +- Alternative considered: write a stdlib-only recursive splitter as part + of Phase 5. Rejected as scope creep — would couple a behaviorally + risky rewrite (chunk boundaries shift, downstream embedding output + changes) to the additive Phase 5 cut, breaking the parity guarantee + the legacy shims rely on. Tracked as a Phase 12 / post-cutover + follow-up: replace the splitter with a stdlib implementation behind + the same `Callable[[str], int]` length-function injection point. +- Scope: limited to `RecursiveCharacterTextSplitter`. No other LangChain + symbol leaks into core; `langchain_core.documents.Document` only + appears inside `Chunk.from_langchain` / `Chunk.to_langchain` / + `Document.from_langchain` / `Document.to_langchain`, all with deferred + imports, exactly as §4.7 prescribes. + +**6. Two file-layout deviations from STRATEGY §3 / §5A / §5B.** +- `core/chunking/markdown_section.py` (§5B, line 1038; §3, line 387) + → renamed to `core/chunking/markdown_utils.py`. The contents are pure + parsing helpers — `MDElement`, `split_md_elements`, `chunk_table`, + `parse_markdown_table`, `get_chunk_page_number` — not a + section-aware chunker strategy. The name `markdown_utils.py` matches + the module's role (utilities consumed by `RecursiveSplitter`); the + separate `markdown_section.py` / `markdown_layout.py` *strategies* + listed in §3's tree aren't built in Phase 5 and remain available + filenames if/when those strategies land. +- `core/retrieval/hydration.py` (§5A, line 1027) → kept as private + `_expand_with_related_chunks` in `retriever.py`. The function is + ~60 LOC, only invoked by `BaseRetriever.expand_search_results`, and + splitting it would add an import + test fixture surface without any + reuse benefit. If a second consumer ever appears (Phase 8 likely), + promoting it to a module-level public function in `hydration.py` is + a one-commit move. +- Why: both deviations make the module names track the actual contents + rather than the strategy doc's pre-write naming guess. Recording so + future readers don't grep for files that aren't there. +- Alternative considered: rename to match the strategy doc verbatim. + Rejected — the strategy filenames anticipated different content + (a section/layout chunker, a standalone hydration entry point) than + what Phase 5 actually produced. + +--- + +## Phase 1 — Registry + Exceptions (2026-04-21) + +**1. Exceptions keep HTTP status_code on the class (OpenRAG style), not in +a separate error handler mapping (mandragora style).** +- Why: Existing code reads `exc.status_code` in multiple places. Switching + to a pure domain exception + API-layer mapping dict would require changing + every consumer now, which is unnecessary churn in Phase 1. +- Alternative considered: mandragora's pattern (bare exceptions in core/, + status code mapping in api/error_handlers.py). Cleaner for hexagonal + purity but rejected for backward compatibility. +- Follow-up: strip status codes from core exceptions in Phase 10 when + `api/error_handlers.py` is built. The error handler will own the mapping. + +--- + +## Phase 5D — Indexing domain logic + parsers (2026-04-30) + +**1. `core/indexing/validators.py` is fully framework-free.** +- All FastAPI types removed (`Form`, `UploadFile`, `HTTPException`, + `status`, `Depends`); validators are pure functions on `str` / `dict` + / `Iterable[str]`. +- `accepted_formats` / `accepted_mimetypes` are passed as args instead + of read from Hydra `config` at module import (legacy module-level reads + of `ACCEPTED_FILE_FORMATS` / `DICT_MIMETYPES` / + `FORBIDDEN_CHARS_IN_FILE_ID` are gone). +- `ValidationError` accepts a `status_code` (and `code`) kwarg. Phase 1 + hardcoded 422; the original validators raised HTTP 400 (invalid + `file_id` / metadata JSON) and HTTP 415 (unsupported format), so a + status-code override is needed to preserve those codes from a + pure-domain exception. Existing precedent in the same module + (`LLMParsingError` overrides `status_code` after `super().__init__`) + shows the pattern is already accepted. +- HTTP translation flows through the existing global + `openrag_exception_handler` (`@app.exception_handler(OpenRAGError)`) + wired in Phase 1, not local `HTTPException` raises in routers. +- Why: A core module that imports FastAPI or reaches into Hydra is not + framework-free, blocks reuse from non-HTTP entry points, and + re-introduces the boundary violation the refactor exists to fix. + Stripping only `Depends()` — the literal task description — would + leave the boundary half-broken. +- Trade-off: error body becomes `{"detail": "[CODE]: msg", "extra": {}}` + instead of FastAPI's `{"detail": "..."}` — matches every other + `OpenRAGError`. +- Alternatives considered: (a) consolidate everything on 422 — rejected, + observable behaviour change; (b) introduce specific subclasses + (`UnsupportedFileFormatError`, etc.) — rejected as premature, only two + call sites need non-default codes today (Phase 10's API error-handler + layer can re-evaluate); (c) keep module-level Hydra reads — rejected, + embeds the infrastructure config object into core; (d) catch and + re-raise as `HTTPException` in the router wrappers — rejected, + duplicates the global handler. + +**2. Exception shims under `utils/exceptions/` use `core.X`, not `openrag.core.X`.** +The legacy shims imported via `openrag.core.utils.exceptions`. With both +`/app` and `/app/openrag/` reachable, Python loads the same file as two +distinct modules, producing two distinct `OpenRAGError` classes — +`isinstance` failed and the global handler never fired. +- Why: Unifying on the bare `core.X` path matches `pythonpath = ./openrag` + and the relative-imports-within-`core/` convention (commit 4528c71). +- Follow-up: ~20 other `from openrag.X` imports across `core/`, `config/`, + and components are latent dual-import traps and should be migrated + in a separate pass. + +**3. Parser layering: native in core, services-backed in services/workers, type-marker bases without vendor names, DI for pools.** +- Native-bytes parsers (PyMuPDF, html_to_markdown, chardet, image) live + in `core/indexing/parsers/`. Service-/Ray-backed parsers (Marker, + LocalWhisper) live in `services/workers/parsers/`. +- Empty marker subclasses `BasePooledParser` / `BaseClientParser` in + `core/indexing/parsers/document_parser.py` categorize parsers by *how* + they get their work done (actor-pool vs HTTP-client) without naming + the implementation. A core base class called `RayPoolParser` or + `OpenaiClientParser` would leak vendor/infrastructure into the + framework-free layer and foreclose swapping the backend. +- Core facades (`MarkerParser`, `LocalWhisperParser`, `ClientPdfParser`, + `ClientAudioParser`) accept any pool/client of the appropriate marker + type via `__init__`; services own the actor lifecycle. +- Why: `@ray.remote` decoration imports infrastructure at + class-definition time and can't be hidden behind a port. DI keeps + facades testable with in-memory fakes. +- Alternatives considered: (a) all parsers in core with Ray injected + via DI — rejected, class-level decoration can't be deferred to + composition; (b) have core facades resolve the actor by name + themselves — rejected, couples core to Ray's named-actor registry. + +**4. Image preprocessing helpers extracted to `core/indexing/image_preprocessor.py`.** +Pure helpers (`ensure_png_compatible_mode`, `pil_to_png_bytes`, +`pil_to_base64`, `is_http_url`, `is_data_uri`, `HTTP_IMAGE_PATTERN`, +`DATA_URI_IMAGE_PATTERN`, `MIN_IMAGE_PIXELS`). Used by the core image +parser and by Marker captioning in services. +- Why: Both layers need PNG normalization and markdown image-reference + detection. Sharing via core (no VLM, no langchain imports) avoids + services depending on `components/indexer/loaders/base.py`. +- Alternative considered: leave helpers in + `components/indexer/loaders/base.py`. Rejected — + services-importing-components is a layering violation, and `base.py` + drags in langchain. + +**5. `services/workers/ray_utils.py` keeps function and decorator forms together; `description=` is a format-string template.** +- `call_ray_actor_with_timeout` / `@with_timeout` and `retry_with_backoff` + / `@with_retry` (with jitter) live in one module — STRATEGY's proposed + `_retry.py` / `_timeout.py` split for `services/inference/` doesn't + apply here because workers need both forms in practice (decorator at + class-definition for static-param call sites, function form for + callsite-resolved values). The decorators delegate to the function + form internally; splitting across two files would duplicate that + wiring. +- `description=` accepts a **format string** like + `"PDF parse ({file_path})"`; `_resolve_description` binds it via + `inspect.signature.bind` against the wrapped call's args at call + time. **Callables (lambdas) are NOT supported** — they fall through + to `if "{" not in template:` and raise `TypeError: argument of type + 'function' is not iterable`. (One outlier in `marker_workers.py` used + a lambda and was fixed in Phase 5E.) +- Inline `call_ray_actor_with_timeout(worker.X.remote(...))` calls in + workers are extracted into one-line `@with_timeout`-decorated helper + methods (`_transcribe_chunk`, `_check_pool_broken`, + `_reset_worker_pool`, `_run_chunk`, `_convert_pdf`) returning the + `ObjectRef`; the decorator awaits it with timeout. Worker files use + only decorator form — no mixed styles. +- Retry-around-timeout semantics preserved: `@with_retry` outer, + `@with_timeout` inner — `TimeoutError` propagates from the inner + helper and the outer decorator re-runs the whole method body (slot + pick, fresh `.remote()`, fresh timeout). +- Alternatives considered: (a) mirror inference's `_retry.py` / + `_timeout.py` split verbatim — rejected, adds files that just import + from each other; (b) keep description static, drop to function form + when dynamic — rejected, re-introduces the verbose + `call_ray_actor_with_timeout(...)` call sites the decorator was meant + to remove; (c) keep function form for the inline cases — rejected, + leaves a mix of styles in the same file with no clear rule. + +**6. `ray_utils` canonical home moved from `components/` to `services/workers/`.** +`components/ray_utils.py` is now a back-compat shim re-exporting from +`services.workers.ray_utils`. +- Why: Ray-actor concurrency primitives belong in the services layer, + not in `components/` (which is on the deprecation path). Routers and + pipeline still import via the components shim during the transition. +- Follow-up: migrate the remaining `components.ray_utils` imports + (pipeline, search router, indexer router, workspaces router, indexer + utils) and delete the shim in Phase 5E. + +**7. Docling and DoclingV2 PDF backends deferred — not migrated in Phase 5D.** +No `core/indexing/parsers/pdf/docling*` modules will be created in this +pass. Legacy `DoclingLoader` and `DoclingLoader2` stay where they are +for now. +- Why: This is a PDF backend we haven't used or tested recently — + porting it now would pin a stale integration into the new layer. We'll + revisit and re-port it (or drop it) in a later pass once the refactor + has shaken out and we know whether Docling is still wanted. +- Alternative considered: port now alongside Marker / OpenAI / DotsOCR + for completeness. Rejected — moves dead-feeling code into the new + layer without verifying it still works. +- Follow-up: revisit during a later parser-coverage sweep. If the + decision is to drop, the legacy modules get deleted in Phase 5E rather + than shimmed. + +**8. `ImageBlock` is the parser↔caption contract — captioning is a downstream stage's job.** +- Every parser (Image, Markdown, Docx, Pptx, Eml, Marker, + `DotsOCRPdfClient`) emits `ImageBlock` with `caption=None`. The + caption stage fills it in. For VLM-PDF specifically, the picture-bbox + crop becomes an `ImageBlock(image_bytes=…, page_number=N)` — the + parser never issues the second VLM call. One uniform contract beats + per-parser carve-outs; the chunker sees the same `ImageBlock` shape + from every parser, including `DotsOCRPdfClient`. +- `ImageBlock.metadata['markdown_ref']` holds the in-text placeholder + (data-URI, `![](pptx-image-N)`, `![](docx-image-N)`, + `![](marker-key)`); the caption stage `str.replace`s it. No + placeholder ⇒ no `markdown_ref` ⇒ caption stage emits a + free-standing `TextBlock`. Contract is documented on `ImageBlock` + itself. +- `ImageBlock` carries `image_bytes` (default `b""`) AND `source_url`. + Locally-extracted images set bytes; HTTP refs (`![alt](https://…)`) + leave bytes empty and set `source_url`. The `image_url` property + returns `data:{mime};base64,…` when bytes are present, else + `source_url` — consumers read `image_url` regardless of shape. +- Why: Refs are per-image-unique and chunk-stable. Legacy + `MarkdownLoader` captioned HTTP images via langchain `ChatOpenAI` + (which accepts URLs natively). The new VLM ABC takes bytes only, so a + fetch stage has to populate them — but the parser still emits one + `ImageBlock` per in-text image, keeping the contract uniform. +- Alternatives considered: (a) positional matching of refs to images — + rejected as fragile; (b) embedding image bytes inside `TextBlock` — + rejected as a heavier model change. + +**9. Paginated parsers emit `list[TextBlock]` with `page_number`; in-band `[PAGE_N]` markers are gone.** +Marker and PPTX previously concatenated all page content into one +`TextBlock` with `[PAGE_N]` markers between pages. They now emit one +`TextBlock` per page with `page_number` set, matching what PyMuPDF +already does. Parsers without natural pagination +(text/html/md/docx/doc/eml/whisper/image) still emit a single +`page_number=1` block. +- Why: Pagination is metadata, not content. Leaking `[PAGE_N]` markers + into chunk text forced every consumer to know the marker syntax; + `TextBlock.page_number` is the canonical channel and was already + half-used. +- Implication for chunking: the chunker must NOT scan for `[PAGE_N]` + markers. Iterate `ProcessedDocument.text_blocks` and carry + `block.page_number` onto every emitted chunk. Page boundaries are + block boundaries. + +**10. Client-backed parsers: generic `Client*Parser` facades; `BaseOpenAIPdfClient` is scaffolding only.** +- Renamed `OpenAIPdfParser` → `ClientPdfParser` + (`core/indexing/parsers/pdf/openai.py` → `pdf/client_based.py`); added + `ClientAudioParser` at `core/indexing/parsers/audio/client_based.py`. + Both accept any `BaseClientParser` and delegate `parse()`. "OpenAI" + was a leaky model-specific label on a class that takes any + HTTP-client-backed parser; whatever DotsOCR / Whisper-vLLM / + Scaleway-Speech is called next quarter, the facade stays the same — + what varies is the injected `BaseClientParser`. +- `BaseOpenAIPdfClient` provides reusable helpers (PDF page rendering, + semaphore-protected `_ocr_one(page_img, prompt) → str | None`, + JSON-fence stripping, JSON loading, picture-bbox cropping). It does + **NOT** define `parse()`, a `PROMPT` class attribute, or abstract + `_caption_images` / `_result_to_md` / `_parse_ocr_response` hooks. + The file was renamed `_openai.py` → `_base_openai_parser.py` to + match the new role. +- Why: The previous abstract pipeline imposed assumptions ("there's one + OCR response per page", "captioning is a parser concern") that didn't + generalise. Treat the base as a toolbox; let each concrete client + (DotsOCR, future variants) drive its own `parse()` and block-emission + strategy. +- Trade-off: more code per concrete subclass. Accepted — + model-specific variation (response schema, block layout, bbox + handling) lives in the subclass anyway. +- Alternative considered: keep one model-specific facade per backend. + Rejected — duplicates the same isinstance + delegate boilerplate. + +**11. DotsOCR response is validated through Pydantic.** +`DotsOCRElement` / `DotsOCRPage(RootModel[list[DotsOCRElement]])` / +`DotsOCRCategory` (Enum) capture the layout-element shape; +`DotsOCRPdfClient._parse_page` runs `model_validate` and returns `None` +on bad payloads. The `{"items": [...]}` envelope is tolerated alongside +a bare list. +- Why: Replaces dict shuffling (`page_res.get("category") == "Picture"`, + `item.get("bbox")`) with typed access (`element.category is + DotsOCRCategory.PICTURE`, `element.bbox`). Bad payloads fail loudly + via `ValidationError` instead of silently returning empty markdown. + +**12. `OpenAIAudioClient` keeps language detection as an injected callable, not a Ray ref-getter.** +Legacy `AudioTranscriber` looked up a `WhisperActor` Ray actor by name. +The new `OpenAIAudioClient` takes `language_detector: Callable[[Path], +Awaitable[str | None]] | None` in its constructor and skips detection +when `None` (vLLM auto-detects). +- Why: Keep the client free of Ray coupling so it can be instantiated + and tested without a Ray cluster. The wiring layer passes a closure + that calls the Whisper actor when `USE_WHISPER_LANG_DETECTOR=true`. +- Alternative considered: keep the Ray actor lookup inside the client + guarded by a config flag. Rejected — pulls Ray into the + `services/inference` layer where the rest of the file is plain HTTP. + +--- + +## Phase 5E — Loader → Parser shims (2026-05-06) + +**1. Legacy loaders are *adapter* shims, not re-export shims.** +The earlier compat-shim pass (commit `93476a6`) used pure `from X +import Y` re-exports because the symbols moved unchanged +(`ray_utils`, `text_sanitizer`, exceptions). The loader→parser move +can't do that: `BaseLoader.aload_document(file_path) → langchain +Document` and `DocumentParser.parse(document) → ProcessedDocument` +have different names *and* different contracts. Each legacy loader +becomes a `BaseLoader` adapter that reads the file into bytes, builds +a `CoreDocument`, calls the new parser, and maps `ProcessedDocument` +back to a langchain `Document`. +- Why: Preserves dynamic loader-discovery + (`BaseLoader.__subclasses__()` in `loaders/__init__.py`) and the + config-string lookup (`file_loaders.pdf: "MarkerLoader"`) without + forcing every consumer to migrate at once. +- Alternative considered: pure re-exports aliasing `*Parser` as + `*Loader`. Rejected — the discovery walk only finds `BaseLoader` + subclasses, so an aliased `DocumentParser` would silently disappear + from the loader registry. + +**2. Shimmed in this pass: text/markdown, image, docx, doc, pptx, pymupdf, marker, local-whisper, openai-audio.** +Each adapter delegates to its core parser and, when the parser emits +`ImageBlock`s with `markdown_ref` set, layers VLM captioning on top +via the existing `BaseLoader` mixin (`self.image_captioning`, +`self.caption_images`, `self.replace_markdown_images_with_captions`). +- Why: Keeps the legacy contract intact (captioned markdown in + `page_content`) while the canonical home is the parser. The + `markdown_ref` substitution path is the same one the future + caption-stage will use. + +**3. `base.py` Stage 1: re-export the four image_preprocessor symbols already in core, leave the captioning mixin in place.** +`ensure_png_compatible_mode`, `HTTP_IMAGE_PATTERN`, +`DATA_URI_IMAGE_PATTERN`, `MIN_IMAGE_PIXELS` now point at the +canonical `core.indexing.image_preprocessor` symbols (class attrs +hold module-level references for `self.X` access). +`_pil_image_to_base64` rewritten on top of `pil_to_png_bytes`. The +VLM endpoint setup, `get_image_description`, `caption_images`, +`replace_markdown_images_with_captions` stay in `base.py` for now. +- Why: Mechanical, behavior-identical change. Stage 2 (move VLM + captioning to `services/inference/captioning`) needs a design call + (where it lives, how the shim acquires it) and is deferred. + +**4. `PyMuPDFParser`: single dedicated thread + retain empty pages for 1-to-1 pagination.** +- PyMuPDF/pymupdf4llm are not thread-safe; concurrent calls raise + `ValueError: not a textpage of this page`. Upstream maintainer + (`pymupdf/PyMuPDF#3771`, closed wontfix) confirms this is documented + behaviour, not a bug. The parser now uses a module-level + `ThreadPoolExecutor(max_workers=1)` instead of `asyncio.to_thread`; + concurrent `parse()` calls queue on the executor, eliminating the + race against the default thread pool. The rest of the indexing + pipeline still parallelizes — only the pymupdf step is serialized. +- Empty pages now produce a `TextBlock` with empty `text` (was + previously dropped while keeping `page_count` accurate). Reverted so + every page produces a `TextBlock`, keeping a 1-to-1 mapping with the + source PDF's pagination — the legacy `\n[PAGE_N]\n` anchor format + the loader-shim emits aligns exactly with the source. + +**5. `TranscriberConfig.direct_upload_suffixes` got lost in the core/config migration; ported to `core/config/indexation.py`.** +The legacy `config/models.py:TranscriberConfig` had the field + +`|`-separated string validator + a default frozenset of audio +extensions. The active `core/config/indexation.py:TranscriberConfig` +(loaded via `openrag.core.config.loader.load_config`) was missing it, +producing `AttributeError: 'TranscriberConfig' object has no attribute +'direct_upload_suffixes'` when the audio shim accessed it. +- Why: `config/models.py` is now vestigial — kept for legacy imports + but no longer drives `load_config()`. Fields added there but not + mirrored to `core/config` are silently inactive at runtime. + +**6. Skipped: eml, `pdf_loaders/openai.py`, `pdf_loaders/dotsocr.py`.** +- `eml_loader.py`: the new `EmlParser` takes `attachment_parsers: + Mapping[str, DocumentParser]`, but the old loader dispatches + attachments through `BaseLoader`-keyed `get_loader_classes` with a + multi-tier PDF fallback chain (`MarkerLoader` → `PyMuPDFLoader` → + `PyMuPDF4LLMLoader` → `DoclingLoader`). The contract bridge isn't + trivial; deferred until services-side attachment-parser composition + lands. +- `pdf_loaders/openai.py` + `pdf_loaders/dotsocr.py`: services-side + `BaseOpenAIPdfClient` / `DotsOCRPdfClient` exist but require a + concrete `core.vlm.VLM` to instantiate, and `vlm_registry` is empty + (no concrete VLM impl exists yet). Both legacy classes are also dead + code on this branch — not in any Hydra config, no external imports. +- Why: Both gaps need new services-side work (attachment-parser DI, + `LangchainOpenAIVLM`-style concrete) before a meaningful shim is + possible. Re-export-only "shims" would relocate the file without + going through the new architecture, defeating the purpose. + +**7. Stale files flagged for deletion (Phase 12 cleanup).** +- `components/indexer/loaders/CustomHTMLLoader.py` and + `components/indexer/loaders/CustomDocLoader.py` — legacy + `BaseLoader` subclasses, not referenced by any Hydra config or + external import. Discoverable via `BaseLoader.__subclasses__()` but + never instantiated. `CustomDocLoader` uses + `UnstructuredWordDocumentLoader` / `UnstructuredODTLoader` — no + clean parser equivalent in core (`DocxParser` uses MarkItDown). +- `config/models.py` (the whole file, incl. its `TranscriberConfig`) + — superseded by `core/config/*`; kept only so legacy imports don't + break. Drift between the two has already caused one runtime bug + (entry 5). +- Why: Out of scope for the loader-shim pass; flagged here so they + don't get re-shimmed by future passes. Removal coordinates with + Phase 12 ("delete old re-export shims"). + +--- + +## Phase 6B — vLLM inference clients + legacy shims (2026-05-07) + +**1. `VLLMVision(VLLMClient, VLM)` — multiple inheritance kept for nominal typing.** +`VLLMClient` provides the full implementation (httpx pool, retry, +circuit breaker, `aclose()`). `VLM` is a pure abstract mixin with no +conflicting methods, so the MRO is linear and clean. Adds only +`_max_tokens`, `caption_image()`, and `caption_images_batch()`. +- Why: VLM and LLM talk to the same vLLM OpenAI-compatible + chat/completions endpoint, so `VLLMClient` is the right concrete + base. Keeping `VLM` in the bases preserves nominal typing — + `isinstance(vision, VLM)` works, and any future code that type-checks + against the VLM ABC will accept `VLLMVision` without a cast. +- Alternative considered: single inheritance `VLLMVision(VLLMClient)` + only, relying on structural/duck typing for registry lookup. Rejected + — the registry is currently structurally typed, but explicit ABC + conformance is cheap here (no diamond, no conflicting methods) and + makes the intent clear to readers. + +**2. `LLM.generate()` and `LLM.chat()` return `dict` (full OpenAI-compatible response body), not `str`.** +The original ABC typed both methods as `→ str`, which forced callers to +re-construct the surrounding OpenAI envelope when building RAG answers +(losing `model`, `usage`, `finish_reason`, etc.). The concrete vLLM +implementation already returned the full `httpx` JSON body; the `str` +annotation was aspirational, not real. +- Why: RAG answers are ultimately forwarded to the client in OpenAI format. + Stripping to plain text at the LLM boundary means the pipeline has to + re-wrap the content into `{"choices": [{"message": {"content": …}}]}` + further up — metadata (token counts, model id, stop reason) is lost in + the process. Returning `dict` preserves the full payload and keeps + back-ends interchangeable without wrapping shims. Using bare `dict` (not + a `TypedDict`) is a deliberate first step: it is backward-compatible with + all current callers and eases compat-shim re-exports while the + refactoring is still ongoing. +- Alternative considered: introduce typed response models (`ChatCompletion`, + `CompletionResponse`, `ChatCompletionChunk`) immediately. Rejected as + premature — Phase 6 adds the concrete client; Phase 10 (API layer + clean-up) is the right time to freeze the contract with typed models. + The `dict` annotation signals intent without coupling every caller to a + model definition that will evolve. +- `stream_chat` stays `AsyncIterator[str]` yielding raw SSE lines + (`data: {…}` strings). Parsing SSE chunks into typed dicts is Phase 10+ + work; the current shape keeps the streaming path consistent with + OpenAI's SDK behaviour. +- Future normalisation — when the typed models land, callers will migrate + to this pattern (TypedDict shown; Pydantic models are equally valid and + would expose `chat_content` / `completion_text` as properties instead): + +```python +from typing import TypedDict + +class _Message(TypedDict): + role: str + content: str + +class _Choice(TypedDict): + index: int + message: _Message # chat completions + finish_reason: str | None + +class _CompletionChoice(TypedDict): + index: int + text: str # text completions + finish_reason: str | None + +class _Usage(TypedDict): + prompt_tokens: int + completion_tokens: int + total_tokens: int + +class ChatCompletion(TypedDict): + id: str + object: str # "chat.completion" + model: str + choices: list[_Choice] + usage: _Usage + +class Completion(TypedDict): + id: str + object: str # "text_completion" + model: str + choices: list[_CompletionChoice] + usage: _Usage + +# Convenience extractors at the pipeline boundary: +def chat_content(resp: ChatCompletion) -> str: + return resp["choices"][0]["message"]["content"] + +def completion_text(resp: Completion) -> str: + return resp["choices"][0]["text"] +``` + + Until then, callers that need the text can use + `resp["choices"][0]["message"]["content"]` directly. + +--- + +## Phase 7A.1 — Connection manager + schema (2026-05-12) + +**1. Pulled the 7A.4 migration directory move forward into the 7A.1 commit set.** +The spec files the directory move (`scripts/migrations/alembic/` → +`services/persistence/migrations/`) and the env.py rewire under a separate +subsection (**7A.4 — Migrations**), distinct from 7A.1 (`connection.py` + +`schema.py`). The move was done in this commit set anyway. +- Why: The 7A.1 work creates `schema.py` whose entire purpose is to be + Alembic's metadata target. Leaving env.py pointing at the legacy + `components.indexer.vectordb.models.Base.metadata` would have created a + short-lived intermediate state where the two metadata definitions had to + stay byte-for-byte identical (or Alembic autogenerate would flag the + schema as drifted). Moving env.py at the same time avoids that risk + window — once schema.py exists, env.py points at it directly. +- Alternative considered: strict 7A.1-only — create `schema.py` but leave + the old `scripts/migrations/alembic/env.py` importing `Base.metadata` + until 7A.4. Rejected for the dual-source-of-truth risk above, and + because the migration move is a pure `git mv` with no code changes + beyond two import lines. + +**2. Extended `RDBConfig` with `database`, `pool_min_size`, `pool_max_size`, `command_timeout`.** +The spec's `ConnectionManager.__init__` pseudocode reads `config.database`, +`config.pool_min_size`, `config.pool_max_size` directly. None of those +fields existed on OpenRAG's `RDBConfig`. Added them (with defaults +`pool_min_size=5`, `pool_max_size=20`, `command_timeout=30`, `database=None`) +plus matching `POSTGRES_DATABASE` / `POSTGRES_POOL_{MIN,MAX}_SIZE` / +`POSTGRES_COMMAND_TIMEOUT` env-var mappings in `core/config/loader.py`. +- Why: 7A.1 doesn't compile otherwise. The spec's "files to create" table + lists only `connection.py` + `schema.py`, but the implementation it shows + has a hard config-shape dependency. Treated as required scaffolding for + 7A.1 rather than as a 7E (DI) concern, since the new fields belong on + the same config object that already carries `host`/`port`/`user`/`password`. +- Alternative considered: pass DSN + pool sizes as bare positional args + to `ConnectionManager.__init__`, leaving `RDBConfig` untouched. Rejected + — the spec's reference implementation accepts a `PostgresConfig` object + and the 7E DI wiring (`create_catalog_store(config)`) hands the whole + config in. Splitting fields across the call site and the config would + diverge from that contract. + +**3. `RDBConfig.database` stays optional; `ConnectionManager.__init__` raises if it's still `None`.** +The legacy code derives the Postgres database name from the Milvus +collection name (`f"partitions_for_collection_{collection_name}"`) at +`MilvusDB` actor startup. The new `RDBConfig` could either (a) require the +caller to set `database` explicitly, or (b) compute the name itself from +`VectorDBConfig.collection_name`. Chose (a) with a None default and a +constructor-time guard. +- Why: Crossing config sections (RDB reading VectorDB) would entangle two + otherwise independent config blocks and make `RDBConfig` non-portable. + The collection→database mapping is an integration concern that belongs + in the 7E DI wiring (`create_catalog_store` will build the database + name from `config.vectordb.collection_name` and inject it). The guard + in `ConnectionManager` makes the missing-database case fail loudly at + construction instead of silently producing a malformed DSN. +- Alternative considered: derive the database name inside + `RDBConfig.model_post_init` from a separately-injected collection name. + Rejected — adds two-way coupling between config sections for no gain; + 7E handles the wiring cleanly in one place. + +**4. Programmatic schema-vs-ORM diff used as the acceptance check.** +After rewriting all 7 tables as `sa.Table(...)` on a shared `MetaData`, +ran a column-by-column / index-by-index / constraint-by-constraint diff +against `components.indexer.vectordb.models.Base.metadata`. Empty diff = +passes. No assertion is shipped — this was a one-time verification, not +runtime behaviour. +- Why: Alembic autogenerate will treat any divergence between the new + metadata target and the live database (which was built from the legacy + ORM) as a pending schema change. The diff confirms that won't happen + on first run, and documents the methodology for the next migration: + any future schema change must update both `schema.py` and the legacy + `models.py` until Phase 9 deletes the latter. +- Alternative considered: ship the diff as a runtime test in 7F. + Deferred — 7F's repo tests already need a live Postgres; a metadata + diff doesn't need one and can live as a one-off check until the + legacy `models.py` goes away in Phase 9. + +--- + +## Phase 7A.2 — Repository implementations (2026-05-12) + +**1. Added a `WorkspaceRepository` port + `Workspace` domain model.** +The Phase-4 ports list had no workspace abstraction even though the +Phase-7 spec explicitly enumerates `workspace_repo.py` as one of the six +"real" repos with ten methods extracted from `PartitionFileManager`. +Added `core/ports/workspace_repo.py`, `core/models/workspace.py`, and +exposed `CatalogStore.workspace_repo`. +- Why: skipping it would leave the `workspaces` + `workspace_files` + tables strictly addressable only through the shim, defeating Phase + 8's point (orchestrators talk to ports, not the legacy actor). The + legacy code's workspace surface is non-trivial (orphan-file + detection, partition-scoped FK resolution) — it needs a first-class + port. +- Alternative considered: fold workspace methods into `DocumentRepository` + or `PartitionRepository`. Rejected — workspaces are a distinct + aggregate (their own row + join table) and conflating them blurs the + responsibility split that the rest of the ports layer enforces. + +**2. Extended `OIDCSession` domain model with three optional `bytes` fields +for the encrypted IdP tokens.** +The Phase-4 model had `id`, `session_token_hash`, `user_id`, `sid`, +`sub`, two timestamps, `last_refresh_at`, `revoked_at` — no token +fields. The port `create_session(session: OIDCSession) -> OIDCSession` +must convey what to store, so omitting them was a port-shape bug. +Added `id_token_encrypted`, `access_token_encrypted`, +`refresh_token_encrypted` as `bytes | None`. +- Why: the auth layer (Phase 6F) encrypts tokens before storage and + decrypts after read; the repo is intentionally byte-blind. Domain + models that hide load-bearing storage fields force callers to + bypass the port (e.g. with separate `set_tokens()` calls), defeating + the point of the typed contract. The "encrypted blobs flow through + the model verbatim" pattern is the same one the legacy + `_oidc_session_to_dict()` already uses. +- Alternative considered: an internal `OIDCSessionWithTokens` model + used only at the repo boundary. Rejected — duplicating models for + the sake of pretending the encrypted bytes aren't part of the + session is structural noise; Phase 8 callers don't read those + fields anyway. + +**3. Concrete repos expose two parallel surfaces: ABC methods + legacy +method names — both writing to the same rows.** +Each `PgRepository` implements the Phase-4 port ABC verbatim +(typed domain models in, typed domain models out) AND carries every +legacy `PartitionFileManager` method name with its original signature +and return shape (positional args, dict returns) as separate methods +marked `# TODO(phase-9): remove`. The legacy methods are NOT on the ABC. +- Why: Phase 8 orchestrators consume the port; Phase 7C shim must + delegate to the existing 76 call sites without rewriting them. A + single typed surface would force the shim to translate at every + call boundary, multiplying the change set and the regression + surface. Two surfaces on the same underlying SQL keeps both clients + happy with zero behavioural drift. Phase 9 deletes the legacy + surface after the shim goes away. +- Alternative considered: legacy method names only, postpone the + typed port to Phase 8. Rejected — the typed port is what makes the + ports/adapters split testable from `core/` unit tests; doing it now + is the cheap moment. + +**4. `users.password_hash`, `users.is_active`, `users.updated_at` exist on +the `User` domain model but not on the schema; mapped to defaults at the +repo boundary.** +The Phase-4 `User` model documents three auth modes (OIDC, token, +password+JWT) so it carries `password_hash`. The current schema only +supports OIDC + token. The user_repo silently drops `password_hash` on +write and synthesises `is_active=True` / `updated_at=created_at` on +read. +- Why: the alternative is dropping fields from the domain model, but + password auth is a planned post-refactoring feature on the roadmap. + Keeping the model field-complete means the orchestrator code can + be written once and only the repo updates when the column lands. +- Alternative considered: add the missing columns now via a new + Alembic migration. Rejected — Phase 7 is a structural refactor, not + a feature expansion. Adding columns expands scope past the spec. + +**5. Api-key repository methods raise `NotImplementedError`; stub repos +raise a `StubRepositoryError(NotImplementedError)` subclass.** +The `UserRepository.create_api_key` / `get_api_keys_by_prefix` / +`list_api_keys_for_user` / `delete_api_key` methods have no backing +table (`users.token` is a single-token field). Six entire ports +(`ChunkRepository`, `JobRepository`, `PromptRepository`, +`ConversationRepository`, `AuditLogRepository`, +`IdempotencyRepository`) plus four extras (`EntityRepository`, +`TopicTagRepository`, `ModelEndpointRepository`, `PresetRepository`) +are full-class stubs. +- Why: a silent fallback (empty list / `None`) is worse than an + exception — an orchestrator that retrieves zero rows from a "doesn't + exist yet" repo behaves indistinguishably from a real empty repo, + hiding bugs. The dedicated `StubRepositoryError` subclass is loudly + grep-findable (`grep -rn stub_not_implemented`) when the + post-refactoring features come online. +- Alternative considered: leave the stub ports unimplemented (no Pg + classes at all). Rejected — `CatalogStore` requires every port via + abstract properties; a partial implementation can't satisfy the + ABC, so Phase 7A.3 (composite store) wouldn't even instantiate. + +**6. Registered `json` / `jsonb` codecs on every connection via the +asyncpg `init` callback so reads return Python dicts.** +The legacy schema stores `files.file_metadata` as `JSON`; asyncpg +returns JSON columns as strings by default. Without the codec every +repo would have to `json.loads()` every row read and `json.dumps()` +every parameter write. +- Why: one place to register, repos stay focused on SQL. The codec + also covers `jsonb` so future migrations from `JSON` to `JSONB` + don't ripple into repo code. +- Alternative considered: per-call `json.loads()` / `json.dumps()` in + each repo. Rejected — duplicate boilerplate per repo, easy to forget + in one spot. The connection-level codec is the asyncpg-recommended + pattern. + +--- + +## Phase 7A.3 — PostgresStore composite (2026-05-13) + +**1. Placed `PostgresStore` in `services/storage/`, not alongside the +repositories in `services/persistence/`.** +The "storage" tier owns the high-level adapters that the rest of the +system depends on (`MilvusStore`, `PostgresStore`); "persistence" owns +the row-level repository implementations. Phase 8 orchestrators only +ever import from `storage` — they never reach into individual repo +modules. Keeping the composite outside `persistence/` makes that +boundary visible in the import graph. +- Why: the phase 7 plan explicitly names this split and it matches the + already-existing `services/storage/milvus_store.py` placeholder. A + future reader can tell the layers apart by directory. +- Alternative considered: `services/persistence/postgres_store.py`. + Rejected — would conflate the composite with its parts and force the + shim/orchestrators to import from `persistence/`, defeating the + point of the directory split. + +**2. Eagerly construct all fifteen repos in `__init__`; share one +`pool_getter` callable across them.** +Repos do not touch the pool until a query runs, so building them at +construction time is free and saves every caller from a lazy-init +dance. Passing a `_pool_getter` bound method (instead of the raw pool +reference) lets the store survive a `shutdown()`/`initialize()` cycle +in tests — repos always see the live pool. +- Why: matches the pool-getter pattern established in 7A.1 and keeps + test fixtures simple (rebuild the store, not the repos). +- Alternative considered: build repos lazily on first property access. + Rejected — extra branching with no measurable win, and the property + getter would lose its read-only character. + +**3. `initialize()` opens the pool *then* runs migrations; both behind +a single entry point.** +The legacy ORM bootstrapped tables synchronously via +`Base.metadata.create_all` before Alembic ever ran, which is why every +Phase 7 migration is idempotent (CLAUDE.md "Alembic Migration +Idempotency"). The composite keeps that ordering so the DI container +calls `await store.initialize()` once and gets a ready-to-query store. +A `run_migrations=False` flag is provided for fast unit tests against +an already-migrated database. +- Why: hiding the two-step lifecycle behind one method keeps the + Phase 7E container wiring identical to the inference adapters and + matches the `CatalogStore` ABC contract (one `initialize`, one + `shutdown`). +- Alternative considered: expose `run_migrations()` separately and + require the container to call both. Rejected — leaks an + implementation detail across the layer boundary and makes every + composition-root harder to write. + +**4. Expose the raw asyncpg pool as a `pool` property on the concrete +class, not on the `CatalogStore` ABC.** +Phase 8 orchestrators need cross-repo transactions +(`async with store.pool.acquire() as conn: async with conn.transaction(): ...`). +That capability is not on the ABC because most consumers do single-repo +calls; adding it to the port would invite leaks of asyncpg specifics +into orchestrator code that doesn't need them. Concrete clients that +truly need transactional escape hatches can depend on `PostgresStore` +directly. +- Why: keeps the ABC minimal while still unlocking the transaction + pattern. Phase 8 will decide whether to formalise a + `UnitOfWork`-style port; until then the escape hatch is explicit and + grep-findable (`grep -rn "store.pool"`). +- Alternative considered: add a `pool` property to `CatalogStore`. + Rejected — turns the ABC into an asyncpg-shaped interface and makes + it harder to swap in a non-Postgres backend (e.g. SQLite for tests). + +--- + +## Phase 7B — Milvus vector store (2026-05-12) + +**1. `hybrid_search(embedding, query_text, …)` is on the `VectorStore` ABC alongside `search`. Initial draft kept it Milvus-specific; reversed before the contract surface settled.** +The Phase 7 plan's `_hybrid_search` example (STRATEGY-adjacent doc `phase 7.md` lines 275–298) references an undeclared `query_text` argument — the doc tacitly admits the embedding-only `search` cannot drive Milvus 2.6's native BM25. Milvus's `Function(FunctionType.BM25)` computes the sparse vector server-side from the `text` field at both insert and query time, so the hybrid path *requires* the raw query text — not a pre-computed sparse vector and not an embedding. +- Why on the ABC: every realistic backend in the SaaS shortlist has a hybrid lexical-plus-dense path (Qdrant via sparse-vector points + fusion, Weaviate hybrid BM25+vector, Pinecone sparse-dense vectors, OpenSearch knn+match). Treating hybrid as a Milvus-only quirk would push an `isinstance(store, MilvusVectorStore)` branch straight to the contract boundary inside the Phase-8 retrieval orchestrator — exactly the leak the narrow store was supposed to prevent. `query_text` is wider than strictly needed for backends that accept a pre-computed sparse vector, but it is the only shape that fits Milvus's server-side BM25 without client-side tokenization, which would force a reindex and contradicts the spec's "Critical: preserve OpenRAG's current Milvus schema". +- Initial draft and reversal: the first pass put `hybrid_search` on `MilvusVectorStore` only and kept the ABC embedding-only — *"BM25 is a Milvus-specific implementation detail; bleeding it into the cross-store contract for one backend's quirk is the wrong direction."* Reversed once the SaaS-end-state backend list made hybrid look like a contract feature rather than a Milvus quirk. Backends without a hybrid path can raise `VDBSearchError` at call time — same shape `MilvusVectorStore` already uses when its backing collection was constructed with `hybrid_search=False`. +- Alternative considered: (a) pre-compute sparse vectors client-side via a tokenizer/IDF table — rejected, abandons Milvus's native BM25 (server-side analyzer + stop words) and would force a reindex; (b) `query_text: str | None = None` added to `search` itself (one method, optional arg) — rejected, conflates dense-only and hybrid call sites in the same method and forces every implementation to inspect both args. + +**2. Embedding dimension comes in via a lazy `await store.initialize(dim)`, not a constructor arg or a config field.** +The Phase 7 plan's example constructor is `MilvusVectorStore(config: MilvusConfig)` and silently elides where `dim` comes from — the schema needs the dim, but the dim lives on the embedder. +- Why: Mirrors `PostgresStore.initialize()` shape so DI wiring (Phase 7E) has one consistent "construct cheap, materialise async" pattern across both stores. Construction stays I/O-free and embedder-free; the DI container resolves the embedder, reads `embedding_dimension`, and passes it to `initialize()`. Idempotent + double-checked-locked so concurrent first-callers don't race. +- Alternative considered: (a) explicit constructor arg `MilvusVectorStore(config, embedding_dimension=…)` — cleaner dependency but forces every test/composition root to resolve the embedder first; (b) `VectorDBConfig.embedding_dimension` — duplicates the embedder's `EmbedderConfig.embedding_dimension` value across two configs, drift risk. + +**3. ABC `collection` arg = Milvus collection name (not partition row-value). Partition lives only in `filters["partition"]`. Added a Milvus-specific `delete_by_filter(filters)`.** +The Phase 7 plan (`phase 7.md` line 300) explicitly maps the ABC's `collection` argument to "the partition row-value", and proposes `drop_collection(name)` deleting rows where `partition == name`. The same word would then mean two different things across the codebase — Milvus's own vocabulary keeps *collection* (top-level container) and *partition* (row tag via `partition_key`) strictly distinct. +- Why: The end-state of this refactor is a multi-tenant SaaS product where each client gets its own Milvus collection (see [[project-saas-collection-per-tenant]] memory). In that world the ABC's `collection` arg is a real per-tenant Milvus collection name; conflating it with partition values would paint the future store-factory/pool into a corner. Strict separation makes the SaaS path a Phase-8+ wrapping layer ("`client_id → MilvusVectorStore`") on top of an unchanged narrow store. +- Concrete shape: `MilvusVectorStore._resolve_collection(name)` accepts only `self._collection_name` or the ABC sentinel `"default"`; anything else raises `ValueError`. `drop_collection(name)` drops the whole Milvus collection (admin/test). Partition-level row deletion (used by the 7C shim's `delete_partition`) goes through a new Milvus-specific public method `delete_by_filter(filters)`, with an explicit guard that refuses empty/wildcard expressions so a typo cannot nuke the entire collection. +- Alternative considered: (a) spec-faithful overload — rejected, conflates two distinct Milvus concepts in code that has to survive the SaaS pivot; (b) ignore the `collection` arg entirely instead of validating — rejected, silently accepting wrong names is the same forward-compat hazard. + +**4. No manual reconnect / retry logic against Milvus — trust pymilvus + gRPC internals.** +The Phase 7 plan's design note recommends `MilvusVectorStore` carry its own retry/reconnect logic, citing `ConnectionNotExistException` and double-checked locking in `_ensure_loaded()`. +- Why: That guidance comes from the pre-2.4 ORM-style `connections.connect(alias=…)` API where named aliases needed explicit re-establishment. Pymilvus 2.6's `MilvusClient(uri=…)` / `AsyncMilvusClient(uri=…)` — per [v2.6.x API reference](https://milvus.io/api-reference/pymilvus/v2.6.x/MilvusClient/Client/MilvusClient.md) and [AsyncMilvusClient v2.6.x](https://milvus.io/api-reference/pymilvus/v2.6.x/MilvusClient/Client/AsyncMilvusClient.md) — expose **no** public retry / reconnect / keepalive knobs and own their gRPC channel internally. The legacy `MilvusDB` does no manual reconnect either. Reintroducing client-side teardown-and-recreate logic risks racing pymilvus's internal channel state for no documented benefit. Documented inline in `MilvusVectorStore.__init__` so the plan's note doesn't get reintroduced later without evidence. +- Alternative considered: (a) lightweight retry without client recreation (sleep + retry-once on connection-error message match) — rejected, no documented gRPC-level guarantee that a fresh call sees a healed channel any sooner than gRPC's own backoff; (b) full client teardown + recreate with double-checked locking — drafted, then dropped after reading the pymilvus 2.6 reference; pure complexity for an unproven failure mode. + +**5. Store surface kept narrow: `VectorStore` ABC + Milvus-only `delete_by_filter`. File/chunk conveniences are 7C shim's job.** +The legacy `MilvusDB` exposes `get_file_chunks`, `get_chunk_by_id`, `get_file_chunk_ids`, `list_all_chunks`, `get_related_chunks`, `get_ancestor_chunks`, `get_surrounding_chunks` — all file-scoped or relationship-scoped reads. +- Why: Each of those is either (a) a thin wrapper over `query_chunks_by_filter` (file-scoped reads) or (b) domain logic that belongs in `core/retrieval/hydration.py` per the spec (surrounding/related/ancestor chunks). Putting them on the store widens the surface only to delete them again in Phase 8. The 7C shim builds the file-scoped variants from `query_chunks_by_filter` (two RPCs vs one — accepted cost for a narrow ABC-aligned store). +- Alternative considered: add the convenience methods directly on the store. Easier 7C shim (one-line delegate per method) but a wider surface to maintain and to migrate again in Phase 8. Rejected. + +--- + +## Phase 7A.4 — Unified migrations namespace (2026-05-12) + +**1. Nest Postgres alembic and Milvus migrations as siblings under `services/persistence/migrations/{alembic,milvus}/`, not as two parallel roots.** +Person A's pulled-forward 7A.4 (commit `91f7078`, logged in [Phase 7A.1 §1](#phase-7a1--connection-manager--schema-2026-05-12)) placed alembic directly at `services/persistence/migrations/`. The phase 7 plan is silent on Milvus migrations — they are an OpenRAG-specific addition (Milvus schema-version property + generic runner under `openrag/scripts/migrations/milvus/`) not contemplated by the spec. The first take on the Milvus side put them at `services/storage/migrations/` next to `milvus_store.py` (commits `b40de00` + `26311f0`, since reset out of history); reshuffled to a unified namespace before push. +- Why: One root with backend subdirs ("where do I run migrations?" → `services/persistence/migrations/`, "for which backend?" → `alembic/` or `milvus/`) is easier to reason about than two roots that split "schema evolution" across two services-layer namespaces purely because their adapters happen to live in different sub-folders. The unified shape also leaves room for additional backends (S3 lifecycle, future tenancy stores) without re-litigating where migrations live each time. Aligns with the SaaS end-state ([[project-saas-collection-per-tenant]] memory) where per-tenant collection lifecycle and per-tenant schema versioning will both grow inside this same namespace. +- Alternative considered: (a) spec-plus-by-symmetry layout — alembic under `services/persistence/migrations/`, Milvus under `services/storage/migrations/`. Rejected: makes the storage layer carry a `migrations/` peer to `milvus_store.py`, conflating "the adapter" with "the adapter's schema-evolution scripts", and forces every reader to know which backend hides where. (b) flatten everything under one dir without backend subdirs. Rejected: alembic's filename conventions (`_.py`) and the Milvus runner's `N.description.py` discovery rules would step on each other. + +**2. Milvus migration runner imports `SCHEMA_VERSION_PROPERTY_KEY` from `services.storage.milvus_store`, not from `components.indexer.vectordb.vectordb`.** +The constant exists identically in both modules right now (the new store copied it byte-for-byte from the legacy MilvusDB during 7B). Either import resolves; we picked the new home. +- Why: The migration set, the constant, and the new adapter all live under `services/` after this move. Importing from the legacy module would create a backwards dependency from the new namespace into the deprecation-path module, pinning `components.indexer.vectordb.vectordb` alive past Phase 9's planned deletion. Updating the import now is a zero-risk follow-up to the move (legacy still defines the constant with the same value, so behaviour is identical) and makes the legacy deletion a single grep-and-delete in Phase 9 rather than "deletion + N migrate.py import updates". +- Alternative considered: leave the import on the legacy module until Phase 9 forces the issue. Rejected — no benefit, and `9` is the wrong phase to be hunting incidental imports. + +--- + +## Phase 7E — DI wiring (2026-05-13) + +**1. Made `ServiceContainer(settings=None)` optional so the pre-Phase-7E +test paths keep working.** +The container's pre-existing job was to populate inference registries +(`ServiceContainer()` with no arguments). Phase 7E adds storage adapter +wiring that needs a `Settings` instance, but rewriting every legacy +test to pass settings is busy-work and risks scope creep. The settings +argument is therefore optional; the storage accessors raise a clear +`RuntimeError` ("ServiceContainer was constructed without a Settings +instance — pass Settings to wire storage adapters") when reached +without one. +- Why: keeps the Phase 7E commit a strict superset of the previous + container behaviour and surfaces the misuse with a message that + points at the fix. +- Alternative considered: split into `ServiceContainer` (registries + only) and `AppContainer(settings)` (registries + storage). Rejected + — Phase 8 orchestrators will pull both layers from one container, + and a hard split now would invite a refactor at the very next phase. + +**2. Centralised the "database name from collection name" idiom in +`create_catalog_store`.** +The legacy `MilvusDB.__init__` derives the Postgres database name from +the Milvus collection name (`partitions_for_collection_`, +vectordb.py:238). That contract is duplicated in `scripts/backup.py`, +`scripts/restore.py`, `scripts/check_file_counts.py`, and the new +Alembic `env.py`. The Phase 7E factory keeps the fallback in one place +so DI wiring code never mentions the prefix; an explicit +`rdb.database` still wins. +- Why: the database name resolution is policy, not orchestrator code + — putting it in the factory keeps `di/container.py` mechanical and + prevents the prefix from drifting into half the call sites. +- Alternative considered: resolve in `PostgresStore.__init__`. + Rejected — pushes a Settings dependency down into the store, which + is happy taking just `RDBConfig` today and shouldn't have to grow a + `vectordb` parameter to learn the collection name. + +**3. `create_vector_store` wires to the real `MilvusVectorStore` (Phase 7B) +in this same branch.** +The factory was first drafted as a fail-loud stub (`raise +NotImplementedError("Phase 7B deliverable")`) so 7E could land +independently of 7B. With 7B already on the branch when 7E was +rebased, the factory's body swapped to its final shape: `return +MilvusVectorStore(settings.vectordb)`. Construction stays I/O-free +and embedder-free; the composition root resolves the embedder and +calls `await store.initialize(embedding_dimension)` later, matching +the lifecycle pattern :class:`PostgresStore` already follows. +- Why: keeps the DI factory shape stable across the 7B handoff — + no caller has to know whether 7B was merged first. The lazy + `initialize(dim)` keeps the construction step out of any async + context and avoids a circular dependency on the embedder factory. +- Alternative considered: pass the embedder dimension into the + factory itself. Rejected — forces every test/composition root to + resolve the embedder before the store, defeating the construct- + cheap / materialise-async split that the Phase 7B store relies on. + +**4. Aligned `services/persistence/` and `services/storage/` imports +to the project's short-form convention (`from core.X`, not +`from openrag.core.X`).** +Pytest sets `pythonpath = ./openrag`, so `openrag/core/foo.py` is +importable as both `core.foo` and `openrag.core.foo` (the editable +install also exposes the `openrag` package). Python treats those as +two distinct modules, so a class defined once but imported via both +paths fails `isinstance` checks — which is exactly how the Phase 7E +container test first surfaced the dual-import bug +(`PgDocumentRepository` not isinstance `DocumentRepository`). Picking +one convention everywhere fixes the bug, and `CLAUDE.md` already +mandates the short form (`from components.ray_utils import ...`). +- Why: matches the rest of the codebase and removes a class of + isinstance bugs that would otherwise dog every Phase 8 orchestrator + test. +- Alternative considered: leave the `openrag.X` prefix in place and + ban the short form in tests. Rejected — the short form is already + load-bearing in `di/`, `core/llm/`, `core/embeddings/`, and the + Phase 6 inference adapters; changing all of those to the long form + would be a much larger and riskier sed pass. + +--- + +## Phase 7F — Integration tests (2026-05-13) + +**1. Integration tests land at `tests/integration/`, not colocated with the SUT and not under `tests/api_tests/`.** +The project today has tests in two places: colocated `openrag/**/test_*.py` (the long-standing convention) and `tests/api_tests/` (HTTP-style black-box tests against a running OpenRAG server). STRATEGY §13C describes a target end-state of `tests/{unit,integration,load}/` with everything under one root. 7F lands integration tests for both the asyncpg repos and `MilvusVectorStore`; the unified layout doesn't arrive until Phase 13C. +- Why: Integration tests for both stores need to escape `pytest.ini`'s `testpaths = openrag` so a bare `uv run pytest` (the default unit run) does not drive real infrastructure. The cleanest interim home is exactly where Phase 13C will park them anyway — `tests/integration/` — so the file lands once and does not need to move during the sweep. Colocating per-repo tests under `openrag/services/persistence/test_*_repo.py` would also force every unit-test invocation to either skip (silently masking failures) or fail (in CI without infra) depending on how the skip is wired, neither of which is a good default. +- Alternative considered: (a) park integration tests under `tests/api_tests/` to match the existing infra-test sibling layout. Rejected — `api_tests/` is HTTP-API-style by convention (httpx against a running server); adapter-level integration tests don't fit that mould, and the strategy doc has already settled the end-state location. (b) keep them colocated and rely on the `integration` pytest marker for deselection. Rejected — relies on every CI invocation remembering `-m "not integration"` and still pulls pymilvus/Postgres imports into the default unit run. + +**2. One ephemeral test database per session, truncate between tests.** +The fixture creates `openrag_phase7_test` on session start with a clean +schema, runs migrations once via `PostgresStore.initialize()`, and lets +every test share the same store. An autouse fixture truncates the seven +user-modifiable tables with `RESTART IDENTITY CASCADE` so each test +starts with primary keys at 1. +- Why: dropping/recreating per test would multiply the session cost by + the test count; a per-test truncate is microseconds. The fresh-DB- + per-session contract gives us migration coverage for free (any + Alembic regression surfaces at fixture setup, not in production). +- Alternative considered: rollback-per-test via savepoints. Rejected + — asyncpg pools cycle connections, so a top-level rollback only + cleans the one connection used; the next test could see partial + state on a different connection. Truncate is straightforward and + matches what Phase 8 orchestrator tests will want. + +**3. Session-scoped loop and fixtures (`loop_scope="session"`).** +`pytest-asyncio`'s default `function`-scope event loop conflicts with a +session-scoped async fixture: the asyncpg pool binds to the loop it was +created on, and each test then gets a different loop. Marking both the +fixtures and the tests with `loop_scope="session"` keeps everything on +one loop, which is also what asyncpg expects. +- Why: a session-scoped pool is the whole reason to run an integration + suite — function-scoped pools would defeat the purpose. The + `loop_scope` knob is the official pytest-asyncio 1.x API for this. +- Alternative considered: function-scoped pools (one per test). + Rejected — every test would pay the connection-establishment cost, + inflating the suite from ~5s to easily 30s+. + +**4. `tests/integration/test_stores.py` is `xfail(strict=True)` until the +cross-store fixture lands in Phase 7C.** +The Phase 7F plan calls for a cross-store full-cycle test +(`create partition → upsert pre-embedded chunks → search → delete`). +:class:`MilvusVectorStore` is in place (Phase 7B) and reachable through +DI, but the existing `postgres_store` fixture does not yet hand out a +companion Milvus store on the same collection — the combined fixture +lands as part of Phase 7C (shim) so the assertion matches the legacy +flow byte-for-byte. The strict marker means the day the fixture is +wired and the body filled in, the unintended pass trips a clear +failure rather than silently turning green. +- Why: the file exists so the eventual diff is just a fixture + body + replacement. Strict xfail prevents "forgot to remove xfail" rot. +- Alternative considered: skip with `pytest.skip(...)` or leave the + file out entirely. Rejected — skips are too easy to forget, and an + absent file means the cross-store test has to remember to be + created. + +**5. Fixed two bugs surfaced by writing the tests, in the same diff.** +The integration suite caught two tz-naive/tz-aware mismatches in code +written during 7A.2: + +* `PgOIDCSessionRepository.delete_expired` constructed + `datetime.now(UTC) - _EXPIRED_RETENTION` (tz-aware) and bound it + against the tz-naive `session_expires_at` column. asyncpg refuses + the cast at bind time. Fix: strip `tzinfo` from the cutoff before + the query. +* `services/persistence/migrations/alembic/env.py` unconditionally clobbered + `sqlalchemy.url` with one derived from `load_config()`, ignoring + the DSN that `ConnectionManager.run_migrations()` had just set — + which is why the first run of the suite tried to resolve the docker + hostname `rdb` from a host process. Fix: defer to the preset URL + unless it's the alembic.ini placeholder. + +Both fixes are tiny but load-bearing for any non-default deployment +(test DBs, named environments, CI overrides). I kept them in this +commit rather than splitting them out — they only show up against a +real database, and splitting would mean landing the new tests broken. +- Why: the alternative (separate fix commits) would have the + integration tests fail on first introduction, which is a bad signal + in `git log`. +- Alternative considered: hard-code a workaround in the test fixture. + Rejected — the underlying bugs would still be there waiting for + Phase 8 orchestrator tests to hit them. + +--- + +## Phase 7C — God-object shim (2026-05-13) + +**1. Kept the "create PG database if missing" bootstrap in the shim, not +in `PostgresStore`.** +The legacy `PartitionFileManager.__init__` auto-created the per-collection +database via `sqlalchemy_utils.create_database` (`utils.py:39-40`) — a +side-effect callers depended on without ever opting in. The Phase 7C +replay needed equivalent behaviour, and the natural architectural home +is `PostgresStore.initialize()` (symmetric with the Alembic-migration +step it already runs, and with `MilvusVectorStore.initialize()` which +materialises its collection). It was instead placed on the shim as +`MilvusDB._ensure_pg_database`, called from `MilvusDB.initialize()` +between the vector-store and catalog-store bootstrap steps. +- Why: Phase 7C's scope contract is "only `components/indexer/vectordb/vectordb.py` + is modified". Pushing the helper down into `PostgresStore` would have + expanded the blast radius into 7A.3 territory while the shim was still + the only caller. Punting the design to Phase 7E (DI wiring, where the + `create_catalog_store` factory already centralises the database-name + derivation) means one place owns both naming and provisioning once + the shim is gone. +- Alternative considered: move it to `PostgresStore.initialize()` now + (cleanest layering). Rejected for the scope reason above — but **flag + for 7E/8 follow-up**: when DI wiring lands the store-creation path + for orchestrators, the database-existence check should move into + `PostgresStore.initialize()` (or `ConnectionManager.initialize()`) + so the shim's `_ensure_pg_database` helper can be deleted in Phase 9 + without leaving a regression behind. Track this together with the + `scripts/*.py` callers that already duplicate the + `partitions_for_collection_` derivation (see Phase 7E entry 2). + +--- +## Phase 8 — Orchestrators (2026-05-19) + +**1. Provider sourcing: routers read `request.app.state.container`.** +- Why: the 8 Phase-8 doc shows providers as one-line accessors; the + ServiceContainer is not attached to the live app until Phase 11. We + attach it best-effort in `main.py` now (not `initialize()`d), so the + thinned routers resolve services but the DB-backed flows stay dormant + until Phase 11. Token-mode auth routes already 400 before any service. +- Alternative considered: a lazy module-level container singleton in + `di/providers.py` (mirrors `components/auth/deps.py`). Rejected — the + one-liner/app.state shape matches the doc and Phase-11 cutover. + +**2. OIDCConfig built from env in the container, not wired into Settings.** +- Why: keeps 8A.1 scoped; `OIDCConfig` exists but is not in root + `Settings`. Added a missing `auto_provision_login` field. +- Alternative: wire it into `Settings` now — deferred (config refactor). + +**3. Thin routers: typed core exceptions propagate to the global +`OpenRAGError` handler; `HTTPException` kept only where the legacy body +must stay byte-identical** (id==1 / 409-exists / file|workspace 404s). +- Why: the legacy `_check_*` guards already surfaced 404s via the global + handler, so propagating `UserNotFoundError`/`PartitionNotFoundError`/ + `ValidationError` reproduces them; only the non-bracketed + `{"detail": ...}` `HTTPException` bodies need to stay verbatim. + +**4. Constructor extensions over the plan's prescribed signatures.** +- Why: to preserve legacy behaviour without reaching into Ray/config, + orchestrators take extra container-supplied args: `collection` + (`settings.vectordb.collection_name` — Partition/Workspace/Retrieval), + `user_repo` (PartitionService, to reproduce `VDBUserNotFound`), + `default_file_quota` (UserService), and orchestrator→orchestrator + injection (UserService←PartitionService for the delete-user cascade). +- Alternative: hold the plan signatures exactly — rejected, would drop + behaviour (cascade) or smuggle config/Ray into the service. + +**5. 8A.2 delete-user owner-partition cascade deferred, then restored in +8B.1.** UserService.delete_user was a plain repo delete in 8A.2 (gap +logged); 8B.1 composes PartitionService so it deletes owned partitions +(vectors + rows) first, matching the legacy Ray `delete_user`. + +**6. 8C searcher backing: inject the Ray-backed `MilvusRayShim` behind +the `RetrievalSearcher` port during the Phase-8 shim.** +- Why: the only `RetrievalSearcher` impl is `MilvusRayShim` (Ray + `Vectordb` actor — embeds + hybrid-searches internally). The 8C plan + text says "retriever calls `vector_store.search()` directly (no Ray)", + but the dev-workflow doc puts Ray cleanup in Phase 9 and allows + orchestrators to call Ray *behind a port* during the shim. Injecting + the shim keeps the orchestrator file Ray-free (8H satisfied: no + Ray remote-call, no Ray import — Ray is behind the port) and avoids a + risky reimplementation of hybrid-search / surrounding-chunks / + similarity behaviour in the hot search path. +- Alternative considered: write a new `VectorStoreSearcher` on the clean + `VectorStore` port now (embed via Embedder factory + `search`). + Rejected for this phase — real regression surface in core search; the + clean adapter lands in Phase 9 when the shim is deleted. **Flag for + Phase 9:** add `VectorStoreSearcher`, swap the container wiring, delete + `MilvusRayShim`. +- Consequence: `RetrievalService.__init__` deviates from the plan's + `(vector_store, embedder_factory, reranker_factory, llm_factory, + document_repo, config)` — with the shim searcher those are unused; + it takes the built `searcher` / `reranker` / `llm` + `config`. + +**7. 8C.2 structured output: core LLM + JSON-mode prompt + parse (no +LangChain).** `RagPipeline.generate_query` and `map_reduce` used a +LangChain structured-output chain for `SearchQueries` / `SummarizedChunk`. +QueryService instead calls the injected core `LLM.chat` with a +JSON-instructed prompt + `response_format={"type":"json_object"}` and +`json.loads` (+ `_json_slice` brace-extraction) into the Pydantic model, +preserving the legacy fallbacks (query-gen: retry once → raw user query; +map-reduce: relevancy=False on any parse error). +- Why: 8H bans LangChain in orchestrators; the plan's explicit "remove + ChatOpenAI — use LLM factory" intent. +- Alternative: wrap the LangChain chain behind a core-LLM-shaped helper + outside orchestrators. Rejected — only partially meets the intent and + keeps a LangChain dependency on the hot path. + +**8. 8C.2 streaming + citations live in QueryService; the router is pure +transport.** `chat_stream` drives the proven +`components.utils.stream_with_source_filtering` (reused verbatim — the +100-char buffer that strips `[Sources: N]` mid-stream is delicate); +`chat`/`complete` return the finalized OpenAI dict with the +citation-filtered `extra`. The router maps partition, wraps +`StreamingResponse`/`JSONResponse`, and passes a request-bound +`prepare_sources` callable so `request.url_for` stays in transport. +- Why: the plan's "service owns streaming" Q&A; keeps the proven SSE + buffer logic intact (no rewrite) while ownership moves to the service. +- Constructor takes `workspace_service` beyond the plan's four (the + legacy `_prepare_for_chat_completion` validated the workspace via the + Ray actor; reusing WorkspaceService.get_workspace keeps it Ray-free). +- Consequence: legacy `components/pipeline.py` (RagPipeline / + RetrieverPipeline) + `map_reduce.py` + `components/retriever.py` are + now dead code (no router imports `RagPipeline`); **flag for Phase 12 + cleanup** to delete them with the other shims. + +**9. 8D.1 indexing dispatch sits behind an `IndexingDispatcher` port + +Ray shim, not direct Ray calls in the service.** The plan's +`IndexingService.__init__(document_repo, workspace_repo, vector_store, +config)` can't reach the `Indexer` / `TaskStateManager` actors. Mirroring +the proven 8C searcher pattern, a new `core/indexing/dispatcher.py` ABC +defines the worker operations the service needs and +`services/storage/indexer_ray_shim.py` adapts the two Ray actors to it; +the container injects it via `from_ray_namespace()` (lazy, like +RetrievalService). `vector_store` and `config` are dropped from the ctor +(the file delete is owned by the Indexer worker; the only config the +legacy router used — data dir, vectordb timeout — is a transport/shim +concern); `dispatcher` is the added arg. +- Why: 8H bans Ray remote calls under `services/orchestrators/` (only + JobService is excepted); the established way to keep an orchestrator + Ray-free during the shim is a core port + `services/storage/` shim. +- Alternative: leave the dispatch in the thin router. Rejected — 8G/8D.1 + explicitly move "inline upload+dispatch → indexing_service.add_file()". +- File save to disk + the byte-identical `HTTPException` guards (409 + exists, 404 not-found, 400 bad workspace_ids, 404 unknown workspace, + 404 no object ref) stay in the router (transport + exact legacy body), + matching the workspaces.py thinning style. +- **Flag for Phase 9:** delete `indexer_ray_shim.py`, have IndexingService + call the pipeline stages directly. + +**10. 8D.2 JobService keeps Ray `.remote` calls (the one 8H exception).** +Per the plan ctor `JobService(task_state_manager)` and the 8H carve-out, +JobService wraps the `TaskStateManager` actor directly — no shim. The +status rollup and `?task_status=` filtering (business logic) move into +the service; `request.url_for` link building stays in the thin queue +router. Container resolves the actor lazily in the cached property +(deferred `utils.dependencies` import) so it is only needed at first +request. +- Why: introducing a shim here would contradict the plan and the + explicit 8H exception; JobService is the documented hook point for the + post-refactor DB-backed job tracking (Phase 9). + +**11. 8D.1 reuses `components.indexer.utils.files.extract_temporal_fields` +as-is, carrying a transitive module-level `load_config()` and an +`HTTPException`.** IndexingService imports that helper for ISO-8601 +metadata parsing. Its module body runs `config = load_config()` at +import time, and the helper raises FastAPI `HTTPException` on a bad +datetime — a transport type leaking into a service. +- Why: it is a pure parsing helper under the shim-exempt `components` + layer (8H greps the orchestrator file itself, which stays clean), and + reusing it verbatim keeps the bad-datetime 400 response byte-identical + with the legacy router. Rewriting it now would risk behavioural drift + for no Phase-8 benefit. +- Alternative: reimplement the parse in `core/` raising `ValidationError` + and map it in the router. Rejected for Phase 8 (byte-identical bodies + are the priority); it is the right move once the shim is gone. +- **Flag for Phase 9:** when `indexer_ray_shim.py` is deleted, move the + temporal-field parsing into core (raise `ValidationError`, drop the + module-level `load_config` and the `HTTPException`). + +**12. 8E ConversionService — serializer behind a `FileSerializer` port + +shim; chunk lookup through the clean `VectorStore`.** Same pattern as +8C/8D.1: new `core/indexing/serializer.py` ABC + `services/storage/ +serializer_ray_shim.py` wrapping the `DocSerializer` actor (it reuses +the legacy `components...serialize_file` helper so the +`call_ray_actor_with_timeout` behaviour the tools router relied on is +preserved). `get_chunk` ports the legacy `get_chunk_by_id` onto +`VectorStore.query_chunks_by_filter({"_id": int(chunk_id)})` and returns +a plain `{"page_content", "metadata"}` dict (no LangChain `Document`), +exactly as PartitionService reads chunks. Ctor deviates from the plan's +`ConversionService(config=config)` → `(serializer, vector_store, +collection)` (the `collection` extra is the established +settings-supplied vector-store name); `config` dropped. +- Why: the plan ctor is underspecified and the 8C/8D shim+port approach + is the established way to keep the orchestrator Ray-free (8H). +- File save + cleanup IO, tool dispatch, and the byte-identical + `HTTPException` bodies (404 not-found, 403 forbidden, the tool-error + 4xx/5xx mapping) stay in the thin routers. +- Note: the legacy router's `task_id` came from + `ray.get_runtime_context().get_task_id()`; outside a Ray task that is + the driver context, so the shim passes a fixed `"tools-extract"` + fallback label (the serializer only uses it for logging / state keys — + no behavioural change). +- **Flag for Phase 9:** delete `serializer_ray_shim.py`, have + ConversionService call the serializer stage directly. + +**13. 8F keeps the lazy-property container (no eager ordered `__init__`); +8H closeout pulls the last Ray call out of a thinned router.** The +plan's 8F snippet builds every orchestrator eagerly in +`ServiceContainer.__init__` in a hand-ordered block. We did *not* adopt +that — the lazy cached-property convention set in 8A.1 (decision 1) +already satisfies "correct instantiation order": dependency resolution +happens on first access (`user_service` → `auth_service` / +`partition_service` / `job_service`; `query_service` → +`retrieval_service` / `workspace_service`), there are no cycles, and the +Ray shims stay un-touched until first request. 8F is therefore a +verification pass, not a rewrite: all nine orchestrators are wired +identically (9 `_x_service` cache slots, 9 properties, 9 providers), +covered by a new `TestPhase8OrchestratorWiring` in `di/test_container.py`. +- During the 8H sweep, `routers/users.py` `/users/info` was found still + calling `task_state_manager.get_user_pending_task_count` directly (a + Ray `.remote` + quota math left in a thinned router by 8A.2). It + passed the literal 8H greps (`vectordb` only) but contradicted "routers + are thin". Fixed per the plan's "Day 3: fix remaining router rewrites": + added `JobService.get_user_pending_task_count`, injected `JobService` + into `UserService` (orchestrator-to-orchestrator, same pattern as + UserService←PartitionService), moved the quota-usage computation into + `UserService.get_current_user_info` (byte-identical response), thinned + the handler. UserService stays Ray-free — the count goes through + JobService, the one 8H-excepted Ray wrapper. +- `routers/actors.py` still imports `utils.dependencies` Ray handles — + intentionally out of scope: it is the Ray actor/health admin router, + not one of the 12 fat routers in the 8G table; it belongs to Phase 9. +- 8H #7 ("no router > 120 lines") is not literally met (indexer 516, + partition 453, …) and is treated as aspirational, consistent with + every prior slice: the verbose OpenAPI `description=` docstrings + dominate the line count; the business-logic extraction — the actual + goal — is complete and grep-verified (#1–#6 clean). + +**14. Phase-8 CI fixes — the §1 "DB-backed flows dormant until Phase 11" +deferral was wrong and is corrected here.** The unit + API CI jobs were +red after Phase 8; four root causes, all behaviour-preserving fixes: +- *Container never initialised (API tests, ~80 failures).* The thinned + routers resolve repos through the container's *own* `PostgresStore`, + a separate instance from the one the legacy Ray `Vectordb` actor owns + and `initialize()`s (`vectordb.py:188`). §1 attached the container but + never opened its asyncpg pool, so every catalog-backed route 500'd + against an uninitialised pool. Fix: `main.py` startup/shutdown hooks + call `container.initialize()/shutdown()` (best-effort, guarded). The + asyncpg layer is idempotent (Phase 7), so a second store against the + same DB alongside the actor's is safe. This pulls a thin slice of + Phase 11 forward — the original deferral broke the live app, which + Phase 8 must not. Phase 11 still folds this into a real lifespan. +- *Auth router 503 instead of 400 in token mode (3 unit failures).* + FastAPI resolves `Depends` in declaration order; `Depends(get_auth_service)` + ran (and 503'd on the absent container) before the in-body + `_require_oidc_mode()` 400 gate. Fix: `_require_oidc_mode` is now a + `Depends` declared *before* the service on login/callback/ + backchannel-logout/logout, so token mode 400s without touching the + container. No behaviour change in oidc mode. +- *`components/utils.py:get_num_tokens` needs an OpenAI key (1 unit + failure, latent keyless-deploy defect).* It built `ChatOpenAI(...)` + just to count tokens; client construction requires a non-empty + api_key (CI mock-vLLM env has none). Fix: fall back to a local + tiktoken `cl100k_base` encoder when the client can't be built. Prod + (key present) behaviour is unchanged; the count is equivalent for the + GPT-3.5/4 family. Also trims LangChain off the QueryService hot path + (aligned with the 8H intent). +- *Stale `routers/test_auth_router.py` (17 collection errors).* The + 882-line file tested the pre-8A.1 fat router (`routers.auth.OIDCClient`, + full OIDC/JWT flow). 8A.1 thinned the router but never updated its + companion test — a Phase-8 omission. All that logic moved to + `AuthService` and is covered by `services/orchestrators/ + test_auth_service.py`. Replaced with a lean transport test (stub + service via `dependency_overrides`: AUTH_MODE gate, delegation, + cookie set/clear, `OIDCFlowError` mapping), matching the phase + principle "logic tests move to the service". +- *Search response lost `metadata.file_id` (5 API `test_search` + filtering failures, second pass).* `Chunk.from_langchain` lifts + `file_id`/`partition`/`page`/`_id` out of the free-form metadata into + typed `Chunk` fields; the thinned `routers/search.py:_documents` + returned only `Chunk.metadata`, so the API contract dropped + `metadata.file_id` (filter tests saw `file_id == None`; `origin` / + temporal fields survived because they stay free-form). Fix: shape the + response via `Chunk.to_langchain().metadata`, which merges the typed + fields back — reproducing the pre-Phase-8 router that returned the raw + Document metadata. Transport-only shaping, stays in the thin router. + +All five CI jobs (Layer guard, Linting, Unit, Integration, API) are +green on the branch after these fixes. + +--- +## Template for future entries + +``` +## Phase N — [short title] ([YYYY-MM-DD]) + +**K. [decision in one line].** +- Why: [what forced the call, what the docs didn't cover]. +- Alternative considered: [what else was on the table, why it was rejected]. +``` diff --git a/docker-compose.yaml b/docker-compose.yaml index e3f52f6ba..8c8fb6d45 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -136,6 +136,7 @@ services: - POSTGRES_USER=${POSTGRES_USER:-root} volumes: - ${DB_VOLUME:-./db}:/var/lib/postgresql/data + - ./scripts/postgres-init:/docker-entrypoint-initdb.d:ro expose: - 5432 @@ -181,4 +182,3 @@ services: # For details see https://github.com/vllm-project/vllm/issues/21179 profiles: - "cpu" - diff --git a/docs/content/docs/documentation/milvus_migration.mdx b/docs/content/docs/documentation/milvus_migration.mdx index 8dd937f0c..cb362d4af 100644 --- a/docs/content/docs/documentation/milvus_migration.mdx +++ b/docs/content/docs/documentation/milvus_migration.mdx @@ -139,13 +139,13 @@ docker compose ps milvus ```bash docker compose run --no-deps --rm --build --entrypoint "" openrag \ - uv run python scripts/migrations/milvus/migrate.py --dry-run + uv run python services/persistence/migrations/milvus/migrate.py --dry-run ``` ```bash docker compose --profile cpu run --no-deps --rm --build --entrypoint "" openrag-cpu \ - uv run python scripts/migrations/milvus/migrate.py --dry-run + uv run python services/persistence/migrations/milvus/migrate.py --dry-run ``` @@ -158,13 +158,13 @@ Review the output to confirm which migrations are pending and what changes they ```bash docker compose run --no-deps --rm --entrypoint "" openrag \ - uv run python scripts/migrations/milvus/migrate.py + uv run python services/persistence/migrations/milvus/migrate.py ``` ```bash docker compose --profile cpu run --no-deps --rm --entrypoint "" openrag-cpu \ - uv run python scripts/migrations/milvus/migrate.py + uv run python services/persistence/migrations/milvus/migrate.py ``` @@ -198,22 +198,22 @@ To upgrade or downgrade to a specific schema version rather than the latest: ```bash # Upgrade to version 2 only docker compose run --no-deps --rm --entrypoint "" openrag \ - uv run python scripts/migrations/milvus/migrate.py --target 2 + uv run python services/persistence/migrations/milvus/migrate.py --target 2 # Downgrade to version 0 (resets version stamp and drops indexes) docker compose run --no-deps --rm --entrypoint "" openrag \ - uv run python scripts/migrations/milvus/migrate.py --downgrade --target 0 + uv run python services/persistence/migrations/milvus/migrate.py --downgrade --target 0 ``` ```bash # Upgrade to version 2 only docker compose --profile cpu run --no-deps --rm --entrypoint "" openrag-cpu \ - uv run python scripts/migrations/milvus/migrate.py --target 2 + uv run python services/persistence/migrations/milvus/migrate.py --target 2 # Downgrade to version 0 (resets version stamp and drops indexes) docker compose --profile cpu run --no-deps --rm --entrypoint "" openrag-cpu \ - uv run python scripts/migrations/milvus/migrate.py --downgrade --target 0 + uv run python services/persistence/migrations/milvus/migrate.py --downgrade --target 0 ``` @@ -226,13 +226,13 @@ Milvus does not support dropping fields. A downgrade only removes indexes and re ```bash docker compose run --no-deps --rm --entrypoint "" openrag \ - uv run python scripts/migrations/milvus/migrate.py --downgrade + uv run python services/persistence/migrations/milvus/migrate.py --downgrade ``` ```bash docker compose --profile cpu run --no-deps --rm --entrypoint "" openrag-cpu \ - uv run python scripts/migrations/milvus/migrate.py --downgrade + uv run python services/persistence/migrations/milvus/migrate.py --downgrade ``` @@ -243,14 +243,14 @@ To fully remove the fields you would need to recreate the collection from scratc ## Adding a New Migration Script -Migration scripts live in `openrag/scripts/migrations/milvus/`. The runner discovers them automatically — no registration step required. +Migration scripts live in `openrag/services/persistence/migrations/milvus/`. The runner discovers them automatically — no registration step required. ### Naming convention Files must follow the pattern `N.short_description.py`, where `N` is the **target schema version** as a positive integer: ``` -openrag/scripts/migrations/milvus/ +openrag/services/persistence/migrations/milvus/ 1.add_temporal_fields.py ← brings the schema to version 1 2.your_new_migration.py ← brings the schema to version 2 migrate.py ← generic runner (do not rename) diff --git a/docs/content/docs/documentation/sql_migration.mdx b/docs/content/docs/documentation/sql_migration.mdx index c2de43959..013d1ee4b 100644 --- a/docs/content/docs/documentation/sql_migration.mdx +++ b/docs/content/docs/documentation/sql_migration.mdx @@ -34,11 +34,11 @@ Replace **``** with your descriptive message before running th docker compose up -d rdb docker compose \ run --no-deps --build --rm \ - --entrypoint "uv run alembic -c /app/openrag/scripts/migrations/alembic/alembic.ini revision --autogenerate -m ''" \ + --entrypoint "uv run alembic -c /app/openrag/services/persistence/migrations/alembic/alembic.ini revision --autogenerate -m ''" \ openrag ``` -It will create a new migration script in the `openrag/scripts/migrations/alembic/versions/` directory. +It will create a new migration script in the `openrag/services/persistence/migrations/alembic/versions/` directory. ### Step 2: Apply the Migration @@ -52,7 +52,7 @@ Rebuilds the image to ensure the migration runs against the exact code being dep docker compose up -d rdb docker compose \ run --no-deps --build --rm \ - --entrypoint "uv run alembic -c /app/openrag/scripts/migrations/alembic/alembic.ini upgrade head" \ + --entrypoint "uv run alembic -c /app/openrag/services/persistence/migrations/alembic/alembic.ini upgrade head" \ openrag; docker compose down ``` @@ -62,7 +62,7 @@ If the stack is already running, skip the rebuild: ```bash title="Apply migrations (exec)" docker compose exec openrag \ - uv run alembic -c /app/openrag/scripts/migrations/alembic/alembic.ini upgrade head + uv run alembic -c /app/openrag/services/persistence/migrations/alembic/alembic.ini upgrade head ``` #### Option C: Local (development) @@ -71,14 +71,14 @@ Run alembic directly with your local venv — no Docker needed: ```bash title="Apply migrations (local)" DATABASE_URL=postgresql://user:pass@localhost:5432/openrag \ - uv run alembic -c openrag/scripts/migrations/alembic/alembic.ini upgrade head + uv run alembic -c openrag/services/persistence/migrations/alembic/alembic.ini upgrade head ``` :::tip[Quick alternative: run inside a running container] If your stack is already running, you can apply migrations directly without rebuilding: ```bash title="Apply migrations via exec" docker compose exec -w /app/openrag openrag \ - /app/.venv/bin/alembic -c scripts/migrations/alembic/alembic.ini upgrade head + /app/.venv/bin/alembic -c services/persistence/migrations/alembic/alembic.ini upgrade head ``` ::: @@ -108,7 +108,7 @@ This usually happens when two or more migration scripts are created independentl ```bash title="Merge Alembic heads" docker compose up -d rdb docker compose run --no-deps --build --rm \ - --entrypoint "uv run alembic -c /app/openrag/scripts/migrations/alembic/alembic.ini merge -m 'merge heads' " \ + --entrypoint "uv run alembic -c /app/openrag/services/persistence/migrations/alembic/alembic.ini merge -m 'merge heads' " \ openrag ``` This will generate a new migration script that merges the two branches. @@ -118,7 +118,7 @@ This usually happens when two or more migration scripts are created independentl ```bash title="Apply migrations after merge" docker compose up -d rdb docker compose run --no-deps --build --rm \ - --entrypoint "uv run alembic -c /app/openrag/scripts/migrations/alembic/alembic.ini upgrade head" \ + --entrypoint "uv run alembic -c /app/openrag/services/persistence/migrations/alembic/alembic.ini upgrade head" \ openrag; docker compose down ``` diff --git a/docs/refactoring/REFACTORING_DEV_WORKFLOW.md b/docs/refactoring/REFACTORING_DEV_WORKFLOW.md new file mode 100644 index 000000000..8b353344f --- /dev/null +++ b/docs/refactoring/REFACTORING_DEV_WORKFLOW.md @@ -0,0 +1,587 @@ +# OpenRAG Hexagonal Refactoring — Development Workflow + +> **Strategy:** Hybrid — merge from dev during safe phases, +> forward-port only critical fixes during transformation phases, +> feature freeze during cutover. + +## Branch Layout + +``` +main (production releases) + | + +-- v1.1.9 (tag, frozen) + | | + | +-- refactor/hexagonal (long-lived refactoring branch) + | | + | +-- refactor/phase-5-retrieval (short-lived, 1-3 days) + | +-- refactor/phase-7-persistence (short-lived, 1-3 days) + | +-- ... + | + +-- dev (continues only on urgent bug fixes ) + | + +-- fix/abc +``` + +**Rules:** + +- `refactor/hexagonal` is branched from `v1.1.9` tag (frozen, tested, deployed) +- Per-task branches are branched from `refactor/hexagonal`, merged back via PR +- `dev` continues independently — no one works on `dev` AND `refactor/hexagonal` simultaneously on the same file +- `refactor/hexagonal` never merges into `dev`. The relationship is one-directional until the final cutover. + +--- + +## Three Modes of Operation + +The refactoring goes through three modes. Each mode has different rules for +how `dev` and `refactor/hexagonal` interact. + +``` +Timeline: + +Phase 0-4 Phase 5-9 Phase 10-12 Phase 13-15 +(Foundation) (Transformation) (Cutover) (Post-cutover) +| | | | +v v v v ++--------+ +-----------+ +--------+ +--------+ +| MODE 1 | | MODE 2 | | MODE 3 | | NORMAL | +| MERGE | | ISOLATE | | FREEZE | | DEV | ++--------+ +-----------+ +--------+ +--------+ + +dev merges into Forward-port Feature freeze refactor/hexagonal +refactor weekly. critical fixes on dev. Final becomes the new +Zero conflicts only. No merges. merge + cutover. dev branch. +(additive only). Features wait. Phases 13-15 as + normal features. +``` + +--- + +## MODE 1 — MERGE (Phases 0-4, ~2 days) + +### What's happening + +Phases 0-4 are purely additive — creating new files in `core/`, `services/`, +`api/`, `di/`. No existing file is modified or deleted. Zero conflict risk. + +### Rules + +| Rule | Detail | +| ----------------- | ----------------------------------- | +| Merge frequency | Once at the end of Mode 1 (day 2) | +| Conflict expected | None (refactor only adds new files) | +| Bug fixes on dev | Continue normally | + +### Workflow + +```bash +# Start of refactoring +git checkout v1.1.9 +git switch -c refactor/hexagonal +git push -u origin refactor/hexagonal + +# Person A works on Phase 0-2 +git switch refactor/hexagonal +git switch -c refactor/phase-0-scaffold +# ... create directories, __init__.py files, import guard ... +# PR into refactor/hexagonal + +# Weekly sync (or after each phase) +git switch refactor/hexagonal +git pull +git merge origin/dev +# No conflicts — our new files don't overlap with dev changes +git push +``` + +### What gets merged from dev + +Everything. Bug fixes, new features, dependency updates — all merge cleanly +because the refactoring hasn't touched any existing files yet. + +### Exit criteria for Mode 1 + +All of these exist and pass: + +- `core/utils/registry.py` — Registry[T] generic +- `core/utils/exceptions.py` — Exception hierarchy +- `core/models/*.py` — All domain models +- `core/config/*.py` — All config schemas + loader +- `core/embeddings/`, `core/rerankers/`, `core/llm/`, `core/vlm/` — ABCs + registries +- `core/vector_stores/`, `core/catalog/` — ABCs +- `core/ports/*.py` — All repository port ABCs +- `core/chunking/`, `core/indexing/parsers/` — ABCs + registries +- `scripts/check_layer_imports.py` passes +- `python -c "from openrag.core.models import Chunk, Document, User"` works +- All existing tests still pass + +--- + +## MODE 2 — ISOLATE (Phases 5-9, ~2 weeks) + +### What's happening + +This is the core transformation. Existing files are being rewritten, gutted, +shimmed, and replaced. File paths change. Imports change. The god object gets +decomposed. This is where merging from `dev` would create painful conflicts. + +### Rules + +| Rule | Detail | +| ------------------- | ----------------------------------------------------------------- | +| Merge frequency | **Never.** No merges from `dev`. | +| Forward-port | Critical bug fixes only (security, data loss, production outages) | +| Features on dev | **No features on dev** | +| Dev changes tracked | Maintain a `FORWARD_PORT_LOG.md` tracking what landed on dev | + +### Forward-porting process + +When a critical fix lands on `dev`: + +```bash +# 1. DON'T merge. Read the diff on dev. +git log origin/dev --oneline -20 # see what landed + +# 2. Understand the fix (read the PR, understand the intent) + +# 3. Re-implement the fix in the new architecture on refactor/hexagonal +git switch refactor/hexagonal +git switch -c fix/forward-port-xyz +# ... write the fix against the new code structure ... +# PR into refactor/hexagonal + +# 4. Log it +echo "- 2026-04-20: Forward-ported fix XYZ (dev commit abc123) -> refactor commit def456" >> FORWARD_PORT_LOG.md +``` + +### FORWARD_PORT_LOG.md + +Keep a running log so nothing is forgotten: + +```markdown +# Forward Port Log + +Tracks dev changes during Mode 2 isolation (Phases 5-9). +Each entry: what changed on dev, whether it was forward-ported or deferred. + +## Forward-ported (critical) + +- 2026-05-05: Security fix — regenerate_token missing auth check + dev: commit abc123, PR #42 + refactor: commit def456 + +- 2026-05-12: Bug fix — Milvus search crash on empty partition + dev: commit ghi789, PR #45 + refactor: commit jkl012 + +## Deferred to cutover (features) + +- 2026-05-08: New endpoint GET /partition/{name}/stats (PR #43) + -> Will re-implement in api/routers/admin/partitions.py + +- 2026-05-15: Added Docling v3 PDF loader (PR #47) + -> Will re-implement in services/inference/ or core/indexing/parsers/ + +- 2026-05-20: Updated pymilvus to 2.5.0 (PR #50) + -> Will update dependencies during Mode 3 +``` + +### Team coordination during Mode 2 + +**All team:** Works exclusively on `refactor/hexagonal`. Does not touch `dev`. + +**Urgent bug fixes:** should land on `dev` but keep in mind: + +- the bug fixes will need to be re-implemented on the new architecture +- only small, well-documented PRs are allowed so forward-porting is easy +- avoid large structural changes to existing files (creates harder forward-ports) + +**Communication:** Daily standup to review: + +- What landed on `dev` today +- What needs forward-porting (critical fixes) +- What's deferred (logged for Mode 3) + +### Parallel work within refactor branch + +During Mode 2, multiple people can work on different phases in parallel: + +``` +refactor/hexagonal + | + +-- Person A: refactor/phase-5-retrieval-core + | (Phases 5A-5C: retrieval, chunking, prompt builders) + | + +-- Person B: refactor/phase-7-persistence + | (Phase 7: god object decomposition) + | + +-- Both merge back into refactor/hexagonal via PR + | + +-- Then Person A + B converge on Phase 8 (orchestrators) +``` + +**Dependency order:** + +``` +Phase 5 (core logic) -----+ +Phase 6 (inference) -----+--> Phase 8 (orchestrators) --> Phase 9 (workers) +Phase 7 (persistence) -----+ +``` + +Phases 5, 6, 7 can run in parallel. Phase 8 needs all three. Phase 9 needs Phase 8. + +### Exit criteria for Mode 2 + +- All core domain logic lives in `core/` (no business logic in `components/`) +- All adapters live in `services/` (inference, storage, persistence) +- All orchestrators live in `services/orchestrators/` +- Ray actors are thin wrappers in `services/workers/` +- Old `components/` files are either deleted or gutted to re-export shims +- Import guard passes +- Integration tests pass (full upload -> search -> chat cycle) +- `FORWARD_PORT_LOG.md` is complete — all dev changes accounted for + +--- + +## MODE 3 — FREEZE (Phases 10-12, ~3-4 days) + +### What's happening + +The API layer is being restructured (routers, middleware, schemas), the DI +container is being wired, and old shims are deleted. This is the final +cutover — after this, the old code is gone. + +### Rules + +| Rule | Detail | +| ------------------------------ | ----------------------------------------------------------------- | +| Feature freeze on dev | **Mandatory.** No new features on dev. Critical bug fixes only. | +| Merge direction | Selected dev changes cherry-picked into refactor (not full merge) | +| Re-implement deferred features | Work through FORWARD_PORT_LOG.md deferred list | +| Duration | 3-4 days maximum — freeze must be time-boxed | + +### The freeze announcement + +``` +Subject: Feature freeze on dev starting [date] — Hexagonal cutover + +Duration: ~3-4 days +What's frozen: New features on dev branch +What's allowed: Critical bug fixes only (security, production outages) +Why: We're cutting over to the new architecture. Parallel changes would + create unmergeable conflicts. +What to do: + - Finish in-progress work before [date] or park it + - Critical fixes: PR into dev as usual, we'll cherry-pick into refactor +``` + +### Phase 10-12 workflow + +```bash +# Phase 10: API layer restructure +git switch refactor/hexagonal +git switch -c refactor/phase-10-api-layer +# Create api/main.py, api/routers/*, api/middleware/*, api/schemas/* +# Mount new routers alongside old ones (parallel operation) +# Test each new router +# Remove old routers one at a time +# PR into refactor/hexagonal + +# Phase 11: Composition root +git switch -c refactor/phase-11-di-container +# Create di/container.py, di/providers.py, di/factories.py +# Wire ServiceContainer into api/main.py lifespan +# Update routers to use Depends() from providers +# Remove global singletons +# PR into refactor/hexagonal + +# Phase 12: Cleanup +git switch -c refactor/phase-12-cleanup +# Delete all old shims, components/, routers/, models/ +# Update Dockerfile, docker-compose, pyproject.toml +# Final verification +# PR into refactor/hexagonal +``` + +### Re-implementing deferred features + +Work through `FORWARD_PORT_LOG.md` deferred list. Each deferred feature +is implemented directly in the new architecture: + +```bash +# Example: re-implement "GET /partition/{name}/stats" from dev PR #43 +git switch refactor/hexagonal +git switch -c feature/partition-stats +# Read the original PR on dev for intent +# Implement in api/routers/admin/partitions.py (new structure) +# Wire through services/orchestrators/partition_service.py +# PR into refactor/hexagonal +``` + +### Dependency sync + +Update all dependencies to match current dev (or newer): + +```bash +git switch refactor/hexagonal +# Compare pyproject.toml between dev and refactor +diff <(git show origin/dev:pyproject.toml) pyproject.toml +# Update versions, add new deps, remove unused +uv sync +uv run pytest -m unit +``` + +### Exit criteria for Mode 3 + +- `openrag/components/` directory does not exist +- `openrag/routers/` directory does not exist (replaced by `openrag/api/routers/`) +- `openrag/models/` directory does not exist (replaced by `openrag/core/models/` + `openrag/api/schemas/`) +- `di/container.py` ServiceContainer is the composition root +- No module-level `config = load_config()` anywhere +- No module-level Ray actor singletons +- Import guard passes with zero violations +- All deferred features from FORWARD_PORT_LOG.md re-implemented +- Full integration test suite passes +- Docker build succeeds +- `docker compose up` + manual smoke test passes + +--- + +## CUTOVER — Replacing dev + +Once Mode 3 is complete and everything passes: + +### Step 1 — Final verification + +```bash +# On refactor/hexagonal +uv run pytest # all tests +python scripts/check_layer_imports.py # layer guard +docker compose -f infra/compose/docker-compose.yaml build # docker build +docker compose -f infra/compose/docker-compose.yaml up -d # full stack +# Run integration tests against running stack +# Manual smoke test: upload document, search, chat, manage users +``` + +### Step 2 — Replace dev + +```bash +# Rename branches +git branch -m dev dev-legacy +git branch -m refactor/hexagonal dev +git push origin dev --force-with-lease +git push origin dev-legacy +``` + +### Step 3 — Communicate + +``` +Subject: Hexagonal refactoring complete — dev branch replaced + +The dev branch now contains the new 3-layer architecture. +dev-legacy preserves the old branch for reference. + +Key changes for developers: +- Import paths changed: components.retriever -> openrag.core.retrieval +- No global singletons: use ServiceContainer + Depends() +- New project layout: see updated README.md and CLAUDE.md +- Tests: uv run pytest -m unit (fast), uv run pytest -m integration (needs stack) + +Please pull fresh and read the updated CLAUDE.md before starting work. +``` + +### Step 4 — Delete the refactoring branch + +```bash +git push origin --delete refactor/hexagonal # remote +git branch -d refactor/hexagonal # local +``` + +--- + +## POST-CUTOVER — Normal Development (Phases 13-15) + +After cutover, development resumes on `dev` with the new architecture. + +Phases 13-15 are developed as normal feature branches on the new `dev`. +They can run in parallel if different people own them — no dependencies +between them. + +```bash +# Phase 13: Project layout restructure +git switch dev +git switch -c feature/phase-13-project-layout +# Move Dockerfiles -> infra/docker/, docker-compose -> infra/compose/ +# Move openrag/scripts/ -> scripts/ +# Unify tests/ (unit + integration + load) +# Move prompts/ -> openrag/prompts/ +# Move extern/indexer-ui -> ui/ +# Update pyproject.toml, CI, README +# PR into dev + +# Phase 14: Per-partition presets +git switch dev +git switch -c feature/phase-14-presets +# Flesh out core/config/{indexation,retrieval,partition,presets}.py +# Alembic migration: presets table + partition config columns +# services/persistence/preset_repo.py +# services/orchestrators/preset_service.py +# Update IndexingService + RetrievalService for per-partition config +# api/routers/admin/presets.py (CRUD endpoints) +# Wire into ServiceContainer, seed defaults +# PR into dev + +# Phase 15: OIDC / Keycloak SSO +git switch dev +git switch -c feature/phase-15-oidc-sso +# Add OIDCConfig to core/config/auth.py +# Create services/auth/{jwt_validator,oidc_mapper,oidc_provisioner}.py +# Add get_user_by_external_id() to UserRepository +# Update api/dependencies/auth.py with dual-auth dispatch +# Wire OIDC into ServiceContainer (conditional on OIDC_ENABLED) +# Update indexer-ui: oidc-client-ts, SSO login button, /auth/callback +# Update .env.example + docker-compose with OIDC env vars +# PR into dev +``` + +No special workflow needed — these are standard features on a clean codebase. + +--- + +## Timeline Summary + +``` +Day 1-2 MODE 1: Phases 0-4 (Foundation) + - 1-2 people, intensive + - Merge dev once at end (zero conflicts) + | +Day 2 Last merge from dev -> refactor/hexagonal + Enter MODE 2 + Announce feature freeze for end of week 2 + | +Day 3-14 MODE 2: Phases 5-9 (Transformation) + - 2-3 people in parallel (5+6 || 7, then converge on 8-9) + - No merges from dev + - Forward-port critical fixes only + - Track all dev changes in FORWARD_PORT_LOG.md + | +Day 12 Enter MODE 3: Feature freeze on dev + | +Day 12-15 MODE 3: Phases 10-12 (Cutover) + - API restructure + DI wiring + cleanup + - Re-implement deferred features (if any) + - Cherry-pick critical fixes from dev + | +Day 15 CUTOVER: refactor/hexagonal replaces dev + Lift feature freeze + | +Week 4+ POST-CUTOVER: Phases 13-15 on new dev (parallel) + - Phase 13: Project layout restructure (infra/, scripts/, tests/, ui/) + - Phase 14: Per-partition presets (indexation + retrieval config) + - Phase 15: OIDC / Keycloak SSO (dual auth, auto-provisioning) + Normal development resumes +``` + +--- + +## CI/CD for refactor/hexagonal + +### Branch protection + +```yaml +# .github/branch-protection for refactor/hexagonal +required_checks: + - unit-tests + - layer-import-guard + - docker-build +``` + +### CI pipeline + +```yaml +# .github/workflows/refactor-ci.yml +name: Refactor CI +on: + push: + branches: [refactor/hexagonal] + pull_request: + branches: [refactor/hexagonal] + +jobs: + unit-tests: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - run: uv sync + - run: uv run pytest -m unit + + layer-guard: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - run: python scripts/check_layer_imports.py + + docker-build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - run: docker build -f Dockerfile -t openrag:refactor . +``` + +### Integration tests (nightly or on-demand) + +## Decision Log Template + +Keep a running decision log during the : + +```markdown +# Refactoring Decision Log + +## 2026-04-21: Branch strategy +Decision: Merge during foundation, isolate during + transformation, freeze during cutover +Reason: dev is active, full merges during Phases 5-9 would create + unmanageable conflicts + +## 2026-04-21: Branch base +Decision: Branch from v1.1.9 tag, not dev +Reason: Frozen, tested baseline. Dev is a moving target. + +## [date]: [decision title] +Decision: [what was decided] +Reason: [why] +Alternative considered: [what else was on the table] +``` + +--- + +## Risk Mitigation Checklist + +Before entering each mode transition: + +### Before Mode 1 -> Mode 2 + +- [ ] All Phase 0-4 deliverables verified +- [ ] Import guard passes +- [ ] Last merge from dev completed and tested +- [ ] FORWARD_PORT_LOG.md created +- [ ] Team informed: "no more merges from dev" + +### Before Mode 2 -> Mode 3 + +- [ ] All Phase 5-9 deliverables verified +- [ ] Integration tests pass on refactor branch +- [ ] FORWARD_PORT_LOG.md reviewed — deferred features catalogued +- [ ] Feature freeze announced to team (2-3 days notice) +- [ ] Freeze start date agreed + +### Before cutover + +- [ ] All Phase 10-12 deliverables verified +- [ ] All deferred features re-implemented +- [ ] Docker build and compose up works +- [ ] Full integration test suite passes +- [ ] Manual smoke test completed +- [ ] README.md and CLAUDE.md updated +- [ ] Cutover communication drafted +- [ ] Rollback plan: dev-legacy branch preserved diff --git a/docs/refactoring/REFACTORING_STRATEGY_v1.md b/docs/refactoring/REFACTORING_STRATEGY_v1.md new file mode 100644 index 000000000..d78cffaa3 --- /dev/null +++ b/docs/refactoring/REFACTORING_STRATEGY_v1.md @@ -0,0 +1,2481 @@ +# OpenRAG → Hexagonal Architecture Refactoring Strategy (v2) + +> **Constraint:** After every commit the system MUST be deployable and pass existing tests. +> **Method:** Strangler Fig — new structure grows alongside old; re-exports keep old +> import paths alive until every consumer has migrated; only then are re-exports removed. + +--- + +## Table of Contents + +1. [Current vs Target Assessment](#1-current-vs-target-assessment) +2. [Key Design Patterns](#2-key-design-patterns) +3. [Target Architecture](#3-target-architecture) +4. [Guiding Principles](#4-guiding-principles) +5. [Phase Overview](#5-phase-overview) +6. [Phase 0 — Scaffold & Import Guard](#phase-0--scaffold--import-guard) +7. [Phase 1 — Generic Registry & Exceptions](#phase-1--generic-registry--exceptions) +8. [Phase 2 — Domain Models](#phase-2--domain-models) +9. [Phase 3 — Configuration Schemas](#phase-3--configuration-schemas) +10. [Phase 4 — ABCs & Ports](#phase-4--interfaces--ports) +11. [Phase 5 — Core Domain Logic](#phase-5--core-domain-logic) +12. [Phase 6 — Inference Adapters](#phase-6--inference-adapters) +13. [Phase 7 — Storage & Persistence Adapters](#phase-7--storage--persistence-adapters) +14. [Phase 8 — Orchestrators](#phase-8--orchestrators) +15. [Phase 9 — Workers (Ray Isolation)](#phase-9--workers-ray-isolation) +16. [Phase 10 — API Layer Restructure](#phase-10--api-layer-restructure) +17. [Phase 11 — Composition Root (DI)](#phase-11--composition-root-di) +18. [Phase 12 — Internal Cleanup & Remove Shims](#phase-12--internal-cleanup--remove-shims) +19. [Phase 13 — Project Layout, Infra, Tests & UI](#phase-13--project-layout-infra-tests--ui) +20. [Phase 14 — Per-Partition Presets (Indexation & Retrieval)](#phase-14--per-partition-presets-indexation--retrieval) +21. [Phase 15 — OIDC / Keycloak SSO Authentication](#phase-15--oidc--keycloak-sso-authentication) +22. [Risk Register](#risk-register) +23. [Migration Utilities](#migration-utilities) + +--- + +## 1. Current vs Target Assessment + +### Current codebase (openrag 1.1.8) + +| Metric | Value | +| ----------------------- | ------------------------------------------------------------------------- | +| Production Python files | 105 | +| Total LOC (non-test) | ~15,824 | +| Package | `openrag/` with `components/`, `routers/`, `config/`, `models/`, `utils/` | +| External infra | Milvus, PostgreSQL, Ray, vLLM/Ollama | + +### Coupling hotspots + +| Hotspot | Symptom | Fix | +| --------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------- | +| **Global `config = load_config()` at module level** (12+ files) | Cannot test/compose without full config | `ServiceContainer` loads config once, passes via constructor | +| **Ray actor singletons at import time** (`utils/dependencies.py`) | Importing any router boots Ray + Milvus + GPU semaphores | Lazy init in `container.initialize()`, actors created only when needed | +| **LangChain `Document` as universal data type** | Domain tied to third-party class | Domain `Chunk` + `Document` models with `from_langchain()` converters | +| **Vectordb actor = God object** (Milvus + PG + embedding + users + partitions + workspaces) | Single point of failure, untestable | Split into `MilvusVectorStore`, `PostgresStore` with 13 repos, `Embedder` port | +| **No interface boundaries** for LLM, VectorDB, persistence | Cannot swap, mock, or test | ABCs co-located in subject folders (`core/embeddings/embedder.py`, etc.) + 13 repository ports in `core/ports/` | +| **Per-component Factory classes** (`RetrieverFactory`, `RerankerFactory`, `ChunkerFactory`, `EmbeddingFactory`) | Duplicated pattern, no DI | Single `Registry[T]` generic + `make_component_factory()` for config-driven instantiation | +| **Router-to-component direct imports** | API fused to business logic | Orchestrator services injected via `Depends(get_service)` | + +### Existing good patterns to preserve + +| Pattern | Location | Action | +| ----------------------------------------------- | ------------------------------------ | ---------------------------------------------------------------------- | +| Retriever strategies (Single, MultiQuery, HyDE) | `components/retriever.py` | Move to `core/retrieval/`, register in `retriever_registry` | +| BaseReranker ABC + factory | `components/reranker/` | Promote to `core/rerankers/reranker.py` + `core/rerankers/registry.py` | +| BaseChunker + RecursiveSplitter | `components/indexer/chunker/` | Move to `core/chunking/` | +| BaseEmbedding + OpenAIEmbedding | `components/indexer/embeddings/` | Promote to `core/embeddings/embedder.py` | +| BaseLoader ABC for file parsers | `components/indexer/loaders/base.py` | Promote to `core/indexing/parsers/document_parser.py` | +| Pydantic config models (frozen) | `config/models.py` | Split into `core/config/` sections | +| Custom exception hierarchy | `utils/exceptions/` | Consolidate in `core/utils/exceptions.py` | +| Web search provider abstraction | `components/websearch/base.py` | Keep as pluggable adapter | +| Disk-loaded prompt templates | `components/prompts/` | Move to `openrag/prompts/` | + +--- + +## 2. Key Design Patterns + +### 2.1 Registry[T] — The plugin pattern core + +```python +class Registry(Generic[T]): + def __init__(self, kind: str) -> None + def register(self, name: str) -> Callable[[Type[T]], Type[T]] # decorator + def create(self, name: str, **kwargs: Any) -> T + def get_class(self, name: str) -> Type[T] + def list_registered(self) -> list[str] + def __contains__(self, name: str) -> bool +``` + +**Key insight:** Implementations register via decorator on class definition: + +```python +@embedder_registry.register("vllm") +class VLLMEmbedder(Embedder): + ... +``` + +Registration happens via side-effect imports in `di/embedders.py`: + +```python +def register_embedders() -> None: + import openrag.services.inference.vllm_client # noqa: F401 + import openrag.services.inference.ollama_client # noqa: F401 +``` + +**Replaces** in OpenRAG: `RetrieverFactory`, `RerankerFactory`, `ChunkerFactory`, +`EmbeddingFactory`, `WebSearchFactory` — all five become `Registry[T]` instances. + +### 2.2 make_component_factory() — Config-driven, cached, thread-safe + +```python +def make_component_factory( + registry: Registry[T], + config_section: dict[str, ModelEndpointConfig], + default_impl: str, + client_caches: list[dict[str, T]], + extra_kwargs_fn: Callable | None = None, +) -> Callable[[str], T]: +``` + +**Key insight:** Returns a `Callable[[str], T]` that orchestrators receive instead of +concrete instances. First call creates + caches; subsequent calls return cached. +Thread-safe with double-checked locking. Cache dict appended to `client_caches` for +lifecycle cleanup in `ServiceContainer.shutdown()`. + +**Replaces** in OpenRAG: Module-level singletons like `ragpipe = RagPipeline()`. + +### 2.3 ServiceContainer — Composition root with sync/async split + +**Lifecycle:** + +1. **`__init__` (sync):** Load config → register impls → create stores → create factories → create services +2. **`initialize()` (async):** Open DB pool → run migrations → seed defaults → mark initialized +3. **`shutdown()` (async):** Close HTTP clients → close DB pool + +**Key insight:** Sync init creates the entire object graph (no I/O). Async init +does all the actual connecting. This means `__init__` can't fail on network issues. + +### 2.4 Co-located ABCs — ABC lives with its subject folder + +Each ABC is co-located inside the subject folder that owns it. This keeps each domain concept +self-contained — the ABC, its registry, and its implementations all discoverable in +one place. + +**Pattern:** Each subject folder owns its ABC + registry + `__init__.py` re-exports: + +``` +core/embeddings/ + embedder.py # Embedder ABC + registry.py # embedder_registry: Registry[Embedder] + __init__.py # re-exports Embedder + embedder_registry +``` + +Consumers import cleanly: + +```python +from openrag.core.embeddings import Embedder, embedder_registry +``` + +**Component ABCs** (co-located in subject folders): + +- `core/embeddings/embedder.py` — Embedder ABC +- `core/rerankers/reranker.py` — Reranker ABC +- `core/llm/llm.py` — LLM ABC +- `core/vlm/vlm.py` — VLM ABC +- `core/vector_stores/vector_store.py` — VectorStore ABC +- `core/chunking/chunking_strategy.py` — ChunkingStrategy ABC +- `core/indexing/parsers/document_parser.py` — DocumentParser ABC + +**Repository ports** (CRUD contracts + CatalogStore aggregate root): + +- `core/ports/catalog_store.py` — CatalogStore ABC (composes all repos) +- `core/ports/` — DocumentRepository, ChunkRepository, UserRepository, etc. + +**Key insight:** Don't put VectorStore in `ports/`. It's a pluggable component +(searchable, swappable backends), not a CRUD repo. + +### 2.5 CatalogStore — Composite repository pattern + +```python +class PostgresStore(CatalogStore): + def __init__(self, config: PostgresConfig): + self._conn = ConnectionManager(config) + pool_getter = lambda: self._conn.pool # lazy: repos don't touch pool until init + self._documents = PgDocumentRepository(pool_getter) + self._users = PgUserRepository(pool_getter) + # ... 13 repos total + + @property + def document_repo(self) -> DocumentRepository: return self._documents +``` + +**Key insight:** Single `ConnectionManager` owns the pool. Repos receive a `pool_getter` +lambda so they can defer access until the pool is live. Composition, not inheritance. + +**Replaces** in OpenRAG: The monolithic `PartitionFileManager` (currently in +`vectordb/utils.py`) which does file + user + partition + workspace + membership ops +through a single class. + +### 2.6 Thin Ray actor wrappers + +```python +@ray.remote(max_concurrency=2, max_restarts=3) +class IndexerActorClass: + def __init__(self): + # Only imports needed for registry registration + import openrag.services.inference.vllm_client # noqa + ... + + async def process_document(self, doc_dict, config_dict, prompt_overrides, job_id): + document = Document.model_validate(doc_dict) + config = Settings.model_validate(config_dict) + result = await ray_data_ingest_documents([document], config, ...) + return {"status": "success", **result.successful[0]} +``` + +**Key insight:** The actor is just serialization boundary + error boundary. All logic +@ray.remote(max_concurrency=2, max_restarts=3) +class IndexerActorClass: + def __init__(self): + # Only imports needed for registry registration + import openrag.services.inference.vllm_client # noqa + ... + + async def process_document(self, doc_dict, config_dict, prompt_overrides, job_id): + document = Document.model_validate(doc_dict) + config = Settings.model_validate(config_dict) + result = await ray_data_ingest_documents([document], config, ...) + return {"status": "success", **result.successful[0]}@ray.remote(max_concurrency=2, max_restarts=3) + class IndexerActorClass: + def __init__(self): + # Only imports needed for registry registration + import openrag.services.inference.vllm_client # noqa + ... + + async def process_document(self, doc_dict, config_dict, prompt_overrides, job_id): + document = Document.model_validate(doc_dict) + config = Settings.model_validate(config_dict) + result = await ray_data_ingest_documents([document], config, ...) + return {"status": "success", **result.successful[0]}lives in the indexing service/pipeline. Models are serialized to dicts by caller, + deserialized by actor. No state between calls. + +**Replaces** in OpenRAG: The 300-line `Indexer` actor that mixes chunking, embedding, +task state management, and file cleanup. + +### 2.8 Pipeline stages as modules + +Each stage is a pure async function. `pipeline_builder.py` chains them sequentially +with per-stage timeout, error marking, and credential scrubbing. + +**Key patterns:** + +- Rows are mutated in-place (no functional pipeline overhead) +- Failed rows get `_error` field (informational, still passed to next stage) +- Credentials scrubbed after the stage that needs them +- Timeout: base + per-chunk scaling (except contextualize — no stage timeout to prevent cascade) + +### 2.9 All-async ABCs + +All component ABCs should be async-native: + +```python +class Embedder(ABC): + async def embed(self, texts: list[str]) -> list[list[float]]: ... + async def embed_single(self, text: str) -> list[float]: ... +``` + +The current `BaseEmbedding` inherits LangChain's +`Embeddings` class (sync `embed_documents` + async wrappers via thread pool). The new +`Embedder` ABC should be async-native. + +### 2.10 Structured error responses + +Error responses should be structured with request tracing: + +```json +{ + "error": { + "message": "Document not found", + "type": "not_found_error", + "code": "NOT_FOUND", + "request_id": "req_a1b2c3d4..." + } +} +``` + +Domain exceptions walk the MRO to find the correct HTTP status code. Request ID +comes from middleware (structlog contextvars). + +### 2.11 Per-partition pipeline configuration + +Each partition can have its own `IndexationPipelineConfig` and `RetrievalPipelineConfig`. +Presets are reusable named configs stored in DB. Partitions reference presets by name. + +Currently OpenRAG has global config only. Adding per-partition +config is a Phase 14 enhancement, not part of the initial refactor. + +# + +### 2.13 BM25 approach + +An alternative to Milvus BM25 is duplicating chunk text in Postgres (with tsvector index) +and using `ts_rank` scoring for full-text search, while keeping Milvus for dense vectors only. +OpenRAG uses Milvus's built-in BM25 (sparse vectors). Both approaches work; OpenRAG's +is more Milvus-native and avoids data duplication. + +--- + +## 3. Target Architecture + +``` ++-----------------------------------------------------------------+ +| API Layer | +| FastAPI routers, middleware, Pydantic request/response | +| schemas, auth dependencies, error handlers | +| | +| openrag/api/ | ++-----------------------------------------------------------------+ +| Services Layer | +| Adapter implementations, orchestrators, Ray workers, | +| inference clients, storage backends, persistence repos | +| | +| openrag/services/ | ++-----------------------------------------------------------------+ +| Core Layer | +| Domain models, interface ABCs, port ABCs, config schemas, | +| retrieval algorithms, chunking strategies, registry, | +| exceptions, observability definitions | +| | +| openrag/core/ (PURE - no infra imports) | ++-----------------------------------------------------------------+ +| Composition Root | +| ServiceContainer, Depends() providers, registry wiring, | +| config-driven factories | +| | +| openrag/di/ | ++-----------------------------------------------------------------+ +``` + +### Dependency rule (STRICT) + +``` +api/ --> di/ --> services/ --> core/ +api/ --> core/ (models, config, ABCs - read-only) +services/ --> core/ +core/ --> (nothing in openrag - only stdlib + pure libs like pydantic) +di/ --> core/ + services/ (the ONLY place that crosses boundaries) +``` + +### Directory structure (full) + +``` +openrag/ +|-- core/ +| |-- config/ # Typed config schemas + YAML loader +| | |-- auth.py # AuthConfig + OIDCConfig (issuer, client_id, claim mapping) +| | |-- endpoints.py # ModelEndpointConfig + ModelsConfig (embedder, reranker, llm, vlm) +| | |-- chunking.py, indexation.py, infrastructure.py, +| | | loader.py, partition.py, presets.py, retrieval.py, root.py +| | +-- __init__.py +| | +| |-- ports/ # Repository contracts + CatalogStore aggregate root +| | |-- catalog_store.py # CatalogStore ABC - composes all repos, owns pool lifecycle +| | |-- document_repo.py, chunk_repo.py, user_repo.py, job_repo.py, +| | | partition_repo.py, conversation_repo.py, prompt_repo.py, +| | | entity_repo.py, topic_tag_repo.py, audit_log_repo.py, +| | | idempotency_repo.py, model_endpoint_repo.py, preset_repo.py +| | +-- __init__.py +| | +| |-- models/ # Domain entities (Pydantic frozen models) +| | |-- catalog.py # DocumentRecord, IndexationJob, status enums +| | |-- chunk.py # Chunk, ChunkType enum +| | |-- contextualization.py # ContextualizedQuery +| | |-- conversation.py # Conversation, Message +| | |-- document.py # Document, ProcessedDocument, TextBlock, ImageBlock +| | |-- prompt.py # Prompt, PromptType enum +| | |-- query.py # RetrievalQuery +| | |-- retrieval_response.py # RetrievalResponse, ScoredChunk +| | |-- retrieval_result.py # RetrievalResult +| | +-- user.py # User, SystemRole, PartitionRole +| | +| |-- chunking/ # ABC + strategies + registry +| | |-- __init__.py # re-exports ChunkingStrategy + chunking_registry +| | |-- chunking_strategy.py # ChunkingStrategy ABC +| | |-- registry.py # chunking_registry: Registry[ChunkingStrategy] +| | |-- recursive.py, fixed.py, sentence.py, +| | | markdown_section.py, markdown_layout.py +| | +| |-- embeddings/ # ABC + registry (impls in services/inference/) +| | |-- __init__.py # re-exports Embedder + embedder_registry +| | |-- embedder.py # Embedder ABC - async embed/embed_single/dimension +| | +-- registry.py # embedder_registry: Registry[Embedder] +| | +| |-- rerankers/ # ABC + registry +| | |-- __init__.py # re-exports Reranker + reranker_registry +| | |-- reranker.py # Reranker ABC - rerank(query, docs, top_k) +| | +-- registry.py # reranker_registry: Registry[Reranker] +| | +| |-- llm/ # ABC + registry only +| | |-- __init__.py # re-exports LLM + llm_registry +| | |-- llm.py # LLM ABC - generate/chat/stream_chat/chat_with_tools +| | +-- registry.py # llm_registry: Registry[LLM] +| | +| |-- vlm/ # ABC + registry +| | |-- __init__.py # re-exports VLM + vlm_registry +| | |-- vlm.py # VLM ABC - caption_image/caption_images_batch +| | +-- registry.py # vlm_registry: Registry[VLM] +| | +| |-- prompts/ # Prompt assembly logic (all builders in one place) +| | |-- __init__.py +| | |-- chat_prompt_builder.py # RAG chat: context + system prompt + query -> messages +| | |-- vlm_prompt_builder.py # VLM: image + caption template -> messages +| | |-- contextualization_builder.py # chunk contextualization prompt assembly +| | |-- query_rewriter.py # multi-query / HyDE prompt building +| | |-- map_reduce_builder.py # map / reduce per-chunk prompts +| | +-- template_loader.py # load_template(name) -> str (reads from openrag/prompts/templates/) +| | +| |-- vector_stores/ # ABC (impls in services/storage/) +| | |-- __init__.py # re-exports VectorStore +| | +-- vector_store.py # VectorStore ABC - upsert/search/delete/ensure_collection +| | +| |-- retrieval/ # Retrieval algorithms (the RAG core) +| | |-- pipeline.py # UnifiedPipeline (dense + BM25 + entity + RRF) +| | |-- retriever.py # Retriever facade +| | |-- entity_retrieval.py +| | |-- hydration.py +| | +-- rrf.py # Reciprocal Rank Fusion (pure math) +| | +| |-- indexing/ # Document ingestion domain logic +| | |-- contextualize.py, text_preprocessor.py, image_preprocessor.py, +| | | validators.py +| | +-- parsers/ +| | |-- __init__.py # re-exports DocumentParser + parser_registry +| | |-- document_parser.py # DocumentParser ABC - parse(document) +| | |-- registry.py # parser_registry: Registry[DocumentParser] +| | |-- text_parser.py, html_parser.py, pdf_parser.py, +| | | image_parser.py, audio_parser.py, video_parser.py +| | +| |-- observability/ +| | +-- metrics.py # Prometheus definitions only +| | +| +-- utils/ +| |-- registry.py # Registry[T] - THE plugin pattern core +| |-- exceptions.py # OpenRAGError hierarchy +| |-- text.py, dates.py, filename.py, mime_validation.py, +| | logging.py, retry.py, streaming.py, tracing.py, +| | debug.py, scrub.py +| +-- __init__.py +| +|-- services/ +| |-- auth/ # Authentication adapters +| | |-- jwt_validator.py # JWKS-based JWT signature verification +| | |-- oidc_mapper.py # JWT claims -> OpenRAG User + partition roles +| | +-- oidc_provisioner.py # Auto-create/sync users from OIDC claims +| | +| |-- inference/ # HTTP clients to inference services +| | |-- vllm_client.py # @embedder_registry.register("vllm"), @llm_registry.register("vllm") +| | |-- ollama_client.py # @embedder_registry.register("ollama"), @llm_registry.register("ollama") +| | |-- infinity_client.py # @reranker_registry.register("infinity") +| | |-- vlm_client.py # @vlm_registry.register("vllm") +| | |-- healthcheck.py +| | |-- _circuit_breaker.py # @with_circuit_breaker decorator (aiobreaker) +| | |-- _retry.py # @with_retry decorator (tenacity + jitter) +| | |-- _timeout.py # @with_timeout decorator (asyncio.timeout) +| | +-- distributed_semaphore.py +| | +| |-- storage/ +| | |-- postgres_store.py # PostgresStore implements CatalogStore (composite of 13 repos) +| | |-- milvus_store.py # MilvusVectorStore implements VectorStore +| | +-- s3_store.py # S3Store - document upload/download +| | +| |-- persistence/ # Postgres repository implementations +| | |-- connection.py # ConnectionManager (asyncpg pool lifecycle + retry) +| | |-- schema.py # SQLAlchemy metadata (all tables) +| | |-- document_repo.py, chunk_repo.py, job_repo.py, user_repo.py, +| | | partition_repo.py, prompt_repo.py, conversation_repo.py, +| | | entity_repo.py, topic_tag_repo.py, audit_log_repo.py, +| | | idempotency_repo.py, model_endpoint_repo.py, preset_repo.py +| | +-- migrations/ +| | |-- env.py +| | +-- versions/ # append-only migration history +| | +| |-- orchestrators/ # Business services (high-level flows) +| | |-- retrieval_service.py, indexing_service.py, document_service.py, +| | | job_service.py, partition_service.py, query_service.py, +| | | query_orchestrator.py, auth_service.py, user_service.py, +| | | prompt_service.py, model_endpoint_service.py, preset_service.py, +| | | cluster_service.py, conversation_service.py, conversion_service.py, +| | | llm_contextualizer.py, research_planner.py +| | +-- __init__.py +| | +| |-- workers/ # Ray-based distributed workers +| | |-- ray_utils.py, pipeline_builder.py, batch_ingest.py, +| | | ray_data_ingest.py, indexer_actor.py, result_aggregation.py +| | +-- stages/ +| | |-- parse.py, caption.py, chunk.py, contextualize.py, +| | | embed.py, store.py +| | +-- __init__.py +| | +| +-- events/ +| +-- job_events.py # In-process SSE event bus +| +|-- api/ +| |-- main.py # FastAPI app, lifespan, middleware, routes +| |-- error_handlers.py # Domain exception -> JSON response +| | +| |-- dependencies/ +| | |-- auth.py # Dual auth: OIDC JWT + API token, RBAC +| | |-- audit.py +| | +-- rate_limit.py +| | +| |-- middleware/ +| | |-- request_id.py, idempotency.py, request_timeout.py, +| | | security_headers.py, instrumentation.py +| | +-- __init__.py +| | +| |-- routers/ +| | |-- auth/login.py +| | |-- user/ +| | | |-- chat.py, retrieve.py, query_plan.py, health.py, +| | | | partitions.py, me.py, account.py, documents.py, +| | | | chat_conversations.py +| | | +-- __init__.py +| | +-- admin/ +| | |-- indexing.py, partitions.py, pipelines.py, documents.py, +| | | jobs.py, prompts.py, model_endpoints.py, presets.py, +| | | system.py, convert.py, users.py, audit_log.py +| | +-- __init__.py +| | +| +-- schemas/ +| |-- user/, admin/, auth/ +| +-- __init__.py +| +|-- di/ +| |-- container.py # ServiceContainer - composition root +| |-- providers.py # FastAPI Depends() accessors +| |-- factories.py # make_component_factory() - config-driven cached factory +| |-- embedders.py # register_embedders() - side-effect imports +| |-- rerankers.py # register_rerankers() +| |-- llms.py # register_llms() +| |-- vlms.py # register_vlms() +| |-- vector_stores.py # create_vector_store() +| +-- repositories.py # create_catalog_store() +| ++-- prompts/ # Disk-loaded prompt templates + +-- templates/ +``` + +--- + +## 4. Guiding Principles + +### 4.1 Strangler Fig migration + +Every module move follows this pattern: + +``` +Commit A: Create new file at target location +Commit B: Update old file to re-export from new location +Commit C: Update consumers to import from new location +... Phase 12: Delete old re-export shim +``` + +### 4.2 Co-located ABCs in subject folders, ports/ for CRUD repos + +Each component ABC lives inside its subject folder alongside its registry. +The folder's `__init__.py` re-exports the public surface for clean imports: + +```python +# core/embeddings/__init__.py +from openrag.core.embeddings.embedder import Embedder +from openrag.core.embeddings.registry import embedder_registry +__all__ = ["Embedder", "embedder_registry"] + +# Consumer code: +from openrag.core.embeddings import Embedder, embedder_registry +``` + +Same pattern for `rerankers/`, `llm/`, `vlm/`, `chunking/`, `vector_stores/`, +and `indexing/parsers/`. + +Ports (`core/ports/`) hold CRUD repository ABCs + `CatalogStore` (the aggregate root +that composes all repos). + +### 4.3 Config injection via ServiceContainer + +Replace `config = load_config()` at module level with constructor injection. +`ServiceContainer.__init__()` is the only place that calls `load_config()`. +Services receive typed config sections, not the entire Settings object. + +### 4.4 Factory callables over concrete instances + +Orchestrators receive `Callable[[str], Embedder]` (factory), not `Embedder` (instance). +This enables lazy creation, model switching, and lifecycle management. + +**Pattern:** + +```python +class RetrievalService: + def __init__( + self, + vector_store: VectorStore, + embedder_factory: Callable[[str], Embedder], # NOT Embedder + reranker_factory: Callable[[str], Reranker], + ... + ): ... +``` + +### 4.5 Ray is an infrastructure detail + +Ray actors in `services/workers/` are thin wrappers. They serialize/deserialize +domain models (dict round-trip), call service methods, and return results. +Core domain logic works without Ray. + +### 4.6 Async-native interfaces + +All interface ABCs are async. No inheriting from LangChain's sync `Embeddings` class. +Sync operations wrapped in `asyncio.to_thread()` at the adapter level. + +### 4.7 Domain models replace LangChain Document + +`core/models/chunk.py:Chunk` and `core/models/document.py:Document` are the domain types. +Boundary converters `from_langchain()` / `to_langchain()` exist during migration. + +--- + +## 5. Phase Overview + +| Phase | Name | Risk | Key deliverable | +| ----- | --------------------------------- | -------- | ------------------------------------------------------------------------- | +| 0 | Scaffold & Import Guard | None | Directory tree + CI guard | +| 1 | Generic Registry & Exceptions | Low | `Registry[T]`, `OpenRAGError` hierarchy | +| 2 | Domain Models | Low | `Chunk`, `Document`, `User`, `RetrievalQuery`, etc. | +| 3 | Configuration Schemas | Low | `core/config/` with typed sections + loader | +| 4 | ABCs & Ports | Low | 8 co-located ABCs + `__init__.py` re-exports + 13 port ABCs | +| 5 | Core Domain Logic | Medium | Retrieval, chunking, indexing in `core/` | +| 6 | Inference Adapters | Medium | vLLM/Ollama/Infinity clients in `services/inference/` | +| 7 | Storage & Persistence | **High** | God object decomposition -> MilvusStore + PostgresStore(13 repos) | +| 8 | Orchestrators | **High** | 15+ business services in `services/orchestrators/` | +| 9 | Workers (Ray) | **High** | Thin actor wrappers + pipeline stages | +| 10 | API Layer | Medium | Routers, middleware, schemas in `api/` | +| 11 | Composition Root | **High** | `ServiceContainer`, `providers.py`, `factories.py` | +| 12 | Internal Cleanup | Medium | Remove shims, delete old `components/`, `routers/`, `models/` | +| 13 | Project Layout, Infra, Tests & UI | Medium | Top-level restructure: `infra/`, `scripts/`, `tests/`, `ui/`, Dockerfiles | +| 14 | Per-Partition Presets | Medium | DB-backed presets for indexation + retrieval per partition | +| 15 | OIDC / Keycloak SSO | Medium | Dual auth (JWT + API tokens), auto-provisioning, role sync | + +--- + +## Phase 0 — Scaffold & Import Guard + +**Goal:** Create directory skeleton + enforce layer boundaries from commit #1. + +### Commits + +**0.1 — Create directory tree** + +```bash +mkdir -p openrag/core/{config,ports,models,chunking,embeddings,rerankers,llm,vlm,prompts,vector_stores,retrieval,indexing/parsers,observability,utils} +mkdir -p openrag/services/{inference,storage,persistence/migrations/versions,orchestrators,workers/stages,events} +mkdir -p openrag/api/{dependencies,middleware,routers/auth,routers/user,routers/admin,schemas/user,schemas/admin,schemas/auth} +mkdir -p openrag/di +# Touch __init__.py in every dir +``` + +**0.2 — Add layer import guard** + +Create `scripts/check_layer_imports.py`: + +```python +""" +CI guard: enforces hexagonal layer dependencies. + +Rules: + core/ -> may NOT import from services/, api/, di/ + services/ -> may NOT import from api/ + api/ -> may NOT import from services/ directly (only via di/) + +Usage: python scripts/check_layer_imports.py +Exit code 0 = pass, 1 = violations found. +""" +``` + +**Verification:** `python -c "import openrag"`, existing tests pass. + +--- + +## Phase 1 — Generic Registry & Exceptions + +**Goal:** Build the two foundational utilities that everything else depends on. + +### 1.1 — Registry[T] + +**File:** `core/utils/registry.py` + +The `Registry[T]` API: + +```python +class Registry(Generic[T]): + def __init__(self, kind: str) -> None + def register(self, name: str) -> Callable[[Type[T]], Type[T]] # decorator + def create(self, name: str, **kwargs: Any) -> T + def get_class(self, name: str) -> Type[T] + def list_registered(self) -> list[str] + def __contains__(self, name: str) -> bool +``` + +Raise `RegistryError` with helpful message listing available implementations. + +### 1.2 — Exception hierarchy + +**File:** `core/utils/exceptions.py` + +Consolidate from `utils/exceptions/{base,vectordb,embeddings}.py`: + +```python +class OpenRAGError(Exception): + """Root. Has code, status_code, to_dict().""" + +# Config & registry +class ConfigError(OpenRAGError): ... +class RegistryError(OpenRAGError): ... + +# Auth +class AuthError(OpenRAGError): ... +class AuthenticationError(AuthError): ... # 401 + +# Validation +class ValidationError(OpenRAGError): ... # 400/422 + +# Infrastructure +class ServiceUnavailableError(OpenRAGError): ... # 503 +class CircuitBreakerOpenError(ServiceUnavailableError): ... + +# Inference +class InferenceError(OpenRAGError): ... # 503 +class LLMParsingError(InferenceError): ... # 502 +class InferenceTimeoutError(InferenceError): ... # 504 +class InferenceConnectionError(InferenceError): # 503 + +# Storage +class StorageError(OpenRAGError): ... # 500 +class MilvusError(StorageError): ... +class PostgresError(StorageError): ... + +# Domain +class NotFoundError(OpenRAGError): ... # 404 +class DocumentNotFoundError(NotFoundError): ... +class PartitionNotFoundError(NotFoundError): ... +class UserNotFoundError(NotFoundError): ... +class QuotaExceededError(OpenRAGError): ... # 429 + +# Pipeline +class PipelineError(OpenRAGError): ... +``` + +### 1.3 — Core utilities + +Move pure utility functions (no infra imports): + +``` +core/utils/text.py <- components/indexer/utils/text_sanitizer.py +core/utils/dates.py <- (new) +core/utils/filename.py <- components/indexer/utils/files.py (filename parts) +core/utils/mime_validation.py <- routers/utils.py (mime validation logic) +core/utils/logging.py <- (structlog setup, no Loguru dependency on Ray) +core/utils/retry.py <- (tenacity config, pure) +core/utils/streaming.py <- (async token stream helpers) +core/utils/tracing.py <- (PipelineTrace - per-request timing) +core/utils/debug.py <- (gated debug file writers) +core/utils/scrub.py <- (secret scrubbing) +``` + +### 1.4 — Update old exceptions to re-export + +```python +# utils/exceptions/__init__.py +from openrag.core.utils.exceptions import * # noqa: F401,F403 +``` + +--- + +## Phase 2 — Domain Models + +**Goal:** Create pure Pydantic models in `core/models/`. No infra imports. + +### Model inventory + +| File | Key types | Source in OpenRAG | +| ----------------------- | -------------------------------------------------------------------------- | ------------------------------------------------- | +| `chunk.py` | `Chunk`, `ChunkType` | Runtime dicts in vectordb | +| `document.py` | `Document`, `ProcessedDocument`, `TextBlock`, `ImageBlock`, `DocumentType` | LangChain `Document` | +| `user.py` | `User`, `PartitionRole`, `UserPartition` | `vectordb/utils.py` User table + `models/user.py` | +| `catalog.py` | `DocumentRecord`, `IndexationJob`, `DocumentStatus`, `JobStatus` | `TaskStateManager` + vectordb File table | +| `query.py` | `RetrievalQuery` | `pipeline.py` SearchQueries | +| `retrieval_result.py` | `RetrievalResult`, `ScoredChunk` | Search result dicts | +| `retrieval_response.py` | `RetrievalResponse` | Pipeline return values | +| `conversation.py` | `Conversation`, `Message` | openai router message handling | +| `contextualization.py` | `ContextualizedQuery` | pipeline.py query generation | +| `prompt.py` | `Prompt`, `PromptType` | prompts system | + +### Key design decisions + +**Chunk model** : + +```python +class Chunk(BaseModel): + id: str # UUID + document_id: str + text: str + chunk_index: int = 0 + chunk_type: ChunkType = ChunkType.TEXT + embedding: list[float] | None = None + metadata: dict[str, Any] = {} + partition: str = "default" + page_number: int | None = None + token_count: int | None = None + context: str | None = None # LLM contextualization text + + def with_embedding(self, embedding: list[float]) -> "Chunk": ... + + # Boundary converters (method body imports only) + @classmethod + def from_langchain(cls, doc: Any) -> "Chunk": ... + def to_langchain(self) -> Any: ... +``` + +**Document model** : + +```python +class Document(BaseModel): + id: str + filename: str + content_type: DocumentType + text: str | None = None + raw_bytes: bytes | None = Field(None, exclude=True) + partition: str = "default" + tags: list[str] = [] + metadata: dict[str, Any] = {} +``` + +### Commits + +``` +2.1 Create core/models/*.py (all domain types) +2.2 Add core/models/__init__.py convenience re-exports +2.3 Add from_langchain/to_langchain converters on Chunk and Document +``` + +--- + +## Phase 3 — Configuration Schemas + +**Goal:** Split monolithic `config/models.py` into typed sections in `core/config/`. + +### File mapping + +| New file | Source section | +| ------------------------------- | -------------------------------------------------- | +| `core/config/root.py` | `Settings` | +| `core/config/auth.py` | env vars (AUTH_TOKEN, SUPER_ADMIN_MODE) | +| `core/config/chunking.py` | `ChunkerConfig` | +| `core/config/endpoints.py` | `ModelEndpointConfig` + `ModelsConfig` (all 4: embedder, reranker, llm, vlm) | +| `core/config/indexation.py` | `LoaderConfig` + `RayConfig` (indexing parts) | +| `core/config/infrastructure.py` | `VectorDBConfig` + `RDBConfig` | +| `core/config/loader.py` | `load_config()` | +| `core/config/partition.py` | (new - per-partition config) | +| `core/config/presets.py` | (new - reusable configs) | +| `core/config/retrieval.py` | `RetrieverConfig` + `RerankerConfig` + `RAGConfig` | + +### Key improvement + +**ModelEndpointConfig pattern** for inference backends: + +```python +class ModelEndpointConfig(BaseModel): + endpoint: str + model_name: str | None = None + batch_size: int = 32 + timeout: float = 30.0 + extra: dict[str, Any] = {} # implementation-specific kwargs + context_window: int = 8192 + +class ModelsConfig(BaseModel): + embedder: dict[str, ModelEndpointConfig] = {} # {"default": ..., "fast": ...} + reranker: dict[str, ModelEndpointConfig] = {} + llm: dict[str, ModelEndpointConfig] = {} + vlm: dict[str, ModelEndpointConfig] = {} +``` + +This enables `make_component_factory()` to look up config by name and create the +right implementation via registry. Future: model endpoints stored in DB, loaded at startup. + +### Commits + +``` +3.1 Create core/config/*.py with typed schemas +3.2 Create core/config/loader.py (YAML + env override logic) +3.3 Update old config/__init__.py to re-export from core/config/ +3.4 Verify all existing imports still resolve +``` + +--- + +## Phase 4 — ABCs & Ports + +**Goal:** Define every contract as an ABC, co-located in its subject folder. +This is the architecturally critical phase. + +### 4A — Component ABCs (co-located in subject folders) + +Each ABC lives inside its subject folder. The folder's `__init__.py` re-exports +the ABC + registry so consumers get clean imports. + +| File | ABC | Key abstract methods | +| ------------------------------------------ | ------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `core/embeddings/embedder.py` | `Embedder` | `async embed(texts) -> list[list[float]]`, `async embed_single(text) -> list[float]`, `dimension: int` (property) | +| `core/llm/llm.py` | `LLM` | `async generate(prompt) -> str`, `async chat(messages) -> str`, `async stream_chat(messages) -> AsyncIterator[str]` (default: non-streaming fallback), `async generate_json(prompt) -> dict` (default: parse generate output), `async chat_with_tools(messages, tools) -> dict` (default: NotImplementedError) | +| `core/rerankers/reranker.py` | `Reranker` | `async rerank(query, documents, top_k) -> list[tuple[int, float]]` | +| `core/vlm/vlm.py` | `VLM` | `async caption_image(image_bytes, prompt) -> str`, `async caption_images_batch(images, prompt) -> list[str]` | +| `core/vector_stores/vector_store.py` | `VectorStore` | `async upsert(chunks, collection)`, `async search(embedding, top_k, collection, filters)`, `async delete(ids, collection)`, `async ensure_collection(name, dimension)`, `async drop_collection(name)`, `async collection_exists(name)`, `async query_ids_by_filter(collection, filters)`, `async query_chunks_by_filter(collection, filters, output_fields)` | +| `core/ports/catalog_store.py` | `CatalogStore` | `async initialize()`, `async shutdown()`, properties: `document_repo`, `job_repo`, `user_repo`, `prompt_repo`, `partition_repo`, ... (13 repos) | +| `core/chunking/chunking_strategy.py` | `ChunkingStrategy` | `chunk(document: ProcessedDocument, partition) -> list[Chunk]` | +| `core/indexing/parsers/document_parser.py` | `DocumentParser` | `async parse(document: Document) -> ProcessedDocument`, `supported_types() -> list[str]` | + +**Each `__init__.py` re-exports the public surface:** + +```python +# core/embeddings/__init__.py +from openrag.core.embeddings.embedder import Embedder +from openrag.core.embeddings.registry import embedder_registry +__all__ = ["Embedder", "embedder_registry"] + +# core/rerankers/__init__.py +from openrag.core.rerankers.reranker import Reranker +from openrag.core.rerankers.registry import reranker_registry +__all__ = ["Reranker", "reranker_registry"] + +# core/llm/__init__.py +from openrag.core.llm.llm import LLM +from openrag.core.llm.registry import llm_registry +__all__ = ["LLM", "llm_registry"] + +# core/vlm/__init__.py +from openrag.core.vlm.vlm import VLM +from openrag.core.vlm.registry import vlm_registry +__all__ = ["VLM", "vlm_registry"] + +# core/vector_stores/__init__.py +from openrag.core.vector_stores.vector_store import VectorStore +__all__ = ["VectorStore"] + +# core/chunking/__init__.py +from openrag.core.chunking.chunking_strategy import ChunkingStrategy +from openrag.core.chunking.registry import chunking_registry +__all__ = ["ChunkingStrategy", "chunking_registry"] + +# core/indexing/parsers/__init__.py +from openrag.core.indexing.parsers.document_parser import DocumentParser +from openrag.core.indexing.parsers.registry import parser_registry +__all__ = ["DocumentParser", "parser_registry"] +``` + +### 4B — Repository ports (core/ports/)### 2.14 Auth upgrade path + +A more complete auth system would use JWT + API keys + bcrypt passwords + +SystemRole (superadmin/admin/user) + PartitionRole (owner/reader). OpenRAG currently +uses SHA-256 token hashing with a simpler role system (viewer/editor/owner). +The refactor should preserve OpenRAG's auth as-is, with the option to upgrade +to JWT/Keycloak later. + +Key examples: + +```python +# core/ports/document_repo.py +class DocumentRepository(ABC): + async def create_document(self, doc: DocumentRecord) -> DocumentRecord: ... + async def get_document(self, document_id: str) -> DocumentRecord | None: ... + async def list_documents(self, partition: str | list[str] | None = None, ...) -> list[DocumentRecord]: ... + async def update_document(self, document_id: str, **fields) -> DocumentRecord | None: ... + async def delete_document(self, document_id: str) -> bool: ... + async def count_documents(self, partition: str | list[str] | None = None, ...) -> int: ... + async def get_by_hash_in_partition(self, partition: str, file_hash: str) -> DocumentRecord | None: ... + +# core/ports/chunk_repo.py +class ChunkRepository(ABC): + async def bulk_insert(self, chunks: list[dict]) -> int: ... + async def get_by_ids(self, chunk_ids: list[str]) -> list[dict]: ... + async def get_by_document_id(self, document_id: str) -> list[dict]: ... + async def delete_by_document_id(self, document_id: str) -> int: ... + async def bm25_search(self, query_text: str, partition: str, top_k: int = 20) -> list[dict]: ... + +# core/ports/user_repo.py +class UserRepository(ABC): + async def create_user(self, user: User) -> User: ... + async def get_user(self, user_id: str) -> User | None: ... + async def get_user_by_token(self, token_hash: str) -> User | None: ... + async def list_users(self, ...) -> list[User]: ... + # + partition assignment methods +``` + +### Commits + +``` +4.1 Create core/embeddings/{embedder.py, registry.py, __init__.py} +4.2 Create core/rerankers/{reranker.py, registry.py, __init__.py} +4.3 Create core/llm/{llm.py, registry.py, __init__.py} +4.4 Create core/vlm/{vlm.py, registry.py, __init__.py} +4.5 Create core/vector_stores/{vector_store.py, __init__.py} +4.6 Create core/chunking/{chunking_strategy.py, registry.py, __init__.py} +4.7 Create core/indexing/parsers/{document_parser.py, registry.py, __init__.py} +4.8 Create all port ABCs in core/ports/ (including catalog_store.py) +``` + +--- + +## Phase 5 — Core Domain Logic + +**Goal:** Move pure business logic into `core/`. First phase that MOVES code. + +### 5A — Retrieval core + +| Target | Source | Key change | +| ----------------------------- | ----------------------------------------------------- | ----------------------------------------------- | +| `core/retrieval/rrf.py` | `components/reranker/base.py` rrf_reranking() | Pure math, no changes | +| `core/retrieval/retriever.py` | `components/retriever.py` ABCRetriever hierarchy | Uses VectorStore port instead of get_vectordb() | +| `core/retrieval/hydration.py` | `components/retriever.py` _expand_with_related_chunks | Uses VectorStore port | +| `core/retrieval/pipeline.py` | `components/pipeline.py` RetrieverPipeline | Orchestrates via injected ports | + +**Critical:** Retriever strategies call `VectorStore.search()` (port method), not +`vectordb.async_search.remote()` (Ray call). The Ray call moves to the adapter. + +### 5B — Chunking strategies + +| Target | Source | +| ----------------------------------- | --------------------------------------------------------- | +| `core/chunking/recursive.py` | `components/indexer/chunker/chunker.py` RecursiveSplitter | +| `core/chunking/markdown_section.py` | `components/indexer/chunker/utils.py` markdown parsing | + +Register via decorator: `@chunking_registry.register("recursive")`. + +### 5C — Prompt builders + +| Target | Source | +| ------------------------------------------- | ----------------------------------------------------------------- | +| `core/prompts/chat_prompt_builder.py` | `components/pipeline.py` context formatting + system prompt | +| `core/prompts/vlm_prompt_builder.py` | `components/indexer/loaders/base.py` VLM prompt logic | +| `core/prompts/contextualization_builder.py` | `components/indexer/chunker/chunker.py` contextualization prompts | +| `core/prompts/query_rewriter.py` | `components/retriever.py` multi-query + HyDE prompt building | +| `core/prompts/map_reduce_builder.py` | `components/map_reduce.py` map/reduce prompts | +| `core/prompts/template_loader.py` | `components/prompts/prompts.py` disk-based template loading | + +### 5D — Indexing domain logic + +| Target | Source | +| ------------------------------------ | ----------------------------------------------------------- | +| `core/indexing/contextualize.py` | `components/indexer/chunker/chunker.py` ChunkContextualizer | +| `core/indexing/text_preprocessor.py` | `components/indexer/utils/text_sanitizer.py` | +| `core/indexing/validators.py` | `routers/utils.py` validation functions | +| `core/indexing/parsers/*.py` | `components/indexer/loaders/*.py` | + +### Commits + +``` +5.1-5.4 Retrieval core (rrf, retriever, hydration, pipeline) +5.5-5.6 Chunking strategies +5.7-5.12 Prompt builders (chat, vlm, contextualization, query rewriter, map-reduce, template loader) +5.13-5.17 Indexing domain logic + parsers +5.15 Update old files to re-export from core/ +``` + +--- + +## Phase 6 — Inference Adapters + +**Goal:** Move all HTTP client code to `services/inference/`. Register with core registries. + +### Files + +| Target | Source | Registers as | +| --------------------------------------------- | --------------------------------------------------------------- | ----------------------------------------------------------------------- | +| `services/inference/vllm_client.py` | `components/llm.py` + `components/indexer/embeddings/openai.py` | `@llm_registry.register("vllm")`, `@embedder_registry.register("vllm")` | +| `services/inference/vlm_client.py` | `components/indexer/loaders/base.py` VLM methods | `@vlm_registry.register("vllm")` | +| `services/inference/infinity_client.py` | `components/reranker/infinity.py` | `@reranker_registry.register("infinity")` | +| `services/inference/distributed_semaphore.py` | `components/utils.py` DistributedSemaphore* | Ray-based cluster-wide limiter | +| `services/inference/_circuit_breaker.py` | (new) | `@with_circuit_breaker` decorator (aiobreaker) | +| `services/inference/_retry.py` | (new) | `@with_retry` decorator (tenacity + jitter) | +| `services/inference/_timeout.py` | (new) | `@with_timeout` decorator (asyncio.timeout) | +| `services/inference/healthcheck.py` | `routers/utils.py` get_openai_models() | Endpoint readiness probes | + +### Key pattern + +**vLLM client with distributed semaphore:** + +```python +@llm_registry.register("vllm") +class VLLMClient(LLM): + def __init__(self, endpoint: str, model_name: str | None = None, timeout: float = 120.0, ...): + self._client = httpx.AsyncClient(...) + self._semaphore = DistributedLLMSemaphore(...) + + async def chat(self, messages, **kwargs): + async with self._slot(): # acquire distributed semaphore + response = await self._client.post(f"{self._endpoint}/v1/chat/completions", ...) + return response.json()["choices"][0]["message"]["content"] +``` + +### Commits + +``` +6.1 Create services/inference/vllm_client.py (LLM + Embedder) +6.2 Create services/inference/infinity_client.py (Reranker) +6.3 Create services/inference/vlm_client.py (VLM) +6.4 Create services/inference/distributed_semaphore.py +6.5 Create services/inference/_circuit_breaker.py + _retry.py + _timeout.py + (shared resilience decorators applied by all inference clients) +6.6 Create services/inference/healthcheck.py +6.7 Update old files to re-export +``` + +--- + +## Phase 7 — Storage & Persistence Adapters + +**This is the highest-risk phase.** The current `Vectordb` Ray actor (god object) must be +decomposed into: + +- `MilvusVectorStore` (implements `VectorStore`) +- `PostgresStore` (implements `CatalogStore`, owns 13 repos) + +### 7A — PostgreSQL persistence layer + +**Follow the composite pattern:** + +**ConnectionManager** (`services/persistence/connection.py`): + +```python +class ConnectionManager: + def __init__(self, config: PostgresConfig): + self._dsn = build_dsn(config) + self._pool: asyncpg.Pool | None = None + + async def initialize(self) -> None: + self._pool = await asyncpg.create_pool(self._dsn, ...) + + @property + def pool(self) -> asyncpg.Pool: + if self._pool is None: + raise RuntimeError("Not initialized") + return self._pool + + async def shutdown(self) -> None: + if self._pool: + await self._pool.close() +``` + +**Repository implementations** — extract from `PartitionFileManager`: + +| New file | Source methods | Implements | +| ------------------------------- | ------------------------------------------------------------------------------ | --------------------- | +| `persistence/document_repo.py` | add_file, remove_file, update_file, list_files, file_exists | `DocumentRepository` | +| `persistence/user_repo.py` | create_user, get_user, get_user_by_token, delete_user, update_user, list_users | `UserRepository` | +| `persistence/partition_repo.py` | create_partition, delete_partition, list_partitions, partition_exists | `PartitionRepository` | +| `persistence/job_repo.py` | TaskStateManager state methods | `JobRepository` | +| `persistence/chunk_repo.py` | (new - for future BM25 in PG) | `ChunkRepository` | +| ... and 8 more repos | | | + +**PostgresStore** (`services/storage/postgres_store.py`): + +```python +class PostgresStore(CatalogStore): + def __init__(self, config: PostgresConfig): + self._conn = ConnectionManager(config) + pool_getter = lambda: self._conn.pool + self._documents = PgDocumentRepository(pool_getter) + self._users = PgUserRepository(pool_getter) + self._partitions = PgPartitionRepository(pool_getter) + self._jobs = PgJobRepository(pool_getter) + # ... 13 total + + @property + def document_repo(self) -> DocumentRepository: + return self._documents + # ... etc +``` + +### 7B — Milvus vector store + +**MilvusVectorStore** (`services/storage/milvus_store.py`): + +```python +class MilvusVectorStore(VectorStore): + def __init__(self, config: VectorDBConfig): + uri = f"http://{config.host}:{config.port}" + self._client = MilvusClient(uri=uri) + self._async_client = AsyncMilvusClient(uri=uri) + self._hybrid = config.hybrid_search + + async def upsert(self, chunks: list[Chunk], collection: str = "default") -> int: ... + async def search(self, embedding, top_k, collection, filters) -> list[dict]: ... + async def delete(self, ids, collection) -> int: ... +``` + +OpenRAG uses Milvus BM25 (sparse vectors), +The alternative (Postgres tsvector) is not used here. Keep OpenRAG's hybrid search approach. + +**Embedding removed from vector store.** Currently `MilvusDB.async_add_documents()` +embeds chunks internally. After refactor, embedding happens in the pipeline stage +BEFORE calling `vector_store.upsert()`. The vector store receives pre-embedded chunks. + +### 7C — Shim the god object + +After creating the new stores, update `vectordb.py` to delegate: + +```python +# Old MilvusDB actor becomes a thin wrapper +class MilvusDB: + def __init__(self): + config = load_config() + self._vector_store = MilvusVectorStore(config.vectordb) + self._catalog_store = PostgresStore(config.rdb) + # ... delegate all methods +``` + +### Commits + +``` +7.1 services/persistence/connection.py +7.2 services/persistence/schema.py (copy SQLAlchemy models) +7.3 services/persistence/document_repo.py +7.4 services/persistence/user_repo.py +7.5 services/persistence/partition_repo.py +7.6 services/persistence/job_repo.py +7.7 services/persistence/chunk_repo.py + remaining repos +7.8 services/storage/milvus_store.py +7.9 services/storage/postgres_store.py (composite) +7.10 Update MilvusDB actor to delegate to new stores (shim) +7.11 Integration test: full upload-search cycle works +``` + +--- + +## Phase 8 — Orchestrators + +**Goal:** Business services in `services/orchestrators/` that coordinate ports. + +### Key pattern + +Orchestrators receive **factory callables**, not instances: + +```python +class RetrievalService: + def __init__( + self, + vector_store: VectorStore, + embedder_factory: Callable[[str], Embedder], + reranker_factory: Callable[[str], Reranker], + llm_factory: Callable[[str], LLM], + config: Settings, + chunk_repo: ChunkRepository | None = None, + ): + self._vector_store = vector_store + self._embedder_factory = embedder_factory + # ... + + async def retrieve(self, query: str, partition: str, **kwargs): + embedder = self._embedder_factory("default") + embedding = await embedder.embed_single(query) + results = await self._vector_store.search(embedding, ...) + if self._reranker_factory: + reranker = self._reranker_factory("default") + results = await reranker.rerank(query, results, ...) + return results +``` + +### Orchestrator inventory + +| Service | Source | Dependencies (injected) | +| ------------------------- | -------------------------------------------- | --------------------------------------------------------- | +| `auth_service.py` | `routers/utils.py` + `api.py` AuthMiddleware | UserRepository, AuthConfig | +| `user_service.py` | `routers/users.py` | UserRepository | +| `partition_service.py` | `routers/partition.py` | PartitionRepository, VectorStore | +| `document_service.py` | vectordb file ops | DocumentRepository, VectorStore | +| `retrieval_service.py` | `pipeline.py` RetrieverPipeline | VectorStore, EmbedderFactory, RerankerFactory, LLMFactory | +| `query_service.py` | `pipeline.py` RagPipeline | RetrievalService, LLMFactory, PromptService | +| `query_orchestrator.py` | `pipeline.py` generate_query | LLMFactory, config | +| `indexing_service.py` | `indexer.py` add_file flow | VectorStore, CatalogStore, config | +| `job_service.py` | TaskStateManager | JobRepository | +| `conversation_service.py` | openai router chat logic | ConversationRepository, LLMFactory | +| `prompt_service.py` | `prompts/prompts.py` | PromptRepository | +| `conversion_service.py` | `loaders/serializer.py` | ParserRegistry, VLMFactory | +| `cluster_service.py` | `routers/actors.py` | Ray introspection | +| `llm_contextualizer.py` | `pipeline.py` query gen | LLMFactory | +| `research_planner.py` | multi-query logic | LLMFactory | + +### Commits + +``` +8.1-8.15 One commit per orchestrator (create + wire) +``` + +--- + +## Phase 9 — Workers (Ray Isolation) + +**Goal:** All Ray code in `services/workers/`. Core and orchestrators are Ray-free. + +### IndexerActor — thin wrapper pattern + +```python +# services/workers/indexer_actor.py +@ray.remote(max_concurrency=2, max_restarts=3) +class IndexerActorClass: + def __init__(self): + # Side-effect imports to register with registries + import openrag.services.inference.vllm_client # noqa + import openrag.services.inference.infinity_client # noqa + + async def process_document(self, doc_dict: dict, config_dict: dict, ...) -> dict: + document = Document.model_validate(doc_dict) + config = Settings.model_validate(config_dict) + result = await ray_data_ingest_documents([document], config, ...) + return {"status": "success", **result} +``` + +### Pipeline stages + +``` +services/workers/stages/parse.py <- loaders/serializer.py +services/workers/stages/caption.py <- loaders/base.py VLM calls +services/workers/stages/chunk.py <- chunker flow +services/workers/stages/contextualize.py <- ChunkContextualizer +services/workers/stages/embed.py <- embedding flow +services/workers/stages/store.py <- vectordb insert flow +``` + +Each stage follows the same pattern: async function, timeout with base + per-chunk +scaling, in-place row mutation, error marking, credential scrubbing after use. + +### Commits + +``` +9.1 services/workers/ray_utils.py +9.2 services/workers/stages/*.py +9.3 services/workers/pipeline_builder.py +9.4 services/workers/indexer_actor.py (thin wrapper) +9.5 Integration test: indexing pipeline e2e +``` + +--- + +## Phase 10 — API Layer Restructure + +**Goal:** Move all FastAPI code to `api/` with clean DI. + +### 10A — Error handlers + + `api/error_handlers.py` maps domain exceptions to HTTP via MRO walk: + +```python +_STATUS_MAP = { + NotFoundError: 404, + AuthenticationError: 401, + AuthError: 403, + ValidationError: 422, + InferenceTimeoutError: 504, + LLMParsingError: 502, + InferenceError: 503, + ServiceUnavailableError: 503, + StorageError: 500, + QuotaExceededError: 429, +} +``` + +Response includes `request_id` from structlog contextvars. + +### 10B — Middleware + +| Target | Source | +| ------------------------------------ | -------------------------------------------- | +| `api/middleware/request_id.py` | (new) | +| `api/middleware/instrumentation.py` | `routers/monitoring.py` MonitoringMiddleware | +| `api/middleware/security_headers.py` | (new) | +| `api/middleware/request_timeout.py` | (new) | +| `api/middleware/idempotency.py` | (new) | + +### 10C — Auth dependencies + +```python +# api/dependencies/auth.py +async def get_current_user( + request: Request, + auth_service: AuthService = Depends(get_auth_service), +) -> User: + token = _extract_token(request) + if not token: + raise AuthenticationError("Missing authentication") + return await auth_service.authenticate(token) +``` + +### 10D — Routers + +All routers use `Depends(get_service)` from `di/providers.py`: + +```python +@router.get("/search") +async def search( + service: RetrievalService = Depends(get_retrieval_service), + query: str = Query(...), + partition: str = Query("default"), +): + return await service.retrieve(query, partition) +``` + +### 10E — main.py + +```python +# api/main.py +from openrag.di.container import ServiceContainer +from openrag.di.providers import set_container + +async def lifespan(app: FastAPI): + container = ServiceContainer() + set_container(container) + await container.initialize() + yield + await container.shutdown() + +app = FastAPI(lifespan=lifespan) +# register middleware, routers, error handlers +``` + +### Commits + +``` +10.1 api/error_handlers.py +10.2 api/middleware/*.py +10.3 api/dependencies/auth.py +10.4 api/schemas/**/*.py +10.5 api/routers/user/health.py (simplest, proof of concept) +10.6-10.15 Remaining routers (one commit each) +10.16 api/main.py (wire everything, keep old routers in parallel) +10.17 Remove old routers from mounts (one at a time) +``` + +--- + +## Phase 11 — Composition Root (DI) + +**Goal:** `ServiceContainer` wires everything. Remove all global singletons. + +### container.py — sync/async lifecycle + +```python +class ServiceContainer: + def __init__(self, config: Settings | None = None): + # 1. Load config + self._config = config or load_config() + setup_logging(self._config) + + # 2. Register implementations (side-effect imports) + register_embedders() # di/embedders.py + register_rerankers() + register_llms() + register_vlms() + + # 3. Create infrastructure stores + self._postgres_store = PostgresStore(self._config.infrastructure.rdb) + self._milvus_store = MilvusVectorStore(self._config.infrastructure.vectordb) + + # 4. Create component factories (thread-safe, cached) + self._client_caches: list[dict] = [] + self._embedder_factory = make_component_factory( + embedder_registry, self._config.models.embedder, "vllm", self._client_caches + ) + self._reranker_factory = make_component_factory( + reranker_registry, self._config.models.reranker, "infinity", self._client_caches + ) + self._llm_factory = make_component_factory( + llm_registry, self._config.models.llm, "vllm", self._client_caches + ) + self._vlm_factory = make_component_factory( + vlm_registry, self._config.models.vlm, "vllm", self._client_caches + ) + + # 5. Create orchestrator services + self._auth_service = AuthService(self._config.auth, self._postgres_store.user_repo) + self._retrieval_service = RetrievalService( + self._milvus_store, self._embedder_factory, self._reranker_factory, + self._llm_factory, self._config, + ) + self._indexing_service = IndexingService( + self._milvus_store, self._config, + ) + # ... all other services + + async def initialize(self): + """Called from FastAPI lifespan. Does all async I/O.""" + await self._postgres_store.initialize() + self._postgres_store.run_migrations() + await self._milvus_store.ensure_collection(...) + # seed defaults... + self._initialized = True + + async def shutdown(self): + for cache in self._client_caches: + for client in cache.values(): + if hasattr(client, "aclose"): + await client.aclose() + await self._postgres_store.shutdown() +``` + +### providers.py + +```python +_container: ServiceContainer | None = None +_lock = threading.Lock() + +def set_container(c: ServiceContainer): + global _container + with _lock: + _container = c + +def _require_initialized() -> ServiceContainer: + if _container is None or not _container.is_initialized: + raise RuntimeError("Container not initialized") + return _container + +def get_retrieval_service() -> RetrievalService: + return _require_initialized().retrieval_service + +def get_config() -> Settings: + return _require_initialized().config +# ... one getter per service +``` + +### factories.py — make_component_factory() + +The make_component_factory() implementation: + +```python +def make_component_factory( + registry: Registry[T], + config_section: dict[str, ModelEndpointConfig], + default_impl: str, + client_caches: list[dict[str, T]], + extra_kwargs_fn: Callable | None = None, +) -> Callable[[str], T]: + cache: dict[str, T] = {} + lock = threading.Lock() + client_caches.append(cache) + + def factory(name: str = "default") -> T: + if name in cache: + return cache[name] + with lock: + if name in cache: + return cache[name] + model_cfg = config_section[name] + impl = model_cfg.extra.get("implementation", default_impl) + kwargs = {"endpoint": model_cfg.endpoint, "model_name": model_cfg.model_name, ...} + if extra_kwargs_fn: + kwargs.update(extra_kwargs_fn(model_cfg)) + instance = registry.create(impl, **kwargs) + cache[name] = instance + return instance + + return factory +``` + +### Registration modules + +```python +# di/embedders.py +def register_embedders() -> None: + import openrag.services.inference.vllm_client # noqa: F401 + +# di/rerankers.py +def register_rerankers() -> None: + import openrag.services.inference.infinity_client # noqa: F401 + +# di/llms.py +def register_llms() -> None: + import openrag.services.inference.vllm_client # noqa: F401 + +# di/vlms.py +def register_vlms() -> None: + import openrag.services.inference.vlm_client # noqa: F401 +``` + +### Commits + +``` +11.1 di/factories.py (make_component_factory) +11.2 di/embedders.py, di/rerankers.py, di/llms.py, di/vlms.py +11.3 di/repositories.py, di/vector_stores.py +11.4 di/container.py (ServiceContainer) +11.5 di/providers.py +11.6 Wire into api/main.py lifespan +11.7 Update routers to use Depends() from providers (one at a time) +11.8 Remove global singletons from utils/dependencies.py +11.9 Remove module-level config = load_config() calls +``` + +--- + +## Phase 12 — Internal Cleanup & Remove Shims + +**Goal:** Delete all backward-compatibility re-exports and old internal code inside +the `openrag/` Python package. After this phase, only the new 3-layer structure +remains inside the package. + +### Commit sequence + +``` +12.1 Run import guard - verify zero violations +12.2 Remove components/retriever.py (-> core/retrieval/) +12.3 Remove components/reranker/ (-> core/rerankers/ + services/inference/) +12.4 Remove components/llm.py (-> core/llm/ + services/inference/) +12.5 Remove components/pipeline.py (-> services/orchestrators/) +12.6 Remove components/map_reduce.py (-> core/prompts/map_reduce_builder.py) +12.7 Remove components/utils.py (-> core/utils/ + services/inference/) +12.8 Remove components/indexer/ (-> core/indexing/ + services/) +12.9 Remove components/websearch/ (-> services/ or keep as adapter) +12.10 Remove routers/ (-> api/routers/) +12.11 Remove models/ (-> core/models/ + api/schemas/) +12.12 Remove utils/dependencies.py (-> di/) +12.13 Remove utils/exceptions/ (-> core/utils/exceptions.py) +12.14 Remove config/ (-> core/config/) +12.15 Delete empty components/, routers/, models/, utils/ directories +12.16 Verify: python -c "import openrag" + all tests pass +``` + +--- + +## Phase 13 — Project Layout, Infra, Tests & UI + +**Goal:** Restructure the top-level project layout from the current flat/scattered +structure to the clean target layout. Move deployment, tests, scripts, and UI +to their proper locations. + +### Current top-level layout (what exists today) + +``` +openrag_1.1.7/ +|-- openrag/ # Python package (mixed: code + scripts + tests + static) +| |-- scripts/ # CLI tools + Alembic migrations (INSIDE package) +| |-- tests/ # Some unit tests (INSIDE package) +| |-- public/ # Static assets (INSIDE package) +| +-- app_front.py # Chainlit frontend (INSIDE package) +| +|-- Dockerfile # Root-level (no infra/ folder) +|-- Dockerfile.ray # Root-level +|-- docker-compose.yaml # Root-level +|-- entrypoint.sh # Root-level +|-- conf/ # Config YAML (OK, stays) +|-- tests/ # Integration tests + Robot Framework +| |-- api_tests/ # pytest integration +| +-- api/ # Robot Framework +|-- extern/ # Submodules: vllm, reranker, indexer-ui +|-- prompts/ # Prompt templates (OK, stays or moves into openrag/) +|-- openrag_metrics/ # Grafana + Prometheus configs +|-- ansible/ # Ansible deployment +|-- charts/ # Helm charts +|-- vdb/ # Milvus config +|-- benchmarks/ # Performance tests +|-- quick_start/ # Getting started examples +|-- utility/ # Misc utility scripts +|-- docs/ # Documentation site (Astro) ++-- pyproject.toml, uv.lock, pytest.ini, etc. +``` + +### Target top-level layout + +``` +openrag/ +|-- openrag/ # Python package (ONLY application code) +| |-- core/ +| |-- services/ +| |-- api/ +| |-- di/ +| +-- prompts/ # Prompt templates (disk-loaded by the app) +| +|-- conf/ # YAML configuration files per environment +|-- infra/ # ALL deployment infrastructure +| |-- docker/ +| | |-- api.Dockerfile # <- was Dockerfile +| | +-- ray.Dockerfile # <- was Dockerfile.ray +| |-- compose/ +| | |-- docker-compose.yaml # <- was root docker-compose.yaml +| | |-- .env.example +| | |-- grafana/ # <- was openrag_metrics/grafana +| | |-- prometheus/ # <- was openrag_metrics/prometheus +| | |-- milvus/ # <- was root vdb/ +| | |-- nginx/ # reverse proxy config +| | +-- postgres/ # Postgres init scripts if any +| |-- scripts/ +| | +-- entrypoint.sh # <- was root entrypoint.sh +| |-- ansible/ # <- was root ansible/ +| +-- charts/ # <- was root charts/ +| +|-- scripts/ # Operational CLI tools +| |-- migrate.py # <- was openrag/scripts/migrations/ +| |-- backup.py # <- was openrag/scripts/backup.py +| |-- restore.py # <- was openrag/scripts/restore.py +| |-- embed.py # <- was openrag/scripts/embed.py +| +-- check_file_counts.py # <- was openrag/scripts/check_file_counts.py +| +|-- tests/ # ALL tests (unified) +| |-- unit/ # Unit tests (mirrors openrag/ package structure) +| | |-- core/ +| | | |-- test_registry.py +| | | |-- test_rrf.py +| | | +-- test_chunk_model.py +| | |-- services/ +| | | |-- test_milvus_store.py +| | | +-- test_vllm_client.py +| | +-- api/ +| | +-- test_error_handlers.py +| |-- integration/ # End-to-end tests (need running services) +| | |-- test_indexer.py # <- was tests/api_tests/test_indexer.py +| | |-- test_search.py # <- was tests/api_tests/test_search.py +| | |-- test_openai_compat.py +| | |-- test_users.py +| | |-- test_partitions.py +| | |-- test_workspaces.py +| | |-- conftest.py # <- was tests/api_tests/conftest.py +| | +-- mock_vllm.py # <- was tests/api_tests/api_run/mock_vllm.py +| |-- load/ # Performance/load tests +| | +-- (from benchmarks/) +| +-- conftest.py # Root conftest with markers: unit, integration, slow +| +|-- docs/ # Human-readable documentation +|-- ui/ # Admin frontend (indexer-ui submodule or standalone) +|-- pyproject.toml ++-- uv.lock +``` + +### 13A — Move deployment infrastructure to infra/ + +``` +13A.1 Create infra/{docker,compose,scripts} directories +13A.2 Move Dockerfile -> infra/docker/api.Dockerfile + Update build context and paths inside the Dockerfile +13A.3 Move Dockerfile.ray -> infra/docker/ray.Dockerfile +13A.4 Move docker-compose.yaml -> infra/compose/docker-compose.yaml + Update build.context, build.dockerfile paths + Update volume mount paths +13A.5 Move entrypoint.sh -> infra/scripts/entrypoint.sh + Update Dockerfile COPY to match new path +13A.6 Move service configs into infra/compose/ (alongside docker-compose): + openrag_metrics/grafana -> infra/compose/grafana/ + openrag_metrics/prometheus -> infra/compose/prometheus/ + vdb/ -> infra/compose/milvus/ + (add nginx/, postgres/ subdirs as needed) +13A.7 Move ansible/ -> infra/ansible/ +13A.8 Move charts/ -> infra/charts/ +13A.9 Verify: docker compose -f infra/compose/docker-compose.yaml build +``` + +### 13B — Move scripts out of the Python package + +Currently `openrag/scripts/` lives inside the Python package, which means CLI tools +are bundled in the Docker image and can import from `openrag.*` but are mixed with +application code. + +``` +13B.1 Create top-level scripts/ directory +13B.2 Move openrag/scripts/backup.py -> scripts/backup.py +13B.3 Move openrag/scripts/restore.py -> scripts/restore.py +13B.4 Move openrag/scripts/embed.py -> scripts/embed.py +13B.5 Move openrag/scripts/check_file_counts.py -> scripts/check_file_counts.py +13B.6 Move openrag/scripts/filter-logs.py -> scripts/filter_logs.py +13B.7 Create scripts/migrate.py wrapping Alembic + Migrations stay in services/persistence/migrations/ (part of the package) + The CLI script just calls alembic programmatically +13B.8 Move shell scripts (backup.sh.example, etc.) -> scripts/ +13B.9 Remove empty openrag/scripts/ directory +13B.10 Update any Docker CMD or documentation referencing old script paths +``` + +### 13C — Restructure tests + +Currently tests are split between `openrag/**/test_*.py` (inside the package) and +`tests/api_tests/` (outside). Unify into a single `tests/` tree. + +``` +13C.1 Create tests/{unit,integration,load} directories +13C.2 Create tests/conftest.py with shared fixtures and markers: + @pytest.mark.unit, @pytest.mark.integration, @pytest.mark.slow +13C.3 Move tests/api_tests/*.py -> tests/integration/ + Update imports in test files +13C.4 Move tests/api_tests/conftest.py -> tests/integration/conftest.py +13C.5 Move tests/api_tests/api_run/mock_vllm.py -> tests/integration/mock_vllm.py +13C.6 Move inline test files from openrag/ to tests/unit/: + openrag/components/indexer/chunker/test_chunking.py -> tests/unit/core/test_chunking.py + openrag/components/reranker/test_rrf_reranking.py -> tests/unit/core/test_rrf.py + openrag/test_token_validation.py -> tests/unit/test_token_validation.py + openrag/test_version.py -> tests/unit/test_version.py + ... (all test_*.py files inside openrag/) +13C.7 Move benchmarks/ -> tests/load/ +13C.8 Update pytest.ini / pyproject.toml [tool.pytest]: + testpaths = ["tests"] + markers: + unit: Unit tests (no external services) + integration: Integration tests (need running services) + slow: Long-running tests +13C.9 Add CI-friendly test commands: + uv run pytest -m unit # fast, no infra needed + uv run pytest -m integration # needs docker services + uv run pytest -m "not slow" # skip load tests +13C.10 Verify: uv run pytest -m unit passes +13C.11 Remove tests/api/ Robot Framework tests or move to tests/robot/ +13C.12 Remove old tests/ subdirectories +``` + +### 13D — Prompts location + +Move prompts inside the Python package so they're bundled with the app: + +``` +13D.1 Move prompts/ -> openrag/prompts/ + (or keep at root and update config paths — choose one) +13D.2 Update core/config paths to reference new prompts location +13D.3 Update Dockerfile COPY to include openrag/prompts/ +``` + +### 13E — UI submodule + +The admin frontend currently lives in `extern/indexer-ui` as a git submodule. + +``` +13E.1 Move extern/indexer-ui -> ui/ + Or: keep as submodule but reference from ui/ symlink +13E.2 Update docker-compose service for indexer-ui to use new path +13E.3 Remove extern/ directory (or keep for vllm/reranker submodules if still needed) +``` + +### 13F — pyproject.toml and root cleanup + +``` +13F.1 Update pyproject.toml: + - Update package name if renaming + - Update [tool.pytest] testpaths + - Update [tool.ruff] src paths + - Verify [project.scripts] entry points +13F.2 Update .github/workflows/ CI: + - Test commands use new paths + - Docker build context uses infra/docker/ + - docker-compose -f infra/compose/docker-compose.yaml +13F.3 Remove root-level files that moved: + - Dockerfile, Dockerfile.ray (now in infra/docker/) + - docker-compose.yaml (now in infra/compose/) + - entrypoint.sh (now in infra/scripts/) + - pytest.ini (config now in pyproject.toml) +13F.4 Remove stale directories: + - quick_start/ (move useful content to docs/ or remove) + - utility/ (merge into scripts/ or remove) + - openrag.egg-info/ (regenerated on build) + - model_weights/, logs/ (runtime dirs, add to .gitignore) +13F.5 Update README.md with new project layout +13F.6 Update CLAUDE.md with new paths and commands +13F.7 Final verification: + - uv sync + - uv run pytest -m unit + - docker compose -f infra/compose/docker-compose.yaml build + - docker compose -f infra/compose/docker-compose.yaml up -d + - Integration tests pass +``` + +--- + +## Phase 14 — Per-Partition Presets (Indexation & Retrieval) + +**Goal:** Add the presetting mechanism that gives fine-grained +indexation and retrieval configuration per partition. This is a feature addition +on top of the clean architecture, not a refactoring step. + +**Key files to create:** + +- `core/config/partition.py` — `PartitionConfig`, `PartitionRow` +- `core/config/presets.py` — `PresetsConfig`, `PresetRow` +- `core/config/indexation.py` — `IndexationPipelineConfig` (25+ fields) +- `core/config/retrieval.py` — `RetrievalPipelineConfig`, `IntentStrategyConfig` +- `services/orchestrators/partition_service.py` — loads partitions into config +- `services/orchestrators/preset_service.py` — preset CRUD + seeding +- `services/persistence/partition_repo.py` — `PartitionRow` persistence +- `services/persistence/preset_repo.py` — `PresetRow` persistence +- `api/routers/admin/presets.py` — preset admin endpoints +- `api/routers/admin/pipelines.py` — dry-run preview endpoint + +### 14A — Config models (already scaffolded in Phase 3) + +Flesh out the config models created in Phase 3 with full pipeline detail: + +**`core/config/indexation.py`** — `IndexationPipelineConfig`: + +```python +class IndexationPipelineConfig(BaseModel): + chunking: ChunkingConfig = Field(default_factory=ChunkingConfig) + parsing_strategy: str = "marker" + vlm: str | None = None + enable_image_captioning: bool = True + enable_contextualization: bool = False + contextualization_llm: str | None = None + contextualization_mode: str = "structured" # none | simple | structured + contextualization_window: int = 1 + contextualization_max_tokens: int = 2048 + enable_metadata_extraction: bool = True + metadata_extraction_llm: str | None = None + enable_entity_extraction: bool = True + entity_labels: list[str] = ["person", "organization", "location", "event"] + enable_topic_tagging: bool = True + max_topic_tags: int = 7 + topic_tagging_llm: str | None = None + # prompt override names (resolved via PromptService) + contextualization_prompt_name: str | None = None + vlm_caption_prompt_name: str | None = None + image_caption_prompt_name: str | None = None +``` + +**`core/config/retrieval.py`** — `RetrievalPipelineConfig`: + +```python +class IntentStrategyConfig(BaseModel): + pipeline: str = "unified" + top_k: int = 10 + top_n: int = 5 + +class RetrievalPipelineConfig(BaseModel): + type: str = "unified" + reranker: str | None = None + llm: str | None = None + top_k: int = 20 + top_n: int = 10 + enable_reranker: bool = True + enable_planner: bool = True + intent_strategies: dict[str, IntentStrategyConfig] = Field(default_factory=_default_intent_strategies) + rrf_k: int = 60 +``` + +**`core/config/partition.py`** — `PartitionConfig`: + +```python +class PartitionConfig(BaseModel): + name: str + description: str = "" + embedder: str = "default" + indexation: IndexationPipelineConfig = Field(default_factory=IndexationPipelineConfig) + retrieval: RetrievalPipelineConfig = Field(default_factory=RetrievalPipelineConfig) + collection_name: str | None = None + chat_history_depth: int = 0 + chat_llm: str | None = None + +class PartitionRow(BaseModel): + """DB representation — references presets by name.""" + name: str + display_name: str | None = None + description: str = "" + embedder: str = "default" + indexation_preset: str = "default" + retrieval_preset: str = "default" + dimension: int = 1024 + collection_name: str | None = None + chat_history_depth: int = 0 + chat_llm: str | None = None +``` + +**`core/config/presets.py`** — `PresetRow`: + +```python +class PresetRow(BaseModel): + name: str + preset_type: str # "indexation" | "retrieval" + config: dict[str, Any] = {} + created_at: datetime + updated_at: datetime +``` + +### 14B — Database layer + +**New port ABCs** (already created as stubs in Phase 4): + +- `core/ports/preset_repo.py` — `PresetRepository`: `get(name, type)`, `list(type)`, `upsert(preset)`, `delete(name, type)` +- `core/ports/partition_repo.py` — extend with `get_partition_config(name)`, `update_partition_config(name, **fields)` + +**New persistence implementations:** + +- `services/persistence/preset_repo.py` — PostgreSQL CRUD for `PresetRow` +- `services/persistence/partition_repo.py` — extend with preset-reference columns + +**Alembic migration:** + +``` +services/persistence/migrations/versions/NNN_add_presets_and_partition_config.py + - CREATE TABLE presets (name, preset_type, config JSONB, created_at, updated_at) + - UNIQUE (name, preset_type) + - ALTER TABLE partitions ADD COLUMN indexation_preset VARCHAR DEFAULT 'default' + - ALTER TABLE partitions ADD COLUMN retrieval_preset VARCHAR DEFAULT 'default' + - ALTER TABLE partitions ADD COLUMN dimension INTEGER DEFAULT 1024 + - ALTER TABLE partitions ADD COLUMN chat_history_depth INTEGER DEFAULT 0 + - ALTER TABLE partitions ADD COLUMN chat_llm VARCHAR NULL +``` + +### 14C — Services + +**`services/orchestrators/preset_service.py`:** + +```python +class PresetService: + def __init__(self, preset_repo: PresetRepository, config: Settings): + ... + + async def seed_defaults(self) -> None: + """Insert default indexation + retrieval presets from YAML if not in DB.""" + + async def load_all(self) -> None: + """Load all presets from DB into config.presets (runtime cache).""" + + async def get_preset(self, name: str, preset_type: str) -> PresetRow: ... + async def list_presets(self, preset_type: str) -> list[PresetRow]: ... + async def upsert_preset(self, preset: PresetRow) -> PresetRow: ... + async def delete_preset(self, name: str, preset_type: str) -> bool: ... +``` + +**`services/orchestrators/partition_service.py`** — extend: + +```python +async def load_partitions(self) -> None: + """Load all partition configs from DB, resolve presets, merge into config.partitions.""" + +async def get_partition_config(self, partition: str) -> PartitionConfig: + """Get resolved config for a partition (preset + overrides).""" +``` + +**Resolution chain** (how a partition gets its full config): + +1. Load `PartitionRow` from DB (has `indexation_preset: str`, `retrieval_preset: str`) +2. Look up `PresetRow` for each preset name +3. Parse preset config JSON into `IndexationPipelineConfig` / `RetrievalPipelineConfig` +4. Build `PartitionConfig` with resolved pipeline configs +5. Cache in `config.partitions[name]` for fast access + +### 14D — Update orchestrators to use per-partition config + +**IndexingService** — currently uses global config. Change to: + +```python +async def index_documents_batch(self, docs: list[Document], partition: str): + partition_config = await self._partition_service.get_partition_config(partition) + idx_config = partition_config.indexation # per-partition indexation settings + # Use idx_config.chunking.strategy, idx_config.enable_image_captioning, etc. +``` + +**RetrievalService** — currently uses global retriever config. Change to: + +```python +async def retrieve(self, query: str, partition: str, **kwargs): + partition_config = await self._partition_service.get_partition_config(partition) + ret_config = partition_config.retrieval # per-partition retrieval settings + top_k = ret_config.top_k + reranker_name = ret_config.reranker or "default" + # ... +``` + +**QueryService** — use `partition_config.chat_llm` and `partition_config.chat_history_depth`. + +### 14E — API endpoints + +**`api/routers/admin/presets.py`:** + +``` +GET /api/v1/admin/presets?type=indexation # list presets +GET /api/v1/admin/presets/{name}?type=indexation +POST /api/v1/admin/presets # create/update preset +DELETE /api/v1/admin/presets/{name}?type=indexation +``` + +**`api/routers/admin/partitions.py`** — extend: + +``` +PATCH /api/v1/admin/partitions/{name}/config # update partition preset assignments +GET /api/v1/admin/partitions/{name}/config # get resolved partition config +``` + +**`api/routers/admin/pipelines.py`** — dry-run preview: + +``` +POST /api/v1/admin/pipelines/preview # preview pipeline config without saving +``` + +### 14F — Default presets (seeded from YAML) + +Add to `conf/`: + +```yaml +# conf/presets/indexation/default.yaml +chunking: + strategy: recursive_splitter + chunk_size: 512 + chunk_overlap: 64 +parsing_strategy: marker +enable_image_captioning: true +enable_contextualization: false +enable_entity_extraction: true +enable_topic_tagging: true + +# conf/presets/retrieval/default.yaml +type: unified +top_k: 20 +top_n: 10 +enable_reranker: true +enable_planner: true +rrf_k: 60 +intent_strategies: + qa: + top_k: 10 + top_n: 5 + summarization: + top_k: 30 + top_n: 15 +``` + +### Commits + +``` +14.1 Flesh out core/config/{indexation,retrieval,partition,presets}.py +14.2 Alembic migration: presets table + partition config columns +14.3 services/persistence/preset_repo.py (implements PresetRepository) +14.4 Extend services/persistence/partition_repo.py with config columns +14.5 services/orchestrators/preset_service.py (CRUD + seed + load) +14.6 Extend services/orchestrators/partition_service.py (load_partitions, resolve presets) +14.7 Update IndexingService to use per-partition config +14.8 Update RetrievalService to use per-partition config +14.9 Update QueryService to use per-partition chat config +14.10 api/routers/admin/presets.py (CRUD endpoints) +14.11 Extend api/routers/admin/partitions.py (config endpoints) +14.12 api/routers/admin/pipelines.py (dry-run preview) +14.13 Add default preset YAML files to conf/presets/ +14.14 Wire PresetService + updated PartitionService into ServiceContainer +14.15 Seed defaults in container.initialize() +14.16 Integration test: create preset -> assign to partition -> index -> retrieve +``` + +### What changes for existing partitions + +- Existing partitions get `indexation_preset = "default"` and `retrieval_preset = "default"` + via the migration (column defaults). +- The `"default"` preset is seeded from YAML on first startup. +- Behavior is **identical** to current OpenRAG until an admin explicitly changes a + partition's preset or creates custom presets. +- **Zero breaking changes** — this is purely additive. + +--- + +## Phase 15 — OIDC / Keycloak SSO Authentication + +**Goal:** Add OIDC-based SSO authentication alongside the existing API token system. +Users can authenticate via Keycloak (JWT) or via API tokens (`or-` prefix) — both +methods coexist. This is a feature addition on the clean architecture. + +### Architecture: Dual Auth + +The auth dependency inspects the token format to dispatch: + +- Starts with `eyJ` (base64 JSON) -> JWT validation path (Keycloak) +- Starts with `or-` -> DB token lookup (existing, unchanged) + +Both paths produce the same `request.state.user` — downstream code is unaware +of which auth method was used. + +### 15A — New files and their locations + +**Core layer** (pure config, no I/O): + +| File | Purpose | +| --------------------- | ----------------------------------------------- | +| `core/config/auth.py` | Add `OIDCConfig` model to existing `AuthConfig` | + +```python +class OIDCConfig(BaseModel): + enabled: bool = False + issuer_url: str = "" # https://keycloak.company.com/realms/corp + client_id: str = "" + client_secret: str = "" # only if confidential client + audience: str | None = None # expected "aud" claim + claim_sub: str = "sub" # claim for user ID + claim_email: str = "email" + claim_name: str = "preferred_username" + claim_groups: str = "groups" # claim containing group memberships + claim_roles: str = "resource_access.openrag.roles" + admin_role: str = "openrag-admin" # role that grants is_admin + group_prefix: str = "/openrag/" # strip from group names + group_pattern: str = r"(.+)/(owner|editor|viewer)" + auto_provision: bool = True # create user on first login + default_quota: int = 10 + jwks_cache_ttl: int = 3600 # cache JWKS keys for 1 hour +``` + +**Services layer** (infrastructure adapters): + +| File | Purpose | +| ----------------------------------- | ----------------------------------------------------------- | +| `services/auth/__init__.py` | Package init | +| `services/auth/jwt_validator.py` | Validates JWT signature against Keycloak's JWKS endpoint | +| `services/auth/oidc_mapper.py` | Extracts user info + partition roles from JWT claims | +| `services/auth/oidc_provisioner.py` | Find-or-create user by `external_user_id`, sync memberships | + +**API layer** (auth dependency update): + +| File | Change | +| -------------------------- | ---------------------------------------------------------------------------------------------- | +| `api/dependencies/auth.py` | Add dual-auth dispatch: `_is_jwt()` detection, JWT path calls validator + mapper + provisioner | + +**DI layer** (wiring): + +| File | Change | +| ----------------- | --------------------------------------------------------------------------------------------------------- | +| `di/container.py` | Create `KeycloakJWTValidator`, `OIDCMapper`, `OIDCProvisioner` if OIDC enabled; inject into `AuthService` | + +**Frontend** (indexer-ui): + +| File | Purpose | +| ------------------------------------------- | -------------------------------------------------------------- | +| `ui/src/lib/auth/oidc.ts` | OIDC client using `oidc-client-ts` (Authorization Code + PKCE) | +| `ui/src/routes/auth/callback/+page.svelte` | OIDC redirect callback handler | +| `ui/src/lib/components/layout/Login.svelte` | Add "Login with SSO" button alongside existing token input | + +### 15B — services/auth/jwt_validator.py + +```python +class KeycloakJWTValidator: + """Validates Keycloak-issued JWTs using OIDC discovery + JWKS.""" + + def __init__(self, config: OIDCConfig): + self._issuer = config.issuer_url + self._audience = config.audience + self._client_id = config.client_id + self._jwks_client = PyJWKClient( + f"{self._issuer}/protocol/openid-connect/certs", + cache_keys=True, + lifespan=config.jwks_cache_ttl, + ) + + def validate(self, token: str) -> dict: + """Validate JWT signature and claims. Returns decoded payload.""" + signing_key = self._jwks_client.get_signing_key_from_jwt(token) + return jwt.decode( + token, signing_key.key, algorithms=["RS256"], + issuer=self._issuer, + audience=self._audience or self._client_id, + ) +``` + +### 15C — services/auth/oidc_mapper.py + +```python +class OIDCUserMapper: + """Maps Keycloak JWT claims to OpenRAG user model.""" + + def extract_user_info(self, claims: dict) -> dict: + return { + "external_user_id": claims[self._config.claim_sub], + "display_name": claims.get(self._config.claim_name) or claims.get(self._config.claim_email), + "is_admin": self._is_admin(claims), + } + + def extract_partitions(self, claims: dict) -> list[dict]: + """Parse Keycloak groups into partition + role pairs.""" + # /openrag/project-alpha/editor -> {"partition": "project-alpha", "role": "editor"} +``` + +### 15D — services/auth/oidc_provisioner.py + +```python +class OIDCUserProvisioner: + """Auto-creates/updates OpenRAG users from Keycloak claims.""" + + def __init__(self, config: OIDCConfig, user_repo: UserRepository): + ... + + async def ensure_user(self, user_info: dict, partitions: list[dict]) -> dict: + """ + 1. Lookup by external_user_id + 2. If not found + auto_provision: create user (no API token) + 3. Sync is_admin from Keycloak roles + 4. Sync partition memberships from Keycloak groups + (Keycloak is source of truth - add missing, update changed, remove stale) + 5. Return full OpenRAG user dict + """ +``` + +### 15E — api/dependencies/auth.py (dual auth dispatch) + +```python +async def get_current_user(request: Request, ...) -> User: + token = _extract_token(request) + + if oidc_enabled and _is_jwt(token): # starts with "eyJ" + claims = jwt_validator.validate(token) + user_info = oidc_mapper.extract_user_info(claims) + partitions = oidc_mapper.extract_partitions(claims) + user = await oidc_provisioner.ensure_user(user_info, partitions) + else: # starts with "or-" + user = await auth_service.authenticate_token(token) + + return user + +def _is_jwt(token: str) -> bool: + return token.startswith("eyJ") +``` + +### 15F — Database changes + +**No schema migration needed.** The `users.external_user_id` column already exists +(nullable, unique, indexed). OIDC users get `token = NULL` (they authenticate via JWT, +not API tokens). The `get_user_by_external_id()` method is added to `UserRepository`. + +### 15G — Environment variables + +```bash +# All optional - if OIDC_ENABLED is not true, OIDC auth is disabled +OIDC_ENABLED=true +OIDC_ISSUER_URL=https://keycloak.company.com/realms/your-realm +OIDC_CLIENT_ID=openrag +OIDC_CLIENT_SECRET= # only if confidential client +OIDC_AUDIENCE=openrag +OIDC_ADMIN_ROLE=openrag-admin +OIDC_GROUP_PREFIX=/openrag/ +OIDC_AUTO_PROVISION=true +OIDC_DEFAULT_QUOTA=10 + +# Frontend +VITE_OIDC_ENABLED=true +VITE_OIDC_ISSUER_URL=https://keycloak.company.com/realms/your-realm +VITE_OIDC_CLIENT_ID=openrag +``` + +### 15H — Frontend OIDC flow + +Add `oidc-client-ts` to indexer-ui. The login page shows two options: + +- "Login with SSO" button (redirects to Keycloak) +- "Login with token" input (existing behavior) + +The choice is driven by `VITE_OIDC_ENABLED`. If OIDC is not configured, +only the token input is shown (no change from current behavior). + +Token refresh is handled automatically by `oidc-client-ts` `automaticSilentRenew`. + +### Commits + +``` +15.1 Add PyJWT[crypto] + oidc-client-ts to dependencies +15.2 Add OIDCConfig to core/config/auth.py +15.3 Create services/auth/jwt_validator.py +15.4 Create services/auth/oidc_mapper.py +15.5 Add get_user_by_external_id() + create_oidc_user() to UserRepository +15.6 Create services/auth/oidc_provisioner.py +15.7 Update api/dependencies/auth.py with dual-auth dispatch +15.8 Wire OIDC components into ServiceContainer (conditional on OIDC_ENABLED) +15.9 Integration test: JWT login + auto-provision + membership sync +15.10 Update indexer-ui: add oidc-client-ts, OIDC login flow, /auth/callback +15.11 Update indexer-ui: conditional login (SSO vs token based on config) +15.12 Update docker-compose + .env.example with OIDC env vars +15.13 Documentation: Keycloak setup guide (realm, client, mappers, groups) +``` + +### What doesn't change + +- API token auth (`or-` tokens) — works exactly as before +- Role hierarchy (viewer/editor/owner) — Keycloak groups map to same roles +- SUPER_ADMIN_MODE — works with OIDC-provisioned admins +- All API endpoints — they see `request.state.user` regardless of auth method +- Database schema — no migration (external_user_id already exists) +- File quota system — applies to OIDC users same as token users + +--- + +## Risk Register + +| Risk | Likelihood | Impact | Mitigation | +| ----------------------------------------------- | ---------- | ------ | ------------------------------------------------------------------------------------------------------------------------ | +| **Ray actor serialization breaks** | High | High | Actors are thin wrappers; domain objects use Pydantic model_dump()/model_validate() for serialization | +| **Import cycles** | Medium | Medium | Layer guard script catches immediately; break via core/models/ | +| **God object decomposition breaks integration** | High | High | Phase 7 uses delegation shim - old actor delegates to new stores. Test full upload-search-delete cycle after each commit | +| **Config loading order** | Medium | Medium | ServiceContainer loads config once, passes down. No module-level load_config() | +| **LangChain Document removal** | Medium | High | Phase 2 adds from_langchain/to_langchain. Remove only in Phase 12 | +| **Async/sync mismatch** | Medium | Medium | All new interfaces are async. Sync operations wrapped in asyncio.to_thread() | +| **Performance regression** | Low | Medium | All layers are zero-cost delegation. Profile RAG pipeline hot path | +| **Test gap during migration** | Medium | High | Run full integration suite after every commit. Add unit tests per core/ module | + +--- + +## Migration Utilities + +### Import rewriter + +```bash +python scripts/rewrite_imports.py --dry-run # preview changes +python scripts/rewrite_imports.py --apply # apply changes +``` + +Reads `scripts/import_mapping.json`: + +```json +{ + "components.retriever": "openrag.core.retrieval.retriever", + "components.reranker": "openrag.core.rerankers", + "config.load_config": "openrag.core.config.loader.load_config" +} +``` + +### Layer import guard + +```bash +python scripts/check_layer_imports.py +``` + +Rules: + +```python +FORBIDDEN = [ + ("openrag/core/", ["openrag.services", "openrag.api", "openrag.di"]), + ("openrag/services/", ["openrag.api"]), +] +``` + +### Test fixtures for new DI + +```python +# tests/conftest.py +@pytest.fixture +def config(): + return load_config(overrides={"vectordb.host": "localhost"}) + +@pytest.fixture +async def container(config): + c = ServiceContainer(config) + await c.initialize() + yield c + await c.shutdown() + +@pytest.fixture +def mock_vector_store(): + return InMemoryVectorStore() # for unit tests +``` + +--- + +## Summary: Execution Order + +``` +Phase 0 Scaffold <- zero risk, directory creation +Phase 1 Registry & Exceptions <- low risk, new files only +Phase 2 Domain Models <- low risk, new files only +Phase 3 Configuration <- low risk, re-export shims +Phase 4 ABCs & Ports <- low risk, new ABCs only +-------------------------------------------------------- foundation complete +Phase 5 Core Domain Logic <- MEDIUM, first code moves +Phase 6 Inference Adapters <- MEDIUM, HTTP clients +Phase 7 Storage & Persistence <- HIGH, god object decomposition +Phase 8 Orchestrators <- HIGH, business logic rewiring +Phase 9 Workers (Ray) <- HIGH, distributed system +-------------------------------------------------------- transformation complete +Phase 10 API Layer <- MEDIUM, router migration +Phase 11 Composition Root <- HIGH, DI wiring +Phase 12 Internal Cleanup <- MEDIUM, delete old code inside openrag/ +-------------------------------------------------------- clean architecture complete +Phase 13 Project Layout & Infra <- MEDIUM, top-level restructure + (infra/, scripts/, tests/, ui/) +Phase 14 Per-Partition Presets <- MEDIUM, feature addition +Phase 15 OIDC / Keycloak SSO <- MEDIUM, dual auth (JWT + API tokens) +``` + +**Phase 0-4** are additive and safe - proceed rapidly. +**Phase 5-9** are the core transformation - one commit at a time, test between each. +**Phase 10-12** are the cutover - clean internal architecture. +**Phase 13** restructures everything outside `openrag/` - deployment, tests, scripts, UI. +**Phase 14** is a feature addition - per-partition presets on the clean architecture. +**Phase 15** adds OIDC SSO - Keycloak JWT alongside existing API tokens. diff --git a/entrypoint.sh b/entrypoint.sh index 6d1e3e509..1d385f093 100644 --- a/entrypoint.sh +++ b/entrypoint.sh @@ -6,8 +6,8 @@ fi if [[ "${ENABLE_RAY_SERVE}" == "true" ]]; then echo "🔁 Starting with Ray Serve..." - uv run $ENV_ARG api.py + uv run $ENV_ARG main.py else echo "🚀 Starting with Uvicorn..." - uv run --no-dev $ENV_ARG uvicorn api:app --host 0.0.0.0 --port ${APP_iPORT:-8080} --reload --workers ${API_NUM_WORKERS:-1} + uv run --no-dev $ENV_ARG uvicorn main:app --host 0.0.0.0 --port ${APP_iPORT:-8080} --reload --workers ${API_NUM_WORKERS:-1} fi diff --git a/openrag/api/__init__.py b/openrag/api/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/api/dependencies/__init__.py b/openrag/api/dependencies/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/api/middleware/__init__.py b/openrag/api/middleware/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/api/routers/__init__.py b/openrag/api/routers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/api/routers/admin/__init__.py b/openrag/api/routers/admin/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/api/routers/auth/__init__.py b/openrag/api/routers/auth/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/api/routers/user/__init__.py b/openrag/api/routers/user/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/api/schemas/__init__.py b/openrag/api/schemas/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/api/schemas/admin/__init__.py b/openrag/api/schemas/admin/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/api/schemas/auth/__init__.py b/openrag/api/schemas/auth/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/api/schemas/user/__init__.py b/openrag/api/schemas/user/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/components/auth/deps.py b/openrag/components/auth/deps.py index c68d50e21..5038ca94b 100644 --- a/openrag/components/auth/deps.py +++ b/openrag/components/auth/deps.py @@ -1,83 +1,24 @@ -"""Lazy, process-local singleton for the OIDCClient. - -Kept in a dedicated module to avoid circular imports between the router -(``openrag/routers/auth.py``) and the application entry point (``openrag/api.py``). - -The OIDC config env vars are resolved here via ``os.getenv`` — the same values -that ``openrag/api.py`` validates at startup. In ``AUTH_MODE=oidc`` mode, these -are guaranteed to be non-empty (api.py refuses to start otherwise), so this -module simply trusts them. -""" +"""Compatibility shim - implementation lives in services.auth.deps.""" from __future__ import annotations -import os -from threading import Lock - -from components.auth.oidc_client import OIDCClient +from services.auth.deps import get_oidc_client as _service_get_oidc_client +from services.auth.deps import reset_oidc_client as _service_reset_oidc_client +from services.auth.oidc_client import OIDCClient _client: OIDCClient | None = None -_lock = Lock() def get_oidc_client() -> OIDCClient: - """Return the shared OIDCClient instance, creating it on first call. - - The instance caches the discovery doc and JWKS, so a single shared client - per worker process is both correct and more efficient than one-per-request. - - Env vars read (all required in AUTH_MODE=oidc): - - OIDC_ENDPOINT - - OIDC_CLIENT_ID - - OIDC_CLIENT_SECRET - - OIDC_REDIRECT_URI - - OIDC_SCOPES (default ``openid email profile offline_access``) - """ - global _client if _client is not None: return _client - with _lock: - if _client is not None: - return _client - issuer = os.environ["OIDC_ENDPOINT"] - client_id = os.environ["OIDC_CLIENT_ID"] - client_secret = os.environ["OIDC_CLIENT_SECRET"] - redirect_uri = os.environ["OIDC_REDIRECT_URI"] - scopes = os.getenv("OIDC_SCOPES", "openid email profile offline_access") - _client = OIDCClient( - issuer=issuer, - client_id=client_id, - client_secret=client_secret, - redirect_uri=redirect_uri, - scopes=scopes, - ) - return _client + return _service_get_oidc_client() def reset_oidc_client() -> None: - """Test hook — drops the cached client so the next call rebuilds from env. - - Best-effort closes the underlying httpx.AsyncClient to avoid "Unclosed - client session" warnings and leaking connections when tests repeatedly - reset the singleton. If no event loop is running we skip the close call - — the GC will eventually reclaim the socket. - """ global _client - with _lock: - old = _client - _client = None - if old is None: - return - try: - import asyncio + _client = None + _service_reset_oidc_client() + - loop = asyncio.get_event_loop_policy().get_event_loop() - if loop.is_running(): - # Schedule close on the running loop without awaiting — caller - # doesn't need to be async. - loop.create_task(old.aclose()) - else: - loop.run_until_complete(old.aclose()) - except Exception: - # Closing is best-effort; never let a reset blow up the caller. - pass +__all__ = ["get_oidc_client", "reset_oidc_client"] diff --git a/openrag/components/auth/middleware.py b/openrag/components/auth/middleware.py index c25cb328f..04903af59 100644 --- a/openrag/components/auth/middleware.py +++ b/openrag/components/auth/middleware.py @@ -27,6 +27,7 @@ import os from collections.abc import Callable +from typing import Any from urllib.parse import quote from components.auth.refresh import refresh_session_if_needed @@ -99,14 +100,13 @@ def is_bypass_path(path: str) -> bool: class AuthMiddleware(BaseHTTPMiddleware): """FastAPI middleware enforcing authentication for both token and oidc modes. - Constructor takes a ``get_vectordb`` callable returning the Ray actor - handle — this indirection keeps the middleware decoupled from - ``utils.dependencies`` so tests can inject a ``MagicMock``. + Constructor takes a ``get_auth_service`` callable so tests can inject a + fake service and the live app can resolve the request-time container. """ - def __init__(self, app, *, get_vectordb: Callable[[], object]): + def __init__(self, app, *, get_auth_service: Callable[[Request], Any]): super().__init__(app) - self._get_vectordb = get_vectordb + self._get_auth_service = get_auth_service async def dispatch(self, request: Request, call_next): # Read env lazily so tests can flip AUTH_MODE per-test. @@ -114,12 +114,11 @@ async def dispatch(self, request: Request, call_next): auth_token = os.getenv("AUTH_TOKEN") enc_key = os.getenv("OIDC_TOKEN_ENCRYPTION_KEY") or "" - vectordb = self._get_vectordb() - # --- Dev mode: AUTH_MODE=token + AUTH_TOKEN unset → user 1 bypass. if auth_mode == "token" and auth_token is None: - user = await vectordb.get_user.remote(1) - user_partitions = await vectordb.list_user_partitions.remote(1) + auth_service = self._get_auth_service(request) + user = await auth_service.get_user_for_request(1) + user_partitions = await auth_service.list_user_partitions_for_request(1) request.state.user = user request.state.user_partitions = user_partitions request.state.oidc_session = None @@ -144,7 +143,8 @@ async def dispatch(self, request: Request, call_next): cookie_token = request.cookies.get(SESSION_COOKIE_NAME) session_valid = False if cookie_token: - session = await vectordb.get_oidc_session_by_token.remote(cookie_token) + auth_service = self._get_auth_service(request) + session = await auth_service.get_oidc_session_by_token_for_request(cookie_token) session_valid = session is not None if not session_valid: next_path = path @@ -158,6 +158,7 @@ async def dispatch(self, request: Request, call_next): user = None session = None + auth_service = self._get_auth_service(request) # --- 1) Cookie session (OIDC UI flow). Gated on oidc mode so the # legacy token-mode contract remains strictly Bearer-only — @@ -165,23 +166,23 @@ async def dispatch(self, request: Request, call_next): # request when AUTH_MODE=token. cookie_token = request.cookies.get(SESSION_COOKIE_NAME) if auth_mode == "oidc" else None if cookie_token: - session = await vectordb.get_oidc_session_by_token.remote(cookie_token) + session = await auth_service.get_oidc_session_by_token_for_request(cookie_token) if session is not None: refreshed = await refresh_session_if_needed( session=session, enc_key=enc_key, - vectordb=vectordb, + auth_service=auth_service, ) if refreshed is None: # Refresh failed or session unusable → revoke and fall through. try: - await vectordb.revoke_oidc_session_by_id.remote(session["id"]) + await auth_service.revoke_oidc_session_by_id_for_request(session["id"]) except Exception as e: logger.bind(error=str(e)).warning("Failed to revoke invalid OIDC session") session = None else: session = refreshed - user = await vectordb.get_user.remote(session["user_id"]) + user = await auth_service.get_user_for_request(session["user_id"]) # --- 2) Fallback: Bearer / ?token= (programmatic clients + internal # callers like Chainlit's header_auth_callback which forwards @@ -200,28 +201,28 @@ async def dispatch(self, request: Request, call_next): # a ``users.token`` hash). Try the session lookup first with # the same lazy-refresh semantics as the cookie branch above. if auth_mode == "oidc": - session = await vectordb.get_oidc_session_by_token.remote(token) + session = await auth_service.get_oidc_session_by_token_for_request(token) if session is not None: refreshed = await refresh_session_if_needed( session=session, enc_key=enc_key, - vectordb=vectordb, + auth_service=auth_service, ) if refreshed is None: try: - await vectordb.revoke_oidc_session_by_id.remote(session["id"]) + await auth_service.revoke_oidc_session_by_id_for_request(session["id"]) except Exception as e: logger.bind(error=str(e)).warning("Failed to revoke invalid OIDC session (bearer path)") session = None else: session = refreshed - user = await vectordb.get_user.remote(session["user_id"]) + user = await auth_service.get_user_for_request(session["user_id"]) if user is None: # Either token mode, or oidc mode with no matching session # — fall back to the long-lived ``users.token`` used by # programmatic clients (CI, scripts, service agents). - user = await vectordb.get_user_by_token.remote(token) + user = await auth_service.get_user_by_token_for_request(token) if not user and auth_mode == "token": # Legacy test contract: robot suite asserts 403 + "Invalid token". return JSONResponse(status_code=403, content={"detail": "Invalid token"}) @@ -243,6 +244,6 @@ async def dispatch(self, request: Request, call_next): # --- Happy path: user resolved. request.state.user = user - request.state.user_partitions = await vectordb.list_user_partitions.remote(user["id"]) + request.state.user_partitions = await auth_service.list_user_partitions_for_request(user["id"]) request.state.oidc_session = session # None when authenticated via Bearer return await call_next(request) diff --git a/openrag/components/auth/oidc_client.py b/openrag/components/auth/oidc_client.py index 50b449b9e..c71f542a5 100644 --- a/openrag/components/auth/oidc_client.py +++ b/openrag/components/auth/oidc_client.py @@ -1,376 +1,5 @@ -"""Lightweight OIDC Relying Party client for OpenRAG. +"""Re-export shim — implementation lives in services.auth.oidc_client.""" -Wraps Authlib's JWT/JWK primitives with: -- Discovery endpoint caching (1 h TTL) -- JWKS caching with automatic refresh on kid-miss -- PKCE pair generation (S256) -- Authorization URL builder -- Code exchange with ID token verification -- Token refresh (lazy, called by middleware when access_token near expiry) -- Userinfo fetch -- Back-channel logout token verification +from services.auth.oidc_client import LogoutTokenClaims, OIDCClient, TokenBundle -One instance per (issuer, client_id, client_secret) tuple. -The instance is not thread-safe for writes but safe for concurrent reads once -the metadata and JWKS caches are populated. -""" - -import base64 -import hashlib -import secrets -import time -from dataclasses import dataclass -from typing import Any -from urllib.parse import urlencode - -import httpx -from authlib.jose import JsonWebKey, JsonWebToken -from authlib.jose.errors import JoseError - - -@dataclass -class TokenBundle: - """Holds the token set returned by the IdP together with verified ID token claims.""" - - id_token: str - access_token: str - refresh_token: str | None - expires_in: int # seconds - token_type: str # usually "Bearer" - claims: dict[str, Any] # verified claims from id_token - - -@dataclass -class LogoutTokenClaims: - """Verified claims from a back-channel logout token.""" - - iss: str - aud: str | list[str] - sub: str | None - sid: str | None - iat: int - jti: str | None - - -class OIDCClient: - """Lightweight OIDC Relying Party client. - - One instance per (issuer, client_id, client_secret) tuple. - """ - - _DISCOVERY_TTL = 3600 # 1 hour - _JWKS_TTL = 3600 # 1 hour - - def __init__( - self, - *, - issuer: str, - client_id: str, - client_secret: str, - redirect_uri: str, - scopes: str, - http_client: httpx.AsyncClient | None = None, - ): - # Keep the issuer string verbatim (including any trailing "/") — the OIDC - # spec mandates strict byte-for-byte equality between ``self.issuer``, the - # issuer advertised by the discovery document, and the ``iss`` claim in - # tokens. Operators must configure ``OIDC_ENDPOINT`` to match EXACTLY - # what the IdP returns. - self.issuer = issuer - self.client_id = client_id - self.client_secret = client_secret - self.redirect_uri = redirect_uri - self.scopes = scopes - self._http = http_client or httpx.AsyncClient(timeout=10.0) - self._metadata: dict | None = None - self._metadata_fetched_at: float = 0.0 - self._jwks: JsonWebKey | None = None - self._jwks_fetched_at: float = 0.0 - - # ------------------------------------------------------------------ - # Discovery - # ------------------------------------------------------------------ - - async def discover(self) -> dict: - """Fetch and cache the OIDC discovery document. - - Returns the cached document if it is less than _DISCOVERY_TTL seconds old. - Raises ValueError if the returned issuer does not match the configured one. - """ - if self._metadata and (time.time() - self._metadata_fetched_at) < self._DISCOVERY_TTL: - return self._metadata - url = f"{self.issuer.rstrip('/')}/.well-known/openid-configuration" - resp = await self._http.get(url) - resp.raise_for_status() - self._metadata = resp.json() - self._metadata_fetched_at = time.time() - if self._metadata.get("issuer") != self.issuer: - raise ValueError(f"Issuer mismatch: configured {self.issuer!r}, got {self._metadata.get('issuer')!r}") - return self._metadata - - # ------------------------------------------------------------------ - # JWKS - # ------------------------------------------------------------------ - - async def _load_jwks(self, force: bool = False) -> JsonWebKey: - meta = await self.discover() - if not force and self._jwks and (time.time() - self._jwks_fetched_at) < self._JWKS_TTL: - return self._jwks - resp = await self._http.get(meta["jwks_uri"]) - resp.raise_for_status() - self._jwks = JsonWebKey.import_key_set(resp.json()) - self._jwks_fetched_at = time.time() - return self._jwks - - # ------------------------------------------------------------------ - # PKCE helpers - # ------------------------------------------------------------------ - - @staticmethod - def generate_pkce_pair() -> tuple[str, str]: - """Generate a PKCE (code_verifier, code_challenge) pair using S256. - - Returns: - (verifier, challenge) — verifier is 128 url-safe chars, - challenge is the base64url-encoded SHA-256 of the verifier. - """ - verifier = secrets.token_urlsafe(96)[:128] - digest = hashlib.sha256(verifier.encode()).digest() - challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode() - return verifier, challenge - - @staticmethod - def generate_state_and_nonce() -> tuple[str, str]: - """Generate cryptographically random state and nonce values.""" - return secrets.token_urlsafe(32), secrets.token_urlsafe(32) - - # ------------------------------------------------------------------ - # Authorization URL - # ------------------------------------------------------------------ - - async def build_authorization_url(self, *, state: str, nonce: str, code_challenge: str) -> str: - """Build the full authorization URL to redirect the browser to.""" - meta = await self.discover() - params = { - "response_type": "code", - "client_id": self.client_id, - "redirect_uri": self.redirect_uri, - "scope": self.scopes, - "state": state, - "nonce": nonce, - "code_challenge": code_challenge, - "code_challenge_method": "S256", - } - return f"{meta['authorization_endpoint']}?{urlencode(params)}" - - # ------------------------------------------------------------------ - # Code exchange - # ------------------------------------------------------------------ - - async def exchange_code(self, *, code: str, code_verifier: str, expected_nonce: str) -> TokenBundle: - """Exchange an authorization code for tokens. - - Verifies the returned id_token (signature, iss, aud, exp, nonce). - - Args: - code: The authorization code from the IdP callback. - code_verifier: The PKCE verifier corresponding to the challenge sent earlier. - expected_nonce: The nonce value that was sent in the authorization request. - - Returns: - A TokenBundle with verified claims. - """ - meta = await self.discover() - data = { - "grant_type": "authorization_code", - "code": code, - "redirect_uri": self.redirect_uri, - "client_id": self.client_id, - "client_secret": self.client_secret, - "code_verifier": code_verifier, - } - resp = await self._http.post(meta["token_endpoint"], data=data, headers={"Accept": "application/json"}) - resp.raise_for_status() - payload = resp.json() - id_token = payload["id_token"] - claims = await self._verify_id_token(id_token, expected_nonce=expected_nonce) - return TokenBundle( - id_token=id_token, - access_token=payload["access_token"], - refresh_token=payload.get("refresh_token"), - expires_in=int(payload.get("expires_in", 0)), - token_type=payload.get("token_type", "Bearer"), - claims=claims, - ) - - # ------------------------------------------------------------------ - # Token refresh - # ------------------------------------------------------------------ - - async def refresh_access_token(self, refresh_token: str) -> TokenBundle: - """Use the refresh_token to obtain a new access_token. - - If the IdP returns a new id_token, it is re-verified (nonce check skipped - per RFC 8252 §8.2 — nonce is only required during the initial code exchange). - If the IdP omits the refresh_token in the response, the caller's existing - refresh_token is preserved. - - Returns: - A new TokenBundle. - """ - meta = await self.discover() - data = { - "grant_type": "refresh_token", - "refresh_token": refresh_token, - "client_id": self.client_id, - "client_secret": self.client_secret, - } - resp = await self._http.post(meta["token_endpoint"], data=data, headers={"Accept": "application/json"}) - resp.raise_for_status() - payload = resp.json() - new_id_token = payload.get("id_token") - claims: dict[str, Any] = {} - if new_id_token: - claims = await self._verify_id_token(new_id_token, expected_nonce=None) - return TokenBundle( - id_token=new_id_token or "", - access_token=payload["access_token"], - # Some IdPs omit the refresh_token on rotation — keep the old one. - refresh_token=payload.get("refresh_token", refresh_token), - expires_in=int(payload.get("expires_in", 0)), - token_type=payload.get("token_type", "Bearer"), - claims=claims, - ) - - # ------------------------------------------------------------------ - # Userinfo - # ------------------------------------------------------------------ - - async def fetch_userinfo(self, access_token: str) -> dict: - """Fetch the userinfo endpoint with the given access token.""" - meta = await self.discover() - resp = await self._http.get( - meta["userinfo_endpoint"], - headers={"Authorization": f"Bearer {access_token}"}, - ) - resp.raise_for_status() - return resp.json() - - # ------------------------------------------------------------------ - # ID token verification - # ------------------------------------------------------------------ - - async def _verify_id_token(self, token: str, *, expected_nonce: str | None) -> dict[str, Any]: - """Verify an ID token's signature and standard claims. - - Retries with a fresh JWKS fetch on kid-miss (covers IdP key rotation). - Raises JoseError / ValueError on any validation failure. - """ - jwks = await self._load_jwks() - jwt = JsonWebToken(["RS256", "ES256", "EdDSA", "RS384", "RS512"]) - try: - claims = jwt.decode(token, jwks) - except JoseError: - # Force JWKS refresh in case of kid rotation; retry once. - jwks = await self._load_jwks(force=True) - claims = jwt.decode(token, jwks) - - # Manual validation — avoids authlib version differences around claims.params - decoded: dict[str, Any] = dict(claims) - now = int(time.time()) - - if decoded.get("iss") != self.issuer: - raise ValueError(f"ID token iss mismatch: expected {self.issuer!r}, got {decoded.get('iss')!r}") - - aud = decoded.get("aud") - if isinstance(aud, list): - if self.client_id not in aud: - raise ValueError(f"ID token aud {aud!r} does not contain client_id {self.client_id!r}") - elif aud != self.client_id: - raise ValueError(f"ID token aud {aud!r} != client_id {self.client_id!r}") - - if "exp" not in decoded: - raise ValueError("ID token missing exp claim") - if int(decoded["exp"]) < now: - raise ValueError("ID token has expired") - - if "iat" not in decoded: - raise ValueError("ID token missing iat claim") - - if expected_nonce is not None: - if decoded.get("nonce") != expected_nonce: - raise ValueError("OIDC nonce mismatch") - - return decoded - - # ------------------------------------------------------------------ - # Back-channel logout token verification - # ------------------------------------------------------------------ - - async def verify_logout_token(self, token: str) -> LogoutTokenClaims: - """Verify an OIDC back-channel logout token. - - Validates: - - Signature (with JWKS kid-miss retry) - - Standard claims (iss, aud, iat) - - events claim contains the back-channel-logout URI key - - nonce must NOT be present (spec requirement) - - At least one of sub or sid must be present - - Returns: - LogoutTokenClaims with the verified values. - Raises: - ValueError: on any spec violation. - """ - jwks = await self._load_jwks() - jwt = JsonWebToken(["RS256", "ES256", "EdDSA", "RS384", "RS512"]) - try: - claims = jwt.decode(token, jwks) - except JoseError: - jwks = await self._load_jwks(force=True) - claims = jwt.decode(token, jwks) - - decoded: dict[str, Any] = dict(claims) - now = int(time.time()) - - if decoded.get("iss") != self.issuer: - raise ValueError(f"logout_token iss mismatch: expected {self.issuer!r}, got {decoded.get('iss')!r}") - - aud = decoded.get("aud") - if isinstance(aud, list): - if self.client_id not in aud: - raise ValueError(f"logout_token aud {aud!r} does not contain client_id {self.client_id!r}") - elif aud != self.client_id: - raise ValueError(f"logout_token aud {aud!r} != client_id {self.client_id!r}") - - if "iat" not in decoded: - raise ValueError("logout_token missing iat claim") - if int(decoded.get("exp", now + 1)) < now: - raise ValueError("logout_token has expired") - - events = decoded.get("events") or {} - if "http://schemas.openid.net/event/backchannel-logout" not in events: - raise ValueError("logout_token missing required back-channel-logout event claim") - - if decoded.get("nonce"): - raise ValueError("logout_token must not contain nonce") - - if not decoded.get("sub") and not decoded.get("sid"): - raise ValueError("logout_token must contain sub or sid") - - return LogoutTokenClaims( - iss=decoded["iss"], - aud=decoded["aud"], - sub=decoded.get("sub"), - sid=decoded.get("sid"), - iat=int(decoded["iat"]), - jti=decoded.get("jti"), - ) - - # ------------------------------------------------------------------ - # Lifecycle - # ------------------------------------------------------------------ - - async def aclose(self) -> None: - """Close the underlying HTTP client.""" - await self._http.aclose() +__all__ = ["OIDCClient", "TokenBundle", "LogoutTokenClaims"] diff --git a/openrag/components/auth/refresh.py b/openrag/components/auth/refresh.py index cd02e80f0..75bcea848 100644 --- a/openrag/components/auth/refresh.py +++ b/openrag/components/auth/refresh.py @@ -1,176 +1,5 @@ -"""Lazy refresh helper for OIDC access tokens. +"""Re-export shim — implementation lives in services.auth.refresh.""" -Extracted from ``AuthMiddleware`` (Phase 5) to keep ``api.py`` small and -independently testable. Called per-request when a valid cookie session is -found; a no-op when the access token is still fresh. +from services.auth.refresh import refresh_session_if_needed -Timezone policy ---------------- -Phase 2 stores all OIDC session timestamps as **naive local time** via -``datetime.now()`` (see ``test_oidc_sessions.py`` and -``PartitionFileManager.get_oidc_session_by_token``). We match that style -everywhere in this module to avoid tz-mismatch bugs when comparing -``access_token_expires_at`` against "now". - -Refresh-token stampede guard (M1) ---------------------------------- -IdPs with refresh_token rotation enabled invalidate the old refresh_token the -first time it is redeemed. Under concurrency, multiple requests can each notice -"my access_token is about to expire" at the same time and race each other to -the token endpoint. The second attempt fails with ``invalid_grant`` and -(without a guard) its session would be revoked mid-flight. - -We mitigate that with two cooperating mechanisms: - -1. A **short-circuit** here: if ``last_refresh_at`` was bumped less than 5 - seconds ago, we assume a sibling request already rotated the tokens, - re-read the row, and reuse those freshly rotated tokens instead of calling - the IdP. -2. A **row-level write lock** in :meth:`PartitionFileManager.update_oidc_session_tokens` - (``SELECT ... FOR UPDATE``) so that only one writer commits at a time on - Postgres. -3. An **error-recovery branch** here: if the IdP does reject our refresh_token - (typically because a sibling raced us and won), we re-read the row once - more and, if the tokens were advanced meanwhile, return the fresh session - rather than giving up. -""" - -from __future__ import annotations - -from datetime import datetime, timedelta -from typing import Any - -from components.auth.deps import get_oidc_client -from components.auth.session_tokens import decrypt_token, encrypt_token -from utils.logger import get_logger - -_REFRESH_BUFFER = timedelta(seconds=60) -_STAMPEDE_WINDOW = timedelta(seconds=5) - -logger = get_logger() - - -def _to_dt(val: Any) -> datetime: - """Coerce a datetime-or-ISO-string into a ``datetime``. - - Ray occasionally ships values across actors in serialised form; accept - either shape so callers never have to care about the transport. - """ - if isinstance(val, datetime): - return val - if isinstance(val, str): - return datetime.fromisoformat(val) - raise TypeError(f"Expected datetime or ISO string, got {type(val).__name__}") - - -async def refresh_session_if_needed( - *, - session: dict[str, Any], - enc_key: str, - vectordb: Any, -) -> dict[str, Any] | None: - """Refresh the IdP access_token if it is within ``_REFRESH_BUFFER`` of expiry. - - Behaviour: - - If the access_token is still valid with the 60s buffer → return ``session`` unchanged. - - Stampede guard: if another request has just refreshed this session - (``last_refresh_at`` within 5s), re-read the row and reuse the fresh - tokens without calling the IdP. - - If near/past expiry AND a ``refresh_token_encrypted`` blob is stored → - call the IdP, persist rotated tokens, return an updated session dict. - - If near/past expiry AND no refresh_token is stored → return ``session`` as-is - when still formally valid, or ``None`` when already expired (caller should - treat as a revoked session). - - If the refresh call raises (typically because a sibling already rotated - the tokens and the IdP now rejects ours) → re-read the row; if a sibling - succeeded, return their fresh session; otherwise ``None``. - - The session dict returned mirrors the DB row shape produced by - ``PartitionFileManager._oidc_session_to_dict``. - """ - now = datetime.now() - access_exp = _to_dt(session["access_token_expires_at"]) - - if access_exp > now + _REFRESH_BUFFER: - return session - - # --- Stampede short-circuit ------------------------------------------- - # If a sibling request just refreshed this same session, re-read the row - # and reuse the freshly rotated tokens. This avoids racing the IdP with a - # refresh_token that the sibling's success has already invalidated. - last_refresh_at = session.get("last_refresh_at") - if last_refresh_at is not None: - try: - last_refresh_at_dt = _to_dt(last_refresh_at) - except TypeError: - last_refresh_at_dt = None - if last_refresh_at_dt is not None and (now - last_refresh_at_dt) < _STAMPEDE_WINDOW: - try: - fresh = await vectordb.get_oidc_session_by_id.remote(session["id"]) - except Exception as e: - logger.bind(session_id=session.get("id"), error=str(e)).warning( - "Stampede-guard re-read failed; falling through to refresh" - ) - fresh = None - if fresh is not None: - fresh_exp = _to_dt(fresh["access_token_expires_at"]) - if fresh_exp > now + _REFRESH_BUFFER: - return fresh - - refresh_enc = session.get("refresh_token_encrypted") - if not refresh_enc: - # No refresh_token available. - # - If still formally valid (within the 60s buffer window but not yet past exp), - # keep using it. - # - If already expired, caller should treat the session as dead. - return session if access_exp > now else None - - try: - refresh_token = decrypt_token(refresh_enc, enc_key) - client = get_oidc_client() - bundle = await client.refresh_access_token(refresh_token) - except Exception as e: - # Maybe a sibling refreshed between our staleness check and the IdP call - # and the IdP has already invalidated our refresh_token. Re-read the - # row once before giving up: if the tokens were rotated meanwhile, - # treat this as a successful refresh (the sibling's). - logger.bind(session_id=session.get("id"), error=str(e)).warning( - "OIDC refresh_token exchange failed — re-reading session for stampede recovery" - ) - try: - fresh = await vectordb.get_oidc_session_by_id.remote(session["id"]) - except Exception as re: - logger.bind(session_id=session.get("id"), error=str(re)).error( - "Post-failure re-read of OIDC session failed — invalidating" - ) - return None - if fresh is not None: - fresh_exp = _to_dt(fresh["access_token_expires_at"]) - if fresh_exp > now + _REFRESH_BUFFER: - return fresh - return None - - new_access_exp = now + timedelta(seconds=max(int(bundle.expires_in or 0), 60)) - new_access_enc = encrypt_token(bundle.access_token, enc_key) - new_refresh_enc = encrypt_token(bundle.refresh_token, enc_key) if bundle.refresh_token else refresh_enc - - try: - await vectordb.update_oidc_session_tokens.remote( - session_id=session["id"], - access_token_encrypted=new_access_enc, - refresh_token_encrypted=new_refresh_enc, - access_token_expires_at=new_access_exp, - ) - except Exception as e: - logger.bind(session_id=session.get("id"), error=str(e)).error( - "Failed to persist refreshed OIDC tokens — invalidating session" - ) - return None - - return { - **session, - "access_token_encrypted": new_access_enc, - "access_token_expires_at": new_access_exp, - "refresh_token_encrypted": new_refresh_enc, - "last_refresh_at": now, - } +__all__ = ["refresh_session_if_needed"] diff --git a/openrag/components/auth/session_tokens.py b/openrag/components/auth/session_tokens.py index 79cc085e1..48a228db8 100644 --- a/openrag/components/auth/session_tokens.py +++ b/openrag/components/auth/session_tokens.py @@ -1,64 +1,5 @@ -"""Session token utilities for OpenRAG OIDC sessions. +"""Re-export shim — implementation lives in services.auth.session_tokens.""" -Opaque session tokens are issued at callback and stored hashed (SHA-256) in the DB. -IdP tokens (access_token, refresh_token) are encrypted with Fernet before storage. +from services.auth.session_tokens import decrypt_token, encrypt_token, hash_session_token, issue_session_token -The Fernet key is provided via the OIDC_TOKEN_ENCRYPTION_KEY environment variable. -Generate one with: - python -c 'from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())' -""" - -import hashlib -import secrets - -from cryptography.fernet import Fernet, InvalidToken - - -def issue_session_token() -> tuple[str, str]: - """Generate a new session token. - - Returns: - (plaintext, sha256_hex) — the plaintext is set in the cookie, - the hash is stored in the database. - """ - plain = secrets.token_urlsafe(32) # 43 chars, >= 256 bits entropy - return plain, hash_session_token(plain) - - -def hash_session_token(token: str) -> str: - """Return the SHA-256 hex digest of the session token.""" - return hashlib.sha256(token.encode("utf-8")).hexdigest() - - -def _fernet(key: str | bytes) -> Fernet: - try: - return Fernet(key.encode("utf-8") if isinstance(key, str) else key) - except Exception as e: - raise ValueError( - "OIDC_TOKEN_ENCRYPTION_KEY is not a valid Fernet key. " - "Generate one with: python -c 'from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())'" - ) from e - - -def encrypt_token(plaintext: str | None, key: str) -> bytes | None: - """Encrypt a plaintext token string. - - Returns None if plaintext is None (refresh_token may be absent). - """ - if plaintext is None: - return None - return _fernet(key).encrypt(plaintext.encode("utf-8")) - - -def decrypt_token(ciphertext: bytes | None, key: str) -> str | None: - """Decrypt a Fernet-encrypted token. - - Returns None if ciphertext is None. - Raises ValueError on key mismatch or data corruption. - """ - if ciphertext is None: - return None - try: - return _fernet(key).decrypt(ciphertext).decode("utf-8") - except InvalidToken as e: - raise ValueError("Failed to decrypt stored OIDC token — key mismatch or corruption") from e +__all__ = ["issue_session_token", "hash_session_token", "encrypt_token", "decrypt_token"] diff --git a/openrag/components/auth/state_cookie.py b/openrag/components/auth/state_cookie.py index c6e5c566f..0d00a0f52 100644 --- a/openrag/components/auth/state_cookie.py +++ b/openrag/components/auth/state_cookie.py @@ -1,55 +1,5 @@ -"""Signed state cookie for OIDC Authorization Code + PKCE flow. +"""Re-export shim — implementation lives in services.auth.state_cookie.""" -The cookie transports state/nonce/code_verifier between /auth/login and /auth/callback. -It is signed (not encrypted) using itsdangerous.URLSafeTimedSerializer with HMAC-SHA1. +from services.auth.state_cookie import StateCookiePayload, StateCookieSerializer -The signing key is the OIDC_TOKEN_ENCRYPTION_KEY (a Fernet base64url key, which is -valid arbitrary bytes for HMAC). The consuming code (phase 4 router) will pass the -key to StateCookieSerializer(key). Using the same key for both Fernet encryption and -HMAC signing is safe since itsdangerous derives separate subkeys via HMAC. - -TTL defaults to 600 s (10 minutes) — long enough for a slow user at the IdP login page. -""" - -from dataclasses import asdict, dataclass - -from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer - - -@dataclass -class StateCookiePayload: - state: str - nonce: str - code_verifier: str - next_url: str = "/" - - -class StateCookieSerializer: - """Signs/verifies the short-lived cookie holding OIDC state/nonce/code_verifier. - - TTL defaults to 600 s (10 minutes) — long enough for a slow user at the IdP. - """ - - COOKIE_NAME = "openrag_oidc_state" - DEFAULT_TTL_SECONDS = 600 - - def __init__(self, secret_key: str, salt: str = "openrag-oidc-state-v1"): - self._serializer = URLSafeTimedSerializer(secret_key, salt=salt) - - def dumps(self, payload: StateCookiePayload) -> str: - """Serialize and sign the payload, returning an opaque cookie value.""" - return self._serializer.dumps(asdict(payload)) - - def loads(self, token: str, max_age: int = DEFAULT_TTL_SECONDS) -> StateCookiePayload: - """Verify and deserialize the cookie value. - - Raises: - ValueError: if the cookie is expired or the signature is invalid. - """ - try: - data = self._serializer.loads(token, max_age=max_age) - except SignatureExpired as e: - raise ValueError("OIDC state cookie expired") from e - except BadSignature as e: - raise ValueError("OIDC state cookie signature invalid") from e - return StateCookiePayload(**data) +__all__ = ["StateCookieSerializer", "StateCookiePayload"] diff --git a/openrag/components/auth/test_middleware.py b/openrag/components/auth/test_middleware.py index 67092d073..1e6d62411 100644 --- a/openrag/components/auth/test_middleware.py +++ b/openrag/components/auth/test_middleware.py @@ -24,45 +24,36 @@ # --------------------------------------------------------------------------- -def _make_vectordb_mock( +def _make_auth_service_mock( *, user=None, user_by_token=None, session=None, partitions=None, ): - """Return a MagicMock exposing the Ray-actor surface used by the middleware. - - Every ``.remote(...)`` call returns a *coroutine* since the middleware - awaits it. - """ + """Return a MagicMock exposing the auth-service surface used by the middleware.""" mock = MagicMock() - mock.get_user = MagicMock() - mock.get_user.remote = AsyncMock(return_value=user or {"id": 1, "display_name": "Admin"}) + mock.get_user_for_request = AsyncMock(return_value=user or {"id": 1, "display_name": "Admin"}) - mock.get_user_by_token = MagicMock() - mock.get_user_by_token.remote = AsyncMock(return_value=user_by_token) + mock.get_user_by_token_for_request = AsyncMock(return_value=user_by_token) - mock.get_oidc_session_by_token = MagicMock() - mock.get_oidc_session_by_token.remote = AsyncMock(return_value=session) + mock.get_oidc_session_by_token_for_request = AsyncMock(return_value=session) + mock.get_oidc_session_by_id_for_request = AsyncMock(return_value=session) - mock.list_user_partitions = MagicMock() - mock.list_user_partitions.remote = AsyncMock(return_value=partitions or []) + mock.list_user_partitions_for_request = AsyncMock(return_value=partitions or []) - mock.revoke_oidc_session_by_id = MagicMock() - mock.revoke_oidc_session_by_id.remote = AsyncMock(return_value=None) + mock.revoke_oidc_session_by_id_for_request = AsyncMock(return_value=None) - mock.update_oidc_session_tokens = MagicMock() - mock.update_oidc_session_tokens.remote = AsyncMock(return_value=None) + mock.update_oidc_session_tokens_for_request = AsyncMock(return_value=None) return mock -def _build_app(vectordb_mock) -> FastAPI: +def _build_app(auth_service_mock) -> FastAPI: """Construct a FastAPI app with the middleware under test.""" app = FastAPI() - app.add_middleware(AuthMiddleware, get_vectordb=lambda: vectordb_mock) + app.add_middleware(AuthMiddleware, get_auth_service=lambda _request: auth_service_mock) @app.get("/") async def root(request: Request): @@ -133,7 +124,7 @@ def _env(self, monkeypatch): monkeypatch.setenv("AUTH_TOKEN", "configured-admin-token") def test_bearer_valid_returns_200(self): - vdb = _make_vectordb_mock(user_by_token={"id": 7, "display_name": "U"}) + vdb = _make_auth_service_mock(user_by_token={"id": 7, "display_name": "U"}) app = _build_app(vdb) with TestClient(app) as client: r = client.get( @@ -144,7 +135,7 @@ def test_bearer_valid_returns_200(self): assert r.json() == {"user": 7} def test_bearer_invalid_returns_403(self): - vdb = _make_vectordb_mock(user_by_token=None) + vdb = _make_auth_service_mock(user_by_token=None) app = _build_app(vdb) with TestClient(app) as client: r = client.get( @@ -155,7 +146,7 @@ def test_bearer_invalid_returns_403(self): assert r.json() == {"detail": "Invalid token"} def test_missing_token_returns_403(self): - vdb = _make_vectordb_mock() + vdb = _make_auth_service_mock() app = _build_app(vdb) with TestClient(app) as client: r = client.get("/v1/chat/completions") @@ -163,7 +154,7 @@ def test_missing_token_returns_403(self): assert r.json() == {"detail": "Missing token"} def test_bypass_path_open(self): - vdb = _make_vectordb_mock() + vdb = _make_auth_service_mock() app = _build_app(vdb) with TestClient(app) as client: r = client.get("/health_check") @@ -176,13 +167,13 @@ class TestTokenModeDevBypass: def test_no_token_resolves_user_1(self, monkeypatch): monkeypatch.setenv("AUTH_MODE", "token") monkeypatch.delenv("AUTH_TOKEN", raising=False) - vdb = _make_vectordb_mock(user={"id": 1, "display_name": "Admin"}) + vdb = _make_auth_service_mock(user={"id": 1, "display_name": "Admin"}) app = _build_app(vdb) with TestClient(app) as client: r = client.get("/v1/chat/completions") assert r.status_code == 200 assert r.json() == {"user": 1} - vdb.get_user.remote.assert_awaited_with(1) + vdb.get_user_for_request.assert_awaited_with(1) # --------------------------------------------------------------------------- @@ -219,15 +210,15 @@ def _fresh_session(self, user_id=42): def test_cookie_valid_and_access_token_fresh_no_refresh(self): session = self._fresh_session(user_id=42) user = {"id": 42, "display_name": "Alice"} - vdb = _make_vectordb_mock(user=user, session=session) + vdb = _make_auth_service_mock(user=user, session=session) app = _build_app(vdb) with TestClient(app) as client: client.cookies.set("openrag_session", "plain-cookie") r = client.get("/v1/chat/completions") assert r.status_code == 200 assert r.json() == {"user": 42} - vdb.update_oidc_session_tokens.remote.assert_not_awaited() - vdb.revoke_oidc_session_by_id.remote.assert_not_awaited() + vdb.update_oidc_session_tokens_for_request.assert_not_awaited() + vdb.revoke_oidc_session_by_id_for_request.assert_not_awaited() def test_cookie_near_expiry_triggers_refresh(self, monkeypatch): """access_token within 60s of expiry AND refresh_token present → refresh.""" @@ -235,13 +226,13 @@ def test_cookie_near_expiry_triggers_refresh(self, monkeypatch): # Force the refresh helper to "see" the token as near-expiry. session["access_token_expires_at"] = datetime.now() + timedelta(seconds=5) user = {"id": 42} - vdb = _make_vectordb_mock(user=user, session=session) + vdb = _make_auth_service_mock(user=user, session=session) # Patch the helper at its import site inside the middleware module # to avoid any dependency on a real OIDC client. - async def fake_refresh(*, session, enc_key, vectordb): + async def fake_refresh(*, session, enc_key, auth_service): new_exp = datetime.now() + timedelta(minutes=30) - await vectordb.update_oidc_session_tokens.remote( + await auth_service.update_oidc_session_tokens_for_request( session_id=session["id"], access_token_encrypted=b"new-enc-access", refresh_token_encrypted=b"new-enc-refresh", @@ -264,15 +255,15 @@ async def fake_refresh(*, session, enc_key, vectordb): r = client.get("/v1/chat/completions") assert r.status_code == 200 - vdb.update_oidc_session_tokens.remote.assert_awaited() + vdb.update_oidc_session_tokens_for_request.assert_awaited() def test_cookie_refresh_fails_session_revoked_and_302(self): """access_token expired + refresh fails → session revoked, UI request → 302.""" session = self._fresh_session(user_id=42) session["access_token_expires_at"] = datetime.now() - timedelta(minutes=1) - vdb = _make_vectordb_mock(user=None, session=session) + vdb = _make_auth_service_mock(user=None, session=session) - async def fake_refresh(*, session, enc_key, vectordb): + async def fake_refresh(*, session, enc_key, auth_service): return None # refresh failed → invalid session with patch( @@ -286,13 +277,13 @@ async def fake_refresh(*, session, enc_key, vectordb): assert r.status_code == 302 assert r.headers["location"].startswith("/auth/login?next=") - vdb.revoke_oidc_session_by_id.remote.assert_awaited_with(1) + vdb.revoke_oidc_session_by_id_for_request.assert_awaited_with(1) # -- bearer fallback ---------------------------------------------------- def test_bearer_fallback_accepted_in_oidc_mode(self): """Programmatic clients keep using ``users.token`` in oidc mode.""" - vdb = _make_vectordb_mock(user_by_token={"id": 9, "display_name": "bot"}) + vdb = _make_auth_service_mock(user_by_token={"id": 9, "display_name": "bot"}) app = _build_app(vdb) with TestClient(app) as client: r = client.get( @@ -305,7 +296,7 @@ def test_bearer_fallback_accepted_in_oidc_mode(self): # -- unauthenticated branching ------------------------------------------ def test_no_creds_api_path_returns_401(self): - vdb = _make_vectordb_mock() + vdb = _make_auth_service_mock() app = _build_app(vdb) with TestClient(app) as client: r = client.get("/indexer/foo") @@ -313,7 +304,7 @@ def test_no_creds_api_path_returns_401(self): assert r.json() == {"detail": "Unauthenticated"} def test_no_creds_root_path_returns_302(self): - vdb = _make_vectordb_mock() + vdb = _make_auth_service_mock() app = _build_app(vdb) with TestClient(app) as client: r = client.get("/", follow_redirects=False) @@ -322,7 +313,7 @@ def test_no_creds_root_path_returns_302(self): assert r.headers["location"] == "/auth/login?next=%2F" def test_no_creds_root_with_query_preserves_next(self): - vdb = _make_vectordb_mock() + vdb = _make_auth_service_mock() app = _build_app(vdb) with TestClient(app) as client: r = client.get("/?foo=bar", follow_redirects=False) @@ -332,14 +323,14 @@ def test_no_creds_root_with_query_preserves_next(self): assert "%2F" in r.headers["location"] def test_no_creds_static_path_returns_302(self): - vdb = _make_vectordb_mock() + vdb = _make_auth_service_mock() app = _build_app(vdb) with TestClient(app) as client: r = client.get("/static/foo.pdf", follow_redirects=False) assert r.status_code == 302 def test_no_creds_v1_chat_returns_401(self): - vdb = _make_vectordb_mock() + vdb = _make_auth_service_mock() app = _build_app(vdb) with TestClient(app) as client: r = client.get("/v1/chat/completions") @@ -362,12 +353,12 @@ async def test_no_refresh_when_token_fresh(self): "refresh_token_encrypted": b"foo", } vdb = MagicMock() - vdb.update_oidc_session_tokens = MagicMock() - vdb.update_oidc_session_tokens.remote = AsyncMock() + vdb.update_oidc_session_tokens_for_request = MagicMock() + vdb.update_oidc_session_tokens_for_request = AsyncMock() - out = await refresh_session_if_needed(session=session, enc_key="k", vectordb=vdb) + out = await refresh_session_if_needed(session=session, enc_key="k", auth_service=vdb) assert out is session - vdb.update_oidc_session_tokens.remote.assert_not_awaited() + vdb.update_oidc_session_tokens_for_request.assert_not_awaited() @pytest.mark.asyncio async def test_expired_no_refresh_token_returns_none(self): @@ -379,7 +370,7 @@ async def test_expired_no_refresh_token_returns_none(self): "refresh_token_encrypted": None, } vdb = MagicMock() - out = await refresh_session_if_needed(session=session, enc_key="k", vectordb=vdb) + out = await refresh_session_if_needed(session=session, enc_key="k", auth_service=vdb) assert out is None # ------------------------------------------------------------------ @@ -390,8 +381,8 @@ async def test_expired_no_refresh_token_returns_none(self): async def test_refresh_short_circuit_when_last_refresh_recent(self): """If another request refreshed <5s ago, reuse the fresh row; do NOT hit the IdP again with a refresh_token that has already been rotated.""" - from components.auth import refresh as refresh_mod - from components.auth.refresh import refresh_session_if_needed + from services.auth import refresh as refresh_mod + from services.auth.refresh import refresh_session_if_needed now = datetime.now() fresh_exp = now + timedelta(minutes=30) @@ -411,10 +402,10 @@ async def test_refresh_short_circuit_when_last_refresh_recent(self): } vdb = MagicMock() - vdb.get_oidc_session_by_id = MagicMock() - vdb.get_oidc_session_by_id.remote = AsyncMock(return_value=fresh_row) - vdb.update_oidc_session_tokens = MagicMock() - vdb.update_oidc_session_tokens.remote = AsyncMock() + vdb.get_oidc_session_by_id_for_request = MagicMock() + vdb.get_oidc_session_by_id_for_request = AsyncMock(return_value=fresh_row) + vdb.update_oidc_session_tokens_for_request = MagicMock() + vdb.update_oidc_session_tokens_for_request = AsyncMock() # Sentinel: the IdP client must NOT be contacted. fake_client = MagicMock() @@ -422,18 +413,18 @@ async def test_refresh_short_circuit_when_last_refresh_recent(self): side_effect=AssertionError("IdP must not be called during stampede short-circuit") ) with patch.object(refresh_mod, "get_oidc_client", return_value=fake_client): - out = await refresh_session_if_needed(session=stale_session, enc_key="k", vectordb=vdb) + out = await refresh_session_if_needed(session=stale_session, enc_key="k", auth_service=vdb) assert out is fresh_row fake_client.refresh_access_token.assert_not_awaited() - vdb.update_oidc_session_tokens.remote.assert_not_awaited() + vdb.update_oidc_session_tokens_for_request.assert_not_awaited() @pytest.mark.asyncio async def test_refresh_recovers_when_idp_rejects_stale_refresh_token(self): """IdP rejects our refresh_token (sibling already rotated it); the helper re-reads the session and returns the sibling's fresh tokens.""" - from components.auth import refresh as refresh_mod - from components.auth.refresh import refresh_session_if_needed + from services.auth import refresh as refresh_mod + from services.auth.refresh import refresh_session_if_needed now = datetime.now() stale_session = { @@ -452,8 +443,8 @@ async def test_refresh_recovers_when_idp_rejects_stale_refresh_token(self): } vdb = MagicMock() - vdb.get_oidc_session_by_id = MagicMock() - vdb.get_oidc_session_by_id.remote = AsyncMock(return_value=fresh_row) + vdb.get_oidc_session_by_id_for_request = MagicMock() + vdb.get_oidc_session_by_id_for_request = AsyncMock(return_value=fresh_row) fake_client = MagicMock() fake_client.refresh_access_token = AsyncMock(side_effect=RuntimeError("invalid_grant")) @@ -461,17 +452,17 @@ async def test_refresh_recovers_when_idp_rejects_stale_refresh_token(self): patch.object(refresh_mod, "get_oidc_client", return_value=fake_client), patch.object(refresh_mod, "decrypt_token", return_value="old-refresh-plain"), ): - out = await refresh_session_if_needed(session=stale_session, enc_key="k", vectordb=vdb) + out = await refresh_session_if_needed(session=stale_session, enc_key="k", auth_service=vdb) assert out is fresh_row fake_client.refresh_access_token.assert_awaited_once() - vdb.get_oidc_session_by_id.remote.assert_awaited_once_with(1) + vdb.get_oidc_session_by_id_for_request.assert_awaited_once_with(1) @pytest.mark.asyncio async def test_refresh_returns_none_when_idp_rejects_and_no_concurrent_refresh(self): """IdP rejects us and no sibling rotated the tokens → invalidate session.""" - from components.auth import refresh as refresh_mod - from components.auth.refresh import refresh_session_if_needed + from services.auth import refresh as refresh_mod + from services.auth.refresh import refresh_session_if_needed now = datetime.now() stale_session = { @@ -484,8 +475,8 @@ async def test_refresh_returns_none_when_idp_rejects_and_no_concurrent_refresh(s stale_row_from_db = dict(stale_session) vdb = MagicMock() - vdb.get_oidc_session_by_id = MagicMock() - vdb.get_oidc_session_by_id.remote = AsyncMock(return_value=stale_row_from_db) + vdb.get_oidc_session_by_id_for_request = MagicMock() + vdb.get_oidc_session_by_id_for_request = AsyncMock(return_value=stale_row_from_db) fake_client = MagicMock() fake_client.refresh_access_token = AsyncMock(side_effect=RuntimeError("invalid_grant")) @@ -493,6 +484,6 @@ async def test_refresh_returns_none_when_idp_rejects_and_no_concurrent_refresh(s patch.object(refresh_mod, "get_oidc_client", return_value=fake_client), patch.object(refresh_mod, "decrypt_token", return_value="old-refresh-plain"), ): - out = await refresh_session_if_needed(session=stale_session, enc_key="k", vectordb=vdb) + out = await refresh_session_if_needed(session=stale_session, enc_key="k", auth_service=vdb) assert out is None diff --git a/openrag/components/files.py b/openrag/components/files.py index e34656cb6..a4abc2520 100644 --- a/openrag/components/files.py +++ b/openrag/components/files.py @@ -1,4 +1,3 @@ -import asyncio import secrets import time from pathlib import Path @@ -42,27 +41,3 @@ async def save_file_to_disk( await buffer.write(chunk) return file_path - - -async def serialize_file(task_id: str, path: str, metadata: dict | None = {}): - import ray - from ray.exceptions import TaskCancelledError - - serializer_queue = ray.get_actor("SerializerQueue", namespace="openrag") - # Kick off the remote task - future = serializer_queue.submit_document.remote(task_id, path, metadata=metadata) - - # Wait for it to complete, with timeout - ready, _ = await asyncio.to_thread(ray.wait, [future]) - - if ready: - try: - doc = await ready[0] - return doc - except TaskCancelledError: - raise - except Exception: - raise - else: - ray.cancel(future, recursive=True) - raise TimeoutError(f"Serialization task {task_id} timed out after seconds") diff --git a/openrag/components/indexer/__init__.py b/openrag/components/indexer/__init__.py index 5fccd9ea7..452a565b5 100644 --- a/openrag/components/indexer/__init__.py +++ b/openrag/components/indexer/__init__.py @@ -1,4 +1,8 @@ -from .indexer import Indexer -from .vectordb import BaseVectorDB, ConnectorFactory +"""Legacy indexer package. -__all__ = [BaseVectorDB, Indexer, ConnectorFactory] +Import concrete classes from their modules directly. Keeping this package +initializer empty avoids import-time Ray/bootstrap side effects when callers +only need nested utility modules such as ``components.indexer.utils.files``. +""" + +__all__: list[str] = [] diff --git a/openrag/components/indexer/chunker/chunker.py b/openrag/components/indexer/chunker/chunker.py index de9f8934c..f4b794889 100644 --- a/openrag/components/indexer/chunker/chunker.py +++ b/openrag/components/indexer/chunker/chunker.py @@ -1,141 +1,92 @@ -from typing import Literal +"""Backward-compatibility shim — chunking primitives delegate to `openrag.core.chunking`. -import openai -from components.indexer.utils.text_sanitizer import sanitize_text +`ChunkerFactory` is config-driven; the new code uses `chunking_registry`. Both +coexist until Phase 8 cutover. + +Scheduled for removal in Phase 12. +""" + +from typing import TYPE_CHECKING, Any, ClassVar, Literal + +# Side-effect import: pre-loads the indexer-utils submodule so the legacy +# circular import between `components.utils` and `components.indexer.utils.files` +# resolves in the correct order. Removing this line breaks chunker collection. +# Slated to disappear when `components.utils` is split (Phase 6+). +from components.indexer.utils import text_sanitizer as _text_sanitizer # noqa: F401 from components.prompts import CHUNK_CONTEXTUALIZER_PROMPT from components.utils import detect_language, get_vlm_semaphore, load_config +from core.chunking.recursive import RecursiveSplitter as _CoreRecursiveSplitter +from core.indexing.contextualize import ChunkContextualizer as _CoreChunkContextualizer +from core.llm.llm import LLM as _CoreLLM +from core.models.chunk import Chunk as _CoreChunk +from core.models.document import ProcessedDocument, TextBlock +from core.prompts.contextualization_builder import wrap_chunk_with_context from langchain_core.documents.base import Document -from langchain_core.messages import HumanMessage, SystemMessage -from langchain_openai import ChatOpenAI -from tqdm.asyncio import tqdm +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage from utils.logger import get_logger -from openrag.consts import IMAGE_PLACEHOLDER - -from ..embeddings import BaseEmbedding -from .utils import MDElement, chunk_table, get_chunk_page_number, split_md_elements +if TYPE_CHECKING: + from components.indexer.embeddings import BaseEmbedding + from langchain_openai import ChatOpenAI logger = get_logger() config = load_config() -# Timeout for individual chunk contextualization LLM calls (in seconds) CONTEXTUALIZATION_TIMEOUT = config.chunker.contextualization_timeout -# Maximum concurrent contextualization tasks to prevent system overload MAX_CONCURRENT_CONTEXTUALIZATION = config.chunker.max_concurrent_contextualization -BASE_CHUNK_FORMAT = "* filename: {filename}\n\n[CHUNK_START]\n\n{content}\n\n[CHUNK_END]" -CHUNK_FORMAT = "[CONTEXT]\n\n{chunk_context}\n\n" + BASE_CHUNK_FORMAT +class _LangChainLLMAdapter(_CoreLLM): + """Wraps a LangChain ``ChatOpenAI`` so it satisfies the core ``LLM`` ABC.""" -class ChunkContextualizer: - """Handles contextualization of document chunks.""" + _ROLE_MAP: ClassVar[dict] = {"user": HumanMessage, "system": SystemMessage, "assistant": AIMessage} - def __init__(self, llm_config: dict): - llm_config: dict = dict(llm_config) - llm_config.update({"timeout": CONTEXTUALIZATION_TIMEOUT}) - self.context_generator = ChatOpenAI(**llm_config) + def __init__(self, lc_llm: "ChatOpenAI") -> None: + self._llm = lc_llm - async def _generate_context( - self, - first_chunks: list[Document], - prev_chunks: list[Document], - current_chunk: Document, - lang: Literal["fr", "en"] = "en", - ) -> str: - """Generate context for a given chunk of text.""" - filename = first_chunks[0].metadata.get("source", "unknown") - - user_msg = f""" - Here is the context to consider for generating the context: - - Filename: {filename} - - First chunks: - {"\n--\n".join(c.page_content for c in first_chunks)} - - - Previous chunks: - {"\n--\n".join(c.page_content for c in prev_chunks)} - - Here is the current chunk to contextualize strictly in this {lang} language: - - Current chunk: - - {current_chunk.page_content} - """ - async with get_vlm_semaphore(): - try: - messages = [ - SystemMessage(content=CHUNK_CONTEXTUALIZER_PROMPT), - HumanMessage(content=user_msg), - ] - output = await self.context_generator.ainvoke(messages) - return output.content - except openai.APITimeoutError: - logger.warning( - f"OpenAI API timeout contextualizing chunk after {CONTEXTUALIZATION_TIMEOUT}s", - filename=filename, - ) - return "" - except Exception as e: - logger.warning( - "Error contextualizing chunk of document", - filename=filename, - error=str(e), - ) - return "" + async def generate(self, prompt: str, **kwargs) -> str: + out = await self._llm.ainvoke(prompt) + return out.content if hasattr(out, "content") else str(out) - async def contextualize_chunks( - self, - chunks: list[Document], - lang: Literal["fr", "en"] = "en", - filename: str = "", - ) -> list[Document]: - """Contextualize a list of document chunks. - - Processes chunks in batches to prevent overwhelming the system with - too many concurrent LLM requests. - """ - try: - first_chunks = chunks[:2] - contexts = [] - batch_size = MAX_CONCURRENT_CONTEXTUALIZATION - - # Process chunks in batches to limit concurrent LLM calls - for batch_start in range(0, len(chunks), batch_size): - batch_end = min(batch_start + batch_size, len(chunks)) - batch_tasks = [ - self._generate_context( - first_chunks=first_chunks, - prev_chunks=chunks[max(0, i - 2) : i] if i > 0 else [], - current_chunk=chunks[i], - lang=lang, - ) - for i in range(batch_start, batch_end) - ] - - batch_contexts = await tqdm.gather( - *batch_tasks, - total=len(batch_tasks), - desc=f"Contextualizing chunks of *{filename}* [{batch_start + 1}-{batch_end}/{len(chunks)}]", - ) - contexts.extend(batch_contexts) + async def chat(self, messages: list[dict[str, str]], **kwargs) -> str: + lc_msgs = [self._ROLE_MAP[m["role"]](content=m["content"]) for m in messages] + out = await self._llm.ainvoke(lc_msgs) + return out.content - return [ - Document( - page_content=CHUNK_FORMAT.format( - content=chunk.page_content, - chunk_context=context, - filename=filename, - ), - metadata=chunk.metadata, - ) - for chunk, context in zip(chunks, contexts, strict=True) - ] + async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> Any: + pass # Not implemented since the contextualizer never streams. - except Exception as e: - logger.warning(f"Error contextualizing chunks from `{filename}`: {e}") - return chunks + +def _chunks_to_documents(chunks: list, base_metadata: dict) -> list[Document]: + """Convert a list of core domain Chunks into legacy LangChain Documents. + + Reproduces the legacy metadata shape: `page`, `chunk_type`, plus the + document/partition keys the legacy code stamps onto every chunk. + """ + out: list[Document] = [] + for c in chunks: + meta = dict(c.metadata) + meta.update( + { + "file_id": c.document_id, + "partition": c.partition, + "page": c.page_number, + "chunk_type": c.chunk_type.value, + } + ) + # Preserve legacy keys that weren't lifted into core fields. + for k, v in base_metadata.items(): + meta.setdefault(k, v) + out.append(Document(page_content=c.text, metadata=meta)) + return out class BaseChunker: - """Base class for document chunkers with built-in contextualization capability.""" + """Legacy chunker shell — markdown-aware splitting delegated to core. + + Subclasses configure ``self._core_splitter`` (a + `core.chunking.recursive.RecursiveSplitter`) in their `__init__`. + """ def __init__( self, @@ -145,6 +96,8 @@ def __init__( contextual_retrieval: bool = False, **kwargs, ): + from langchain_openai import ChatOpenAI + self.chunk_size = chunk_size self.chunk_overlap_rate = chunk_overlap_rate self.chunk_overlap = int(self.chunk_size * self.chunk_overlap_rate) @@ -152,12 +105,20 @@ def __init__( self.llm = ChatOpenAI(**llm_config) self._length_function = self.llm.get_num_tokens - self.text_splitter = None + self._core_splitter: _CoreRecursiveSplitter | None = None self.contextual_retrieval = contextual_retrieval - - # Initialize contextualizer only if needed - self.contextualizer = ChunkContextualizer(llm_config) if contextual_retrieval else None + if contextual_retrieval: + _lc_llm = ChatOpenAI(**{**llm_config, "timeout": CONTEXTUALIZATION_TIMEOUT}) + self.contextualizer: _CoreChunkContextualizer | None = _CoreChunkContextualizer( + llm=_LangChainLLMAdapter(_lc_llm), + system_prompt=CHUNK_CONTEXTUALIZER_PROMPT, + timeout_seconds=CONTEXTUALIZATION_TIMEOUT, + max_concurrent=MAX_CONCURRENT_CONTEXTUALIZATION, + semaphore=get_vlm_semaphore(), + ) + else: + self.contextualizer = None async def _apply_contextualization( self, @@ -169,120 +130,31 @@ async def _apply_contextualization( if not self.contextual_retrieval or len(chunks) < 2: return [ Document( - page_content=BASE_CHUNK_FORMAT.format(chunk_context="", filename=filename, content=c.page_content), + page_content=wrap_chunk_with_context(c.page_content, filename), metadata=c.metadata, ) for c in chunks ] - return await self.contextualizer.contextualize_chunks(chunks, lang=lang, filename=filename) - - def _prepare_md_elements(self, content: str) -> tuple[list[MDElement], list[MDElement]]: - """Prepare and combine markdown elements from raw content.""" - md_elements: list[MDElement] = split_md_elements(content) - - tables_and_images, texts = [], [] - - for e in md_elements: - if e.type in ("table", "image"): - if e.type == "image" and IMAGE_PLACEHOLDER.lower() in e.content.lower(): # skip placeholder images - continue - - if self._length_function(e.content) <= 100: # do not isolate small tables/images - texts.append(e) - else: - tables_and_images.append(e) - else: - texts.append(e) - - return texts, tables_and_images - - def split_text(self, text: str) -> list[str]: - """Split text into chunks using the text splitter.""" - if not self.text_splitter: - logger.warning("Text splitter not initialized. Initializing with default RecursiveCharacterTextSplitter.") - from langchain.text_splitter import RecursiveCharacterTextSplitter - - self.text_splitter = RecursiveCharacterTextSplitter( - chunk_size=self.chunk_size, - chunk_overlap=self.chunk_overlap, - length_function=self._length_function, - ) - - return self.text_splitter.split_text(text) + core_chunks = [_CoreChunk.from_langchain(c) for c in chunks] + contextualized = await self.contextualizer.contextualize(core_chunks, filename=filename, lang=lang) + return [c.to_langchain(with_id=False) for c in contextualized] def _get_chunks(self, content: str, metadata: dict | None = None, log=None) -> list[Document]: log = log or logger - texts, tables_and_images = self._prepare_md_elements(content=content) - combined_texts = "\n".join([e.content for e in texts]) - - # Sanitize the combined text before chunking to remove excessive whitespace - # and useless characters, which saves tokens and improves quality - sanitized_texts = sanitize_text( - combined_texts, - normalize_whitespace=True, - remove_control_chars=True, - remove_zero_width_chars=True, - max_consecutive_newlines=2, - normalize_unicode=True, - ) + metadata = metadata or {} + partition = metadata.get("partition", "default") - text_chunks = self.split_text(sanitized_texts) - - # Manage tables and images as separate chunks - chunks = [] - for e in tables_and_images: - if e.type == "table" and self._length_function(e.content) > self.chunk_size: - # Chunk large tables separately - subtables = chunk_table( - table_element=e, - chunk_size=self.chunk_size, - length_function=self._length_function, - ) - - s = [ - Document( - page_content=subtable.content.strip(), - metadata={ - **metadata, - "page": subtable.page_number, - "chunk_type": "table", - }, - ) - for subtable in subtables - ] - - else: - s = [ - Document( - page_content=e.content.strip(), - metadata={ - **metadata, - "page": e.page_number, - "chunk_type": e.type, - }, - ) - ] - chunks.extend(s) - - prev_page_num = 1 - for c in text_chunks: - page_info = get_chunk_page_number(chunk_str=c, previous_chunk_ending_page=prev_page_num) - start_page = page_info["start_page"] - prev_page_num = page_info["end_page"] - chunks.append( - Document( - page_content=c.strip(), - metadata={**metadata, "page": start_page, "chunk_type": "text"}, - ) - ) - - if chunks: - chunks.sort(key=lambda d: d.metadata.get("page")) - return chunks - else: + doc = ProcessedDocument( + document_id=metadata.get("file_id", ""), + text_blocks=[TextBlock(text=content)], + metadata=metadata, + ) + chunks = self._core_splitter.chunk(doc, partition=partition) + if not chunks: log.warning("No chunks created. Content is empty or image is not informative.") return [] + return _chunks_to_documents(chunks, base_metadata=metadata) async def split_document(self, doc: Document, task_id: str | None = None) -> list[Document]: """Split document into chunks with optional contextualization.""" @@ -297,11 +169,9 @@ async def split_document(self, doc: Document, task_id: str | None = None) -> lis detected_lang = detect_language(text=doc.page_content) - # Process document through pipeline chunks = self._get_chunks(doc.page_content.strip(), metadata, log=log) if chunks: - # Apply contextualization if enabled log.info( "Contextualizing chunks", apply_contextualization=self.contextual_retrieval, @@ -323,15 +193,10 @@ def __init__( **kwargs, ): super().__init__(chunk_size, chunk_overlap_rate, llm_config, contextual_retrieval, **kwargs) - - from langchain.text_splitter import RecursiveCharacterTextSplitter - - self.text_splitter = RecursiveCharacterTextSplitter( + self._core_splitter = _CoreRecursiveSplitter( chunk_size=self.chunk_size, - chunk_overlap=self.chunk_overlap, + chunk_overlap_rate=self.chunk_overlap_rate, length_function=self._length_function, - is_separator_regex=True, - separators=["\n", r"(?<=[\.\?\!])"], ) @@ -343,13 +208,11 @@ class ChunkerFactory: @staticmethod def create_chunker( config, - embedder: BaseEmbedding | None = None, + embedder: "BaseEmbedding | None" = None, ) -> BaseChunker: - # Extract parameters chunker_params = config.chunker.model_dump() name = chunker_params.pop("name") - # Initialize and return the chunker chunker_cls: BaseChunker = ChunkerFactory.CHUNKERS.get(name) if not chunker_cls: diff --git a/openrag/components/indexer/chunker/utils.py b/openrag/components/indexer/chunker/utils.py index 2529b616d..b07753233 100644 --- a/openrag/components/indexer/chunker/utils.py +++ b/openrag/components/indexer/chunker/utils.py @@ -1,252 +1,33 @@ -import re -from collections.abc import Callable -from typing import Literal - -from components.indexer.utils.text_sanitizer import clean_markdown_table_spacing - -# Regex to match a Markdown table (header + delimiter + at least one row) -TABLE_RE = re.compile( - r"((?:^|\n)\|.*?\|\r?\n\|\s*[:-]+(?:\s*\|[:-]+)*\|\r?\n(?:\|.*?\|\r?\n)+)", - re.DOTALL | re.MULTILINE, +"""Backward-compatibility shim — re-exports from `openrag.core.chunking.markdown_utils`. + +The implementation moved to `openrag/core/chunking/markdown_utils.py` in +Phase 5B. New code should import from there directly. This file is kept +so existing legacy imports keep working until the consumers migrate; +scheduled for removal in Phase 12. +""" + +from core.chunking.markdown_utils import ( + IMAGE_RE, + PAGE_RE, + TABLE_RE, + MDElement, + chunk_table, + get_chunk_page_number, + get_page_number, + parse_markdown_table, + span_inside, + split_md_elements, ) -# Regex to match image descriptions -IMAGE_RE = re.compile(r"((.*?))", re.DOTALL) - -# Regex to match page markers -PAGE_RE = re.compile(r"\[PAGE_(\d+)\]") - - -class MDElement: - """Class representing a segment of markdown content.""" - - def __init__( - self, - type: Literal["text", "table", "image"], - content: str, - page_number: int | None = None, - ): - self.type = type # 'text', 'table', 'image' - self.content = content - self.page_number = page_number - - def __repr__(self): - return f"Element(type={self.type}, page_number={self.page_number}, content={self.content[:100]}...)" - - -def span_inside(span: tuple[int, int], container: tuple[int, int]) -> bool: - return container[0] <= span[0] and span[1] <= container[1] - - -def get_page_number(position, page_markers): - """ - Given a position in the text and list of (position, page_number) tuples, - return the page number for that position. - Content after [PAGE_N] marker belongs to page N+1. - """ - current_page = 1 # Default to page 1 if before any markers - for marker_pos, page_num in page_markers: - if position >= marker_pos: - current_page = page_num + 1 # Content after [PAGE_N] is on page N+1 - else: - break - return current_page - - -def split_md_elements(md_text: str) -> list[MDElement]: - """ - Split markdown text into segments of text, tables, and images. - Returns a list of tuples: - - ('text', content) for text segments - - ('table', content, page_number) for tables - - ('image', content, page_number) for images - """ - # Find all page markers - page_markers = [] - for match in PAGE_RE.finditer(md_text): - page_markers.append((match.start(), int(match.group(1)))) - page_markers.sort() # Ensure they're in order - - all_matches = [] - - # Find image matches first and record their spans - image_spans = [] - for match in IMAGE_RE.finditer(md_text): - span = match.span() - page_num = get_page_number(span[0], page_markers) - all_matches.append((span, "image", match.group(1).strip(), page_num)) - image_spans.append(span) - - # Find table matches, but skip those that are fully inside an image description - for match in TABLE_RE.finditer(md_text): - span = match.span() - if not any(span_inside(span, image_span) for image_span in image_spans): - page_num = get_page_number(span[0], page_markers) - all_matches.append((span, "table", match.group(1).strip(), page_num)) - - # Sort matches by start position - all_matches.sort(key=lambda x: x[0][0]) - - parts = [] - last = 0 - - for (start, end), match_type, content, page_num in all_matches: - # Add text segment before this match if there is any - if start > last: - text_segment = md_text[last:start] - if text_segment.strip(): # Only add non-empty text segments - parts.append(("text", text_segment.strip())) - - # Add the matched segment with page number - parts.append((match_type, content, page_num)) - last = end - - # Add remaining text after the last match - if last < len(md_text): - remaining_text = md_text[last:] - if remaining_text.strip(): # Only add non-empty text segments - parts.append(("text", remaining_text.strip())) - - return [MDElement(*p) for p in parts] - - -def get_chunk_page_number(chunk_str: str, previous_chunk_ending_page=1): - """ - Determine the start and end pages for a text chunk containing [PAGE_N] separators. - PAGE_N marks the end of page N - text before separator is on page N. - """ - # Find all page separator matches in the chunk - matches = list(PAGE_RE.finditer(chunk_str)) - - if not matches: - # No separators found - entire chunk is on previous page - return { - "start_page": previous_chunk_ending_page, - "end_page": previous_chunk_ending_page, - } - - first_match = matches[0] - last_match = matches[-1] - last_char_idx = len(chunk_str) - 1 - - # Determine start page - if first_match.start() == 0: - # Chunk starts with a separator - begins on next page - start_page = int(first_match.group(1)) + 1 - else: - # Text precedes first separator - starts on previous page - start_page = previous_chunk_ending_page - - # Determine end page - if last_match.end() - 1 == last_char_idx: - # Chunk ends exactly at a separator - ends on that page - end_page = int(last_match.group(1)) - else: - # Chunk ends after separator - ends on next page - end_page = int(last_match.group(1)) + 1 - - return {"start_page": start_page, "end_page": end_page} - - -def parse_markdown_table(markdown_table): - """ - Parse a markdown table and extract header and groups based on Domain column. - - Returns: - tuple: (header_lines, groups) - - header_lines: list of [header_row, separator_row] - - groups: list of lists, each containing rows belonging to one domain - """ - lines = markdown_table.strip().split("\n") - - # Extract header (first 2 lines) - header_lines = lines[:2] - data_rows = lines[2:] - - # Group rows by Domain (first column) - groups = [] - current_group = [] - - for row in data_rows: - # Parse first column (Domain) - cells = [cell.strip() for cell in row.split("|")[1:-1]] - if not cells: - continue # skip malformed rows - - domain = cells[0] - - # If Domain is not empty, start a new group - if domain: - if current_group: # Save previous group - groups.append(current_group) - current_group = [row] # Start new group - else: - # Domain is empty, continue current group - current_group.append(row) - - # Don't forget the last group - if current_group: - groups.append(current_group) - - return header_lines, groups - - -def chunk_table( - table_element: MDElement, - chunk_size: int = 512, - length_function: Callable[[str], int] | None = None, -) -> list[MDElement]: - txt = clean_markdown_table_spacing(table_element.content) - header_lines, groups = parse_markdown_table(txt) - - # Convert header lines → text block - header_text = "\n".join(header_lines) - - # Convert group lists → text blocks - group_texts = ["\n".join(g) for g in groups] - - # Precompute token length - header_ntoks = length_function(header_text) - groups_ntoks = [length_function(g) for g in group_texts] - - subtables = [] - current_rows = [header_text] - current_size = header_ntoks - - prev_last_row = None # for overlap - - for group_txt, g_ntoks in zip(group_texts, groups_ntoks, strict=True): - # If adding this group exceeds the chunk limit - if current_size + g_ntoks > chunk_size: - # ---- finalize current subtable ---- - subtables.append("\n".join(current_rows)) - - # ---- start new subtable with OVERLAP ---- - current_rows = [header_text] # always restart headers - if prev_last_row: - current_rows.append(prev_last_row) # add overlapping row - - current_rows.append(group_txt) - current_size = header_ntoks + (length_function(prev_last_row) if prev_last_row else 0) + g_ntoks - - else: - # fits → just append normally - current_rows.append(group_txt) - current_size += g_ntoks - - # track last row for overlap - prev_last_row = group_txt - - # finalize last subtable - if current_rows: - subtables.append("\n".join(current_rows)) - - # wrap into MDElement list - return [ - MDElement( - type="table", - content=subtable, - page_number=table_element.page_number, - ) - for subtable in subtables - ] +__all__ = [ + "IMAGE_RE", + "MDElement", + "PAGE_RE", + "TABLE_RE", + "chunk_table", + "get_chunk_page_number", + "get_page_number", + "parse_markdown_table", + "span_inside", + "split_md_elements", +] diff --git a/openrag/components/indexer/embeddings/__init__.py b/openrag/components/indexer/embeddings/__init__.py index 69850e74f..a1d35852c 100644 --- a/openrag/components/indexer/embeddings/__init__.py +++ b/openrag/components/indexer/embeddings/__init__.py @@ -1,5 +1,7 @@ +from services.inference.vllm_client import VLLMEmbedder # noqa: F401 + from .base import BaseEmbedding -from .openai import OpenAIEmbedding +from .openai import _ShimOpenAIEmbedding as OpenAIEmbedding EMBEDDER_MAPPING = { "openai": OpenAIEmbedding, diff --git a/openrag/components/indexer/embeddings/openai.py b/openrag/components/indexer/embeddings/openai.py index 23aafd904..2c23e6f90 100644 --- a/openrag/components/indexer/embeddings/openai.py +++ b/openrag/components/indexer/embeddings/openai.py @@ -1,6 +1,17 @@ +"""Backward-compatibility shim — delegates to services.inference.vllm_client. + +All new code should import directly from ``services.inference.vllm_client``. +""" + +import asyncio +from concurrent.futures import ThreadPoolExecutor + +import httpx import openai +from core.config.endpoints import EmbedderConfig from langchain_core.documents.base import Document from openai import OpenAI +from services.inference.vllm_client import VLLMEmbedder # noqa: F401 from utils.exceptions.embeddings import * from utils.logger import get_logger @@ -9,7 +20,75 @@ logger = get_logger() +_SYNC_POOL = ThreadPoolExecutor(max_workers=1) + + +def _run_sync(coro): + """Run an async coroutine from sync code, safe inside a running event loop (e.g. Ray).""" + return _SYNC_POOL.submit(asyncio.run, coro).result() + + +def _normalize_texts(texts: list[str | Document]) -> list[str]: + return [item.page_content if isinstance(item, Document) else item for item in texts] + + +class _ShimOpenAIEmbedding(BaseEmbedding): + """Legacy shim — delegates to ``VLLMEmbedder`` for actual HTTP transport. + + Preserves the sync ``embed_documents``/``embed_query`` contract expected by + ``vectordb.py`` (via LangChain's ``aembed_documents`` thread wrapper) while + using VLLMEmbedder's long-lived async httpx pool under the hood. + """ + + def __init__(self, embeddings_config: EmbedderConfig): + self._delegate = VLLMEmbedder( + endpoint=embeddings_config.base_url, + model_name=embeddings_config.model_name, + max_model_len=embeddings_config.max_model_len, + api_key=embeddings_config.api_key, + ) + + @property + def embedding_dimension(self) -> int: + # Probe once if unknown — legacy callers (e.g. MilvusDB schema creation) read + # this before any embed() call, but VLLMEmbedder only learns its dimension + # from a real response. The probe must run on a one-off sync httpx.Client: + # asyncio.run() here would tear down the loop and leave the delegate's + # long-lived AsyncClient pool with stale connections, breaking the next + # real async call with "Event loop is closed". + try: + return self._delegate.dimension + except RuntimeError: + pass + body: dict = {"model": self._delegate._model, "input": ["dim-probe"]} + if self._delegate._max_model_len is not None: + body["truncate_prompt_tokens"] = self._delegate._max_model_len + with httpx.Client(timeout=30.0, headers=dict(self._delegate._client.headers)) as client: + resp = client.post(f"{self._delegate._endpoint}/embeddings", json=body) + resp.raise_for_status() + self._delegate._dimension = len(resp.json()["data"][0]["embedding"]) + return self._delegate._dimension + + def embed_documents(self, texts: list[str | Document]) -> list[list[float]]: + if not texts: + return [] + return _run_sync(self._delegate.embed(_normalize_texts(texts))) + + async def aembed_documents(self, texts: list[str | Document]) -> list[list[float]]: + if not texts: + return [] + return await self._delegate.embed(_normalize_texts(texts)) + + def embed_query(self, text: str) -> list[float]: + return _run_sync(self._delegate.embed_single(text)) + + async def aembed_query(self, text: str) -> list[float]: + return await self._delegate.embed_single(text) + + class OpenAIEmbedding(BaseEmbedding): + """Legacy OpenAI embedding wrapper. New code should use VLLMEmbedder (via DI).""" + def __init__(self, embeddings_config): self.embedding_model = embeddings_config.model_name self.base_url = embeddings_config.base_url @@ -20,16 +99,12 @@ def __init__(self, embeddings_config): @property def embedding_dimension(self) -> int: try: - # Test call to get embedding dimension output = self.embed_documents([Document(page_content="test")]) return len(output[0]) except Exception: raise def embed_documents(self, texts: list[str | Document]) -> list[list[float]]: - """ - Embed documents using the configured embedder. - """ if isinstance(texts[0], Document): texts = [doc.page_content for doc in texts] @@ -69,9 +144,6 @@ def embed_documents(self, texts: list[str | Document]) -> list[list[float]]: ) def embed_query(self, text: str) -> list[float]: - """ - Embed a query using the configured embedder. - """ try: output = self.embed_documents([Document(page_content=text)]) return output[0] diff --git a/openrag/components/indexer/indexer.py b/openrag/components/indexer/indexer.py deleted file mode 100644 index 434353126..000000000 --- a/openrag/components/indexer/indexer.py +++ /dev/null @@ -1,437 +0,0 @@ -import asyncio -import gc -import os -import traceback -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any - -import ray -import torch -from config import load_config -from langchain_core.documents.base import Document - -from .chunker import BaseChunker, ChunkerFactory -from .utils import serialize_file - -config = load_config() -save_uploaded_files = os.environ.get("SAVE_UPLOADED_FILES", "true").lower() == "true" - -POOL_SIZE = config.ray.pool_size -MAX_TASKS_PER_WORKER = config.ray.max_tasks_per_worker - - -@ray.remote( - max_concurrency=config.ray.indexer.concurrency_groups.default, - max_task_retries=config.ray.indexer.max_task_retries, - concurrency_groups={ - "update": config.ray.indexer.concurrency_groups.update, - "search": config.ray.indexer.concurrency_groups.search, - "delete": config.ray.indexer.concurrency_groups.delete, - "insert": config.ray.indexer.concurrency_groups.insert, - "chunk": config.ray.indexer.concurrency_groups.chunk, - "serialize": config.ray.indexer.concurrency_groups.serialize, - }, -) -class Indexer: - def __init__(self): - from utils.logger import get_logger - - self.config = load_config() - self.logger = get_logger() - - # Initialize chunker - self.chunker: BaseChunker = ChunkerFactory.create_chunker(self.config) - - self.default_partition = "_default" - self.enable_insertion = self.config.vectordb.enable - self.handle = ray.get_actor("Indexer", namespace="openrag") - - self.logger.info("Indexer actor initialized.") - - @ray.method(concurrency_group="chunk") - async def chunk(self, doc: Document, file_path: str, task_id: str = None) -> list[Document]: - chunks = await self.chunker.split_document(doc, task_id) - return chunks - - @ray.method(concurrency_group="serialize") - async def serialize_file( - self, - path: str, - metadata: dict = {}, - task_id: str = None, - ): - # Serialize - doc = await serialize_file(task_id, path, metadata=metadata) - return doc - - async def add_file( - self, - path: str, - metadata: dict | None = None, - partition: str | None = None, - user: dict | None = None, - workspace_ids: list[str] | None = None, - replace: bool = False, - ): - task_state_manager = ray.get_actor("TaskStateManager", namespace="openrag") - task_id = ray.get_runtime_context().get_task_id() - metadata = metadata or {} - - file_id = metadata.get("file_id", None) - log = self.logger.bind(file_id=file_id, partition=partition, task_id=task_id) - log.info("Queued file for indexing.") - try: - # Set task details - user_metadata = {k: v for k, v in metadata.items() if k not in {"file_id", "source"}} - - await task_state_manager.set_details.remote( - task_id, - file_id=metadata.get("file_id"), - partition=partition, - metadata=user_metadata, - user_id=user.get("id"), - ) - - # Check/normalize partition - partition = self._check_partition_str(partition) - metadata = {**metadata, "partition": partition} - - # Serialize - doc = await self.handle.serialize_file.remote(path=path, metadata=metadata, task_id=task_id) - - # Chunk - if doc: - await task_state_manager.set_state.remote(task_id, "CHUNKING") - chunks = await self.handle.chunk.remote(doc, str(path), task_id) - else: - log.warning("No document returned from serialization; skipping indexing.") - chunks = [] - - if self.enable_insertion: - if chunks: - await task_state_manager.set_state.remote(task_id, "INSERTING") - if replace: - # PUT flow: PG File row already exists; insert new Milvus chunks - # and update PG metadata in-place (no File row creation). - await self.handle.replace_file_documents.remote(chunks, user=user) - else: - await self.handle.insert_documents.remote(chunks, user=user) - log.info(f"Document {path} indexed successfully") - else: - log.debug("No chunks to insert !!! Potentially the uploaded file is empty") - else: - log.info(f"Vectordb insertion skipped (enable_insertion={self.enable_insertion}).") - - # Mark task as completed before workspace association so the file - # record exists in the DB before we reference it from workspace_files. - await task_state_manager.set_state.remote(task_id, "COMPLETED") - - # Associate with workspaces only after successful indexing (best-effort). - # Not needed for replace=True since the PG row (and its workspace FKs) is preserved. - if workspace_ids and not replace: - vectordb = ray.get_actor("Vectordb", namespace="openrag") - try: - await asyncio.gather( - *[vectordb.add_files_to_workspace.remote(ws_id, [file_id]) for ws_id in workspace_ids] - ) - except Exception as ws_err: - log.warning( - "Failed to associate file with workspaces; file is indexed but workspace links may be incomplete", - error=str(ws_err), - workspace_ids=workspace_ids, - ) - - except Exception as e: - tb = "".join(traceback.format_exception(type(e), e, e.__traceback__)) - log.error(f"Task {task_id} failed in add_file\n{tb}") - await task_state_manager.set_failed_if_not_cancelled.remote(task_id, tb) - raise - - finally: - if torch.cuda.is_available(): - gc.collect() - torch.cuda.empty_cache() - torch.cuda.ipc_collect() - try: - # Cleanup input file - if not save_uploaded_files: - Path(path).unlink(missing_ok=True) - log.debug(f"Deleted input file: {path}") - except Exception as cleanup_err: - log.warning(f"Failed to delete input file {path}: {cleanup_err}") - return True - - @ray.method(concurrency_group="insert") - async def insert_documents(self, chunks, user): - vectordb = ray.get_actor("Vectordb", namespace="openrag") - await vectordb.async_add_documents.remote(chunks, user) - - @ray.method(concurrency_group="insert") - async def replace_file_documents(self, chunks, user): - """Insert chunks for an existing file after its old Milvus chunks have been deleted. - - Unlike insert_documents, this calls add_documents_for_existing_file which - updates the PostgreSQL File row in-place instead of creating a new one. - """ - vectordb = ray.get_actor("Vectordb", namespace="openrag") - await vectordb.add_documents_for_existing_file.remote(chunks, user) - - @ray.method(concurrency_group="delete") - async def delete_file(self, file_id: str, partition: str) -> bool: - log = self.logger.bind(file_id=file_id, partition=partition) - vectordb = ray.get_actor("Vectordb", namespace="openrag") - if not self.enable_insertion: - log.error("Vector database is not enabled, but delete_file was called.") - return False - - try: - await vectordb.delete_file.remote(file_id, partition) - log.info("Deleted file from partition.", file_id=file_id, partition=partition) - - except Exception as e: - log.error("Error in delete_file", error=str(e)) - raise - - @ray.method(concurrency_group="update") - async def update_file_metadata( - self, - file_id: str, - metadata: dict, - partition: str, - user: dict | None = None, - ): - log = self.logger.bind(file_id=file_id, partition=partition) - vectordb = ray.get_actor("Vectordb", namespace="openrag") - if not self.enable_insertion: - log.error("Vector database is not enabled, but update_file_metadata was called.") - return - - try: - # Upsert metadata in-place: updates Milvus chunks (preserving vectors, - # no re-embedding) and the PostgreSQL file record. No delete step, so - # workspace FK references and file_count are never disturbed. - await vectordb.upsert_file_metadata.remote(file_id, partition, metadata) - log.info("Metadata updated for file.") - except Exception as e: - log.error("Error in update_file_metadata", error=str(e)) - raise - - @ray.method(concurrency_group="update") - async def copy_file( - self, - file_id: str, - metadata: dict, - partition: str, - user: dict | None = None, - ): - log = self.logger.bind(file_id=file_id, partition=partition) - vectordb = ray.get_actor("Vectordb", namespace="openrag") - if not self.enable_insertion: - log.error("Vector database is not enabled, but copy_file was called.") - return - - try: - docs = await vectordb.get_file_chunks.remote(file_id, partition) - for doc in docs: - doc.metadata.update(metadata) - - await vectordb.async_add_documents.remote(docs, user=user) - - log.info( - "File copy completed", - file_id=file_id, - partition=partition, - new_file_id=metadata.get("file_id"), - new_partition=metadata.get("partition"), - ) - except Exception as e: - log.error("Error in copy_file", error=str(e)) - raise - - @ray.method(concurrency_group="search") - async def asearch( - self, - query: str, - top_k: int = 5, - similarity_threshold: float = 0.60, - partition: str | list[str] | None = None, - filter: str | None = None, - filter_params: dict | None = None, - ) -> list[Document]: - partition_list = self._check_partition_list(partition) - vectordb = ray.get_actor("Vectordb", namespace="openrag") - return await vectordb.async_search.remote( - query=query, - partition=partition_list, - top_k=top_k, - similarity_threshold=similarity_threshold, - filter=filter, - filter_params=filter_params, - ) - - def _check_partition_str(self, partition: str | None) -> str: - if partition is None: - self.logger.warning("partition not provided; using default.") - return self.default_partition - if not isinstance(partition, str): - raise ValueError("Partition must be a string.") - return partition - - def _check_partition_list(self, partition: str | list[str] | None) -> list[str]: - if partition is None: - self.logger.warning("partition not provided; using default.") - return [self.default_partition] - if isinstance(partition, str): - return [partition] - if isinstance(partition, list) and all(isinstance(p, str) for p in partition): - return partition - raise ValueError("Partition must be a string or a list of strings.") - - -@dataclass -class TaskInfo: - state: str | None = None - error: str | None = None - details: dict[str, Any] = field(default_factory=dict) - object_ref: ray.ObjectRef | None = None - - -@ray.remote(concurrency_groups={"set": 1000, "get": 1000, "queue_info": 1000}) -class TaskStateManager: - def __init__(self): - self.tasks: dict[str, TaskInfo] = {} - self.user_index: dict[int, set[str]] = {} - self.lock = asyncio.Lock() - - async def _ensure_task(self, task_id: str) -> TaskInfo: - """Helper to get-or-create the TaskInfo object under lock.""" - if task_id not in self.tasks: - self.tasks[task_id] = TaskInfo() - return self.tasks[task_id] - - @ray.method(concurrency_group="set") - async def set_state(self, task_id: str, state: str): - async with self.lock: - info = await self._ensure_task(task_id) - info.state = state - - @ray.method(concurrency_group="set") - async def set_error(self, task_id: str, tb_str: str): - async with self.lock: - info = await self._ensure_task(task_id) - info.error = tb_str - - @ray.method(concurrency_group="set") - async def set_failed_if_not_cancelled(self, task_id: str, tb_str: str) -> bool: - """Atomically set state to FAILED and record the traceback, unless the - task is already CANCELLED. Returns True if the state was set to FAILED.""" - async with self.lock: - info = self.tasks.get(task_id) - if info is None or info.state == "CANCELLED": - return False - info.state = "FAILED" - info.error = tb_str - return True - - @ray.method(concurrency_group="set") - async def set_details( - self, - task_id: str, - *, - file_id: str, - partition: int, - metadata: dict, - user_id: int, - ): - async with self.lock: - info = await self._ensure_task(task_id) - info.details = { - "file_id": file_id, - "partition": partition, - "metadata": metadata, - "user_id": user_id, - } - self.user_index.setdefault(user_id, set()).add(task_id) - - @ray.method(concurrency_group="set") - async def set_object_ref(self, task_id: str, object_ref: dict): - async with self.lock: - info = await self._ensure_task(task_id) - info.object_ref = object_ref - - @ray.method(concurrency_group="get") - async def get_state(self, task_id: str) -> str | None: - async with self.lock: - info = self.tasks.get(task_id) - return info.state if info else None - - @ray.method(concurrency_group="get") - async def get_error(self, task_id: str) -> str | None: - async with self.lock: - info = self.tasks.get(task_id) - return info.error if info else None - - @ray.method(concurrency_group="get") - async def get_details(self, task_id: str) -> dict | None: - async with self.lock: - info = self.tasks.get(task_id) - return info.details if info else None - - @ray.method(concurrency_group="get") - async def get_object_ref(self, task_id: str) -> dict | None: - async with self.lock: - info = self.tasks.get(task_id) - return info.object_ref if info else None - - @ray.method(concurrency_group="queue_info") - async def get_all_states(self) -> dict[str, str]: - async with self.lock: - return {tid: info.state for tid, info in self.tasks.items()} - - @ray.method(concurrency_group="queue_info") - async def get_all_info(self) -> dict[str, dict]: - async with self.lock: - return { - task_id: { - "state": info.state, - "error": info.error, - "details": info.details, - } - for task_id, info in self.tasks.items() - } - - @ray.method(concurrency_group="queue_info") - async def get_all_user_info(self, user_id: int) -> dict[str, dict]: - async with self.lock: - task_ids = self.user_index.get(user_id, set()) - return { - tid: { - "state": self.tasks[tid].state, - "error": self.tasks[tid].error, - "details": self.tasks[tid].details, - } - for tid in task_ids - if tid in self.tasks - } - - @ray.method(concurrency_group="queue_info") - async def get_pool_info(self) -> dict[str, int]: - return { - "pool_size": POOL_SIZE, - "max_tasks_per_worker": MAX_TASKS_PER_WORKER, - "total_capacity": POOL_SIZE * MAX_TASKS_PER_WORKER, - } - - @ray.method(concurrency_group="queue_info") - async def get_user_pending_task_count(self, user_id: int) -> int: - """Count tasks for a user that are not yet COMPLETED or FAILED.""" - async with self.lock: - task_ids = self.user_index.get(user_id, set()) - pending_states = {"QUEUED", "SERIALIZING", "CHUNKING", "INSERTING"} - count = 0 - for tid in task_ids: - info = self.tasks.get(tid) - if info and info.state in pending_states: - count += 1 - return count diff --git a/openrag/components/indexer/loaders/audio/local_whisper.py b/openrag/components/indexer/loaders/audio/local_whisper.py index e8eb51e3d..ec1c70823 100644 --- a/openrag/components/indexer/loaders/audio/local_whisper.py +++ b/openrag/components/indexer/loaders/audio/local_whisper.py @@ -1,121 +1,67 @@ +""" +Local Whisper-backed audio loader. + +The Ray actor + pool that drive ``faster-whisper`` (``WhisperActor``, +``WhisperPool``) and the services-side :class:`BasePooledParser` +implementation now live in +``services/workers/parsers/whisper_workers.py``; this module re-exports +``WhisperActor`` and ``WhisperPool`` for legacy import paths +(``components.indexer.loaders.audio.local_whisper.WhisperActor`` is +still used by the OpenAI audio loader for language detection, and by +``services/workers/bootstrap.py`` for the actor bootstrap). + +``LocalWhisperLoader`` is now a thin :class:`BaseLoader` adapter that +delegates to +:class:`core.indexing.parsers.audio.local_whisper.LocalWhisperParser`, +which itself wraps the services-side pool. New code should call the +core parser directly; this shim keeps the legacy loader-discovery path +alive until consumers migrate. +""" + import asyncio from pathlib import Path -import ray -import torch -from config import load_config -from faster_whisper import WhisperModel +from core.indexing.parsers.audio.local_whisper import LocalWhisperParser +from core.models.document import Document as CoreDocument from langchain_core.documents.base import Document +from services.workers.parsers.whisper_workers import ( # noqa: F401 (re-exported for legacy import paths) + LocalWhisperLoader as _ServicesWhisperPool, +) +from services.workers.parsers.whisper_workers import ( # noqa: F401 + WhisperActor, + WhisperPool, +) from utils.logger import get_logger from ..base import BaseLoader logger = get_logger() -config = load_config() - - -if torch.cuda.is_available(): - WHISPER_NUM_GPUS = config.loader.local_whisper.whisper_num_gpus -else: # On CPU - WHISPER_NUM_GPUS = 0 - -WHISPER_CONCURRENCY_PER_WORKER = config.loader.local_whisper.whisper_concurrency_per_worker - - -@ray.remote( - num_gpus=WHISPER_NUM_GPUS, max_restarts=5, max_concurrency=WHISPER_CONCURRENCY_PER_WORKER -) # Ensure each worker processes one file at a time -class WhisperActor: - def __init__(self): - import torch - from config import load_config - from utils.logger import get_logger - - self.logger = get_logger() - self.config = load_config() - - device = "cuda" if torch.cuda.is_available() else "cpu" - compute_type = "float16" if device == "cuda" else "int8" - model_name = self.config.loader.local_whisper.model - - self.logger.info("Loading Whisper model", model_name=model_name, device=device, compute_type=compute_type) - self.model = WhisperModel(model_name, device=device, compute_type=compute_type) - self.logger.info("Whisper model loaded successfully", model_name=model_name, device=device) - - async def transcribe(self, wav_path: str | Path) -> str: - self.logger.info("Transcribing audio file", file_path=Path(wav_path).name) - - def _transcribe_sync() -> str: - segments, _ = self.model.transcribe(str(wav_path)) - return "".join(segment.text for segment in segments) - - return await asyncio.to_thread(_transcribe_sync) - - async def detect_language(self, wav_path: str | Path, fallback_language="en") -> str: - try: - self.logger.info("Detecting language for audio file", file_path=Path(wav_path).name) - - def _detect_language_sync() -> str: - # beam_size=1 + max_new_tokens=1 runs only language detection, no full transcription - _, info = self.model.transcribe(str(wav_path), beam_size=1, max_new_tokens=1) - return info.language - - return await asyncio.to_thread(_detect_language_sync) - - except Exception as e: - self.logger.error("Error detecting language", error=str(e)) - return fallback_language - - -@ray.remote -class WhisperPool: - def __init__(self): - from utils.logger import get_logger - - self.logger = get_logger() - - n_workers = config.loader.local_whisper.whisper_n_workers - self.logger.info(f"Starting WhisperPool with {n_workers} workers") - self.workers = [WhisperActor.remote() for _ in range(n_workers)] - self._pending = [0] * n_workers - - async def transcribe(self, path): - from components.ray_utils import call_ray_actor_with_timeout, retry_with_backoff - - timeout = config.loader.local_whisper.whisper_timeout - - async def attempt(i: int): - idx = min(range(len(self._pending)), key=lambda j: self._pending[j]) - self._pending[idx] += 1 - try: - return await call_ray_actor_with_timeout( - self.workers[idx].transcribe.remote(path), - timeout=timeout, - task_description=f"WhisperPool transcribe ({path})", - ) - finally: - self._pending[idx] -= 1 - - return await retry_with_backoff( - attempt, - max_retries=config.loader.local_whisper.whisper_max_task_retry, - base_delay=config.loader.local_whisper.whisper_retry_base_delay, - task_description=f"WhisperPool transcribe ({path})", - ) class LocalWhisperLoader(BaseLoader): + """Adapter shim — delegates to ``LocalWhisperParser`` via the services-side pool.""" + def __init__(self, **kwargs): super().__init__(**kwargs) - self.whisper_actor: WhisperPool = ray.get_actor("WhisperPool", namespace="openrag") + self._parser = LocalWhisperParser(pool=_ServicesWhisperPool()) async def aload_document(self, file_path, metadata: dict = None, save_markdown=False): + path = Path(file_path) + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=CoreDocument.detect_content_type(path.name), + raw_bytes=raw_bytes, + metadata=dict(metadata) if metadata else {}, + ) try: - content = await self.whisper_actor.transcribe.remote(file_path) - doc = Document(page_content=content, metadata=metadata) - if save_markdown: - self.save_content(content, str(file_path)) - return doc + processed = await self._parser.parse(core_doc) except Exception as e: logger.error("Error loading document", error=str(e)) raise + + content = "".join(b.text for b in processed.text_blocks) + doc = Document(page_content=content, metadata=dict(metadata) if metadata else {}) + if save_markdown: + self.save_content(content, str(file_path)) + return doc diff --git a/openrag/components/indexer/loaders/audio/openai.py b/openrag/components/indexer/loaders/audio/openai.py index 86d235d08..695e99528 100644 --- a/openrag/components/indexer/loaders/audio/openai.py +++ b/openrag/components/indexer/loaders/audio/openai.py @@ -1,124 +1,73 @@ +""" +OpenAI-compatible audio loader. + +The transcription client now lives in +``services/inference/parsers/openai_audio.py`` as +:class:`OpenAIAudioClient` (a :class:`BaseClientParser`). +``OpenAIAudioLoader`` is a thin :class:`BaseLoader` adapter that +constructs the services-side client (with a Whisper-actor-backed +language detector when ``transcriber.use_whisper_lang_detector`` is +enabled) and wraps it in +:class:`core.indexing.parsers.audio.client_based.ClientAudioParser`. +New code should call the core parser directly; this shim keeps the +legacy loader-discovery path alive until consumers migrate. +""" + import asyncio from pathlib import Path -import ray -from components.utils import get_audio_semaphore +from core.indexing.parsers.audio.client_based import ClientAudioParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document -from openai import AsyncOpenAI -from pydub import AudioSegment +from services.inference.parsers.openai_audio import OpenAIAudioClient +from services.workers.parsers.whisper_workers import detect_language_via_actor from utils.logger import get_logger from ..base import BaseLoader -from .local_whisper import WhisperActor logger = get_logger() -# Duration of the audio sample used for language detection -LANG_DETECT_SAMPLE_MS = 30_000 # 30 s - - -class AudioTranscriber: - """Transcribes audio in a single request (no chunking). - - Language detection is handled locally by WhisperActor (faster-whisper). - vLLM's native language detection fix is not yet merged (PR #34342) missed the v0.16.0 branch - cut (Feb 8) — it was merged Feb 21 and will ship in v0.17.0. - """ - - def __init__(self, config): - self.client = AsyncOpenAI( - base_url=config.loader.transcriber.base_url, - api_key=config.loader.transcriber.api_key, - timeout=config.loader.transcriber.timeout, - ) - self.model_name = config.loader.transcriber.model_name - self.use_whisper_lang_detector = config.loader.transcriber.use_whisper_lang_detector - self.direct_upload_suffixes = config.loader.transcriber.direct_upload_suffixes - async def transcribe(self, file_path: Path) -> str: - # Formats in self.direct_upload_suffixes (configurable via - # TRANSCRIBER_DIRECT_UPLOAD_SUFFIXES) are sent as-is to avoid the ~10x - # size inflation from WAV conversion (Scaleway cap: 100 MB; OpenAI: 25 MB). - # Everything else falls back to WAV for vLLM/libsndfile deployments. - - tmp_wav = None - try: - logger.bind(file=file_path.name) - suffix = file_path.suffix.lower() - if suffix in self.direct_upload_suffixes: - wav_path = file_path - # We still need to load the audio so language detection can - # extract its 30-second sample. ``AudioSegment.from_file`` - # uses ffmpeg under the hood, so it handles every format. - sound = await asyncio.to_thread(AudioSegment.from_file, file_path) - else: - sound = await asyncio.to_thread(AudioSegment.from_file, file_path) - logger.info("Converting audio to WAV (unsupported container)", duration_s=f"{len(sound) / 1000:.1f}") - tmp_wav = file_path.with_suffix(".wav") - await asyncio.to_thread(sound.export, tmp_wav, format="wav") - wav_path = tmp_wav - - language = await self._detect_language(sound, wav_path) if self.use_whisper_lang_detector else None - logger.info("Transcribing audio as a single request", language=language) - - async with get_audio_semaphore(): - return await self._transcribe_file(wav_path, language) - except Exception as e: - logger.exception("Error in transcribe", error=str(e)) - raise e - finally: - if tmp_wav: - await asyncio.to_thread(tmp_wav.unlink, True) - - async def _detect_language(self, sound: AudioSegment, wav_path: Path, fallback: str = "en") -> str: - """Detect language via local WhisperActor from a short audio sample.""" - sample = sound[:LANG_DETECT_SAMPLE_MS] - tmp_path = wav_path.parent / f"{wav_path.stem}_langdetect.wav" - await asyncio.to_thread(sample.export, tmp_path, format="wav") - try: - whisper_actor = self._get_whisper_actor() - return await whisper_actor.detect_language.remote(tmp_path, fallback) - except Exception as e: - logger.exception("Language detection failed", error=str(e)) - return fallback - finally: - await asyncio.to_thread(tmp_path.unlink, True) - - def _get_whisper_actor(self): - actor_name = "WhisperActor" - try: - return ray.get_actor(actor_name, namespace="openrag") - except ValueError: - return WhisperActor.options(name=actor_name, namespace="openrag").remote() - except Exception as e: - logger.error("Error getting WhisperActor", error=str(e)) - raise - - async def _transcribe_file(self, wav_path: Path, language: str = None) -> str: - """Send a single file to the transcription endpoint.""" - try: - kwargs = {"model": self.model_name, "file": wav_path} - if language: - kwargs["language"] = language - result = await self.client.audio.transcriptions.create(**kwargs) - return result.text - except Exception as e: - logger.exception("Error transcribing file", file=wav_path.name, error=str(e)) - raise e +async def _whisper_language_detector(file_path: Path) -> str | None: + """Detect language via the singleton ``WhisperActor`` (worker-side helper).""" + return await detect_language_via_actor(file_path) class OpenAIAudioLoader(BaseLoader): + """Adapter shim — delegates to ``OpenAIAudioClient`` via ``ClientAudioParser``.""" + def __init__(self, **kwargs): super().__init__(**kwargs) - self.transcriber = AudioTranscriber(config=self.config) + cfg = self.config.loader.transcriber + _client = OpenAIAudioClient( + base_url=cfg.base_url, + api_key=cfg.api_key, + model=cfg.model_name, + timeout=cfg.timeout, + direct_upload_suffixes=cfg.direct_upload_suffixes, + language_detector=_whisper_language_detector if cfg.use_whisper_lang_detector else None, + ) + self._parser = ClientAudioParser(client=_client) async def aload_document(self, file_path, metadata: dict = None, save_markdown=False): + if metadata is None: + metadata = {} + path = Path(file_path) + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.AUDIO, + raw_bytes=raw_bytes, + metadata=dict(metadata), + ) try: - content = await self.transcriber.transcribe(Path(file_path)) - doc = Document(page_content=content, metadata=metadata) - if save_markdown: - self.save_content(content, str(file_path)) - return doc - except Exception as e: - logger.exception("Error in OpenAIAudioLoader", path=file_path, error=str(e)) - raise e + processed = await self._parser.parse(core_doc) + except Exception: + logger.exception("Error in OpenAIAudioLoader", path=str(file_path)) + raise + content = "\n\n".join(b.text for b in processed.text_blocks) + doc = Document(page_content=content, metadata=metadata) + if save_markdown: + self.save_content(content, str(file_path)) + return doc diff --git a/openrag/components/indexer/loaders/base.py b/openrag/components/indexer/loaders/base.py index 70f05e394..647e06fd8 100644 --- a/openrag/components/indexer/loaders/base.py +++ b/openrag/components/indexer/loaders/base.py @@ -2,11 +2,23 @@ import base64 import re from abc import ABC, abstractmethod -from io import BytesIO from pathlib import Path from components.prompts import IMAGE_DESCRIBER from components.utils import get_vlm_semaphore, load_config +from core.indexing.image_preprocessor import ( + DATA_URI_IMAGE_PATTERN as _CORE_DATA_URI_IMAGE_PATTERN, +) +from core.indexing.image_preprocessor import ( + HTTP_IMAGE_PATTERN as _CORE_HTTP_IMAGE_PATTERN, +) +from core.indexing.image_preprocessor import ( + MIN_IMAGE_PIXELS as _CORE_MIN_IMAGE_PIXELS, +) +from core.indexing.image_preprocessor import ( + ensure_png_compatible_mode, # noqa: F401 (re-exported for legacy import path) + pil_to_png_bytes, +) from langchain_core.messages import HumanMessage from langchain_openai import ChatOpenAI from openai import BadRequestError @@ -19,20 +31,13 @@ config = load_config() -def ensure_png_compatible_mode(image: Image.Image) -> Image.Image: - """Convert incompatible PIL image modes to PNG-saveable modes.""" - if image.mode in ("CMYK", "YCbCr", "LAB"): - return image.convert("RGB") - if image.mode in ("P", "LA", "PA"): - return image.convert("RGBA") - return image - - class BaseLoader(ABC): - # Class-level compiled regex patterns (shared across all instances) - HTTP_IMAGE_PATTERN = re.compile(r"!\[(.*?)\]\((https?://[^)]+)\)") - DATA_URI_IMAGE_PATTERN = re.compile(r"!\[(.*?)\]\((data:image/[^;]+;base64,[^)]+)\)") - MIN_IMAGE_PIXELS = 784 # Qwen2.5-VL min_pixels threshold + # Class-level compiled regex patterns and constants — single source of truth + # lives in ``core.indexing.image_preprocessor``; pinned here as class attrs so + # subclasses keep working via ``self.X``. + HTTP_IMAGE_PATTERN = _CORE_HTTP_IMAGE_PATTERN + DATA_URI_IMAGE_PATTERN = _CORE_DATA_URI_IMAGE_PATTERN + MIN_IMAGE_PIXELS = _CORE_MIN_IMAGE_PIXELS def __init__(self, **kwargs) -> None: self.page_sep = "[PAGE_SEP]" @@ -68,14 +73,12 @@ def save_content(self, text_content: str, path: str): def _pil_image_to_base64(self, image: Image.Image) -> str: """Convert PIL Image to base64 string.""" - buffered = BytesIO() try: - image = ensure_png_compatible_mode(image) - image.save(buffered, format="PNG") + png_bytes = pil_to_png_bytes(image) except Exception as e: logger.warning("Failed to convert image to PNG", error=str(e), mode=getattr(image, "mode", "unknown")) return "" - return base64.b64encode(buffered.getvalue()).decode() + return base64.b64encode(png_bytes).decode() def _is_http_url(self, data: str) -> bool: """Check if string is an HTTP/HTTPS URL.""" diff --git a/openrag/components/indexer/loaders/doc.py b/openrag/components/indexer/loaders/doc.py index 387b46543..46b4e1cb4 100644 --- a/openrag/components/indexer/loaders/doc.py +++ b/openrag/components/indexer/loaders/doc.py @@ -1,46 +1,78 @@ +""" +Legacy ``.doc`` file loader implementation. + +``DocLoader`` is now a thin :class:`BaseLoader` adapter that delegates +to :class:`core.indexing.parsers.doc_parser.DocParser` (which itself +runs Spire.Doc → .docx conversion and then ``DocxParser``) and layers +VLM captioning of embedded images on top. New code should call the +core parser directly; this shim keeps the legacy loader-discovery path +alive until consumers migrate. +""" + +import asyncio import os -import tempfile +from io import BytesIO +from pathlib import Path +from core.indexing.parsers.doc_parser import DocParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document as LCDocument -from spire.doc import Document, FileFormat +from PIL import Image from utils.logger import get_logger from .base import BaseLoader -from .docx import DocxLoader -os.environ["DOTNET_SYSTEM_GLOBALIZATION_INVARIANT"] = "1" # Disable Globalization +os.environ["DOTNET_SYSTEM_GLOBALIZATION_INVARIANT"] = "1" logger = get_logger() class DocLoader(BaseLoader): + """Adapter shim — delegates to ``DocParser``; layers image captioning on top.""" + def __init__(self, **kwargs) -> None: super().__init__(**kwargs) - self.MDLoader = DocxLoader(**kwargs) + self._parser = DocParser() async def aload_document(self, file_path, metadata, save_markdown=False): - """Convert .doc to .docx format, then use DocxLoader to convert to markdown. - Falls back to plain text extraction if the .docx conversion fails.""" - temp_path = None - document = Document() - try: - document.LoadFromFile(str(file_path)) - with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as temp_file: - temp_path = temp_file.name - document.SaveToFile(temp_path, FileFormat.Docx2016) - except Exception as e: - logger.bind(file_id=metadata.get("file_id"), partition=metadata.get("partition")).warning( - f"Spire.Doc conversion to .docx failed, falling back to text extraction: {e}" - ) - text = document.GetText() - doc = LCDocument(page_content=text, metadata=metadata) - if save_markdown: - self.save_content(text, str(file_path)) - return doc + path = Path(file_path) + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.DOC, + raw_bytes=raw_bytes, + metadata=dict(metadata) if metadata else {}, + ) + processed = await self._parser.parse(core_doc) + result = "\n\n".join(b.text for b in processed.text_blocks).strip() + + if not self.image_captioning: + logger.info("Image captioning disabled. Ignoring images.") + for block in processed.images: + ref = (block.metadata or {}).get("markdown_ref") + if ref: + result = result.replace(ref, "") else: - result = await self.MDLoader.aload_document(temp_path, metadata, save_markdown) - return result - finally: - document.Close() - if temp_path and os.path.exists(temp_path): - os.remove(temp_path) + if processed.images: + pil_images: list[Image.Image] = [] + for block in processed.images: + img = Image.open(BytesIO(block.image_bytes)) + img.load() + pil_images.append(img) + captions = await self.caption_images(pil_images, desc="Captioning embedded images") + for block, caption in zip(processed.images, captions): + ref = (block.metadata or {}).get("markdown_ref") + if ref: + result = result.replace(ref, caption.replace("\\", "/")) + + result = await self.replace_markdown_images_with_captions( + result, + caption_data_uris=False, + desc="Captioning linked images", + ) + + doc = LCDocument(page_content=result, metadata=dict(metadata) if metadata else {}) + if save_markdown: + self.save_content(result, str(file_path)) + return doc diff --git a/openrag/components/indexer/loaders/docx.py b/openrag/components/indexer/loaders/docx.py index 615f58abf..6692cac20 100644 --- a/openrag/components/indexer/loaders/docx.py +++ b/openrag/components/indexer/loaders/docx.py @@ -1,10 +1,24 @@ -import re +""" +DOCX file loader implementation. + +``DocxLoader`` is now a thin :class:`BaseLoader` adapter that delegates +extraction to :class:`core.indexing.parsers.docx_parser.DocxParser` and +layers VLM captioning of embedded images on top via the ``BaseLoader`` +mixin. The legacy ``convert_to_png_image`` helper and the +``get_images_from_zip`` instance method are preserved for backward +compatibility with existing test consumers; new code should use the +core parser directly. +""" + +import asyncio import zipfile from io import BytesIO +from pathlib import Path -from docx import Document as DocxDocument +from core.indexing.parsers.docx_parser import DocxParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document -from markitdown import MarkItDown from PIL import Image from utils.logger import get_logger @@ -18,68 +32,59 @@ def convert_to_png_image(image: Image.Image) -> Image.Image: with BytesIO() as buffer: image.save(buffer, format="PNG") buffer.seek(0) - # Reload the image from the buffer as a PNG png_image = Image.open(buffer).convert("RGBA") return png_image class DocxLoader(BaseLoader): + """Adapter shim — delegates to ``DocxParser``; layers image captioning on top.""" + def __init__(self, **kwargs): super().__init__(**kwargs) - self.converter = MarkItDown() + self._parser = DocxParser() async def aload_document(self, file_path, metadata, save_markdown=False): - try: - result = self.converter.convert(file_path).text_content - except Exception as markitdown_err: - logger.warning( - "MarkItDown conversion failed, falling back to python-docx plain text extraction", - path=str(file_path), - error=str(markitdown_err), - ) - try: - result = self._fallback_extract_text(file_path) - except Exception as docx_err: - raise RuntimeError( - f"DOCX conversion failed with both MarkItDown ({markitdown_err}) and python-docx ({docx_err})" - ) from docx_err - - if self.image_captioning: - # Handle embedded images (extracted from docx zip) - # images may contain None entries for unsupported formats (e.g. EMF, WMF) - images = self.get_images_from_zip(file_path) - valid_images = [img for img in images if img is not None] - captions = await self.caption_images(valid_images, desc="Captioning embedded images") - - # Rebuild caption list preserving positional alignment with markdown refs - caption_iter = iter(captions) - for img in images: - caption = next(caption_iter) if img is not None else "" - result = re.sub( - r"!\[.*?\]\(data:image/.*?\)", - caption.replace("\\", "/") if caption else "", - string=result, - count=1, - ) - - # Handle linked images (HTTP URLs) using shared method - # Only caption HTTP URLs, data URIs are already handled above + path = Path(file_path) + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.DOCX, + raw_bytes=raw_bytes, + metadata=dict(metadata) if metadata else {}, + ) + processed = await self._parser.parse(core_doc) + result = "\n\n".join(b.text for b in processed.text_blocks).strip() + + if not self.image_captioning: + logger.info("Image captioning disabled. Ignoring images.") + for block in processed.images: + ref = (block.metadata or {}).get("markdown_ref") + if ref: + result = result.replace(ref, "") + else: + pil_images: list[Image.Image] = [] + for block in processed.images: + img = Image.open(BytesIO(block.image_bytes)) + img.load() + pil_images.append(img) + captions = await self.caption_images(pil_images, desc="Captioning embedded images") + for block, caption in zip(processed.images, captions): + ref = (block.metadata or {}).get("markdown_ref") + if ref: + result = result.replace(ref, caption.replace("\\", "/")) + result = await self.replace_markdown_images_with_captions( result, caption_data_uris=False, desc="Captioning linked images", ) - else: - logger.info("Image captioning disabled. Ignoring images.") - doc = Document(page_content=result, metadata=metadata) + doc = Document(page_content=result, metadata=dict(metadata) if metadata else {}) if save_markdown: self.save_content(result, str(file_path)) return doc - def _fallback_extract_text(self, file_path) -> str: - doc = DocxDocument(file_path) - return "\n\n".join(p.text for p in doc.paragraphs if p.text.strip()) + # ----- legacy helpers retained for test_docx_loader.py compatibility ----- def get_images_from_zip(self, input_file): try: @@ -89,16 +94,11 @@ def get_images_from_zip(self, input_file): return [] with docx: file_names = docx.namelist() - # word/media/ may also contain non-image files (e.g. oleObject, hdphoto, ink) image_files = [f for f in file_names if f.startswith("word/media/")] if not image_files: return [] images_not_in_order, order = [], [] - - # the images got from the original file is not in the right order - # but the target_ref contains the position of the image in the document - for image_file in image_files: image_data = docx.read(image_file) image_extension = image_file.split(".")[-1].lower() @@ -116,7 +116,6 @@ def get_images_from_zip(self, input_file): if not images_not_in_order: return [] - # Reorder images by their original position in the document max_order = max(order) images = [None] * max_order for i, pos in enumerate(order): diff --git a/openrag/components/indexer/loaders/image.py b/openrag/components/indexer/loaders/image.py index a1356c802..6f4280692 100644 --- a/openrag/components/indexer/loaders/image.py +++ b/openrag/components/indexer/loaders/image.py @@ -1,7 +1,23 @@ +""" +Image file loader implementation. + +``ImageLoader`` is now a thin :class:`BaseLoader` adapter that delegates +decode (raster + SVG) to +:class:`core.indexing.parsers.image_parser.ImageParser` and then layers +VLM captioning on top via the ``BaseLoader`` mixin. The legacy +``ImageLoadError`` contract is preserved on decode failure. New code +should call the core parser directly; this shim keeps the legacy +loader-discovery path alive until consumers migrate. +""" + +import asyncio from io import BytesIO from pathlib import Path -import cairosvg +from core.indexing.parsers.image_parser import ImageParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType +from core.utils.exceptions import OpenRAGError from langchain_core.documents import Document from PIL import Image from utils.logger import get_logger @@ -11,34 +27,57 @@ log = get_logger() -class ImageLoadError(Exception): +class ImageLoadError(OpenRAGError): """Raised when an image file cannot be loaded or converted.""" + def __init__(self, message: str, **kwargs): + super().__init__(message, code="IMAGE_LOAD_ERROR", status_code=500, **kwargs) + class ImageLoader(BaseLoader): def __init__(self, **kwargs): super().__init__(**kwargs) + # ``min_pixels=0`` so the parser does not drop small images; the + # size threshold is enforced by ``get_image_description`` (which + # returns the legacy "Image too small for captioning" marker). + self._parser = ImageParser(min_pixels=0) async def aload_document(self, file_path, metadata=None, save_markdown=False): - path = Path(file_path) + if metadata is None: + metadata = {} + path = Path(file_path) try: - # Handle SVG files by converting to PNG first - if path.suffix.lower() == ".svg": - png_data = cairosvg.svg2png(url=str(path)) - img = Image.open(BytesIO(png_data)) - else: - img = Image.open(path) + raw_bytes = await asyncio.to_thread(path.read_bytes) except Exception as e: log.error( - "Failed to load image file", + "Failed to read image file", file_path=str(path), error_type=type(e).__name__, error=str(e), ) raise ImageLoadError(f"Cannot load image '{path.name}': {type(e).__name__}") from e + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.IMAGE, + raw_bytes=raw_bytes, + metadata=metadata, + ) + try: + processed = await self._parser.parse(core_doc) + if not processed.images or not processed.images[0].image_bytes: + raise ImageLoadError(f"Cannot load image '{path.name}': failed to decode") + + img = Image.open(BytesIO(processed.images[0].image_bytes)) + img.load() + except ImageLoadError: + raise + except Exception as e: + log.error("Failed to decode image file", file_path=str(path), error_type=type(e).__name__, error=str(e)) + raise ImageLoadError(f"Cannot load image '{path.name}': {type(e).__name__}") from e description = await self.get_image_description(image_data=img) + doc = Document(page_content=description, metadata=metadata) if save_markdown: self.save_content(description, str(path)) diff --git a/openrag/components/indexer/loaders/pdf_loaders/docling2.py b/openrag/components/indexer/loaders/pdf_loaders/docling2.py index a686465c2..22049a155 100644 --- a/openrag/components/indexer/loaders/pdf_loaders/docling2.py +++ b/openrag/components/indexer/loaders/pdf_loaders/docling2.py @@ -1,150 +1,84 @@ -import asyncio - -import ray -import torch -from config import load_config -from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend -from docling.datamodel.base_models import InputFormat -from docling.datamodel.document import ConversionResult -from docling.datamodel.pipeline_options import ( - AcceleratorDevice, - AcceleratorOptions, - PdfPipelineOptions, - TableFormerMode, - TableStructureOptions, -) -from docling.document_converter import DocumentConverter, PdfFormatOption -from docling_core.types.doc.document import PictureItem +"""Docling-backed PDF loader. + +The Ray actor + pool (``DoclingWorker``, ``DoclingPool``) and the +services-side :class:`BasePooledParser` implementation now live in +``services/workers/parsers/docling_workers.py``; this module re-exports +them for legacy import paths (``services.workers.bootstrap`` constructs the +named ``DoclingPool`` actor at startup via ``get_or_create_actor``). + +``DoclingLoader2`` is a thin :class:`BaseLoader` adapter that delegates +to :class:`core.indexing.parsers.pdf.docling.DoclingParser`, which +wraps the services-side pool. New code should call the core parser +directly; this shim keeps the legacy loader-discovery path alive until +consumers migrate. +""" + +from __future__ import annotations + from langchain_core.documents.base import Document -from tqdm.asyncio import tqdm +from services.workers.parsers.docling_workers import ( # noqa: F401 (re-exported for legacy paths) + DoclingLoader, + DoclingPool, + DoclingWorker, +) from utils.logger import get_logger from ..base import BaseLoader logger = get_logger() -config = load_config() - -if torch.cuda.is_available(): - DOCLING_NUM_GPUS = config.loader.docling_num_gpus -else: # On CPU - DOCLING_NUM_GPUS = 0 -DOCLING_MAX_TASKS_PER_WORKER = config.loader.docling_max_tasks_per_worker +class DoclingLoader2(BaseLoader): + """Adapter shim — delegates to ``DoclingParser`` via the services-side pool.""" + def __init__(self, **kwargs): + super().__init__(**kwargs) + from core.indexing.parsers.pdf.docling import DoclingParser + from services.workers.parsers.docling_workers import DoclingLoader as _DoclingLoader -@ray.remote(num_gpus=DOCLING_NUM_GPUS) -class DoclingWorker: - def __init__(self): - img_scale = 2 - pipeline_options = PdfPipelineOptions( - do_ocr=True, - do_table_structure=True, - generate_picture_images=True, - images_scale=img_scale, - # generate_table_images=True, - # generate_page_images=True - ) - pipeline_options.table_structure_options = TableStructureOptions( - do_cell_matching=True, mode=TableFormerMode.ACCURATE - ) + self._parser = DoclingParser(pool=_DoclingLoader()) - pipeline_options.accelerator_options = AcceleratorOptions(device=AcceleratorDevice.AUTO) - self.converter = DocumentConverter( - format_options={ - InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options, backend=PyPdfiumDocumentBackend) - } + async def aload_document(self, file_path, metadata, save_markdown=False): + import asyncio + from pathlib import Path + + from core.models.document import Document as CoreDocument + from core.models.document import DocumentType + + path = Path(file_path) + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.PDF, + raw_bytes=raw_bytes, + metadata=dict(metadata or {}), ) + processed = await self._parser.parse(core_doc) - async def convert(self, file_path) -> ConversionResult: - with torch.no_grad(): - o = await asyncio.to_thread(self.converter.convert, str(file_path)) - return o + markdown = "" + for block in processed.text_blocks: + markdown += block.text + f"\n[PAGE_{block.page_number}]\n" + if self.image_captioning and processed.images: + import io -@ray.remote -class DoclingPool: - def __init__(self): - from config import load_config - from utils.logger import get_logger + from PIL import Image - self.logger = get_logger() - self.config = load_config() - self.pool_size = self.config.loader.docling_pool_size + pil_images = [] + for img_block in processed.images: + if img_block.image_bytes: + pil_images.append(Image.open(io.BytesIO(img_block.image_bytes))) - self.actors = [DoclingWorker.remote() for _ in range(self.pool_size)] - self._queue: asyncio.Queue[ray.actor.ActorHandle] = asyncio.Queue() - - for _ in range(DOCLING_MAX_TASKS_PER_WORKER): - for actor in self.actors: - self._queue.put_nowait(actor) - - total_slots = self.pool_size * DOCLING_MAX_TASKS_PER_WORKER - self.logger.info( - f"Docling pool: {self.pool_size} actors × {DOCLING_MAX_TASKS_PER_WORKER} slots = " - f"{total_slots} PDF concurrency" - ) - - async def process_pdf(self, file_path: str) -> ConversionResult: - from components.ray_utils import call_ray_actor_with_timeout, retry_with_backoff - - timeout = self.config.loader.docling_timeout - - async def attempt(i: int): - actor: DoclingWorker = await self._queue.get() - try: - return await call_ray_actor_with_timeout( - actor.convert.remote(file_path), - timeout=timeout, - task_description=f"DoclingPool PDF ({file_path})", - ) - finally: - await self._queue.put(actor) - - return await retry_with_backoff( - attempt, - max_retries=self.config.loader.docling_max_task_retry, - base_delay=self.config.loader.docling_retry_base_delay, - task_description=f"DoclingPool PDF ({file_path})", - ) - - -class DoclingLoader2(BaseLoader): - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.docling_actor: DoclingPool = ray.get_actor("DoclingPool", namespace="openrag") - - async def aload_document(self, file_path, metadata, save_markdown=False): - result: ConversionResult = await self.docling_actor.process_pdf.remote(file_path) - n_pages = len(result.pages) - - s = "" - for i in range(1, n_pages + 1): - s += result.document.export_to_markdown(page_no=i) - s += f"\n[PAGE_{i}]\n" - - enriched_content = s - if self.image_captioning: - pictures = result.document.pictures - descriptions = await self.get_captions(pictures) - for description in descriptions: - enriched_content = enriched_content.replace("", description, 1) + if pil_images: + captions = await self.caption_images(pil_images) + for img_block, caption in zip(processed.images, captions): + ref = (img_block.metadata or {}).get("markdown_ref") + if ref: + markdown = markdown.replace("", caption, 1) else: logger.debug("Image captioning disabled. Ignoring images.") - doc = Document(page_content=enriched_content, metadata=metadata) + doc = Document(page_content=markdown, metadata=metadata) if save_markdown: - self.save_document(Document(page_content=enriched_content), str(file_path)) + self.save_document(Document(page_content=markdown), str(file_path)) return doc - - async def get_captions(self, pictures: list[PictureItem]): - tasks = [] - for picture in pictures: - tasks.append(self.get_image_description(picture.image.pil_image)) - try: - results = await tqdm.gather(*tasks, desc="Captioning imgs") - except asyncio.CancelledError: - for task in tasks: - task.cancel() - raise - return results diff --git a/openrag/components/indexer/loaders/pdf_loaders/marker.py b/openrag/components/indexer/loaders/pdf_loaders/marker.py index 62e136dc1..fce7b3116 100644 --- a/openrag/components/indexer/loaders/pdf_loaders/marker.py +++ b/openrag/components/indexer/loaders/pdf_loaders/marker.py @@ -1,306 +1,52 @@ +""" +Marker-backed PDF loader. + +The Ray actor + pool that drive Marker (``MarkerWorker``, +``MarkerPool``) and the services-side :class:`BasePooledParser` +implementation now live in +``services/workers/parsers/marker_workers.py``; this module re-exports +``MarkerWorker`` and ``MarkerPool`` for legacy import paths +(``services.workers.bootstrap`` constructs the named ``MarkerPool`` actor at +startup via ``get_or_create_actor``). + +``MarkerLoader`` is now a thin :class:`BaseLoader` adapter that +delegates to :class:`core.indexing.parsers.pdf.marker.MarkerParser`, +which itself wraps the services-side pool. New code should call the +core parser directly; this shim keeps the legacy loader-discovery path +alive until consumers migrate. +""" + import asyncio -import gc -import re import time +from io import BytesIO from pathlib import Path -import pypdfium2 -import ray -import torch -from config import load_config +from core.indexing.parsers.pdf.marker import MarkerParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document -from marker.converters.pdf import PdfConverter +from PIL import Image +from services.workers.parsers.marker_workers import ( # noqa: F401 (re-exported for legacy import paths) + MarkerLoader as _ServicesMarkerPool, +) +from services.workers.parsers.marker_workers import ( # noqa: F401 + MarkerPool, + MarkerWorker, +) from utils.logger import get_logger from ..base import BaseLoader logger = get_logger() -config = load_config() - -if torch.cuda.is_available(): - MARKER_NUM_GPUS = config.loader.marker_num_gpus -else: # On CPU - MARKER_NUM_GPUS = 0 - - -@ray.remote(num_gpus=MARKER_NUM_GPUS, max_restarts=5) -class MarkerWorker: - def __init__(self): - import os - - from config import load_config - from utils.logger import get_logger - - self.logger = get_logger() - self.config = load_config() - self.page_sep = "[PAGE_SEP]" - - self._workers = self.config.loader.marker_max_processes - - self.converter_config = { - "output_format": "markdown", - "paginate_output": True, - "page_separator": self.page_sep, - "pdftext_workers": self.config.loader.marker_pdftext_workers, - "disable_multiprocessing": False, - } - os.environ["RAY_ADDRESS"] = "auto" - - self.executor = None - self.init_resources() - - def init_resources(self): - from marker.models import create_model_dict - - self.model_dict = create_model_dict() - for v in self.model_dict.values(): - if hasattr(v.model, "share_memory"): - v.model.share_memory() - - self.setup_mp() - - def setup_mp(self): - """Initialize ProcessPoolExecutor for PDF processing. - - We use ProcessPoolExecutor instead of multiprocessing.Pool because: - - Ray actors run as daemon processes - - Pool workers are daemonic by default and cannot spawn children - - The pdftext library (used by Marker) internally spawns processes - - ProcessPoolExecutor workers are non-daemon, allowing nested process creation - """ - from concurrent.futures import ProcessPoolExecutor - - import torch.multiprocessing as mp - - if self.executor: - self.logger.warning("Resetting ProcessPoolExecutor") - self.executor.shutdown(wait=False, cancel_futures=True) - self.executor = None - - # Ensure spawn method for CUDA compatibility - try: - if mp.get_start_method(allow_none=True) != "spawn": - mp.set_start_method("spawn", force=True) - except RuntimeError: - self.logger.warning("Process start method already set, using existing method") - - self.logger.info(f"Initializing MarkerWorker with {self._workers} workers") - self.executor = ProcessPoolExecutor( - max_workers=self._workers, - initializer=self._worker_init, - initargs=(self.model_dict,), - mp_context=mp.get_context("spawn"), - max_tasks_per_child=self.config.loader.marker_max_tasks_per_child, - ) - self.logger.info("MarkerWorker initialized with ProcessPoolExecutor") - - @staticmethod - def _worker_init(model_dict): - global worker_model_dict - worker_model_dict = model_dict - logger.debug("Worker initialized with model dictionary") - - @staticmethod - def _process_pdf(file_path, config): - global worker_model_dict - - page_range = config.get("page_range") - if page_range is not None: - label = f"[p{page_range[0]}-{page_range[-1]}]" - else: - label = "(all pages)" - - try: - logger.debug("Processing PDF", path=file_path, label=label) - converter = PdfConverter( - artifact_dict=worker_model_dict, - config=config, - ) - render = converter(file_path) - return render - except Exception as e: - logger.exception("Error processing PDF", path=file_path, label=label, error=str(e)) - raise - finally: - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.ipc_collect() - - async def process_pdf(self, file_path: str, page_range: list[int] | None = None): - from concurrent.futures import TimeoutError as FuturesTimeoutError - - converter_config = self.converter_config.copy() - if page_range is not None: - converter_config["page_range"] = page_range - - loop = asyncio.get_event_loop() - timeout = self.config.loader.marker_timeout - - def run_with_timeout(): - future = self.executor.submit(self._process_pdf, file_path, converter_config) - try: - result = future.result(timeout=timeout) - return result - except FuturesTimeoutError: - self.logger.exception("MarkerWorker child process timed out", path=file_path) - raise - except Exception: - self.logger.exception("Error processing with MarkerWorker", path=file_path) - raise - - result = await loop.run_in_executor(None, run_with_timeout) - return result.markdown, result.images - - def is_pool_broken(self): - # ProcessPoolExecutor auto-replaces dead/finished workers on next - # submit(), so counting live processes is unreliable and unnecessary. - # Only a None or shut-down executor requires reinitialization. - return self.executor is None or bool(getattr(self.executor, "_broken", False)) - - def __del__(self): - """Clean up ProcessPoolExecutor on actor destruction""" - if self.executor: - try: - self.executor.shutdown(wait=False, cancel_futures=True) - except Exception: - pass # Best effort cleanup - - -@ray.remote(max_restarts=5) -class MarkerPool: - def __init__(self): - from config import load_config - from utils.logger import get_logger - - self.logger = get_logger() - self.config = load_config() - self.max_processes = self.config.loader.marker_max_processes - self.pool_size = self.config.loader.marker_pool_size - self.actors = [MarkerWorker.remote() for _ in range(self.pool_size)] - self._queue: asyncio.Queue[ray.actor.ActorHandle] = asyncio.Queue() - - for _ in range(self.max_processes): - for actor in self.actors: - self._queue.put_nowait(actor) - - self.logger.info( - f"Marker pool: {self.pool_size} actors × {self.max_processes} slots = " - f"{self.pool_size * self.max_processes} PDF concurrency" - ) - - @staticmethod - def _get_page_count(file_path: str) -> int: - pdf = pypdfium2.PdfDocument(file_path) - try: - return len(pdf) - finally: - pdf.close() - - @staticmethod - def _create_chunks(page_count: int, chunk_size: int) -> list[tuple[list[int], str]]: - if page_count <= chunk_size: - return [(list(range(page_count)), f"({page_count}p)")] - chunks = [] - for start in range(0, page_count, chunk_size): - end = min(start + chunk_size, page_count) - page_range = list(range(start, end)) - label = f"[p{start}-{end - 1}]" - chunks.append((page_range, label)) - return chunks - - async def ensure_worker_pool_healthy(self, worker): - from components.ray_utils import call_ray_actor_with_timeout - - timeout = self.config.loader.marker_timeout - broken = await call_ray_actor_with_timeout( - worker.is_pool_broken.remote(), - timeout=timeout, - task_description="MarkerWorker pool health check", - ) - if broken: - self.logger.warning("Worker ProcessPoolExecutor is broken. Reinitializing pool...") - await call_ray_actor_with_timeout( - worker.setup_mp.remote(), - timeout=timeout, - task_description="MarkerWorker pool reset", - ) - - async def _process_chunk(self, file_path: str, page_range: list[int] | None, label: str): - """Acquire a worker slot, process a PDF chunk, and release the slot. - - Retries on failure with exponential backoff up to marker_max_task_retry times. - A fresh worker is acquired per attempt so a flaky worker can be sidestepped - and ensure_worker_pool_healthy re-runs each time. - """ - from components.ray_utils import call_ray_actor_with_timeout, retry_with_backoff - - timeout = self.config.loader.marker_timeout - - async def attempt(i: int): - worker = await self._queue.get() - try: - self.logger.info(f"MarkerWorker allocated for {label} (attempt {i + 1})") - await self.ensure_worker_pool_healthy(worker) - future = worker.process_pdf.remote(file_path, page_range=page_range) - return await call_ray_actor_with_timeout( - future, - timeout=timeout, - task_description=f"MarkerPool PDF {label} ({file_path})", - ) - finally: - await self._queue.put(worker) - self.logger.debug(f"MarkerWorker returned to pool for {label}") - - return await retry_with_backoff( - attempt, - max_retries=self.config.loader.marker_max_task_retry, - base_delay=self.config.loader.marker_retry_base_delay, - task_description=f"MarkerPool PDF {label} ({file_path})", - ) - - async def process_pdf(self, file_path: str): - chunk_size = self.config.loader.marker_chunk_size - - if chunk_size <= 0: - return await self._process_chunk(file_path, page_range=None, label="(all pages)") - - page_count = self._get_page_count(file_path) - chunks = self._create_chunks(page_count, chunk_size) - - if len(chunks) == 1: - page_range, label = chunks[0] - return await self._process_chunk(file_path, page_range=None, label=label) - - self.logger.info( - f"Splitting {page_count}-page PDF into {len(chunks)} chunks of ~{chunk_size} pages for parallel processing" - ) - - tasks = [asyncio.create_task(self._process_chunk(file_path, page_range, label)) for page_range, label in chunks] - try: - results = await asyncio.gather(*tasks) - except Exception: - for task in tasks: - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - raise - - # Reassemble: concatenate markdown in order, merge image dicts - all_markdown = [] - all_images = {} - for markdown, images in results: - all_markdown.append(markdown) - all_images.update(images) - - combined_markdown = "\n\n".join(all_markdown) - return combined_markdown, all_images class MarkerLoader(BaseLoader): + """Adapter shim — delegates to ``MarkerParser`` via the services-side pool.""" + def __init__(self, **kwargs): super().__init__(**kwargs) self.page_sep = "[PAGE_SEP]" - self.worker = ray.get_actor("MarkerPool", namespace="openrag") + self._parser = MarkerParser(pool=_ServicesMarkerPool()) async def aload_document( self, @@ -308,41 +54,46 @@ async def aload_document( metadata: dict | None = None, save_markdown: bool = False, ) -> Document: - from components.ray_utils import call_ray_actor_with_timeout - if metadata is None: metadata = {} + path = Path(file_path) file_path_str = str(file_path) start = time.time() try: - timeout = self.config.loader.marker_timeout - future = self.worker.process_pdf.remote(file_path_str) - markdown, images = await call_ray_actor_with_timeout( - future, - timeout=timeout, - task_description=f"MarkerLoader PDF loading ({file_path_str})", + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.PDF, + raw_bytes=raw_bytes, + metadata=dict(metadata), ) + processed = await self._parser.parse(core_doc) + markdown = "".join(f"{b.text}\n[PAGE_{b.page_number}]\n" for b in processed.text_blocks) if not markdown: raise RuntimeError(f"Conversion failed for {file_path_str}") - if self.image_captioning: - keys = list(images.keys()) - captions = await self.caption_images(list(images.values())) - for key, caption in zip(keys, captions): - markdown = markdown.replace(f"![]({key})", caption) - - else: + if not self.image_captioning: logger.debug("Image captioning disabled.") - - markdown = markdown.split(self.page_sep, 1)[1] - markdown = re.sub(r"\{(\d+)\}" + re.escape(self.page_sep), r"[PAGE_\1]", markdown) - markdown = markdown.replace("
", "").strip() + for block in processed.images: + ref = (block.metadata or {}).get("markdown_ref") + if ref: + markdown = markdown.replace(ref, "") + elif processed.images: + pil_images: list[Image.Image] = [] + for block in processed.images: + img = Image.open(BytesIO(block.image_bytes)) + img.load() + pil_images.append(img) + captions = await self.caption_images(pil_images) + for block, caption in zip(processed.images, captions): + ref = (block.metadata or {}).get("markdown_ref") + if ref: + markdown = markdown.replace(ref, caption) doc = Document(page_content=markdown, metadata=metadata) - if save_markdown: self.save_content(markdown, file_path_str) diff --git a/openrag/components/indexer/loaders/pdf_loaders/pymupdf.py b/openrag/components/indexer/loaders/pdf_loaders/pymupdf.py index f00018b4f..c856c9523 100644 --- a/openrag/components/indexer/loaders/pdf_loaders/pymupdf.py +++ b/openrag/components/indexer/loaders/pdf_loaders/pymupdf.py @@ -1,25 +1,60 @@ +""" +PyMuPDF-backed PDF loader implementation. + +``PyMuPDFLoader`` and ``PyMuPDF4LLMLoader`` are now thin +:class:`BaseLoader` adapters that delegate to +:class:`core.indexing.parsers.pdf.pymupdf.PyMuPDFParser` (text and +markdown modes respectively). The markdown adapter additionally layers +VLM captioning of embedded images on top via the ``BaseLoader`` mixin. +New code should call the core parser directly; this shim keeps the +legacy loader-discovery path alive until consumers migrate. +""" + +import asyncio +from io import BytesIO from pathlib import Path -import pymupdf4llm -from langchain_community.document_loaders import PyMuPDFLoader as pymupdf_loader +from core.indexing.parsers.pdf.pymupdf import PyMuPDFParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document +from PIL import Image +from utils.logger import get_logger from ..base import BaseLoader +logger = get_logger() + + +def _join_pages_with_anchors(text_blocks) -> str: + """Join one ``TextBlock`` per page with the legacy ``\\n[PAGE_N]\\n`` anchors.""" + return "".join(f"{b.text}\n[PAGE_{b.page_number}]\n" for b in text_blocks) + + +async def _read_pdf_bytes(file_path) -> tuple[Path, bytes]: + path = Path(file_path) + raw_bytes = await asyncio.to_thread(path.read_bytes) + return path, raw_bytes + class PyMuPDFLoader(BaseLoader): + """Adapter shim — delegates to ``PyMuPDFParser(mode='text')``.""" + def __init__(self, **kwargs): super().__init__(**kwargs) + self._parser = PyMuPDFParser(mode="text") async def aload_document(self, file_path, metadata: dict = None, save_markdown=False): - loader = pymupdf_loader( - file_path=Path(file_path), + metadata = {} if metadata is None else dict(metadata) + path, raw_bytes = await _read_pdf_bytes(file_path) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.PDF, + raw_bytes=raw_bytes, + metadata=metadata, ) - pages = await loader.aload() - - s = "" - for page_num, segment in enumerate(pages, start=1): - s += segment.page_content.strip() + f"\n[PAGE_{page_num}]\n" + processed = await self._parser.parse(core_doc) + s = _join_pages_with_anchors(processed.text_blocks) doc = Document(page_content=s, metadata=metadata) if save_markdown: @@ -28,15 +63,43 @@ async def aload_document(self, file_path, metadata: dict = None, save_markdown=F class PyMuPDF4LLMLoader(BaseLoader): + """Adapter shim — delegates to ``PyMuPDFParser(mode='markdown')``; layers image captioning on top.""" + def __init__(self, **kwargs) -> None: super().__init__(**kwargs) + self._parser = PyMuPDFParser(mode="markdown") async def aload_document(self, file_path, metadata: dict = None, save_markdown=False): - pages = pymupdf4llm.to_markdown(file_path, write_images=False, page_chunks=True) + metadata = {} if metadata is None else dict(metadata) + path, raw_bytes = await _read_pdf_bytes(file_path) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.PDF, + raw_bytes=raw_bytes, + metadata=metadata, + ) + processed = await self._parser.parse(core_doc) + s = _join_pages_with_anchors(processed.text_blocks) - s = "" - for page_num, segment in enumerate(pages, start=1): - s += segment.get("text").strip() + f"\n[PAGE_{page_num}]\n" + if not self.image_captioning: + # Legacy parity: the old loader called ``pymupdf4llm`` with the + # default ``embed_images=False`` and surfaced no images. The new + # parser embeds them as data URIs; strip those refs to match. + for block in processed.images: + ref = (block.metadata or {}).get("markdown_ref") + if ref: + s = s.replace(ref, "") + elif processed.images: + pil_images: list[Image.Image] = [] + for block in processed.images: + img = Image.open(BytesIO(block.image_bytes)) + img.load() + pil_images.append(img) + captions = await self.caption_images(pil_images, desc="Captioning embedded images") + for block, caption in zip(processed.images, captions): + ref = (block.metadata or {}).get("markdown_ref") + if ref: + s = s.replace(ref, caption.replace("\\", "/")) doc = Document(page_content=s, metadata=metadata) if save_markdown: diff --git a/openrag/components/indexer/loaders/pptx_loader.py b/openrag/components/indexer/loaders/pptx_loader.py index 69fa45be5..f91c75832 100644 --- a/openrag/components/indexer/loaders/pptx_loader.py +++ b/openrag/components/indexer/loaders/pptx_loader.py @@ -1,9 +1,20 @@ -import html -import re +""" +PPTX file loader implementation. + +``PPTXLoader`` is now a thin :class:`BaseLoader` adapter that delegates +extraction to :class:`core.indexing.parsers.pptx_parser.PptxParser` and +layers VLM captioning of slide pictures on top via the ``BaseLoader`` +mixin. New code should call the core parser directly; this shim keeps +the legacy loader-discovery path alive until consumers migrate. +""" + +import asyncio from io import BytesIO +from pathlib import Path -import pptx -from html_to_markdown import convert +from core.indexing.parsers.pptx_parser import PptxParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document from PIL import Image from utils.logger import get_logger @@ -13,149 +24,46 @@ logger = get_logger() -class PPTXConverter: - """Implementation based on PPTX converter in MarkItDown library. - - https://github.com/microsoft/markitdown/blob/main/packages/markitdown/src/markitdown/converters/_pptx_converter.py - """ - - def __init__(self, image_placeholder=r"", page_separator: str = "[PAGE_SEP]"): - self.image_placeholder = image_placeholder - self.page_separator = page_separator - - def convert(self, local_path): - md_content = "" - presentation = pptx.Presentation(local_path) - slide_num = 0 - images_list = [] - for slide in presentation.slides: - slide_num += 1 - - title = slide.shapes.title - for shape in slide.shapes: - if self._is_picture(shape): - images_list.append(Image.open(BytesIO(shape.image.blob))) - md_content += self.image_placeholder - - # Tables - if self._is_table(shape): - html_table = "" - first_row = True - for row in shape.table.rows: - html_table += "" - for cell in row.cells: - if first_row: - html_table += "" - else: - html_table += "" - html_table += "" - first_row = False - html_table += "
" + html.escape(cell.text) + "" + html.escape(cell.text) + "
" - md_content += "\n" + convert(html_table).strip() + "\n" - - # Charts - if shape.has_chart: - md_content += self._convert_chart_to_markdown(shape.chart) - - # Text areas - elif shape.has_text_frame: - if shape == title: - md_content += "# " + shape.text.lstrip() + "\n" - else: - md_content += shape.text + "\n" - - md_content = md_content.strip() - - if slide.has_notes_slide: - md_content += "\n\n### Notes:\n" - notes_frame = slide.notes_slide.notes_text_frame - if notes_frame is not None: - md_content += notes_frame.text - md_content = md_content.strip() - - md_content += f"\n[PAGE_{slide_num}]\n" - - return md_content, images_list - - def _is_picture(self, shape): - try: - if shape.shape_type == pptx.enum.shapes.MSO_SHAPE_TYPE.PICTURE: - return True - if shape.shape_type == pptx.enum.shapes.MSO_SHAPE_TYPE.PLACEHOLDER: - if hasattr(shape, "image"): - return True - except NotImplementedError: - # https://python-pptx.readthedocs.io/en/latest/_modules/pptx/shapes/autoshape.html - # Not all shape types are implemented in python-pptx - logger.warning("Encountered an unimplemented shape type.") - - return False - - def _is_table(self, shape): - try: - if shape.shape_type == pptx.enum.shapes.MSO_SHAPE_TYPE.TABLE: - return True - except NotImplementedError: - # # https://python-pptx.readthedocs.io/en/latest/_modules/pptx/shapes/autoshape.html - # Not all shape types are implemented in python-pptx - logger.warning("Encountered an unimplemented shape type.") - return False - - def _convert_chart_to_markdown(self, chart): - try: - md = "\n\n### Chart" - if chart.has_title: - md += f": {chart.chart_title.text_frame.text}" - md += "\n\n" - data = [] - category_names = [c.label for c in chart.plots[0].categories] - series_names = [s.name for s in chart.series] - data.append(["Category"] + series_names) - - for idx, category in enumerate(category_names): - row = [category] - for series in chart.series: - row.append(series.values[idx]) - data.append(row) - - markdown_table = [] - for row in data: - markdown_table.append("| " + " | ".join(map(str, row)) + " |") - header = markdown_table[0] - separator = "|" + "|".join(["---"] * len(data[0])) + "|" - return md + "\n".join([header, separator] + markdown_table[1:]) - except ValueError as e: - # Handle the specific error for unsupported chart types - if "unsupported plot type" in str(e): - return "\n\n[unsupported chart]\n\n" - except Exception: - # Catch any other exceptions that might occur - return "\n\n[unsupported chart]\n\n" - - class PPTXLoader(BaseLoader): + """Adapter shim — delegates to ``PptxParser``; layers image captioning on top.""" + def __init__(self, **kwargs) -> None: super().__init__(**kwargs) - self.image_placeholder = r"" - self.converter = PPTXConverter(image_placeholder=self.image_placeholder, page_separator=self.page_sep) + self._parser = PptxParser() async def aload_document(self, file_path, metadata=None, save_markdown=False): - md_content, imgs = self.converter.convert(local_path=file_path) - - if self.image_captioning: - images_captions = await self.caption_images(imgs, desc="Generating captions") - - for caption in images_captions: - md_content = re.sub( - self.image_placeholder, - caption.replace("\\", "/"), - md_content, - count=1, - ) - else: + metadata = {} if metadata is None else dict(metadata) + path = Path(file_path) + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.PPTX, + raw_bytes=raw_bytes, + metadata=dict(metadata) if metadata else {}, + ) + processed = await self._parser.parse(core_doc) + + # Reconstitute the legacy ``\n[PAGE_N]\n`` page-anchored layout. + slides = [f"{b.text}\n[PAGE_{b.page_number}]" for b in processed.text_blocks] + md_content = ("\n".join(slides) + "\n") if slides else "" + + if not self.image_captioning: logger.info("Image captioning disabled. Ignoring images.") - # Remove image placeholders when captioning is disabled - md_content = md_content.replace(self.image_placeholder, "") + for block in processed.images: + ref = (block.metadata or {}).get("markdown_ref") + if ref: + md_content = md_content.replace(ref, "") + elif processed.images: + pil_images: list[Image.Image] = [] + for block in processed.images: + img = Image.open(BytesIO(block.image_bytes)) + img.load() + pil_images.append(img) + captions = await self.caption_images(pil_images, desc="Generating captions") + for block, caption in zip(processed.images, captions): + ref = (block.metadata or {}).get("markdown_ref") + if ref: + md_content = md_content.replace(ref, caption.replace("\\", "/")) doc = Document(page_content=md_content, metadata=metadata) if save_markdown: diff --git a/openrag/components/indexer/loaders/serializer.py b/openrag/components/indexer/loaders/serializer.py index 28e278892..0f9c5d6c5 100644 --- a/openrag/components/indexer/loaders/serializer.py +++ b/openrag/components/indexer/loaders/serializer.py @@ -1,89 +1,11 @@ -import gc -from pathlib import Path +"""DocSerializer Ray actor — legacy re-export shim. -import ray -import torch -from config import load_config -from langchain_core.documents.base import Document +The implementation now lives in +``services/workers/parsers/doc_serializer.py``; this module re-exports +``DocSerializer`` so existing import paths (``services/workers/bootstrap.py``, +``services/storage/serializer_ray_shim.py``) are unaffected. +""" -from . import get_loader_classes +from services.workers.parsers.doc_serializer import DocSerializer # noqa: F401 -config = load_config() - -# Set ray resources -if torch.cuda.is_available(): - NUM_GPUS = config.ray.num_gpus -else: # On CPU - NUM_GPUS = 0 - -DICT_MIMETYPES = config.loader.mimetypes.to_dict() - - -@ray.remote(max_restarts=5) -class DocSerializer: - def __init__(self, data_dir=None, **kwargs) -> None: - from config import load_config - from utils.logger import get_logger - - self.logger = get_logger() - self.config = load_config() - self.data_dir = data_dir - self.kwargs = kwargs - self.kwargs["config"] = self.config - self.save_markdown = self.config.loader.save_markdown - - # Initialize loader classes: - self.loader_classes = get_loader_classes(config=self.config) - self.logger.info("DocSerializer initialized.") - - async def serialize_document( - self, - task_id: str, - path: str | Path, - metadata: dict | None = None, - ) -> Document: - metadata = metadata or {} - # Set task state - log = self.logger.bind( - file_id=metadata.get("file_id"), - partition=metadata.get("partition"), - task_id=task_id, - ) - task_state_manager = ray.get_actor("TaskStateManager", namespace="openrag") - await task_state_manager.set_state.remote(task_id, "SERIALIZING") - - log.info("Starting document serialization") - - p = Path(path) - file_ext = p.suffix.lower() - mimetype = metadata.get("mimetype", None) - # Get appropriate loader for the file type - if mimetype is None: - loader_cls = self.loader_classes.get(file_ext) - else: - loader_cls = self.loader_classes.get(DICT_MIMETYPES.get(mimetype)) - - if loader_cls is None: - log.warning(f"No loader available for {p.name}") - raise ValueError(f"No loader available for file type {file_ext}.") - - log.debug(f"Loading document: {p.name} with loader {loader_cls.__name__}") - loader = loader_cls(**self.kwargs) - - try: - # Load the doc - doc: Document = await loader.aload_document( - file_path=path, metadata=metadata, save_markdown=self.save_markdown - ) - - # Clean up resources - del loader - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.ipc_collect() - log.info("Document serialized successfully") - return doc - except Exception as e: - log.exception("Failed to serialize document", error=str(e)) - raise +__all__ = ["DocSerializer"] diff --git a/openrag/components/indexer/loaders/test_doc_loader.py b/openrag/components/indexer/loaders/test_doc_loader.py index 68d2f2b30..a74e01b24 100644 --- a/openrag/components/indexer/loaders/test_doc_loader.py +++ b/openrag/components/indexer/loaders/test_doc_loader.py @@ -1,21 +1,25 @@ """ -Unit tests for DocLoader .doc to .docx conversion with fallback. - -Mocks spire.doc.Document entirely since it's a native .NET library -that cannot run without real .doc files. +Unit tests for the legacy ``DocLoader`` shim. + +The .doc → .docx → markdown conversion logic itself is tested in +``core/indexing/parsers/test_doc_parser.py``. These tests cover only +shim-level concerns: the langchain ``Document`` round-trip, the +``save_markdown=True`` integration with ``BaseLoader.save_content``, +and that errors raised by the underlying ``DocParser`` propagate +without being swallowed. """ -import os from unittest.mock import AsyncMock, MagicMock, patch import pytest from config.models import LoaderConfig, VLMConfig +from core.models.document import ProcessedDocument, TextBlock from langchain_core.documents.base import Document as LCDocument @pytest.fixture def mock_config(): - """Create a minimal mock config for BaseLoader.""" + """Minimal mock config for BaseLoader.""" config = MagicMock() config.vlm = VLMConfig(model="mock", base_url="http://mock", api_key="mock") config.loader = LoaderConfig(image_captioning=False, image_captioning_url=False) @@ -27,21 +31,17 @@ def metadata(): return {"file_id": "test-file-id", "partition": "test-partition"} -# All patches needed to import and instantiate DocLoader without real dependencies _PATCHES = [ - patch("components.indexer.loaders.doc.Document"), - patch("components.indexer.loaders.doc.DocxLoader"), + patch("components.indexer.loaders.doc.DocParser"), patch("components.indexer.loaders.base.ChatOpenAI"), patch("components.indexer.loaders.base.load_config"), ] def _start_patches(mock_config): - """Start all patches and return (MockSpireDoc, MockDocxLoader).""" mocks = [p.start() for p in _PATCHES] - mock_spire_doc, mock_docx_loader_cls, mock_chat, mock_load_config = mocks + _mock_doc_parser_cls, _mock_chat, mock_load_config = mocks mock_load_config.return_value = mock_config - return mock_spire_doc, mock_docx_loader_cls def _stop_patches(): @@ -58,150 +58,88 @@ def _patch_cleanup(): _stop_patches() -class TestDocLoader: - """Test DocLoader .doc to .docx conversion and fallback logic.""" - - def _make_loader(self, mock_config): - """Create a DocLoader with all dependencies mocked. Patches must be active.""" - from components.indexer.loaders.doc import DocLoader +def _make_loader(mock_config): + from components.indexer.loaders.doc import DocLoader - loader = DocLoader(config=mock_config) - return loader - - @pytest.mark.asyncio - async def test_successful_conversion(self, mock_config, metadata): - """Test happy path: .doc converts to .docx successfully.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) + return DocLoader(config=mock_config) - expected_doc = LCDocument(page_content="converted markdown", metadata=metadata) - loader.MDLoader.aload_document = AsyncMock(return_value=expected_doc) - mock_doc_instance = MagicMock() - mock_spire_doc.return_value = mock_doc_instance +def _processed(text: str = "markdown") -> ProcessedDocument: + return ProcessedDocument( + document_id="test", + text_blocks=[TextBlock(text=text, page_number=1)] if text else [], + metadata={}, + page_count=1 if text else 0, + ) - result = await loader.aload_document("/fake/path.doc", metadata) - mock_doc_instance.LoadFromFile.assert_called_once_with("/fake/path.doc") - mock_doc_instance.SaveToFile.assert_called_once() - mock_doc_instance.Close.assert_called_once() - - # DocxLoader was called with a temp .docx path - loader.MDLoader.aload_document.assert_called_once() - call_args = loader.MDLoader.aload_document.call_args - assert call_args[0][0].endswith(".docx") - - assert result == expected_doc - mock_doc_instance.GetText.assert_not_called() +class TestDocLoaderShim: + """Shim-level integration: ``DocParser`` ↔ langchain ``Document`` ↔ ``BaseLoader``.""" @pytest.mark.asyncio - async def test_fallback_on_spire_exception(self, mock_config, metadata): - """Test fallback to text extraction when SaveToFile crashes.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) - - mock_doc_instance = MagicMock() - mock_doc_instance.SaveToFile.side_effect = Exception("TypeInitialization_Type_NoTypeAvailable") - mock_doc_instance.GetText.return_value = "Plain text content from .doc" - mock_spire_doc.return_value = mock_doc_instance + async def test_happy_path_returns_langchain_document(self, mock_config, metadata, tmp_path): + """Parser output is joined into ``page_content``; ``metadata`` is passed through.""" + _start_patches(mock_config) + loader = _make_loader(mock_config) + loader._parser = MagicMock() + loader._parser.parse = AsyncMock(return_value=_processed("converted markdown")) - result = await loader.aload_document("/fake/path.doc", metadata) + file_path = tmp_path / "x.doc" + file_path.write_bytes(b"\xd0\xcf\x11\xe0fake-doc") - mock_doc_instance.SaveToFile.assert_called_once() - mock_doc_instance.GetText.assert_called_once() - mock_doc_instance.Close.assert_called_once() + result = await loader.aload_document(str(file_path), metadata) - # DocxLoader should NOT have been called - loader.MDLoader.aload_document.assert_not_called() - assert result.page_content == "Plain text content from .doc" + assert isinstance(result, LCDocument) + assert result.page_content == "converted markdown" assert result.metadata == metadata - @pytest.mark.asyncio - async def test_temp_file_cleaned_up_on_success(self, mock_config, metadata): - """Test temp file is removed after successful conversion.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) - - created_temp_files = [] - - mock_doc_instance = MagicMock() - - def capture_temp_path(path, fmt): - created_temp_files.append(path) - - mock_doc_instance.SaveToFile.side_effect = capture_temp_path - mock_spire_doc.return_value = mock_doc_instance - - expected_doc = LCDocument(page_content="content", metadata=metadata) - loader.MDLoader.aload_document = AsyncMock(return_value=expected_doc) - - await loader.aload_document("/fake/path.doc", metadata) - - # The temp file should have been cleaned up by the finally block - for path in created_temp_files: - assert not os.path.exists(path) + loader._parser.parse.assert_awaited_once() + forwarded = loader._parser.parse.await_args.args[0] + assert forwarded.raw_bytes == b"\xd0\xcf\x11\xe0fake-doc" + assert forwarded.filename == "x.doc" @pytest.mark.asyncio - async def test_temp_file_cleaned_up_on_failure(self, mock_config, metadata): - """Test temp file is removed even when conversion fails.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) - - mock_doc_instance = MagicMock() - created_temp_files = [] - - def save_then_fail(path, fmt): - created_temp_files.append(path) - # Create the file so we can verify it's cleaned up - with open(path, "w") as f: - f.write("partial") - raise Exception("Spire crash") - - mock_doc_instance.SaveToFile.side_effect = save_then_fail - mock_doc_instance.GetText.return_value = "fallback text" - mock_spire_doc.return_value = mock_doc_instance - - result = await loader.aload_document("/fake/path.doc", metadata) - - assert result.page_content == "fallback text" - for path in created_temp_files: - assert not os.path.exists(path), f"Temp file was not cleaned up: {path}" + async def test_empty_parser_result_yields_empty_content(self, mock_config, metadata, tmp_path): + """An empty ``ProcessedDocument`` yields an empty langchain document.""" + _start_patches(mock_config) + loader = _make_loader(mock_config) + loader._parser = MagicMock() + loader._parser.parse = AsyncMock(return_value=_processed(text="")) + + file_path = tmp_path / "x.doc" + file_path.write_bytes(b"") + + result = await loader.aload_document(str(file_path), metadata) + assert result.page_content == "" + assert result.metadata == metadata @pytest.mark.asyncio - async def test_fallback_with_save_markdown(self, mock_config, metadata, tmp_path): - """Test fallback path respects save_markdown flag.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) - - mock_doc_instance = MagicMock() - mock_doc_instance.SaveToFile.side_effect = Exception("Spire crash") - mock_doc_instance.GetText.return_value = "Extracted text" - mock_spire_doc.return_value = mock_doc_instance + async def test_save_markdown_writes_extracted_content(self, mock_config, metadata, tmp_path): + """``save_markdown=True`` calls ``BaseLoader.save_content`` with the extracted text and source path.""" + _start_patches(mock_config) + loader = _make_loader(mock_config) + loader._parser = MagicMock() + loader._parser.parse = AsyncMock(return_value=_processed("Extracted text")) - file_path = str(tmp_path / "test.doc") + file_path = tmp_path / "x.doc" + file_path.write_bytes(b"x") with patch.object(loader, "save_content") as mock_save: - result = await loader.aload_document(file_path, metadata, save_markdown=True) - mock_save.assert_called_once_with("Extracted text", file_path) + result = await loader.aload_document(str(file_path), metadata, save_markdown=True) + mock_save.assert_called_once_with("Extracted text", str(file_path)) assert result.page_content == "Extracted text" @pytest.mark.asyncio - async def test_docx_loader_error_propagates(self, mock_config, metadata): - """Test that MDLoader errors are NOT caught by the Spire fallback.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) - - mock_doc_instance = MagicMock() - mock_spire_doc.return_value = mock_doc_instance - - # Spire conversion succeeds, but DocxLoader fails - loader.MDLoader.aload_document = AsyncMock(side_effect=ValueError("DocxLoader broke")) - - with pytest.raises(ValueError, match="DocxLoader broke"): - await loader.aload_document("/fake/path.doc", metadata) - - # GetText fallback should NOT have been used - mock_doc_instance.GetText.assert_not_called() - # But Close should still be called (via finally) - mock_doc_instance.Close.assert_called_once() + async def test_parser_error_propagates(self, mock_config, metadata, tmp_path): + """Exceptions from the underlying ``DocParser`` are not swallowed.""" + _start_patches(mock_config) + loader = _make_loader(mock_config) + loader._parser = MagicMock() + loader._parser.parse = AsyncMock(side_effect=ValueError("DocParser broke")) + + file_path = tmp_path / "x.doc" + file_path.write_bytes(b"x") + + with pytest.raises(ValueError, match="DocParser broke"): + await loader.aload_document(str(file_path), metadata) diff --git a/openrag/components/indexer/loaders/txt_loader.py b/openrag/components/indexer/loaders/txt_loader.py index 33e272663..775b798c0 100644 --- a/openrag/components/indexer/loaders/txt_loader.py +++ b/openrag/components/indexer/loaders/txt_loader.py @@ -1,11 +1,24 @@ """ Text and Markdown file loader implementation. + +``TextLoader`` and ``MarkdownLoader`` are now thin :class:`BaseLoader` +adapters that delegate extraction to the corresponding core parsers +(:class:`core.indexing.parsers.text_parser.TextParser`, +:class:`core.indexing.parsers.markdown_parser.MarkdownParser`). Image +captioning is layered on top of the markdown adapter via the +``BaseLoader`` mixin. New code should call the core parsers directly; +these shims keep the legacy loader-discovery path alive until consumers +migrate. """ +import asyncio from pathlib import Path from components.indexer.loaders.base import BaseLoader -from langchain_community.document_loaders import TextLoader as LangchainTextLoader +from core.indexing.parsers.markdown_parser import MarkdownParser +from core.indexing.parsers.text_parser import TextParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document from utils.logger import get_logger @@ -13,12 +26,11 @@ class TextLoader(BaseLoader): - """ - Loader for plain text files (.txt). - """ + """Adapter shim — delegates to ``TextParser`` and returns a LangChain ``Document``.""" def __init__(self, **kwargs) -> None: super().__init__(**kwargs) + self._parser = TextParser() async def aload_document( self, @@ -30,13 +42,15 @@ async def aload_document( metadata = {} path = Path(file_path) - loader = LangchainTextLoader(file_path=str(path), autodetect_encoding=True) - - # Load document segments asynchronously - doc_segments = await loader.aload() - - # Create final document - content = doc_segments[0].page_content.strip() + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.TEXT, + raw_bytes=raw_bytes, + metadata=metadata, + ) + processed = await self._parser.parse(core_doc) + content = "\n\n".join(block.text for block in processed.text_blocks).strip() doc = Document(page_content=content, metadata=metadata) if save_markdown: @@ -46,12 +60,11 @@ async def aload_document( class MarkdownLoader(BaseLoader): - """ - Loader for markdown files (.md). - """ + """Adapter shim — delegates to ``MarkdownParser`` and layers image captioning on top.""" def __init__(self, **kwargs) -> None: super().__init__(**kwargs) + self._parser = MarkdownParser() async def aload_document( self, @@ -63,15 +76,16 @@ async def aload_document( metadata = {} path = Path(file_path) - loader = LangchainTextLoader(file_path=str(path), autodetect_encoding=True) - - # Load document segments asynchronously - doc_segments = await loader.aload() - - # Create final document - content = doc_segments[0].page_content.strip() + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.MARKDOWN, + raw_bytes=raw_bytes, + metadata=metadata, + ) + processed = await self._parser.parse(core_doc) + content = "\n\n".join(block.text for block in processed.text_blocks).strip() - # Caption any images in the markdown content = await self.replace_markdown_images_with_captions(content) doc = Document(page_content=content, metadata=metadata) diff --git a/openrag/components/indexer/utils/files.py b/openrag/components/indexer/utils/files.py index e64ed100e..73b407fc3 100644 --- a/openrag/components/indexer/utils/files.py +++ b/openrag/components/indexer/utils/files.py @@ -6,12 +6,8 @@ import aiofiles import consts -from components.utils import load_config from fastapi import HTTPException, UploadFile, status -config = load_config() -SERIALIZE_TIMEOUT = config.ray.indexer.serialize_timeout - def sanitize_filename(filename: str) -> str: # Split filename into name and extension @@ -71,22 +67,6 @@ async def save_file_to_disk( return file_path -async def serialize_file(task_id: str, path: str, metadata: dict | None = None): - import ray - from components.ray_utils import call_ray_actor_with_timeout - - metadata = metadata or {} - - serializer = ray.get_actor("DocSerializer", namespace="openrag") - future = serializer.serialize_document.remote(task_id, path, metadata=metadata) - - return await call_ray_actor_with_timeout( - future, - timeout=SERIALIZE_TIMEOUT, - task_description=f"Serialization task {task_id}", - ) - - def extract_temporal_fields(metadata: dict, temporal_fields: list) -> dict: result = {} for field in temporal_fields: diff --git a/openrag/components/indexer/utils/text_sanitizer.py b/openrag/components/indexer/utils/text_sanitizer.py index 304a094cb..ef14f6f8d 100644 --- a/openrag/components/indexer/utils/text_sanitizer.py +++ b/openrag/components/indexer/utils/text_sanitizer.py @@ -1,149 +1,7 @@ -""" -Text sanitization utilities for cleaning extracted text and improving quality. - -This module provides functions to clean and normalize text extracted from various -document sources (PDFs, Office files, etc.) by removing excessive whitespace, -special characters, and other artifacts that don't add value. -""" - -import re -import unicodedata - - -def sanitize_text( - text: str, - normalize_whitespace: bool = True, - remove_control_chars: bool = True, - remove_zero_width_chars: bool = True, - max_consecutive_newlines: int = 2, - normalize_unicode: bool = True, -) -> str: - """ - Sanitize text by removing useless characters and normalizing whitespace. - - This function performs comprehensive text cleaning including: - - Removing or normalizing control characters - - Removing zero-width spaces and invisible characters - - Normalizing excessive whitespace (spaces, tabs) - - Limiting consecutive newlines - - Unicode normalization - - Args: - text: The input text to sanitize - normalize_whitespace: If True, normalize spaces and tabs to single spaces - remove_control_chars: If True, remove control characters (except \n, \r, \t) - remove_zero_width_chars: If True, remove zero-width spaces and similar chars - max_consecutive_newlines: Maximum number of consecutive newlines to keep (0 = unlimited) - normalize_unicode: If True, normalize unicode to NFC form - - Returns: - Sanitized text string - - Examples: - >>> sanitize_text("Hello world\\n\\n\\n\\nTest") - 'Hello world\\n\\nTest' - >>> sanitize_text("Text with\\t\\ttabs") - 'Text with tabs' - """ - if not text: - return text - - # Normalize unicode to NFC form (composed form) - if normalize_unicode: - text = unicodedata.normalize("NFC", text) - - # Remove zero-width spaces and similar invisible characters - if remove_zero_width_chars: - # Zero-width space (U+200B), zero-width non-joiner (U+200C), - # zero-width joiner (U+200D), word joiner (U+2060), - # zero-width no-break space (U+FEFF) - text = re.sub(r"[\u200B-\u200D\u2060\uFEFF]", "", text) - - # Remove control characters except newline, carriage return, and tab - if remove_control_chars: - # Remove C0 control characters (0x00-0x1F) except \t (0x09), \n (0x0A), \r (0x0D) - # and C1 control characters (0x80-0x9F) - text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F-\x9F]", "", text) - - # Normalize whitespace - if normalize_whitespace: - # Convert multiple spaces to single space - text = re.sub(r" {2,}", " ", text) - - # Convert tabs to single space - text = re.sub(r"\t+", " ", text) - - # Remove spaces at the beginning of lines - text = re.sub(r"(?m)^ +", "", text) - - # Remove spaces at the end of lines - text = re.sub(r"(?m) +$", "", text) - - # Normalize line breaks - # First, normalize different line break styles to \n - text = re.sub(r"\r\n", "\n", text) - text = re.sub(r"\r", "\n", text) - - # Limit consecutive newlines - if max_consecutive_newlines > 0: - pattern = r"\n{" + str(max_consecutive_newlines + 1) + r",}" - replacement = "\n" * max_consecutive_newlines - text = re.sub(pattern, replacement, text) - - # Remove leading/trailing whitespace - text = text.strip() - return text - - -def clean_markdown_table_spacing(markdown_table: str) -> str: - """ - Normalize spacing inside a markdown table: - - trims each cell - - keeps table shape intact - - Args: - markdown_table: Markdown table text to clean - - Returns: - Cleaned markdown table with normalized spacing - """ - cleaned_lines = [] - - for line in markdown_table.strip().split("\n"): - if "|" not in line: - cleaned_lines.append(line.strip()) - continue - - # Split row into cells (preserve leading/trailing pipes) - parts = line.split("|") - - # Strip each cell except the outer empty ones - cleaned_cells = [cell.strip() for cell in parts] - - # Rebuild with a single space around each cell - new_line = "| " + " | ".join(cleaned_cells[1:-1]) + " |" - cleaned_lines.append(new_line) - - return "\n".join(cleaned_lines) - - -def sanitize_extracted_text(text: str) -> str: - """ - Convenience function for sanitizing text extracted from documents. - - This applies a standard set of cleaning operations suitable for - text extraction endpoints and general document processing. - Uses the default sanitization settings which include: - - Normalize whitespace - - Remove control characters - - Remove zero-width characters - - Limit consecutive newlines to 2 - - Normalize Unicode - - Args: - text: The extracted text to sanitize - - Returns: - Sanitized text - """ - return sanitize_text(text) +# Re-export from canonical location for backward compatibility. +# New code should import from `core.utils.text` directly. +from core.utils.text import ( # noqa: F401 + clean_markdown_table_spacing, + sanitize_extracted_text, + sanitize_text, +) diff --git a/openrag/components/indexer/vectordb/__init__.py b/openrag/components/indexer/vectordb/__init__.py deleted file mode 100644 index bf9de8be6..000000000 --- a/openrag/components/indexer/vectordb/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .vectordb import * diff --git a/openrag/components/indexer/vectordb/models.py b/openrag/components/indexer/vectordb/models.py deleted file mode 100644 index bcc2c5e76..000000000 --- a/openrag/components/indexer/vectordb/models.py +++ /dev/null @@ -1,197 +0,0 @@ -from datetime import datetime - -from sqlalchemy import ( - JSON, - Boolean, - CheckConstraint, - Column, - DateTime, - ForeignKey, - Index, - Integer, - LargeBinary, - String, - UniqueConstraint, -) -from sqlalchemy.orm import ( - declarative_base, - relationship, -) - -Base = declarative_base() - - -class File(Base): - __tablename__ = "files" - - id = Column(Integer, primary_key=True) - file_id = Column(String, nullable=False, index=True) # Added index for file_id lookups - # Foreign key points directly to the partition string - partition_name = Column(String, ForeignKey("partitions.partition"), nullable=False, index=True) # Added index - file_metadata = Column(JSON, nullable=True, default=dict) - - created_by = Column(Integer, ForeignKey("users.id", ondelete="SET NULL"), nullable=True, index=True) - - # Document relationship fields - relationship_id = Column( - String, nullable=True, index=True - ) # Groups related documents (e.g., email thread ID, folder path) - parent_id = Column( - String, nullable=True, index=True - ) # Hierarchical parent reference (e.g., parent email, parent folder) - - # relationship to the Partition object - partition = relationship("Partition", back_populates="files") - - # Enforce uniqueness of (file_id, partition_name) - this also creates an index - __table_args__ = ( - UniqueConstraint("file_id", "partition_name", name="uix_file_id_partition"), - # Additional composite index for common query patterns (partition first for better selectivity) - Index("ix_partition_file", "partition_name", "file_id"), - # Indexes for relationship queries - Index("ix_relationship_partition", "relationship_id", "partition_name"), - Index("ix_parent_partition", "parent_id", "partition_name"), - ) - - def to_dict(self): - metadata = self.file_metadata or {} - d = { - "partition": self.partition_name, - "file_id": self.file_id, - "relationship_id": self.relationship_id, - "parent_id": self.parent_id, - **metadata, - } - return d - - def __repr__(self): - return f"" - - -class Partition(Base): - __tablename__ = "partitions" - - id = Column(Integer, primary_key=True) - partition = Column(String, unique=True, nullable=False, index=True) # Index already exists due to unique constraint - created_at = Column( - DateTime, default=datetime.now, nullable=False, index=True - ) # Added index for time-based queries - files = relationship("File", back_populates="partition", cascade="all, delete-orphan") - memberships = relationship("PartitionMembership", back_populates="partition", cascade="all, delete-orphan") - workspaces = relationship( - "Workspace", - cascade="all, delete-orphan", - backref="partition_ref", - foreign_keys="Workspace.partition_name", - primaryjoin="Partition.partition == Workspace.partition_name", - ) - - def to_dict(self): - d = { - "partition": self.partition, - "created_at": self.created_at.isoformat(), - } - return d - - def __repr__(self): - return f"" - - -class User(Base): - __tablename__ = "users" - - id = Column(Integer, primary_key=True) - external_user_id = Column(String, unique=True, nullable=True, index=True) - display_name = Column(String, nullable=True) - email = Column(String, unique=True, nullable=True, index=True) - token = Column(String, unique=True, nullable=True, index=True) - is_admin = Column(Boolean, default=False, nullable=False) - created_at = Column(DateTime, default=datetime.now, nullable=False) - file_quota = Column(Integer, nullable=True, default=None) - file_count = Column(Integer, nullable=False, default=0) - memberships = relationship("PartitionMembership", back_populates="user", cascade="all, delete-orphan") - oidc_sessions = relationship("OIDCSession", back_populates="user", cascade="all, delete-orphan") - - -class OIDCSession(Base): - __tablename__ = "oidc_sessions" - - id = Column(Integer, primary_key=True) - session_token_hash = Column(String(64), unique=True, nullable=False, index=True) - user_id = Column( - Integer, - ForeignKey("users.id", ondelete="CASCADE"), - nullable=False, - index=True, - ) - sid = Column(String, nullable=True, index=True) # OIDC session id claim (for back-channel logout) - sub = Column(String, nullable=False) # OIDC subject claim - id_token_encrypted = Column(LargeBinary, nullable=True) - access_token_encrypted = Column(LargeBinary, nullable=True) - refresh_token_encrypted = Column(LargeBinary, nullable=True) - access_token_expires_at = Column(DateTime, nullable=False) - session_expires_at = Column(DateTime, nullable=False) - created_at = Column(DateTime, default=datetime.now, nullable=False) - last_refresh_at = Column(DateTime, nullable=True) - revoked_at = Column(DateTime, nullable=True) - - user = relationship("User", back_populates="oidc_sessions") - - __table_args__ = (Index("ix_oidc_sessions_user_sub", "user_id", "sub"),) - - -class PartitionMembership(Base): - __tablename__ = "partition_memberships" - - id = Column(Integer, primary_key=True) - partition_name = Column( - String, - ForeignKey("partitions.partition", ondelete="CASCADE"), - nullable=False, - index=True, - ) - user_id = Column(Integer, ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True) - role = Column(String, nullable=False) # 'owner' | 'editor' | 'viewer' - added_at = Column(DateTime, default=datetime.now, nullable=False) - - __table_args__ = ( - UniqueConstraint("partition_name", "user_id", name="uix_partition_user"), - CheckConstraint("role IN ('owner','editor','viewer')", name="ck_membership_role"), - Index("ix_user_partition", "user_id", "partition_name"), - ) - - partition = relationship("Partition", back_populates="memberships") - user = relationship("User", back_populates="memberships") - - -class Workspace(Base): - __tablename__ = "workspaces" - - id = Column(Integer, primary_key=True) - workspace_id = Column(String, unique=True, nullable=False, index=True) - partition_name = Column( - String, - ForeignKey("partitions.partition", ondelete="CASCADE"), - nullable=False, - index=True, - ) - created_by = Column(Integer, ForeignKey("users.id", ondelete="SET NULL"), nullable=True) - display_name = Column(String, nullable=True) - created_at = Column(DateTime, default=datetime.now) - - files = relationship("WorkspaceFile", cascade="all, delete-orphan", backref="workspace") - - -class WorkspaceFile(Base): - __tablename__ = "workspace_files" - - id = Column(Integer, primary_key=True) - workspace_id = Column( - String, - ForeignKey("workspaces.workspace_id", ondelete="CASCADE"), - nullable=False, - index=True, - ) - file_id = Column(Integer, ForeignKey("files.id", ondelete="CASCADE"), nullable=False, index=True) - - __table_args__ = (UniqueConstraint("workspace_id", "file_id", name="uix_workspace_file"),) diff --git a/openrag/components/indexer/vectordb/test_oidc_sessions.py b/openrag/components/indexer/vectordb/test_oidc_sessions.py deleted file mode 100644 index d63090ede..000000000 --- a/openrag/components/indexer/vectordb/test_oidc_sessions.py +++ /dev/null @@ -1,398 +0,0 @@ -"""Unit tests for OIDC user/session methods on ``PartitionFileManager``. - -These tests exercise the real ORM methods added in Phase 2 of the OIDC -integration (see ``.omc/plans/oidc-auth/plan.md`` §4, §6.3). They run -against an in-memory SQLite database — no Ray, no Postgres, no Milvus -required. - -A ``PartitionFileManager`` instance is created without invoking -``__init__`` (which assumes Postgres + ``sqlalchemy_utils.database_exists`` -semantics); we instead wire up a SQLite engine via ``Base.metadata.create_all`` -and attach it to the object. This mirrors how ``PartitionFileManager`` -itself initialises the schema in production (see ``utils.py``:``__init__`` -which does the same ``Base.metadata.create_all`` call against Postgres). -""" - -from datetime import datetime, timedelta - -import pytest -from components.indexer.vectordb.models import Base, User -from components.indexer.vectordb.utils import PartitionFileManager -from models.user import UserCreate -from sqlalchemy import create_engine -from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import sessionmaker -from utils.logger import get_logger - - -@pytest.fixture() -def pfm(): - """In-memory ``PartitionFileManager`` with a clean schema per test.""" - engine = create_engine("sqlite:///:memory:") - Base.metadata.create_all(engine) - Session = sessionmaker(bind=engine, expire_on_commit=False) - - # Bypass __init__ — it targets Postgres and uses sqlalchemy_utils' - # create_database() which misbehaves with in-memory SQLite. Instead, - # construct a bare PartitionFileManager and attach the schema-ready - # engine ourselves. This keeps the method bodies under test untouched. - mgr = PartitionFileManager.__new__(PartitionFileManager) - mgr.engine = engine - mgr.Session = Session - mgr.logger = get_logger() - mgr.file_quota_per_user = -1 # unlimited for tests - - yield mgr - engine.dispose() - - -def _make_user( - pfm, - *, - display_name: str = "Test User", - email: str | None = None, - external_user_id: str | None = None, -) -> int: - """Helper — insert a user row directly (bypasses create_user's token - hashing since we don't need the API token path here) and return the id.""" - with pfm.Session() as s: - u = User( - display_name=display_name, - email=email, - external_user_id=external_user_id, - is_admin=False, - ) - s.add(u) - s.commit() - s.refresh(u) - return u.id - - -# --------------------------------------------------------------------------- -# user lookup by external_user_id -# --------------------------------------------------------------------------- - - -def test_get_user_by_external_id_returns_user(pfm): - user_id = _make_user(pfm, external_user_id="sub-abc-123") - found = pfm.get_user_by_external_id("sub-abc-123") - assert found is not None - assert found["id"] == user_id - assert found["external_user_id"] == "sub-abc-123" - - -def test_get_user_by_external_id_returns_none_for_unknown(pfm): - _make_user(pfm, external_user_id="sub-abc-123") - assert pfm.get_user_by_external_id("sub-other") is None - - -# --------------------------------------------------------------------------- -# create_user — email uniqueness vs. null emails -# --------------------------------------------------------------------------- - - -def test_create_user_allows_multiple_null_emails(pfm): - """No email provided → the user is still created, and a second null-email - user is allowed too. NULLs are distinct under the unique email index, so - email-less accounts never collide with each other.""" - u1 = pfm.create_user(UserCreate(display_name="A", external_user_id="sub-a", email=None)) - u2 = pfm.create_user(UserCreate(display_name="B", external_user_id="sub-b", email=None)) - # Empty-string email is treated as "no email" too. - u3 = pfm.create_user(UserCreate(display_name="C", external_user_id="sub-c", email="")) - assert u1["email"] is None - assert u2["email"] is None - assert u3["email"] is None - assert len({u1["id"], u2["id"], u3["id"]}) == 3 - - -def test_create_user_rejects_duplicate_email_case_insensitive(pfm): - """A real (non-null) email is unique and matched case-insensitively, since - create_user lowercases before storing.""" - pfm.create_user(UserCreate(display_name="A", external_user_id="sub-a", email="dup@example.com")) - with pytest.raises(IntegrityError): - pfm.create_user(UserCreate(display_name="B", external_user_id="sub-b", email="DUP@Example.com")) - - -def test_get_user_by_email_is_case_insensitive_and_none_safe(pfm): - pfm.create_user(UserCreate(display_name="A", external_user_id="sub-a", email="Found@Example.com")) - assert pfm.get_user_by_email("found@example.com") is not None - assert pfm.get_user_by_email("FOUND@EXAMPLE.COM") is not None - assert pfm.get_user_by_email("missing@example.com") is None - assert pfm.get_user_by_email("") is None - assert pfm.get_user_by_email(None) is None - - -# --------------------------------------------------------------------------- -# update_user_fields — OIDC claim-mapping write path -# --------------------------------------------------------------------------- - - -def test_update_user_fields_updates_display_name_and_lowercases_email(pfm): - user_id = _make_user(pfm, display_name="Old", email="old@example.com") - pfm.update_user_fields(user_id, {"display_name": "New Name", "email": "NEW@Example.COM"}) - refreshed = pfm.get_user_by_external_id("sub-missing") # None lookup, use session - # Re-read via direct ORM since the user has no external_user_id set. - with pfm.Session() as s: - row = s.query(User).filter_by(id=user_id).first() - assert row.display_name == "New Name" - assert row.email == "new@example.com" # normalised - # `refreshed` is unrelated to the assertions — silences unused warning. - assert refreshed is None - - -def test_update_user_fields_raises_on_non_whitelisted_field(pfm): - user_id = _make_user(pfm) - with pytest.raises(ValueError, match="non-whitelisted"): - pfm.update_user_fields(user_id, {"is_admin": True}) - - -def test_update_user_fields_raises_for_unknown_user(pfm): - with pytest.raises(ValueError, match="not found"): - pfm.update_user_fields(999999, {"display_name": "Ghost"}) - - -def test_update_user_fields_drops_none_values(pfm): - user_id = _make_user(pfm, display_name="Keep", email="keep@example.com") - # None values are silently dropped — here every value is None so nothing - # should be written and the row must remain intact. - pfm.update_user_fields(user_id, {"display_name": None, "email": None}) - with pfm.Session() as s: - row = s.query(User).filter_by(id=user_id).first() - assert row.display_name == "Keep" - assert row.email == "keep@example.com" - - -def test_update_user_fields_empty_dict_is_noop(pfm): - """Empty dict must short-circuit before opening a DB session.""" - user_id = _make_user(pfm, display_name="Stable") - # Monkey-patch Session to detect unexpected opens. - original_session = pfm.Session - opened = {"count": 0} - - class _SpySession: - def __call__(self, *a, **kw): - opened["count"] += 1 - return original_session(*a, **kw) - - pfm.Session = _SpySession() - try: - pfm.update_user_fields(user_id, {}) - assert opened["count"] == 0, "update_user_fields({}) must short-circuit before touching the DB" - finally: - pfm.Session = original_session - - -# --------------------------------------------------------------------------- -# create_oidc_session / get_oidc_session_by_token — round-trip -# --------------------------------------------------------------------------- - - -def _session_kwargs(user_id, *, sid="sid-xyz", session_token_plain="plain-token-aaaa"): - now = datetime.now() - return { - "user_id": user_id, - "sub": "sub-abc-123", - "sid": sid, - "session_token_plain": session_token_plain, - "id_token_encrypted": b"\x01\x02\x03", - "access_token_encrypted": b"\xaa\xbb\xcc", - "refresh_token_encrypted": b"\xdd\xee\xff", - "access_token_expires_at": now + timedelta(minutes=5), - "session_expires_at": now + timedelta(hours=8), - } - - -def test_create_and_get_oidc_session_round_trip(pfm): - user_id = _make_user(pfm) - kwargs = _session_kwargs(user_id, session_token_plain="tok-roundtrip-01") - created = pfm.create_oidc_session(**kwargs) - assert created["id"] is not None - assert created["user_id"] == user_id - assert created["sub"] == kwargs["sub"] - assert created["sid"] == kwargs["sid"] - # encrypted blobs passed through untouched - assert created["access_token_encrypted"] == kwargs["access_token_encrypted"] - assert created["revoked_at"] is None - - fetched = pfm.get_oidc_session_by_token("tok-roundtrip-01") - assert fetched is not None - assert fetched["id"] == created["id"] - assert fetched["user_id"] == user_id - - -def test_get_oidc_session_by_token_returns_none_for_unknown(pfm): - user_id = _make_user(pfm) - pfm.create_oidc_session(**_session_kwargs(user_id, session_token_plain="tok-A")) - assert pfm.get_oidc_session_by_token("tok-does-not-exist") is None - - -# --------------------------------------------------------------------------- -# revocation + expiry visibility -# --------------------------------------------------------------------------- - - -def test_get_oidc_session_returns_none_when_revoked(pfm): - user_id = _make_user(pfm) - created = pfm.create_oidc_session(**_session_kwargs(user_id, session_token_plain="tok-revoke")) - pfm.revoke_oidc_session_by_id(created["id"]) - assert pfm.get_oidc_session_by_token("tok-revoke") is None - - -def test_get_oidc_session_returns_none_when_session_expired(pfm): - user_id = _make_user(pfm) - past = datetime.now() - timedelta(hours=1) - pfm.create_oidc_session( - user_id=user_id, - sub="sub-abc-123", - sid="sid-expired", - session_token_plain="tok-expired", - id_token_encrypted=None, - access_token_encrypted=None, - refresh_token_encrypted=None, - access_token_expires_at=past, - session_expires_at=past, # already expired at insert time - ) - assert pfm.get_oidc_session_by_token("tok-expired") is None - - -def test_revoke_oidc_sessions_by_sid_revokes_all_matching(pfm): - user_id = _make_user(pfm) - # Two sessions sharing one sid, one with a different sid. - pfm.create_oidc_session(**_session_kwargs(user_id, sid="sid-shared", session_token_plain="tok-1")) - pfm.create_oidc_session(**_session_kwargs(user_id, sid="sid-shared", session_token_plain="tok-2")) - pfm.create_oidc_session(**_session_kwargs(user_id, sid="sid-other", session_token_plain="tok-3")) - - count = pfm.revoke_oidc_sessions_by_sid("sid-shared") - assert count == 2 - - # The revoked sessions must no longer be retrievable by token. - assert pfm.get_oidc_session_by_token("tok-1") is None - assert pfm.get_oidc_session_by_token("tok-2") is None - # The other-sid session must still be live. - assert pfm.get_oidc_session_by_token("tok-3") is not None - - -def test_revoke_oidc_sessions_by_sid_idempotent(pfm): - """Calling twice with the same sid must only revoke non-revoked rows.""" - user_id = _make_user(pfm) - pfm.create_oidc_session(**_session_kwargs(user_id, sid="sid-X", session_token_plain="tok-X")) - first = pfm.revoke_oidc_sessions_by_sid("sid-X") - second = pfm.revoke_oidc_sessions_by_sid("sid-X") - assert first == 1 - assert second == 0 # already revoked - - -# --------------------------------------------------------------------------- -# update_oidc_session_tokens — post-refresh update -# --------------------------------------------------------------------------- - - -def test_update_oidc_session_tokens_updates_fields_and_last_refresh(pfm): - user_id = _make_user(pfm) - created = pfm.create_oidc_session(**_session_kwargs(user_id, session_token_plain="tok-refresh")) - original_session_expiry = created["session_expires_at"] - new_expiry = datetime.now() + timedelta(minutes=10) - - pfm.update_oidc_session_tokens( - session_id=created["id"], - access_token_encrypted=b"\x11\x22\x33", - refresh_token_encrypted=b"\x44\x55\x66", - access_token_expires_at=new_expiry, - ) - - fetched = pfm.get_oidc_session_by_token("tok-refresh") - assert fetched is not None - assert fetched["access_token_encrypted"] == b"\x11\x22\x33" - assert fetched["refresh_token_encrypted"] == b"\x44\x55\x66" - # datetime comparison — tolerate microsecond differences from DB round-trip - assert abs((fetched["access_token_expires_at"] - new_expiry).total_seconds()) < 1 - assert fetched["last_refresh_at"] is not None - # session_expires_at (the hard cap) is untouched - assert fetched["session_expires_at"] == original_session_expiry - - -def test_update_oidc_session_tokens_accepts_none_refresh(pfm): - """Some IdPs don't rotate refresh_token on refresh (omit refresh_token - in the response). We must keep the old encrypted value.""" - user_id = _make_user(pfm) - created = pfm.create_oidc_session(**_session_kwargs(user_id, session_token_plain="tok-nrr")) - new_expiry = datetime.now() + timedelta(minutes=10) - pfm.update_oidc_session_tokens( - session_id=created["id"], - access_token_encrypted=b"\x99\x88\x77", - refresh_token_encrypted=None, - access_token_expires_at=new_expiry, - ) - fetched = pfm.get_oidc_session_by_token("tok-nrr") - assert fetched["refresh_token_encrypted"] == b"\xdd\xee\xff" # original - - -def test_update_oidc_session_tokens_raises_for_unknown_id(pfm): - with pytest.raises(ValueError, match="does not exist"): - pfm.update_oidc_session_tokens( - session_id=424242, - access_token_encrypted=b"x", - refresh_token_encrypted=None, - access_token_expires_at=datetime.now(), - ) - - -# --------------------------------------------------------------------------- -# cleanup_expired_oidc_sessions -# --------------------------------------------------------------------------- - - -def test_cleanup_deletes_only_rows_older_than_retention(pfm): - """Rows are purged only once ``session_expires_at`` is older than - the 7-day retention window. Still-live and recently-expired rows stay.""" - user_id = _make_user(pfm) - now = datetime.now() - - # (1) Live — must stay - pfm.create_oidc_session( - user_id=user_id, - sub="sub", - sid="sid-live", - session_token_plain="tok-live", - id_token_encrypted=None, - access_token_encrypted=None, - refresh_token_encrypted=None, - access_token_expires_at=now + timedelta(minutes=5), - session_expires_at=now + timedelta(hours=1), - ) - - # (2) Recently expired (within 7-day retention) — must stay - pfm.create_oidc_session( - user_id=user_id, - sub="sub", - sid="sid-recent", - session_token_plain="tok-recent", - id_token_encrypted=None, - access_token_encrypted=None, - refresh_token_encrypted=None, - access_token_expires_at=now - timedelta(hours=2), - session_expires_at=now - timedelta(days=1), - ) - - # (3) Past retention — must be deleted - pfm.create_oidc_session( - user_id=user_id, - sub="sub", - sid="sid-stale", - session_token_plain="tok-stale", - id_token_encrypted=None, - access_token_encrypted=None, - refresh_token_encrypted=None, - access_token_expires_at=now - timedelta(days=30), - session_expires_at=now - timedelta(days=10), - ) - - deleted = pfm.cleanup_expired_oidc_sessions() - assert deleted == 1 - - # Live session still retrievable. - assert pfm.get_oidc_session_by_token("tok-live") is not None - # Recently expired: row kept, but still_masked as expired by get_by_token. - assert pfm.get_oidc_session_by_token("tok-recent") is None - # Stale: row gone entirely. - assert pfm.get_oidc_session_by_token("tok-stale") is None diff --git a/openrag/components/indexer/vectordb/utils.py b/openrag/components/indexer/vectordb/utils.py deleted file mode 100644 index 56f515f08..000000000 --- a/openrag/components/indexer/vectordb/utils.py +++ /dev/null @@ -1,1111 +0,0 @@ -import hashlib -import os -import secrets -from datetime import datetime, timedelta - -from config import load_config -from models.user import UserCreate, UserUpdate -from sqlalchemy import create_engine, delete, func, select, text, update -from sqlalchemy.dialects.postgresql import insert as pg_insert -from sqlalchemy.orm import sessionmaker -from sqlalchemy_utils import ( - create_database, - database_exists, -) -from utils.exceptions.vectordb import * -from utils.logger import get_logger - -from .models import ( - Base, - File, - OIDCSession, - Partition, - PartitionMembership, - User, - Workspace, - WorkspaceFile, -) - -logger = get_logger() -config = load_config() - -DEFAULT_FILE_QUOTA = config.rdb.default_file_quota - - -class PartitionFileManager: - def __init__(self, database_url: str, logger=logger): - try: - self.engine = create_engine(database_url) - if not database_exists(database_url): - create_database(database_url) - - Base.metadata.create_all(self.engine) - self.logger = logger - self.Session = sessionmaker(bind=self.engine) - AUTH_TOKEN = os.getenv("AUTH_TOKEN") - self._ensure_admin_user(AUTH_TOKEN) - self.file_quota_per_user = DEFAULT_FILE_QUOTA - - except Exception as e: - raise VDBConnectionError( - f"Failed to connect to database: {e!s}", - db_url=database_url, - db_type="SQLAlchemy", - ) - - def _ensure_admin_user(self, admin_token: str): - if not admin_token: - admin_token = f"or-{secrets.token_hex(16)}" - hashed_token = self.hash_token(admin_token) - with self.Session() as s: - admin = s.query(User).filter_by(id=1).first() - if not admin: - admin = User( - display_name="Admin", - token=hashed_token, - is_admin=True, - ) - s.add(admin) - s.commit() - self.logger.info("Created admin user") - else: - admin.is_admin = True - admin.token = hashed_token - s.commit() - self.logger.info("Upgraded existing user to admin") - - def list_partition_files(self, partition: str, limit: int | None = None): - """List files in a partition with optional limit - Optimized by querying File table directly""" - log = self.logger.bind(partition=partition) - with self.Session() as session: - log.debug("Listing partition files") - - # Query files directly - if partition doesn't exist, files will be empty - files_query = session.query(File).filter(File.partition_name == partition) - if limit is not None: - files_query = files_query.limit(limit) - - files = files_query.all() - - # If no files found - if not files: - log.warning("Partition doesn't exist or has no files") - return {} - - result = { - "files": [file.to_dict() for file in files], - } - - log.info(f"Listed {len(files)} files from partition") - return result - - def add_file_to_partition( - self, - file_id: str, - partition: str, - file_metadata: dict | None = None, - user_id: int | None = None, - relationship_id: None | str = None, - parent_id: None | str = None, - ): - """Add a file to a partition with optional relationship fields. - - Args: - file_id: Unique identifier for the file - partition: Partition name - file_metadata: Additional metadata as JSON - user_id: User ID for ownership (creates partition membership) - relationship_id: Groups related documents (e.g., email thread ID, folder path) - parent_id: Hierarchical parent reference (e.g., parent email file_id) - """ - log = self.logger.bind(file_id=file_id, partition=partition) - with self.Session() as session: - try: - existing_file = ( - session.query(File.id).filter(File.file_id == file_id, File.partition_name == partition).first() - ) - if existing_file: - log.warning("File already exists") - return False - - partition_obj = session.query(Partition).filter(Partition.partition == partition).first() - if not partition_obj: - partition_obj = Partition(partition=partition) - session.add(partition_obj) - log.info("Created new partition") - - membership = PartitionMembership(partition_name=partition, user_id=user_id, role="owner") - session.add(membership) - - # Add file to partition - file = File( - file_id=file_id, - partition_name=partition, - file_metadata=file_metadata, - relationship_id=relationship_id, - parent_id=parent_id, - created_by=user_id, - ) - - session.add(file) - # Increment uploader's file_count - if user_id: - session.query(User).filter(User.id == user_id).update( - {User.file_count: User.file_count + 1}, synchronize_session=False - ) - session.commit() - log.info("Added file successfully") - return True - except Exception: - session.rollback() - log.exception("Error adding file to partition") - raise - - def remove_file_from_partition(self, file_id: str, partition: str): - """Remove a file from its partition - Optimized without join""" - log = self.logger.bind(file_id=file_id, partition=partition) - with self.Session() as session: - try: - # Direct filter without join (uses composite index) - file = session.query(File).filter(File.file_id == file_id, File.partition_name == partition).first() - if file: - uploader_id = file.created_by - session.delete(file) - if uploader_id: - session.query(User).filter(User.id == uploader_id).update( - {User.file_count: func.greatest(User.file_count - 1, 0)}, - synchronize_session=False, - ) - session.commit() - log.info(f"Removed file {file_id} from partition {partition}") - return True - log.warning("File not found in partition") - return False - except Exception as e: - session.rollback() - log.error(f"Error removing file: {e}") - raise e - - def update_file_metadata_in_db(self, file_id: str, partition: str, file_metadata: dict) -> bool: - """Update the file_metadata JSON column and structured fields for an existing file. - - Returns True if the file was found and updated, False otherwise. - Unlike remove_file_from_partition + add_file_to_partition, this preserves - the files.id primary key so that workspace FK references remain valid. - - If file_metadata contains keys that correspond to structured File columns - (relationship_id, parent_id), those columns are updated too so they stay - in sync with the JSON blob. - """ - log = self.logger.bind(file_id=file_id, partition=partition) - with self.Session() as session: - try: - file = session.query(File).filter(File.file_id == file_id, File.partition_name == partition).first() - if not file: - log.warning("File not found for metadata update") - return False - file.file_metadata = file_metadata - # Sync structured columns when the corresponding keys are present - # in the metadata payload, so PG columns never diverge from the JSON. - if "relationship_id" in file_metadata: - file.relationship_id = file_metadata["relationship_id"] - if "parent_id" in file_metadata: - file.parent_id = file_metadata["parent_id"] - session.commit() - log.info("Updated file metadata in-place") - return True - except Exception: - session.rollback() - log.exception("Error updating file metadata") - raise - - # Sentinel object to distinguish "not provided" from explicit None. - _UNSET = object() - - def update_file_in_partition( - self, - file_id: str, - partition: str, - file_metadata: dict | None = None, - relationship_id: str | None | object = _UNSET, - parent_id: str | None | object = _UNSET, - ) -> bool: - """Update an existing file record in-place (for PUT: new content, same file_id). - - Preserves files.id so workspace FK references stay intact. - Unlike delete+re-add, this never touches file_count or created_by. - - Pass relationship_id=None or parent_id=None explicitly to clear a stale - link. Omit the argument entirely to leave the column unchanged. - """ - log = self.logger.bind(file_id=file_id, partition=partition) - with self.Session() as session: - try: - file = session.query(File).filter(File.file_id == file_id, File.partition_name == partition).first() - if not file: - log.warning("File not found for update") - return False - if file_metadata is not None: - file.file_metadata = file_metadata - if relationship_id is not self._UNSET: - file.relationship_id = relationship_id - if parent_id is not self._UNSET: - file.parent_id = parent_id - session.commit() - log.info("Updated file record in-place") - return True - except Exception: - session.rollback() - log.exception("Error updating file in partition") - raise - - def delete_partition(self, partition: str): - """Delete a partition and all its files""" - with self.Session() as session: - partition_obj = session.query(Partition).filter_by(partition=partition).first() - if partition_obj: - # Count files per uploader before cascade deletes them - uploader_counts = ( - session.query(File.created_by, func.count(File.id)) - .filter(File.partition_name == partition, File.created_by.isnot(None)) - .group_by(File.created_by) - .all() - ) - session.delete(partition_obj) # Cascades to files and memberships - for uploader_id, count in uploader_counts: - session.query(User).filter(User.id == uploader_id).update( - {User.file_count: func.greatest(User.file_count - count, 0)}, - synchronize_session=False, - ) - session.commit() - self.logger.info("Deleted partition", partition=partition) - return True - else: - self.logger.info("Partition does not exist", partition=partition) - return False - - def list_partitions(self): - """List all existing partitions""" - with self.Session() as session: - partitions = session.query(Partition).all() - return [partition.to_dict() for partition in partitions] - - def get_partition_file_count(self, partition: str): - """Get the count of files in a partition - Optimized with direct count""" - with self.Session() as session: - # Optimized: Direct count query instead of loading partition and files - return session.query(File).filter(File.partition_name == partition).count() - - def get_total_file_count(self): - """Get the total count of files across all partitions""" - with self.Session() as session: - return session.query(File).count() - - def partition_exists(self, partition: str): - """Check if a partition exists by its key - Optimized with exists()""" - with self.Session() as session: - # Optimized: Use exists() for better performance - return session.query(session.query(Partition).filter(Partition.partition == partition).exists()).scalar() - - def file_exists_in_partition(self, file_id: str, partition: str): - """Check if a file exists in a specific partition - Optimized without join""" - with self.Session() as session: - # Optimized: Direct filter without join, use exists() for better performance - return session.query( - session.query(File).filter(File.file_id == file_id, File.partition_name == partition).exists() - ).scalar() - - # Users - - def create_user(self, body: UserCreate) -> dict: - """Create a user and generate an API token for them.""" - with self.Session() as s: - token = f"or-{secrets.token_hex(16)}" - hashed_token = self.hash_token(token) - file_quota = body.file_quota - if self.file_quota_per_user > 0 and file_quota is None: - file_quota = self.file_quota_per_user # default to default quota - - user = User( - display_name=body.display_name, - external_user_id=body.external_user_id, - email=(body.email.strip().lower() if body.email else None), - token=hashed_token, - is_admin=body.is_admin, - file_quota=file_quota, - ) - s.add(user) - s.commit() - s.refresh(user) - - return { - "id": user.id, - "display_name": user.display_name, - "external_user_id": user.external_user_id, - "email": user.email, - "token": token, - "is_admin": user.is_admin, - "file_quota": user.file_quota, - "file_count": user.file_count, - } - - def list_users(self) -> list[dict]: - with self.Session() as s: - users = s.query(User).all() - return [ - { - "id": u.id, - "display_name": u.display_name, - "external_user_id": u.external_user_id, - "is_admin": u.is_admin, - "file_quota": u.file_quota, - "file_count": u.file_count, - "created_at": u.created_at.isoformat(), - } - for u in users - ] - - def get_user_by_token(self, token: str) -> dict | None: - with self.Session() as s: - hashed_token = self.hash_token(token) - user = s.query(User).filter(User.token == hashed_token).first() - if not user: - return None - - memberships = [ - { - "partition": m.partition_name, - "role": m.role, - "added_at": m.added_at.isoformat(), - } - for m in user.memberships - ] - - return { - "id": user.id, - "display_name": user.display_name, - "external_user_id": user.external_user_id, - "is_admin": user.is_admin, - "file_quota": user.file_quota, - "file_count": user.file_count, - "memberships": memberships, - } - - def get_user_by_id(self, user_id: int) -> dict | None: - with self.Session() as s: - user = s.query(User).filter(User.id == user_id).first() - if not user: - return None - - memberships = [ - { - "partition": m.partition_name, - "role": m.role, - "added_at": m.added_at.isoformat(), - } - for m in user.memberships - ] - - return { - "id": user.id, - "display_name": user.display_name, - "external_user_id": user.external_user_id, - "is_admin": user.is_admin, - "file_quota": user.file_quota, - "file_count": user.file_count, - "memberships": memberships, - } - - def delete_user(self, user_id: int) -> bool: - with self.Session() as s: - user = s.query(User).filter(User.id == user_id).first() - if not user: - return False - s.delete(user) - s.commit() - return True - - def regenerate_user_token(self, user_id: int) -> dict | None: - with self.Session() as s: - user = s.query(User).filter(User.id == user_id).first() - if not user: - return None - new_token = f"or-{secrets.token_hex(16)}" - hashed_token = self.hash_token(new_token) - user.token = hashed_token - s.commit() - s.refresh(user) - - return { - "id": user.id, - "display_name": user.display_name, - "external_user_id": user.external_user_id, - "token": new_token, - "is_admin": user.is_admin, - "file_quota": user.file_quota, - "file_count": user.file_count, - } - - # Memberships - def list_partition_members(self, partition: str) -> list[dict]: - with self.Session() as s: - if not s.query(Partition).filter(Partition.partition == partition).first(): - self.logger.warning(f"Partition '{partition}' does not exist.") - return [] - ms = s.query(PartitionMembership).filter_by(partition_name=partition).all() - return [ - { - "user_id": m.user_id, - "role": m.role, - "added_at": m.added_at.isoformat(), - } - for m in ms - ] - - def add_partition_member(self, partition: str, user_id: int, role: str) -> bool: - with self.Session() as s: - if not s.query(Partition).filter(Partition.partition == partition).first(): - s.add(Partition(partition=partition)) - m = s.query(PartitionMembership).filter_by(partition_name=partition, user_id=user_id).first() - if m: - m.role = role - else: - s.add(PartitionMembership(partition_name=partition, user_id=user_id, role=role)) - s.commit() - return True - - def remove_partition_member(self, partition: str, user_id: int) -> bool: - with self.Session() as s: - m = s.query(PartitionMembership).filter_by(partition_name=partition, user_id=user_id).first() - if not m: - return False - s.delete(m) - s.commit() - return True - - def update_partition_member_role(self, partition: str, user_id: int, new_role: str) -> bool: - with self.Session() as s: - m = s.query(PartitionMembership).filter_by(partition_name=partition, user_id=user_id).first() - if not m: - return False - m.role = new_role - s.commit() - return True - - def create_partition(self, partition: str, user_id: int): - with self.Session() as s: - if s.query(Partition).filter(Partition.partition == partition).first(): - self.logger.warning(f"Partition '{partition}' already exists.") - return - p = Partition(partition=partition) - s.add(p) - # Add creator as owner - m = PartitionMembership(partition_name=partition, user_id=user_id, role="owner") - s.add(m) - s.commit() - self.logger.info(f"Partition '{partition}' created by user_id {user_id}.") - - def list_user_partitions(self, user_id: int): - """Return full partition objects (to_dict) with role for a given user.""" - with self.Session() as s: - # Join Partition and PartitionMembership - results = ( - s.query(Partition, PartitionMembership.role) - .join( - PartitionMembership, - Partition.partition == PartitionMembership.partition_name, - ) - .filter(PartitionMembership.user_id == user_id) - .all() - ) - - partitions = [] - for partition_obj, role in results: - d = partition_obj.to_dict() - d["role"] = role - partitions.append(d) - - return partitions - - def user_exists(self, user_id: int) -> bool: - with self.Session() as s: - return s.query(User).filter(User.id == user_id).first() is not None - - def user_is_partition_member(self, user_id: int, partition: str) -> bool: - with self.Session() as s: - return s.query(PartitionMembership).filter_by(user_id=user_id, partition_name=partition).first() is not None - - def update_user(self, user_id: int, body: UserUpdate) -> dict: - """Update user's profile fields. Only provided (non-None) fields are updated.""" - with self.Session() as s: - user = s.query(User).filter(User.id == user_id).first() - for field, value in body.model_dump(exclude_unset=True).items(): - setattr(user, field, value) - - s.commit() - s.refresh(user) - return { - "id": user.id, - "display_name": user.display_name, - "external_user_id": user.external_user_id, - "is_admin": user.is_admin, - "created_at": user.created_at.isoformat(), - "file_quota": user.file_quota, - "file_count": user.file_count, - } - - def hash_token(self, token: str) -> str: - """Return a SHA-256 hash of a token string.""" - return hashlib.sha256(token.encode("utf-8")).hexdigest() - - # Document relationship methods - - def get_files_by_relationship(self, partition: str, relationship_id: str) -> list[dict]: - """Get all files sharing a relationship_id within a partition. - - Args: - partition: Partition name - relationship_id: The relationship group identifier - - Returns: - List of file dictionaries - """ - with self.Session() as session: - files = ( - session.query(File) - .filter( - File.partition_name == partition, - File.relationship_id == relationship_id, - ) - .all() - ) - return [f.to_dict() for f in files] - - def get_file_ids_by_relationship(self, partition: str, relationship_id: str) -> list[str]: - """Get file_ids for all files sharing a relationship_id. - - Args: - partition: Partition name - relationship_id: The relationship group identifier - - Returns: - List of file_id strings - """ - with self.Session() as session: - results = ( - session.query(File.file_id) - .filter( - File.partition_name == partition, - File.relationship_id == relationship_id, - ) - .all() - ) - return [r[0] for r in results] - - def get_file_ancestors(self, partition: str, file_id: str, max_ancestor_depth: int | None = None) -> list[dict]: - """Get all ancestors of a file using recursive CTE. - - Returns ordered list from root to the specified file (direct path only). - - Args: - partition: Partition name - file_id: The file identifier to find ancestors for - max_ancestor_depth: Maximum depth to traverse (None = unlimited) - - Returns: - List of file dictionaries ordered from root to the specified file - """ - - with self.Session() as session: - # Recursive CTE for ancestor traversal with optional max depth - depth_condition = "WHERE a.depth < :max_ancestor_depth" if max_ancestor_depth is not None else "" - query = text(f""" - WITH RECURSIVE ancestors AS ( - -- Base case: start with the target file - SELECT id, file_id, partition_name, parent_id, file_metadata, - relationship_id, 0 as depth - FROM files - WHERE file_id = :file_id AND partition_name = :partition - AND relationship_id IS NOT NULL - - UNION ALL - - -- Recursive case: get parent - SELECT f.id, f.file_id, f.partition_name, f.parent_id, - f.file_metadata, f.relationship_id, a.depth + 1 - FROM files f - INNER JOIN ancestors a ON f.file_id = a.parent_id - AND f.partition_name = a.partition_name - AND f.relationship_id IS NOT NULL - {depth_condition} - ) - SELECT * FROM ancestors ORDER BY depth DESC - """) - - params = {"file_id": file_id, "partition": partition} - if max_ancestor_depth is not None: - params["max_ancestor_depth"] = max_ancestor_depth - - result = session.execute(query, params) - - return [ - { - "file_id": row.file_id, - "partition": row.partition_name, - "parent_id": row.parent_id, - "relationship_id": row.relationship_id, - "depth": row.depth, - **(row.file_metadata or {}), - } - for row in result - ] - - def get_ancestor_file_ids(self, partition: str, file_id: str, max_ancestor_depth: int | None = None) -> list[str]: - """Get file_ids for all ancestors of a file. - - Returns ordered list from root to the specified file (direct path only). - - Args: - partition: Partition name - file_id: The file identifier to find ancestors for - max_ancestor_depth: Maximum depth to traverse (None = unlimited) - Returns: - List of file_id strings ordered from root to the specified file - """ - ancestors = self.get_file_ancestors(partition, file_id, max_ancestor_depth) - return [a["file_id"] for a in ancestors] - - # --- Workspace methods --- - - def create_workspace(self, workspace_id: str, partition: str, user_id: int | None, display_name: str | None = None): - with self.Session() as session: - ws = Workspace( - workspace_id=workspace_id, partition_name=partition, created_by=user_id, display_name=display_name - ) - session.add(ws) - session.commit() - - def list_workspaces(self, partition: str) -> list[dict]: - with self.Session() as session: - result = session.execute(select(Workspace).where(Workspace.partition_name == partition)) - return [ - { - "workspace_id": w.workspace_id, - "partition_name": w.partition_name, - "display_name": w.display_name, - "created_by": w.created_by, - "created_at": str(w.created_at), - } - for w in result.scalars() - ] - - def get_workspace(self, workspace_id: str) -> dict | None: - with self.Session() as session: - result = session.execute(select(Workspace).where(Workspace.workspace_id == workspace_id)) - w = result.scalar_one_or_none() - if not w: - return None - return { - "workspace_id": w.workspace_id, - "partition_name": w.partition_name, - "display_name": w.display_name, - "created_by": w.created_by, - "created_at": str(w.created_at), - } - - def delete_workspace(self, workspace_id: str) -> list[str]: - """Delete workspace, return list of orphaned file_ids. - - A file is orphaned if it belongs to this workspace and no other. - Since WorkspaceFile.file_id is now an integer FK to files.id, every - workspace file has a backing files row, so the only condition is - "not present in any other workspace". - """ - with self.Session() as session: - # File PKs present in at least one other workspace - subq_other_ws = select(WorkspaceFile.file_id).where(WorkspaceFile.workspace_id != workspace_id) - # Orphaned = in this workspace, not in any other - result = session.execute( - select(File.file_id) - .join(WorkspaceFile, WorkspaceFile.file_id == File.id) - .where(WorkspaceFile.workspace_id == workspace_id) - .where(WorkspaceFile.file_id.notin_(subq_other_ws)) - ) - orphaned_file_ids = [r[0] for r in result.all()] - - # Delete workspace (cascades workspace_files) - session.execute(delete(Workspace).where(Workspace.workspace_id == workspace_id)) - session.commit() - return orphaned_file_ids - - def get_existing_file_ids(self, partition: str, file_ids: list[str]) -> set[str]: - """Return the subset of *file_ids* that actually exist in *partition*.""" - with self.Session() as session: - result = session.execute( - select(File.file_id).where( - File.partition_name == partition, - File.file_id.in_(file_ids), - ) - ) - return {r[0] for r in result.all()} - - def add_files_to_workspace(self, workspace_id: str, file_ids: list[str]) -> list[str]: - """Add files to a workspace. Returns list of file_ids that could not be resolved.""" - with self.Session() as session: - # Resolve the workspace's partition to scope the File lookup - workspace = session.execute( - select(Workspace).where(Workspace.workspace_id == workspace_id) - ).scalar_one_or_none() - if workspace is None: - return file_ids - partition = workspace.partition_name - - # Bulk-resolve all file_ids → File.id in a single query - rows = session.execute( - select(File.file_id, File.id).where(File.file_id.in_(file_ids), File.partition_name == partition) - ).all() - id_map = {r[0]: r[1] for r in rows} - missing = [fid for fid in file_ids if fid not in id_map] - - for fid, file_pk in id_map.items(): - stmt = pg_insert(WorkspaceFile).values(workspace_id=workspace_id, file_id=file_pk) - stmt = stmt.on_conflict_do_nothing(constraint="uix_workspace_file") - session.execute(stmt) - session.commit() - return missing - - def remove_file_from_workspace(self, workspace_id: str, file_id: str) -> bool: - """Remove a file from a workspace. Returns True if the association existed, False otherwise.""" - with self.Session() as session: - workspace = session.execute( - select(Workspace).where(Workspace.workspace_id == workspace_id) - ).scalar_one_or_none() - if workspace is None: - return False - file_pk = session.execute( - select(File.id).where(File.file_id == file_id, File.partition_name == workspace.partition_name) - ).scalar_one_or_none() - if file_pk is None: - return False - result = session.execute( - delete(WorkspaceFile).where( - WorkspaceFile.workspace_id == workspace_id, - WorkspaceFile.file_id == file_pk, - ) - ) - session.commit() - return result.rowcount > 0 - - def list_workspace_files(self, workspace_id: str) -> list[str]: - with self.Session() as session: - result = session.execute( - select(File.file_id) - .join(WorkspaceFile, WorkspaceFile.file_id == File.id) - .where(WorkspaceFile.workspace_id == workspace_id) - ) - return [r[0] for r in result.all()] - - def get_file_workspaces(self, file_id: str, partition: str) -> list[str]: - """Return the workspace IDs that contain the given file, scoped to the given partition.""" - with self.Session() as session: - file_pk = session.execute( - select(File.id).where(File.file_id == file_id, File.partition_name == partition) - ).scalar_one_or_none() - if file_pk is None: - return [] - ws_ids = select(Workspace.workspace_id).where(Workspace.partition_name == partition) - result = session.execute( - select(WorkspaceFile.workspace_id).where( - WorkspaceFile.file_id == file_pk, - WorkspaceFile.workspace_id.in_(ws_ids), - ) - ) - return [r[0] for r in result.all()] - - def remove_file_from_all_workspaces(self, file_id: str, partition: str): - """Remove file from all workspaces in the given partition — called during file deletion.""" - with self.Session() as session: - file_pk = session.execute( - select(File.id).where(File.file_id == file_id, File.partition_name == partition) - ).scalar_one_or_none() - if file_pk is None: - return - ws_ids = select(Workspace.workspace_id).where(Workspace.partition_name == partition) - session.execute( - delete(WorkspaceFile).where( - WorkspaceFile.file_id == file_pk, - WorkspaceFile.workspace_id.in_(ws_ids), - ) - ) - session.commit() - - # ------------------------------------------------------------------ - # OIDC — user lookup by sub + optional claim-mapping update - # ------------------------------------------------------------------ - - # Whitelist mirrored in ``api._OIDC_CLAIM_MAPPING_ALLOWED_FIELDS`` and in - # ``routers/auth.py`` — three locations on purpose: the startup validator - # and request-time parser both filter their own inputs, and the DB method - # refuses writes outside this set as a last line of defence against a - # caller bypassing the upstream checks. - _OIDC_WRITABLE_USER_FIELDS = {"display_name", "email"} - - def _user_to_dict(self, user: User) -> dict: - """Serialize a User ORM object to the dict shape used elsewhere.""" - memberships = [ - { - "partition": m.partition_name, - "role": m.role, - "added_at": m.added_at.isoformat(), - } - for m in user.memberships - ] - return { - "id": user.id, - "display_name": user.display_name, - "email": user.email, - "external_user_id": user.external_user_id, - "is_admin": user.is_admin, - "file_quota": user.file_quota, - "file_count": user.file_count, - "memberships": memberships, - } - - def get_user_by_external_id(self, external_user_id: str) -> dict | None: - """Return the user (as dict) whose ``external_user_id`` matches, or None. - - ``external_user_id`` stores the stable OIDC ``sub`` claim in OIDC mode. - """ - with self.Session() as s: - user = s.query(User).filter(User.external_user_id == external_user_id).first() - if not user: - return None - return self._user_to_dict(user) - - def get_user_by_email(self, email: str) -> dict | None: - """Return the user (as dict) whose ``email`` matches (case-insensitive), or None. - - Email is normalized the same way ``create_user`` stores it (stripped, - lowercased). This is used only to explain an auto-provisioning collision - on the unique ``email`` index — never as a login-matching path, which - stays ``external_user_id``-only. - """ - if not isinstance(email, str) or not email.strip(): - return None - with self.Session() as s: - user = s.query(User).filter(User.email == email.strip().lower()).first() - if not user: - return None - return self._user_to_dict(user) - - def update_user_fields(self, user_id: int, fields: dict[str, object]) -> None: - """Update whitelisted scalar fields on the users table. - - Enforces the same whitelist as the OIDC claim-mapping parser - (``display_name``, ``email``) — writing to any other field raises - ``ValueError``. ``None`` values are silently dropped (defensive: the - claim was missing upstream). Empty mapping is a no-op and does not - open a DB session. - """ - if not fields: - return - bad = set(fields) - self._OIDC_WRITABLE_USER_FIELDS - if bad: - raise ValueError(f"Cannot update non-whitelisted user fields: {sorted(bad)}") - cleaned = {k: v for k, v in fields.items() if v is not None} - if not cleaned: - return - # Normalize email to lowercase if present (consistent with create_user). - if "email" in cleaned and isinstance(cleaned["email"], str): - cleaned["email"] = cleaned["email"].strip().lower() - with self.Session() as s: - user = s.query(User).filter_by(id=user_id).first() - if user is None: - raise ValueError(f"User {user_id} not found") - for k, v in cleaned.items(): - setattr(user, k, v) - s.commit() - - # ------------------------------------------------------------------ - # OIDC — sessions - # ------------------------------------------------------------------ - - def _oidc_session_to_dict(self, session_row: OIDCSession) -> dict: - """Serialize an OIDCSession ORM row to a dict. Encrypted blobs are - passed through untouched — the caller (middleware) decrypts them.""" - return { - "id": session_row.id, - "user_id": session_row.user_id, - "sub": session_row.sub, - "sid": session_row.sid, - "id_token_encrypted": session_row.id_token_encrypted, - "access_token_encrypted": session_row.access_token_encrypted, - "refresh_token_encrypted": session_row.refresh_token_encrypted, - "access_token_expires_at": session_row.access_token_expires_at, - "session_expires_at": session_row.session_expires_at, - "created_at": session_row.created_at, - "last_refresh_at": session_row.last_refresh_at, - "revoked_at": session_row.revoked_at, - } - - def create_oidc_session( - self, - *, - user_id: int, - sub: str, - sid: str | None, - session_token_plain: str, - id_token_encrypted: bytes | None, - access_token_encrypted: bytes | None, - refresh_token_encrypted: bytes | None, - access_token_expires_at: datetime, - session_expires_at: datetime, - ) -> dict: - """Insert a new OIDC session row. ``session_token_plain`` is hashed - (SHA-256) before storage — the plaintext is never persisted. - - Returns the row as a dict. Caller is responsible for setting the - ``openrag_session`` cookie with ``session_token_plain``. - """ - session_token_hash = self.hash_token(session_token_plain) - with self.Session() as s: - row = OIDCSession( - session_token_hash=session_token_hash, - user_id=user_id, - sub=sub, - sid=sid, - id_token_encrypted=id_token_encrypted, - access_token_encrypted=access_token_encrypted, - refresh_token_encrypted=refresh_token_encrypted, - access_token_expires_at=access_token_expires_at, - session_expires_at=session_expires_at, - ) - s.add(row) - s.commit() - s.refresh(row) - self.logger.bind(user_id=user_id, sid=sid).info("Created OIDC session") - return self._oidc_session_to_dict(row) - - def get_oidc_session_by_token(self, session_token_plain: str) -> dict | None: - """Look up an OIDC session by its plaintext cookie token. - - Returns None if: - - no row with matching ``session_token_hash`` - - the row is revoked (``revoked_at IS NOT NULL``) - - the session has expired (``session_expires_at < now()``) - """ - session_token_hash = self.hash_token(session_token_plain) - now = datetime.now() - with self.Session() as s: - row = s.query(OIDCSession).filter(OIDCSession.session_token_hash == session_token_hash).first() - if row is None: - return None - if row.revoked_at is not None: - return None - if row.session_expires_at < now: - return None - return self._oidc_session_to_dict(row) - - def get_oidc_session_by_id(self, session_id: int) -> dict | None: - """Look up an OIDC session by primary key. - - Used by the refresh-token stampede guard in - ``components.auth.refresh.refresh_session_if_needed``: when a concurrent - request may have already rotated the tokens, the helper re-reads the row - to see whether it can reuse the fresh tokens instead of calling the IdP - with the (now-invalidated) old refresh_token. - - Returns the row as a dict, or ``None`` if the row does not exist, is - revoked, or the hard session cap has elapsed. - """ - now = datetime.now() - with self.Session() as s: - row = s.query(OIDCSession).filter(OIDCSession.id == session_id).first() - if row is None: - return None - if row.revoked_at is not None: - return None - if row.session_expires_at < now: - return None - return self._oidc_session_to_dict(row) - - def update_oidc_session_tokens( - self, - *, - session_id: int, - access_token_encrypted: bytes, - refresh_token_encrypted: bytes | None, - access_token_expires_at: datetime, - ) -> None: - """Persist refreshed tokens after a successful refresh_token exchange. - - Also bumps ``last_refresh_at`` to ``now()``. Does NOT extend - ``session_expires_at`` — the hard session cap is set at creation time - and is unaffected by access-token rotation. - - The row is locked with ``SELECT ... FOR UPDATE`` so that concurrent - refresh calls on the same session serialize at the DB level (Postgres). - SQLite silently ignores the lock hint, which is fine for tests — the - ``last_refresh_at`` short-circuit in :mod:`components.auth.refresh` - already handles the common stampede case without needing a real lock. - """ - with self.Session() as s: - row = s.query(OIDCSession).filter(OIDCSession.id == session_id).with_for_update().first() - if row is None: - raise ValueError(f"oidc_session id={session_id} does not exist") - row.access_token_encrypted = access_token_encrypted - if refresh_token_encrypted is not None: - row.refresh_token_encrypted = refresh_token_encrypted - row.access_token_expires_at = access_token_expires_at - row.last_refresh_at = datetime.now() - s.commit() - - def revoke_oidc_sessions_by_sid(self, sid: str) -> int: - """Revoke all non-revoked sessions matching the given OIDC ``sid``. - - Used by the OIDC Back-Channel Logout flow — the IdP POSTs a signed - logout_token with a ``sid`` claim; we mark every session with that - sid as revoked. Returns the count affected. - """ - now = datetime.now() - with self.Session() as s: - stmt = ( - update(OIDCSession) - .where(OIDCSession.sid == sid) - .where(OIDCSession.revoked_at.is_(None)) - .values(revoked_at=now) - ) - result = s.execute(stmt) - s.commit() - count = result.rowcount or 0 - self.logger.bind(sid=sid, count=count).info("Revoked OIDC sessions by sid") - return count - - def revoke_oidc_session_by_id(self, session_id: int) -> None: - """Revoke a single session by primary key. Used by RP-initiated logout.""" - now = datetime.now() - with self.Session() as s: - row = s.query(OIDCSession).filter(OIDCSession.id == session_id).first() - if row is None: - return - if row.revoked_at is None: - row.revoked_at = now - s.commit() - self.logger.bind(session_id=session_id).info("Revoked OIDC session") - - def cleanup_expired_oidc_sessions(self) -> int: - """Delete rows whose ``session_expires_at`` is older than 7 days. - - A retention window after expiry helps post-mortem debugging while - keeping the table bounded. Intended to be called by a future cron. - Returns the number of rows deleted. - """ - cutoff = datetime.now() - timedelta(days=7) - with self.Session() as s: - stmt = delete(OIDCSession).where(OIDCSession.session_expires_at < cutoff) - result = s.execute(stmt) - s.commit() - count = result.rowcount or 0 - if count: - self.logger.info(f"Cleaned up {count} expired OIDC sessions") - return count diff --git a/openrag/components/indexer/vectordb/vectordb.py b/openrag/components/indexer/vectordb/vectordb.py deleted file mode 100644 index 3c600ad30..000000000 --- a/openrag/components/indexer/vectordb/vectordb.py +++ /dev/null @@ -1,1551 +0,0 @@ -import asyncio -import time -from abc import ABC, abstractmethod -from datetime import UTC, datetime - -import numpy as np -import ray -from config import load_config -from langchain_core.documents.base import Document -from models.user import UserCreate, UserUpdate -from pymilvus import ( - AnnSearchRequest, - AsyncMilvusClient, - DataType, - Function, - FunctionType, - MilvusClient, - MilvusException, - RRFRanker, -) -from sqlalchemy import URL -from utils.exceptions.base import EmbeddingError -from utils.exceptions.vectordb import * -from utils.logger import get_logger - -from ..embeddings import BaseEmbedding, EmbeddingFactory -from .utils import PartitionFileManager - -logger = get_logger() -config = load_config() - - -class BaseVectorDB(ABC): - """ - Abstract base class for a Vector Database. - This class defines the interface for a vector database connector. - """ - - @abstractmethod - async def list_collections(self): - pass - - @abstractmethod - def collection_exists(self, collection_name: str): - pass - - @abstractmethod - def list_partitions(self): - pass - - @abstractmethod - def partition_exists(self, partition: str) -> bool: - pass - - @abstractmethod - async def delete_partition(self, partition: str): - pass - - @abstractmethod - def list_partition_files(self, partition: str, limit: int | None = None): - pass - - @abstractmethod - async def delete_file(self, file_id: str, partition: str): - pass - - @abstractmethod - async def async_add_documents(self, chunks: list[Document], user: dict): - pass - - @abstractmethod - async def async_search( - self, - query: str, - top_k: int = 5, - similarity_threshold: float = 0.60, - partition: list[str] = None, - filter: str | None = None, - filter_params: dict | None = None, - with_surrounding_chunks: bool = False, - ) -> list[Document]: - pass - - @abstractmethod - async def async_multi_query_search( - self, - partition: list[str], - queries: list[str], - top_k_per_query: int = 5, - similarity_threshold: float = 0.6, - filter: str | None = None, - filter_params: dict | None = None, - with_surrounding_chunks: bool = False, - ) -> list[Document]: - pass - - @abstractmethod - async def list_all_chunk(self, partition: str, include_embedding: bool = True) -> list[Document]: - pass - - @abstractmethod - async def get_file_chunks(self, file_id: str, partition: str, include_id: bool = False, limit: int = 2000): - pass - - @abstractmethod - async def get_chunk_by_id(self, chunk_id: str): - pass - - -SCHEMA_VERSION_PROPERTY_KEY = "openrag.schema_version" -INDEXED_TIME_FIELDS = ["created_at"] - -MAX_LENGTH = 65_535 - -analyzer_params = { - "tokenizer": "standard", - "filter": [ - { - "type": "stop", # Specifies the filter type as stop - "stop_words": [ - "", - "", - "[Image Placeholder]", - "_english_", - "_french_", - "[CHUNK_START]", - "[CHUNK_END]", - "[CONTEXT]", - ], # Defines custom stop words and includes the English and French stop word list - } - ], -} - - -@ray.remote -class MilvusDB(BaseVectorDB): - def __init__(self): - try: - from config import load_config - from utils.logger import get_logger - - self.config = load_config() - self.logger = get_logger() - - # init milvus clients - self.port = self.config.vectordb.port - self.host = self.config.vectordb.host - uri = f"http://{self.host}:{self.port}" - self.uri = uri - try: - self._client = MilvusClient(uri=uri) - self._async_client = AsyncMilvusClient(uri=uri) - except MilvusException as e: - raise VDBConnectionError( - f"Failed to connect to Milvus: {e!s}", - db_url=uri, - db_type="Milvus", - ) - - # embedder - self.embedder: BaseEmbedding = EmbeddingFactory.get_embedder(embeddings_config=self.config.embedder) - - self.hybrid_search = self.config.vectordb.hybrid_search - # partition related params - self.rdb_host = self.config.rdb.host - self.rdb_port = self.config.rdb.port - self.rdb_user = self.config.rdb.user - self.rdb_password = self.config.rdb.password - self.partition_file_manager: PartitionFileManager = None - - # Initialize collection-related attributes - self.collection_name = self.config.vectordb.collection_name - self.collection_loaded = False - self.load_collection() - - except VDBError: - raise - - except Exception as e: - self.logger.exception("Unexpected error initializing Milvus clients", error=str(e)) - raise VDBConnectionError( - f"Unexpected error initializing Milvus clients: {e!s}", - db_url=uri, - db_type="Milvus", - ) - - def load_collection(self): - if not self.collection_loaded: - self.logger = self.logger.bind(collection=self.collection_name, database="Milvus") - try: - if self._client.has_collection(self.collection_name): - self.logger.warning(f"Collection `{self.collection_name}` already exists. Loading it.") - self._check_schema_version() - else: - self.logger.info("Creating empty collection") - index_params = self._create_index() - schema = self._create_schema() - consistency_level = "Strong" - try: - self._client.create_collection( - collection_name=self.collection_name, - schema=schema, - consistency_level=consistency_level, - index_params=index_params, - enable_dynamic_field=True, - ) - except MilvusException as e: - self.logger.exception( - f"Failed to create collection `{self.collection_name}`", - error=str(e), - ) - raise VDBCreateOrLoadCollectionError( - f"Failed to create collection `{self.collection_name}`: {e!s}", - collection_name=self.collection_name, - operation="create_collection", - ) - self._store_schema_version() - try: - self._client.load_collection(self.collection_name) - self.collection_loaded = True - except MilvusException as e: - self.logger.exception( - f"Failed to load collection `{self.collection_name}`", - error=str(e), - ) - raise VDBCreateOrLoadCollectionError( - f"Failed to load existing collection `{self.collection_name}`: {e!s}", - collection_name=self.collection_name, - operation="load_collection", - ) - - database_url = URL.create( - drivername="postgresql", - username=self.rdb_user, - password=self.rdb_password, - host=self.rdb_host, - port=self.rdb_port, - database=f"partitions_for_collection_{self.collection_name}", - ) - self.partition_file_manager = PartitionFileManager( - database_url=database_url.render_as_string(hide_password=False), - logger=self.logger, - ) - self.logger.info("Milvus collection loaded.") - except VDBError: - raise - except Exception as e: - self.logger.exception( - f"Unexpected error setting collection name `{self.collection_name}`", - error=str(e), - ) - raise UnexpectedVDBError( - f"Unexpected error setting collection name `{self.collection_name}`: {e!s}", - collection_name=self.collection_name, - ) - - def _create_schema(self): - self.logger.info("Creating Schema") - schema = self._client.create_schema(enable_dynamic_field=True) - schema.add_field(field_name="_id", datatype=DataType.INT64, is_primary=True, auto_id=True) - schema.add_field( - field_name="text", - datatype=DataType.VARCHAR, - enable_analyzer=True, - enable_match=True, - max_length=MAX_LENGTH, - analyzer_params=analyzer_params, - ) - - schema.add_field( - field_name="partition", - datatype=DataType.VARCHAR, - max_length=MAX_LENGTH, - is_partition_key=True, - ) - - schema.add_field( - field_name="file_id", - datatype=DataType.VARCHAR, - max_length=MAX_LENGTH, - ) - - schema.add_field( - field_name="vector", - datatype=DataType.FLOAT_VECTOR, - dim=self.embedder.embedding_dimension, - ) - - for time_field in INDEXED_TIME_FIELDS: - schema.add_field(field_name=time_field, datatype=DataType.TIMESTAMPTZ, nullable=True) - - if self.hybrid_search: - # Add sparse field for BM25 - this will be auto-generated - schema.add_field( - field_name="sparse", - datatype=DataType.SPARSE_FLOAT_VECTOR, - index_type="SPARSE_INVERTED_INDEX", - ) - - # BM25 function to auto-generate sparse embeddings - bm25_function = Function( - name="text_bm25_emb", - function_type=FunctionType.BM25, - input_field_names=["text"], - output_field_names=["sparse"], - ) - - # Add the function to our schema - schema.add_function(bm25_function) - return schema - - def _create_index(self): - self.logger.info("Creating Index") - index_params = self._client.prepare_index_params() - # Add index for file_id field - index_params.add_index( - field_name="file_id", - index_type="INVERTED", - index_name="file_id_idx", - ) - - # ADD index for partition field - index_params.add_index(field_name="partition", index_type="INVERTED", index_name="partition_idx") - - # Add index for vector field - index_params.add_index( - field_name="vector", - index_type="HNSW", - metric_type="COSINE", - index_params={"M": 128, "efConstruction": 256, "metric_type": "COSINE"}, - ) - - # Add index for sparase field - index_params.add_index( - field_name="sparse", - index_name="sparse_idx", - index_type="SPARSE_INVERTED_INDEX", - index_params={ - "metric_type": "BM25", - "inverted_index_algo": "DAAT_MAXSCORE", - "bm25_k1": 1.2, - "bm25_b": 0.75, - }, - ) - # indexes for dates TIMESTAMPTZ field - for time_field in INDEXED_TIME_FIELDS: - index_params.add_index( - field_name=time_field, - index_type="STL_SORT", # Index for TIMESTAMPTZ - index_name=f"{time_field}_idx", - ) - - return index_params - - def _store_schema_version(self) -> None: - """Persist the configured schema_version as a collection property after collection creation.""" - schema_version = self.config.vectordb.schema_version - self._client.alter_collection_properties( - collection_name=self.collection_name, - properties={SCHEMA_VERSION_PROPERTY_KEY: str(schema_version)}, - ) - self.logger.info(f"Schema version {schema_version} stored on collection `{self.collection_name}`.") - - def _check_schema_version(self) -> None: - """ - Read the stored schema version from collection properties and compare it - against the configured schema_version. Raises VDBSchemaMigrationRequiredError - if they diverge so the application fails fast instead of silently working on a - stale schema. - """ - expected_version = self.config.vectordb.schema_version - desc = self._client.describe_collection(self.collection_name) - props = desc.get("properties", {}) - raw = props.get(SCHEMA_VERSION_PROPERTY_KEY) - - try: - stored_version = int(raw) if raw is not None else 0 - except (ValueError, TypeError): - stored_version = 0 - - if stored_version != expected_version: - raise VDBSchemaMigrationRequiredError( - f"Collection `{self.collection_name}` is at schema version {stored_version} " - f"but the application requires version {expected_version}. " - "Please perform the migration script.", - collection_name=self.collection_name, - stored_version=stored_version, - expected_version=expected_version, - ) - - self.logger.info(f"Collection `{self.collection_name}` schema version {stored_version} — OK.") - - async def list_collections(self) -> list[str]: - return self._client.list_collections() - - async def async_add_documents(self, chunks: list[Document], user: dict) -> None: - """Asynchronously add documents to the vector store.""" - - try: - file_metadata = dict(chunks[0].metadata) - file_metadata.pop("page") - file_id, partition = ( - file_metadata.get("file_id"), - file_metadata.get("partition"), - ) - - # Extract relationship fields (will be stored in both Milvus and PostgreSQL) - relationship_id = file_metadata.get("relationship_id") - parent_id = file_metadata.get("parent_id") - - self.logger.bind( - partition=partition, - file_id=file_id, - filename=file_metadata.get("filename"), - ) - - # check if this file_id exists - res = self.partition_file_manager.file_exists_in_partition(file_id=file_id, partition=partition) - if res: - error_msg = f"This File Id ({file_id}) already exists in Partition ({partition})" - self.logger.error(error_msg) - raise VDBInsertError( - error_msg, - status_code=409, - collection_name=self.collection_name, - partition=partition, - file_id=file_id, - ) - - entities = [] - vectors = await self.embedder.aembed_documents(chunks) - order_metadata_l: list[dict] = _gen_chunk_order_metadata(n=len(chunks)) - indexed_at = datetime.now(UTC).isoformat() - - for chunk, vector, order_metadata in zip(chunks, vectors, order_metadata_l): - entities.append( - { - "text": chunk.page_content, - "vector": vector, - "indexed_at": indexed_at, - **order_metadata, - **chunk.metadata, - } - ) - - await self._async_client.insert( - collection_name=self.collection_name, - data=entities, - ) - - # insert file_id and partition into partition_file_manager - file_metadata.update({"indexed_at": indexed_at}) - self.partition_file_manager.add_file_to_partition( - file_id=file_id, - partition=partition, - file_metadata=file_metadata, - user_id=user.get("id"), - relationship_id=relationship_id, - parent_id=parent_id, - ) - self.logger.info(f"File '{file_id}' added to partition '{partition}'") - except EmbeddingError as e: - self.logger.exception("Embedding failed", error=str(e)) - raise - except VDBError as e: - self.logger.exception("VectorDB operation failed", error=str(e)) - raise - - except Exception as e: - self.logger.exception("Unexpected error while adding a document", error=str(e)) - raise UnexpectedVDBError( - f"Unexpected error while adding a document: {e!s}", - collection_name=self.collection_name, - ) - - async def async_multi_query_search( - self, - partition, - queries, - top_k_per_query=5, - similarity_threshold=0.6, - filter=None, - filter_params=None, - with_surrounding_chunks=False, - ) -> list[Document]: - # Gather all search tasks concurrently - search_tasks = [ - self.async_search( - query=query, - top_k=top_k_per_query, - similarity_threshold=similarity_threshold, - partition=partition, - filter=filter, - filter_params=filter_params, - with_surrounding_chunks=with_surrounding_chunks, - ) - for query in queries - ] - retrieved_results = await asyncio.gather(*search_tasks) - retrieved_chunks = {} - # Process the retrieved documents - for retrieved in retrieved_results: - if retrieved: - for document in retrieved: - retrieved_chunks[document.metadata["_id"]] = document - return list(retrieved_chunks.values()) - - async def async_search( - self, - query: str, - top_k: int = 5, - similarity_threshold: float = 0.60, - partition: list[str] = None, - filter: str | None = None, - filter_params: dict | None = None, - with_surrounding_chunks: bool = False, - ) -> list[Document]: - expr_parts = [] - if partition != ["all"]: - expr_parts.append(f"partition in {partition}") - - if filter: - expr_parts.append(filter) - - if filter_params: - # Don't mutate the caller's dict — concurrent calls may share it - filter_params = dict(filter_params) - if "workspace_id" in filter_params: - workspace_id = filter_params.pop( - "workspace_id" - ) # workspace_id is only used for filtering in the partition_file_manager, not in Milvus directly - ws = self.partition_file_manager.get_workspace(workspace_id) - if not ws: - return [] # Workspace not found → no results - - file_ids = self.partition_file_manager.list_workspace_files(workspace_id) - if not file_ids: - return [] # Empty workspace → no results - - # Pin to the workspace's own partition regardless of the requested - # partition set — file_id is only unique per (file_id, partition_name) - # so a cross-partition search could otherwise return chunks from a - # different partition that reuses the same file_id. - ws_partition = ws["partition_name"] - # Replace any outer partition filter with the workspace's partition - expr_parts = [p for p in expr_parts if not p.startswith("partition in ")] - expr_parts.append(f'partition == "{ws_partition}"') - - id_list = ", ".join(f'"{fid}"' for fid in file_ids) - expr_parts.append(f"file_id IN [{id_list}]") - - # Join all parts with " and " only if there are multiple conditions - expr = " and ".join(expr_parts) if expr_parts else "" - - try: - query_vector = await self.embedder.aembed_query(query) - vector_param = { - "data": [query_vector], - "anns_field": "vector", - "param": { - "metric_type": "COSINE", - "params": { - "ef": 64, - "radius": similarity_threshold, - "range_filter": 1.0, - }, - }, - "limit": top_k, - "expr": expr, - } - if self.hybrid_search: - sparse_param = { - "data": [query], - "anns_field": "sparse", - "param": { - "metric_type": "BM25", - "params": {"drop_ratio_build": 0.2}, - }, - "limit": top_k, - "expr": expr, - } - reqs = [ - AnnSearchRequest(**vector_param), - AnnSearchRequest(**sparse_param), - ] - response = await self._async_client.hybrid_search( - collection_name=self.collection_name, - reqs=reqs, - ranker=RRFRanker(100), - output_fields=["*"], - limit=top_k, - ) - else: - vector_param = { - "data": [query_vector], - "anns_field": "vector", - "search_params": { - "metric_type": "COSINE", - "params": { - "ef": 64, - "radius": similarity_threshold, - "range_filter": 1.0, - }, - }, - "limit": top_k, - } - response = await self._async_client.search( - collection_name=self.collection_name, - output_fields=["*"], - filter=expr, - **vector_param, - ) - - docs = _parse_documents_from_search_results(response) - if with_surrounding_chunks: - self.logger.debug("Fetching surrounding chunks") - surrounding_chunks = await self.get_surrounding_chunks(docs) - self.logger.debug("Fetched surrounding chunks", count=len(surrounding_chunks)) - docs.extend(surrounding_chunks) - - return docs - - except MilvusException as e: - self.logger.exception("Search failed in Milvus", error=str(e)) - raise VDBSearchError( - f"Search failed in Milvus: {e!s}", - collection_name=self.collection_name, - partition=partition, - ) - except EmbeddingError as e: - self.logger.exception("Embedding failed while processing the query", error=str(e)) - raise - - except Exception as e: - self.logger.exception("Unexpected error occurred", error=str(e)) - raise UnexpectedVDBError( - f"Unexpected error occurred: {e!s}", - collection_name=self.collection_name, - partition=partition, - ) - - async def get_surrounding_chunks(self, docs: list[Document]) -> list[Document]: - existant_ids = {doc.metadata.get("_id") for doc in docs} - - # Collect all prev/next section IDs - section_ids = [ - section_id - for doc in docs - for section_id in [ - doc.metadata.get("prev_section_id"), - doc.metadata.get("next_section_id"), - ] - if section_id is not None - ] - - if not section_ids: - return [] - - # Query all sections in parallel - tasks = [ - self._async_client.query( - collection_name=self.collection_name, - filter=f"section_id == {section_id}", - limit=1, - ) - for section_id in section_ids - ] - responses = await asyncio.gather(*tasks) - - # Build output, skipping duplicates - output_docs = [] - for response in responses: - if not response: - continue - doc_id = response[0].get("_id") - if doc_id not in existant_ids: - existant_ids.add(doc_id) - output_docs.append( - Document( - page_content=response[0]["text"], - metadata={key: value for key, value in response[0].items() if key not in ["text", "vector"]}, - ) - ) - - return output_docs - - async def delete_file(self, file_id: str, partition: str): - log = self.logger.bind(file_id=file_id, partition=partition) - try: - res = await self._async_client.delete( - collection_name=self.collection_name, - filter=f"partition == '{partition}' and file_id == '{file_id}'", - ) - - self.partition_file_manager.remove_file_from_all_workspaces(file_id, partition) - self.partition_file_manager.remove_file_from_partition(file_id=file_id, partition=partition) - log.info("Deleted file chunks from partition.", count=res.get("delete_count", 0)) - - except MilvusException as e: - log.exception(f"Couldn't delete file chunks for file_id {file_id}", error=str(e)) - raise VDBDeleteError( - f"Couldn't delete file chunks for file_id {file_id}: {e!s}", - collection_name=self.collection_name, - partition=partition, - file_id=file_id, - ) - except VDBError: - raise - except Exception as e: - log.exception("Unexpected error while deleting file chunks", error=str(e)) - raise UnexpectedVDBError( - f"Unexpected error while deleting file chunks {file_id}: {e!s}", - collection_name=self.collection_name, - partition=partition, - file_id=file_id, - ) - - async def delete_chunks_by_ids(self, chunk_ids: list[int]): - """Delete specific Milvus chunks by their _id primary keys.""" - if not chunk_ids: - return - try: - await self._async_client.delete( - collection_name=self.collection_name, - ids=chunk_ids, - ) - self.logger.info("Deleted old chunks by ID.", count=len(chunk_ids)) - except MilvusException as e: - self.logger.exception("Failed to delete old chunks by ID", error=str(e)) - raise VDBDeleteError( - f"Failed to delete old chunks by ID: {e!s}", - collection_name=self.collection_name, - ) - except Exception as e: - self.logger.exception("Unexpected error while deleting chunks by ID", error=str(e)) - raise UnexpectedVDBError( - f"Unexpected error while deleting chunks by ID: {e!s}", - collection_name=self.collection_name, - ) - - async def get_file_chunk_ids(self, file_id: str, partition: str) -> list[int]: - """Return the Milvus _id values for all chunks of a file.""" - log = self.logger.bind(file_id=file_id, partition=partition) - try: - results = [] - offset = 0 - limit = 100 - while True: - response = await self._async_client.query( - collection_name=self.collection_name, - filter="partition == {partition} and file_id == {file_id}", - filter_params={"partition": partition, "file_id": file_id}, - output_fields=["_id"], - limit=limit, - offset=offset, - ) - if not response: - break - results.extend(r["_id"] for r in response) - offset += len(response) - return results - except MilvusException as e: - log.exception("Failed to get file chunk IDs", error=str(e)) - raise VDBSearchError( - f"Failed to get file chunk IDs for {file_id}: {e!s}", - collection_name=self.collection_name, - partition=partition, - file_id=file_id, - ) - except Exception as e: - log.exception("Unexpected error while getting file chunk IDs", error=str(e)) - raise UnexpectedVDBError( - f"Unexpected error while getting file chunk IDs for {file_id}: {e!s}", - collection_name=self.collection_name, - partition=partition, - file_id=file_id, - ) - - async def upsert_file_metadata(self, file_id: str, partition: str, metadata: dict): - """Update metadata on all chunks of a file in-place via Milvus upsert. - - Fetches existing chunks (with _id and vectors), merges new metadata, - then upserts back into Milvus. No re-embedding is performed. - Also updates the PostgreSQL file record metadata in-place. - """ - log = self.logger.bind(file_id=file_id, partition=partition) - try: - # Fetch all chunks with their _id and vector so we can upsert without re-embedding. - docs = await self.get_file_chunks(file_id, partition, include_id=True, include_vectors=True) - if not docs: - log.warning("No chunks found for metadata upsert") - return - - entities = [] - for doc in docs: - chunk_metadata = dict(doc.metadata) - # Merge new metadata into the chunk metadata. - # _id and vector are already in chunk_metadata (via include_id/include_vectors). - chunk_metadata.update(metadata) - entities.append( - { - "text": doc.page_content, - **chunk_metadata, - } - ) - - await self._async_client.upsert( - collection_name=self.collection_name, - data=entities, - ) - - # Build file-level metadata from the first chunk (same as async_add_documents). - # Strip per-chunk fields that don't belong in the file-level PG record. - file_metadata = dict(docs[0].metadata) - for key in ("_id", "vector", "page", "section_id", "prev_section_id", "next_section_id"): - file_metadata.pop(key, None) - file_metadata.update(metadata) - if not self.partition_file_manager.update_file_metadata_in_db(file_id, partition, file_metadata): - # PG row was concurrently deleted; Milvus upsert already succeeded. - # Log warning but don't fail — Milvus data will be orphaned until - # next cleanup, but the user-facing operation should still succeed. - log.warning("PG file row not found during metadata upsert; Milvus updated but PG skipped") - - log.info("Upserted file metadata in-place.", chunk_count=len(entities)) - - except MilvusException as e: - log.exception("Milvus upsert failed", error=str(e)) - raise VDBInsertError( - f"Couldn't upsert metadata for file {file_id}: {e!s}", - collection_name=self.collection_name, - partition=partition, - file_id=file_id, - ) - except VDBError: - raise - except Exception as e: - log.exception("Unexpected error during metadata upsert", error=str(e)) - raise UnexpectedVDBError( - f"Unexpected error during metadata upsert for {file_id}: {e!s}", - collection_name=self.collection_name, - ) - - async def add_documents_for_existing_file(self, chunks: list[Document], user: dict) -> None: - """Replace Milvus chunks for a file that already exists in PostgreSQL. - - Used by PUT (file replace). The flow is insert-before-delete so the file - is never left in a half-replaced state: - 1. Snapshot old chunk _id values - 2. Embed and insert new chunks - 3. Delete old chunks by _id - 4. Update the PostgreSQL File row in-place - - If step 2 fails, old chunks remain intact. If step 3 fails, we have - duplicates temporarily but no data loss — a retry or manual cleanup - can resolve it. - - Note: this implements strict PUT semantics — the new chunk metadata - fully replaces the old. Fields like ``relationship_id`` and ``parent_id`` - are taken from the new chunks' metadata; if the caller omits them, the - PG columns are cleared. To preserve old values across a PUT, the caller - must re-supply them in the request metadata. - """ - log = self.logger # Fallback; rebound with context below - try: - file_metadata = dict(chunks[0].metadata) - file_metadata.pop("page") - file_id, partition = file_metadata.get("file_id"), file_metadata.get("partition") - relationship_id = file_metadata.get("relationship_id") - parent_id = file_metadata.get("parent_id") - - log = self.logger.bind(partition=partition, file_id=file_id, filename=file_metadata.get("filename")) - - # 1. Snapshot old chunk _id values before inserting new ones. - old_chunk_ids = await self.get_file_chunk_ids(file_id, partition) - - # 2. Embed and insert new chunks. - entities = [] - vectors = await self.embedder.aembed_documents(chunks) - order_metadata_l: list[dict] = _gen_chunk_order_metadata(n=len(chunks)) - for chunk, vector, order_metadata in zip(chunks, vectors, order_metadata_l): - entities.append( - { - "text": chunk.page_content, - "vector": vector, - **order_metadata, - **chunk.metadata, - } - ) - - await self._async_client.insert( - collection_name=self.collection_name, - data=entities, - ) - - # 3. Delete old chunks by _id (new ones are already durable). - await self.delete_chunks_by_ids(old_chunk_ids) - - # 4. Update existing PostgreSQL file record in-place (preserves files.id PK) - if not self.partition_file_manager.update_file_in_partition( - file_id=file_id, - partition=partition, - file_metadata=file_metadata, - relationship_id=relationship_id, - parent_id=parent_id, - ): - # PG row was concurrently deleted after we inserted new Milvus chunks. - # Log warning — Milvus has new orphaned chunks but data is consistent - # (old chunks deleted, new chunks inserted, no PG record). - log.warning("PG file row not found during replace; Milvus updated but PG skipped") - log.info(f"File '{file_id}' chunks replaced in partition '{partition}'") - - except EmbeddingError as e: - log.exception("Embedding failed", error=str(e)) - raise - except VDBError as e: - log.exception("VectorDB operation failed", error=str(e)) - raise - except Exception as e: - log.exception("Unexpected error while adding chunks for existing file", error=str(e)) - raise UnexpectedVDBError( - f"Unexpected error while adding chunks for existing file: {e!s}", - collection_name=self.collection_name, - ) - - async def get_file_chunks( - self, - file_id: str, - partition: str, - include_id: bool = False, - include_vectors: bool = False, - limit: int = 2000, - ): - log = self.logger.bind(file_id=file_id, partition=partition) - try: - self._check_file_exists(file_id, partition) - filter_expr = f'partition == "{partition}" and file_id == "{file_id}"' - excluded_keys = {"text"} - if not include_id: - excluded_keys.add("_id") - if not include_vectors: - excluded_keys.add("vector") - - # Milvus query with output_fields=["*"] returns all scalar fields - # but excludes vector fields. To include vectors, request them explicitly. - output_fields = ["*", "vector"] if include_vectors else ["*"] - - results = [] - iterator = self._client.query_iterator( - collection_name=self.collection_name, - filter=filter_expr, - limit=limit, - batch_size=min(limit, 16000), - output_fields=output_fields, - ) - try: - while True: - batch = iterator.next() - if not batch: - break - results.extend(batch) - finally: - iterator.close() - - docs = [ - Document( - page_content=res["text"], - metadata={key: value for key, value in res.items() if key not in excluded_keys}, - ) - for res in results - ] - log.info("Fetched file chunks.", count=len(results)) - return docs - - except MilvusException as e: - log.exception(f"Couldn't get file chunks for file_id {file_id}", error=str(e)) - raise VDBSearchError( - f"Couldn't get file chunks for file_id {file_id}: {e!s}", - collection_name=self.collection_name, - partition=partition, - file_id=file_id, - ) - except VDBError: - raise - - except Exception as e: - log.exception("Unexpected error while getting file chunks", error=str(e)) - raise VDBSearchError( - f"Unexpected error while getting file chunks {file_id}: {e!s}", - collection_name=self.collection_name, - partition=partition, - file_id=file_id, - ) - - async def get_chunk_by_id(self, chunk_id: str): - """ - Retrieve a chunk by its ID. - Args: - chunk_id (str): The ID of the chunk to retrieve (Milvus Int64 _id as string). - Returns: - Document: The retrieved chunk, or None if not found or invalid ID format. - """ - log = self.logger.bind(chunk_id=chunk_id) - # Milvus _id is Int64, so we need to convert the string to int - try: - chunk_id_int = int(chunk_id) - except (ValueError, TypeError): - log.warning("Invalid chunk_id format - must be an integer") - return None - - try: - response = await self._async_client.query( - collection_name=self.collection_name, - filter=f"_id == {chunk_id_int}", - limit=1, - ) - if response: - return Document( - page_content=response[0]["text"], - metadata={key: value for key, value in response[0].items() if key not in ["text", "vector"]}, - ) - return None - except MilvusException as e: - log.exception("Milvus query failed", error=str(e)) - raise VDBSearchError( - f"Milvus query failed: {e!s}", - collection_name=self.collection_name, - ) - - except Exception as e: - log.exception("Unexpected error while retrieving chunk", error=str(e)) - raise UnexpectedVDBError( - f"Unexpected error while retrieving chunk {chunk_id}: {e!s}", - collection_name=self.collection_name, - ) - - def file_exists(self, file_id: str, partition: str): - """ - Check if a file exists in Milvus - """ - try: - return self.partition_file_manager.file_exists_in_partition(file_id=file_id, partition=partition) - except Exception as e: - self.logger.exception( - "File existence check failed.", - file_id=file_id, - partition=partition, - error=str(e), - ) - return False - - def list_partition_files(self, partition: str, limit: int | None = None): - try: - self._check_partition_exists(partition) - return self.partition_file_manager.list_partition_files(partition=partition, limit=limit) - - except VDBError: - raise - - except Exception as e: - self.logger.exception( - f"Unexpected error while listing files in partition {partition}", - error=str(e), - ) - raise UnexpectedVDBError( - f"Unexpected error while listing files in partition {partition}: {e!s}", - collection_name=self.collection_name, - partition=partition, - ) - - def list_partitions(self): - try: - return self.partition_file_manager.list_partitions() - except Exception as e: - self.logger.exception("Failed to list partitions", error=str(e)) - raise - - def collection_exists(self, collection_name: str): - """ - Check if a collection exists in Milvus - """ - return self._client.has_collection(collection_name=collection_name) - - async def delete_partition(self, partition: str): - self._check_partition_exists(partition) - log = self.logger.bind(partition=partition) - - try: - count = self._client.delete( - collection_name=self.collection_name, - filter=f"partition == '{partition}'", - ) - - self.partition_file_manager.delete_partition(partition) - log.info("Deleted points from partition", count=count.get("delete_count")) - - except MilvusException as e: - log.exception("Failed to delete partition", error=str(e)) - raise VDBDeleteError( - f"Failed to delete partition `{partition}`: {e!s}", - collection_name=self.collection_name, - partition=partition, - ) - except VDBError as e: - log.exception("VectorDB operation failed", error=str(e)) - raise e - except Exception as e: - log.exception("Unexpected error while deleting partition", error=str(e)) - raise UnexpectedVDBError( - f"Unexpected error while deleting partition {partition}: {e!s}", - collection_name=self.collection_name, - partition=partition, - ) - - def partition_exists(self, partition: str): - """ - Check if a partition exists in Milvus - """ - log = self.logger.bind(partition=partition) - try: - return self.partition_file_manager.partition_exists(partition=partition) - except Exception as e: - log.exception("Partition existence check failed.", error=str(e)) - return False - - async def list_all_chunk(self, partition: str, include_embedding: bool = True): - """ - List all chunk from a given partition. - """ - try: - self._check_partition_exists(partition) - - # Create a filter expression for the query - filter_expression = "partition == {partition}" - expr_params = {"partition": partition} - - excluded_keys = ["text"] - if not include_embedding: - excluded_keys.append("vector") - - def prepare_metadata(res: dict): - metadata = {} - for k, v in res.items(): - if k not in excluded_keys: - if k == "vector": - v = str(np.array(v).flatten().tolist()) - metadata[k] = v - return metadata - - chunks = [] - iterator = self._client.query_iterator( - collection_name=self.collection_name, - filter=filter_expression, - expr_params=expr_params, - batch_size=16000, - output_fields=["*"], - ) - - try: - while True: - result = iterator.next() - if not result: - break - chunks.extend( - [ - Document( - page_content=res["text"], - metadata=prepare_metadata(res), - ) - for res in result - ] - ) - finally: - iterator.close() - - return chunks - - except MilvusException as e: - self.logger.exception("Milvus query failed", error=str(e)) - raise VDBSearchError( - f"Milvus query failed: {e!s}", - collection_name=self.collection_name, - partition=partition, - ) - except VDBError: - raise - - except Exception as e: - self.logger.exception( - f"Unexpected error while listing all chunks in partition {partition}", - error=str(e), - ) - raise UnexpectedVDBError( - f"Unexpected error while listing all chunks in partition {partition}: {e!s}", - collection_name=self.collection_name, - partition=partition, - ) - - async def create_user(self, body: UserCreate): - return self.partition_file_manager.create_user(body) - - async def get_user(self, user_id: int): - self._check_user_exists(user_id) - return self.partition_file_manager.get_user_by_id(user_id) - - async def delete_user(self, user_id: int): - self._check_user_exists(user_id) - user_partitions = [ - p["partition"] for p in self.partition_file_manager.list_user_partitions(user_id) if p["role"] == "owner" - ] - for partition in user_partitions: - await self.delete_partition(partition) - self.partition_file_manager.delete_user(user_id) - - async def list_users(self): - return self.partition_file_manager.list_users() - - async def get_user_by_token(self, token: str): - return self.partition_file_manager.get_user_by_token(token) - - async def regenerate_user_token(self, user_id: int): - self._check_user_exists(user_id) - return self.partition_file_manager.regenerate_user_token(user_id) - - async def update_user(self, user_id: int, body: UserUpdate): - self._check_user_exists(user_id) - return self.partition_file_manager.update_user(user_id, body) - - async def list_user_partitions(self, user_id: int): - self._check_user_exists(user_id) - return self.partition_file_manager.list_user_partitions(user_id) - - # ------------------------------------------------------------------ - # OIDC — exposed on the Ray actor (thin delegations) - # ------------------------------------------------------------------ - - async def get_user_by_external_id(self, external_user_id: str): - return self.partition_file_manager.get_user_by_external_id(external_user_id) - - async def update_user_fields(self, user_id: int, fields: dict): - self._check_user_exists(user_id) - return self.partition_file_manager.update_user_fields(user_id, fields) - - async def create_oidc_session( - self, - *, - user_id: int, - sub: str, - sid: str | None, - session_token_plain: str, - id_token_encrypted: bytes | None, - access_token_encrypted: bytes | None, - refresh_token_encrypted: bytes | None, - access_token_expires_at, - session_expires_at, - ): - self._check_user_exists(user_id) - return self.partition_file_manager.create_oidc_session( - user_id=user_id, - sub=sub, - sid=sid, - session_token_plain=session_token_plain, - id_token_encrypted=id_token_encrypted, - access_token_encrypted=access_token_encrypted, - refresh_token_encrypted=refresh_token_encrypted, - access_token_expires_at=access_token_expires_at, - session_expires_at=session_expires_at, - ) - - async def get_oidc_session_by_token(self, session_token_plain: str): - return self.partition_file_manager.get_oidc_session_by_token(session_token_plain) - - async def get_oidc_session_by_id(self, session_id: int): - return self.partition_file_manager.get_oidc_session_by_id(session_id) - - async def update_oidc_session_tokens( - self, - *, - session_id: int, - access_token_encrypted: bytes, - refresh_token_encrypted: bytes | None, - access_token_expires_at, - ): - return self.partition_file_manager.update_oidc_session_tokens( - session_id=session_id, - access_token_encrypted=access_token_encrypted, - refresh_token_encrypted=refresh_token_encrypted, - access_token_expires_at=access_token_expires_at, - ) - - async def revoke_oidc_sessions_by_sid(self, sid: str) -> int: - return self.partition_file_manager.revoke_oidc_sessions_by_sid(sid) - - async def revoke_oidc_session_by_id(self, session_id: int) -> None: - return self.partition_file_manager.revoke_oidc_session_by_id(session_id) - - async def cleanup_expired_oidc_sessions(self) -> int: - return self.partition_file_manager.cleanup_expired_oidc_sessions() - - async def list_partition_members(self, partition: str) -> list[dict]: - self._check_partition_exists(partition) - return self.partition_file_manager.list_partition_members(partition) - - async def update_partition_member_role(self, partition: str, user_id: int, new_role: str): - self._check_membership_exists(partition, user_id) - self.partition_file_manager.update_partition_member_role(partition, user_id, new_role) - self.logger.info(f"User_id {user_id} role updated to '{new_role}' in partition '{partition}'.") - - async def create_partition(self, partition: str, user_id: int): - self._check_user_exists(user_id) - self.partition_file_manager.create_partition(partition, user_id) - self.logger.info(f"Partition '{partition}' created by user_id {user_id}.") - - async def add_partition_member(self, partition: str, user_id: int, role: str): - self._check_partition_exists(partition) - self._check_user_exists(user_id) - self.partition_file_manager.add_partition_member(partition, user_id, role) - self.logger.info(f"User_id {user_id} added to partition '{partition}'.") - - async def remove_partition_member(self, partition: str, user_id: int) -> bool: - self._check_membership_exists(partition, user_id) - self.partition_file_manager.remove_partition_member(partition, user_id) - self.logger.info(f"User_id {user_id} removed from partition '{partition}'.") - - def _check_user_exists(self, user_id: int): - if not self.partition_file_manager.user_exists(user_id): - self.logger.warning(f"User with ID {user_id} does not exist.") - raise VDBUserNotFound( - f"User with ID {user_id} does not exist.", - collection_name=self.collection_name, - user_id=user_id, - ) - - def _check_partition_exists(self, partition: str): - if not self.partition_file_manager.partition_exists(partition): - self.logger.warning(f"Partition '{partition}' does not exist.") - raise VDBPartitionNotFound( - f"Partition '{partition}' does not exist.", - collection_name=self.collection_name, - partition=partition, - ) - - def _check_membership_exists(self, partition: str, user_id: int): - self._check_partition_exists(partition) - self._check_user_exists(user_id) - if not self.partition_file_manager.user_is_partition_member(user_id, partition): - raise VDBMembershipNotFound( - f"User with ID {user_id} is not a member of partition '{partition}'.", - collection_name=self.collection_name, - user_id=user_id, - partition=partition, - ) - - def _check_file_exists(self, file_id, partition: str): - if not self.partition_file_manager.file_exists_in_partition(file_id=file_id, partition=partition): - raise VDBFileNotFoundError( - f"File ID '{file_id}' does not exist in partition '{partition}'", - collection_name=self.collection_name, - partition=partition, - file_id=file_id, - ) - - # Document relationship methods - - def get_files_by_relationship(self, partition: str, relationship_id: str) -> list[dict]: - """Get all files sharing a relationship_id within a partition. - - Args: - partition: Partition name - relationship_id: The relationship group identifier - - Returns: - List of file dictionaries - """ - return self.partition_file_manager.get_files_by_relationship( - partition=partition, relationship_id=relationship_id - ) - - def get_file_ancestors(self, partition: str, file_id: str, max_ancestor_depth: int | None = None) -> list[dict]: - """Get all ancestors of a file (direct path from root to file). - - Args: - partition: Partition name - file_id: The file identifier - - Returns: - List of file dictionaries ordered from root to the specified file - """ - return self.partition_file_manager.get_file_ancestors( - partition=partition, file_id=file_id, max_ancestor_depth=max_ancestor_depth - ) - - async def get_related_chunks(self, partition: str, relationship_id: str, limit: int = 100) -> list[Document]: - """Get all chunks for files in a relationship group. - - Args: - partition: Partition name - relationship_id: The relationship group identifier - limit: Maximum number of chunks to return - - Returns: - List of Document objects - """ - file_ids = self.partition_file_manager.get_file_ids_by_relationship( - partition=partition, relationship_id=relationship_id - ) - - if not file_ids: - return [] - - # Build filter expression for Milvus query - file_id_list = ", ".join(f'"{fid}"' for fid in file_ids) - filter_expr = f'partition == "{partition}" and file_id in [{file_id_list}]' - - results = await self._async_client.query( - collection_name=self.collection_name, - filter=filter_expr, - limit=limit, - output_fields=["*"], - ) - - return [ - Document( - page_content=res["text"], - metadata={k: v for k, v in res.items() if k not in ["text", "vector"]}, - ) - for res in results - ] - - async def get_ancestor_chunks( - self, partition: str, file_id: str, limit: int = 100, max_ancestor_depth: int | None = None - ) -> list[Document]: - """Get all chunks for ancestor files (direct path from root to file). - - Args: - partition: Partition name - file_id: The file identifier - limit: Maximum number of chunks to return - - Returns: - List of Document objects ordered by ancestry - """ - ancestor_file_ids = self.partition_file_manager.get_ancestor_file_ids( - partition=partition, file_id=file_id, max_ancestor_depth=max_ancestor_depth - ) - - if not ancestor_file_ids: - return [] - - # Build filter expression for Milvus query - file_id_list = ", ".join(f'"{fid}"' for fid in ancestor_file_ids) - filter_expr = f'partition == "{partition}" and file_id in [{file_id_list}]' - - results = await self._async_client.query( - collection_name=self.collection_name, - filter=filter_expr, - limit=limit, - output_fields=["*"], - ) - - return [ - Document( - page_content=res["text"], - metadata={k: v for k, v in res.items() if k not in ["text", "vector"]}, - ) - for res in results - ] - - # --- Workspace methods --- - - async def create_workspace( - self, workspace_id: str, partition: str, user_id: int | None = None, display_name: str | None = None - ): - self.partition_file_manager.create_workspace(workspace_id, partition, user_id, display_name) - - async def list_workspaces(self, partition: str) -> list[dict]: - return self.partition_file_manager.list_workspaces(partition) - - async def get_workspace(self, workspace_id: str) -> dict | None: - return self.partition_file_manager.get_workspace(workspace_id) - - async def delete_workspace(self, workspace_id: str) -> list[str]: - """Delete workspace and return orphaned file_ids. Caller must delete those files from Milvus.""" - return self.partition_file_manager.delete_workspace(workspace_id) - - async def get_existing_file_ids(self, partition: str, file_ids: list[str]) -> list[str]: - """Return the subset of file_ids that exist in the given partition.""" - return list(self.partition_file_manager.get_existing_file_ids(partition, file_ids)) - - async def add_files_to_workspace(self, workspace_id: str, file_ids: list[str]) -> list[str]: - return self.partition_file_manager.add_files_to_workspace(workspace_id, file_ids) - - async def remove_file_from_workspace(self, workspace_id: str, file_id: str) -> bool: - return self.partition_file_manager.remove_file_from_workspace(workspace_id, file_id) - - async def list_workspace_files(self, workspace_id: str) -> list[str]: - return self.partition_file_manager.list_workspace_files(workspace_id) - - async def get_file_workspaces(self, file_id: str, partition: str) -> list[str]: - """Return workspace IDs that contain the given file, scoped to the partition.""" - return self.partition_file_manager.get_file_workspaces(file_id, partition) - - -def _gen_chunk_order_metadata(n: int = 20) -> list[dict]: - # Use base timestamp + index to ensure uniqueness - base_ts = int(time.time_ns()) - ids: list[int] = [base_ts + i for i in range(n)] - L = [] - for i in range(n): - prev_chunk_id = ids[i - 1] if i > 0 else None - next_chunk_id = ids[i + 1] if i < n - 1 else None - L.append( - { - "prev_section_id": prev_chunk_id, - "section_id": ids[i], - "next_section_id": next_chunk_id, - } - ) - return L - - -def _parse_documents_from_search_results(search_results) -> list[Document]: - if not search_results: - return [] - - ret = [] - excluded_keys = ["text", "vector"] - for result in search_results[0]: - entity = result.get("entity", {}) - metadata = {k: v for k, v in entity.items() if k not in excluded_keys} - doc = Document( - page_content=entity["text"], - metadata=metadata, - ) - ret.append(doc) - - return ret - - -class ConnectorFactory: - CONNECTORS: dict[BaseVectorDB] = { - "milvus": MilvusDB, - # "qdrant": QdrantDB, - } - - @staticmethod - def get_vectordb_cls(): - name = config.vectordb.connector_name - vdb_cls = ConnectorFactory.CONNECTORS.get(name) - if not vdb_cls: - raise ValueError(f"VECTORDB '{name}' is not supported.") - return vdb_cls diff --git a/openrag/components/llm.py b/openrag/components/llm.py index bfaf4ed76..e4151879f 100644 --- a/openrag/components/llm.py +++ b/openrag/components/llm.py @@ -1,13 +1,59 @@ +"""Backward-compatibility shim — delegates to services.inference.vllm_client. + +All new code should import directly from ``services.inference.vllm_client``. +""" + import copy import json +import warnings import httpx +from config.models import LLMConfig +from services.inference.vllm_client import VLLMClient # noqa: F401 from utils.logger import get_logger logger = get_logger() +class _LLMShim: + """Legacy shim — delegates to ``VLLMClient`` for retry, circuit breaker, + and connection pooling while preserving the generator-based interface.""" + + def __init__(self, llm_config: LLMConfig, logger=None): + warnings.warn( + "components.llm.LLM is deprecated — use services.inference.vllm_client.VLLMClient", + DeprecationWarning, + stacklevel=2, + ) + self.logger = logger + config_kwargs = {k: v for k, v in llm_config.model_dump().items() if k not in ("api_key", "base_url", "model")} + self._delegate = VLLMClient( + endpoint=llm_config.base_url, + model_name=llm_config.model, + api_key=llm_config.api_key, + **config_kwargs, + ) + + async def completions(self, request: dict): + prompt = request.pop("prompt") + response = await self._delegate.generate(prompt, **request) + yield response + + async def chat_completion(self, request: dict): + messages = request.pop("messages") + stream = request.pop("stream", False) + + if stream: + async for line in self._delegate.stream_chat(messages, **request): + yield line + else: + resp_dict = await self._delegate.chat(messages, **request) + yield resp_dict + + class LLM: + """Legacy LLM wrapper. New code should use VLLMClient (via DI) instead.""" + def __init__(self, llm_config, logger=None): self.logger = logger default_llm_config = llm_config.model_dump() @@ -21,7 +67,6 @@ def __init__(self, llm_config, logger=None): } def _extract_llm_overrides(self, request: dict): - """Extract and apply LLM overrides from metadata.llm_override.""" metadata = request.get("metadata") or {} llm_override = metadata.pop("llm_override", None) or {} @@ -87,7 +132,7 @@ async def chat_completion(self, request: dict): logger.error(f"Error while streaming chat completion: {str(e)}") raise - else: # Handle non-streaming response + else: try: response = await client.post( url=f"{base_url}/chat/completions", diff --git a/openrag/components/map_reduce.py b/openrag/components/map_reduce.py deleted file mode 100644 index af2d4ce47..000000000 --- a/openrag/components/map_reduce.py +++ /dev/null @@ -1,157 +0,0 @@ -from pathlib import Path - -from config import load_config -from langchain_core.documents.base import Document -from langchain_openai import ChatOpenAI -from pydantic import BaseModel, Field -from tqdm.asyncio import tqdm -from utils.logger import get_logger - -from .utils import get_llm_semaphore - -logger = get_logger() -config = load_config() - -LOG_DIR = Path(config.paths.log_dir) - - -system_prompt_map = """You are an AI assistant specialized in extracting and synthesizing relevant information from text. - -Your task: -1. Analyze the provided text in relation to the user's question -2. Extract only the essential information that directly addresses the query -3. Preserve necessary context (Key words, project names or initiatives, dates, etc.) to maintain accuracy and clarity of the summary for it to be self-understandable - -Guidelines: -- Present information clearly and concisely without unnecessary rephrasing or commentary -- Focus on precision: include what matters, exclude what doesn't. -- If a document does have any relevant content with respect to a query, classify it irrelevant such without providing a `synthesis`. -""" - - -class SummarizedChunk(BaseModel): - relevancy: bool = Field(..., description="Indicates if the chunk is relevant to the query") - summary: str = Field( - "", - description="The summarized content of the chunk. The field should be empty if relevancy is False.", - ) - - -user_prompt = """ -Here is a text: -{content} - -From this document, identify and comprehensively summarize the information useful for answering the following question: -{query} -""" - - -class RAGMapReduce: - def __init__(self, config): - self.config = config - self.slm: ChatOpenAI = ChatOpenAI(**config.llm.model_dump()).with_structured_output(SummarizedChunk) - map_reduce_config = self.config.map_reduce - self.initial_batch_size = map_reduce_config.initial_batch_size - self.expansion_batch_size = map_reduce_config.expansion_batch_size - self.max_total_documents = map_reduce_config.max_total_documents - - self.debug = map_reduce_config.debug - - assert self.max_total_documents >= self.initial_batch_size, ( - "`max_total_documents` must be greater than or equal to `initial_batch_size`" - ) - - async def infer_chunk_relevancy(self, query, chunk: Document) -> SummarizedChunk: - async with get_llm_semaphore(): - try: - params = { - "max_tokens": 512, - "temperature": 0.3, - } - output_chunk: SummarizedChunk = await self.slm.ainvoke( - [ - {"role": "system", "content": system_prompt_map}, - { - "role": "user", - "content": user_prompt.format(query=query, content=chunk.page_content), - }, - ], - **params, - ) - return output_chunk - except Exception as e: - logger.error("Error during chunk relevancy inference", error=str(e)) - return SummarizedChunk(relevancy=False, summary="") - - async def map_batch( - self, - query: str, - chunks: list[Document], - summaries: list[SummarizedChunk], - kth_batch=1, - ): - """Process a batch of chunks""" - logger.debug(f"Processing {kth_batch}-th batch of chunks", batch_size=len(chunks)) - tasks = [self.infer_chunk_relevancy(query, chunk) for chunk in chunks] - outputs: list[SummarizedChunk] = await tqdm.gather( - *tasks, desc="Map & Reduce processing chunks", total=len(chunks) - ) - - # if the last 'expansion_batch_size' chunks are all irrelevant, we can terminate - terminate = all(not o.relevancy for o in outputs[-self.expansion_batch_size :]) - - for o, chunk in zip(outputs, chunks): - if o.relevancy: - summaries.append(Document(page_content=o.summary, metadata=chunk.metadata)) - - if self.debug: - with open(LOG_DIR / "map_reduce.md", "a") as f: - f.write(f"### Query: \n{query}\n") - f.write(f"### Chunk Content: \n* Relevancy: {o.relevancy} \n\n {chunk.page_content}\n") - f.write(f"### Summary: \n{o.summary}\n") - f.write("\n-------\n\n") - - return outputs, terminate - - async def map(self, query: str, chunks: list[Document]): - """Perform the map phase of map-reduce on the provided chunks. - Initally processes `initial_batch_size` number of documents to identify relevant ones. If they are all found to be relevant, - it continues to process additional documents in batches of `expansion_batch_size` until a - Args: - query (str): The user's query. - chunks (list[Document]): The list of document chunks (the `RETRIEVER_TOP_K` documents from the retreiver) to process. - - Returns: list[Document]: A list of relevant document summaries. - """ - - summaries: list[Document] = [] - - initial_batch, remaining_chunks = ( - chunks[: self.initial_batch_size], - chunks[self.initial_batch_size :], - ) - _, terminate = await self.map_batch(query, initial_batch, summaries=summaries, kth_batch=1) - - if terminate or not remaining_chunks or len(summaries) >= self.max_total_documents: - return summaries - - for jth_batch, i in enumerate(range(0, len(remaining_chunks), self.expansion_batch_size), start=2): - n = min(self.expansion_batch_size, self.max_total_documents - len(summaries)) - if n <= 0: - break - - logger.debug( - f"Expanding map phase: processing batch {jth_batch} with size {n}", - summaries_count=len(summaries), - ) - - next_batch = remaining_chunks[i : i + n] - _, terminate = await self.map_batch( - query=query, chunks=next_batch, summaries=summaries, kth_batch=jth_batch - ) - - if terminate or len(summaries) >= self.max_total_documents: - break - - logger.debug("Map reduce completed", relevant_chunks_count=len(summaries), query=query) - return summaries diff --git a/openrag/components/pipeline.py b/openrag/components/pipeline.py deleted file mode 100644 index 8187392d0..000000000 --- a/openrag/components/pipeline.py +++ /dev/null @@ -1,456 +0,0 @@ -import asyncio -import copy -from datetime import datetime -from enum import Enum -from typing import Literal - -import openai -import ray -from components.prompts import ( - QUERY_CONTEXTUALIZER_PROMPT, - SPOKEN_STYLE_ANSWER_PROMPT, - SYS_PROMPT_TMPLT, -) -from components.ray_utils import call_ray_actor_with_timeout -from components.utils import detect_language, format_context, format_web_context -from components.websearch import WebSearchFactory -from config import load_config -from langchain_core.documents.base import Document -from langchain_core.exceptions import OutputParserException -from langchain_openai import ChatOpenAI -from pydantic import BaseModel, Field, ValidationError -from utils.logger import get_logger - -from .llm import LLM -from .map_reduce import RAGMapReduce -from .reranker import BaseReranker, RerankerFactory -from .retriever import BaseRetriever, RetrieverFactory -from .utils import SOURCE_SEPARATOR - -logger = get_logger() -config = load_config() -VECTORDB_TIMEOUT = config.ray.indexer.vectordb_timeout - - -class RAGMODE(Enum): - SIMPLERAG = "SimpleRag" - CHATBOTRAG = "ChatBotRag" - - -class TemporalPredicate(BaseModel): - """A single constraint on a document's creation date. - - Multiple predicates on the same `Query` are combined with logical AND. - Use two predicates to express a closed range (e.g. last month): - [{op: ">=", value: "2026-03-01..."}, {op: "<=", value: "2026-03-31..."}] - """ - - field: Literal["created_at"] = Field( - default="created_at", - description="Document metadata field to filter on. Always `created_at` for now.", - ) - operator: Literal[">", "<", ">=", "<="] = Field( - description="Comparison operator applied to the date field.", - ) - value: str = Field( - description='ISO 8601 datetime with timezone, e.g. "2026-03-15T00:00:00+00:00".', - ) - - -class Query(BaseModel): - """A single vector database search query with optional temporal filters on document creation date. - - Predicates in `temporal_filters` are AND-combined. To express an exclusion - (e.g. "last year except March"), emit TWO `Query` objects, each with its own - AND-combined predicates covering one side of the gap. - """ - - query: str = Field(description="A semantically enriched, descriptive query for vector similarity search.") - temporal_filters: list[TemporalPredicate] | None = Field( - default=None, - description="Date predicates on `created_at`, AND-combined. Null when no temporal reference in the query.", - ) - - def to_milvus_filter(self) -> str | None: - """The temporal_filters attributes are already checked through the Pydantic types, except for date value that is kept as string, - as LLM sometimes give correct but not entirely complete date - """ - - if not self.temporal_filters: - return None - parts = [] - for p in self.temporal_filters: - try: - datetime.fromisoformat(p.value) - except (TypeError, ValueError): - logger.warning( - "Dropping temporal predicate with non-ISO value", - field=p.field, - operator=p.operator, - value=p.value, - ) - continue - parts.append(f'{p.field} {p.operator} ISO "{p.value}"') - if not parts: - return None - return " and ".join(parts) - - def __str__(self) -> str: - return f"Query: {self.query}, Filter: {self.to_milvus_filter()}" - - -class SearchQueries(BaseModel): - query_list: list[Query] = Field(..., description="Search sub-queries to retrieve relevant documents.") - - def __str__(self) -> str: - return " --- ".join(str(q) for q in self.query_list) - - -class RetrieverPipeline: - def __init__(self) -> None: - # retriever - self.retriever: BaseRetriever = RetrieverFactory.create_retriever(config=config) - self.allow_filterless_fallback = config.retriever.allow_filterless_fallback - - # reranker - self.reranker_enabled = config.reranker.enabled - self.reranker: BaseReranker = RerankerFactory.get_reranker(config) - logger.debug("Reranker", enabled=self.reranker_enabled, provider=config.reranker.provider) - self.reranker_top_k = config.reranker.top_k - - async def retrieve_docs( - self, - partition: list[str], - query: Query, - top_k: int | None = None, - filter_params: dict | None = None, - ) -> list[Document]: - milvus_filter = query.to_milvus_filter() - docs = await self.retriever.retrieve( - partition=partition, query=query.query, filter=milvus_filter, filter_params=filter_params - ) - - # Fallback: drop temporal filter if it wiped out all candidates. - # Gated by `retriever.allow_filterless_fallback` so deployments that - # prefer strict temporal retrieval can opt out (returns no docs - # rather than temporally-incorrect ones). - if not docs and milvus_filter and self.allow_filterless_fallback: - logger.warning( - "Temporal filter dropped: no documents matched, retrying without filter", - query=str(query.query), - filter=milvus_filter, - partition=partition, - ) - docs = await self.retriever.retrieve( - partition=partition, query=query.query, filter=None, filter_params=filter_params - ) - - logger.debug("Documents retreived", document_count=len(docs)) - - if docs: - # 1. rerank all the docs - if self.reranker_enabled: - docs = await self.reranker.rerank(query=query.query, documents=docs, top_k=None) - logger.debug("Documents reranked", document_count=len(docs)) - - # 2. expand the docs with related documents - if self.retriever.expansion_enabled: - # Limit the number of docs to expand - top_k = max(self.reranker_top_k, top_k) if top_k else self.reranker_top_k - docs2expand = copy.deepcopy(docs[:top_k]) - - logger.debug("Documents to expand", document_count=len(docs2expand)) - expanded_docs = await self.retriever.expand_search_results(results=docs2expand) - if len(docs2expand) == len(expanded_docs): # no expansion found, keep the original docs - return docs - - logger.debug("Documents expanded", document_count=len(expanded_docs)) - docs = expanded_docs - - # rerank again after expansion if reranker is enabled - if self.reranker_enabled: - docs = await self.reranker.rerank(query=query.query, documents=docs, top_k=None) - logger.debug("Documents after expansion and reranking", document_count=len(docs)) - - return docs - - async def get_relevant_docs( - self, - partition: list[str], - search_queries: SearchQueries, - top_k: int | None = None, - filter_params: dict | None = None, - ) -> list[Document]: - tasks = [ - self.retrieve_docs(partition=partition, query=q, top_k=top_k, filter_params=filter_params) - for q in search_queries.query_list - ] - results = await asyncio.gather(*tasks) - results = self.reranker.rrf_reranking(doc_lists=results) - if top_k is not None: - results = results[:top_k] - logger.debug("Final relevant documents after RRF reranking", document_count=len(results)) - return results - - -class RagPipeline: - def __init__(self) -> None: - # retriever pipeline - self.retriever_pipeline = RetrieverPipeline() - - # RAG - self.rag_mode = config.rag.mode - self.chat_history_depth = config.rag.chat_history_depth - self.max_context_tokens = config.reranker.top_k * config.chunker.chunk_size - - self.llm_client = LLM(config.llm, logger) - - llm = ChatOpenAI( - base_url=config.llm.base_url, - api_key=config.llm.api_key, - model=config.llm.model, - temperature=config.llm.temperature, - ) - - primary = llm.with_structured_output(SearchQueries, method="json_schema", strict=True) - fallback = llm.with_structured_output(SearchQueries, method="function_calling", strict=False) - self.query_generator = primary.with_fallbacks( - [fallback], - exceptions_to_handle=(openai.BadRequestError,), - ) - - self.max_contextualized_query_len = config.rag.max_contextualized_query_len - - # map reduce - self.map_reduce: RAGMapReduce = RAGMapReduce(config=config) - - # Web search - self.web_search_service = WebSearchFactory.create_service(config) - if self.web_search_service.provider: - logger.info("Web search enabled", provider=config.websearch.provider) - else: - logger.info("Web search disabled (WEBSEARCH_API_TOKEN not set)") - - async def generate_query(self, messages: list[dict]) -> SearchQueries: - match RAGMODE(self.rag_mode): - case RAGMODE.SIMPLERAG: - # For SimpleRag, we don't need to contextualize the query as the chat history is not taken into account - last_msg = messages[-1] - return SearchQueries(query_list=[Query(query=last_msg["content"])]) - - case RAGMODE.CHATBOTRAG: - # Contextualize the query based on the chat history - chat_history = "" - for m in messages: - chat_history += f"{m['role']}: {m['content']}\n" - - last_user_query = messages[-1]["content"] - query_language = detect_language(last_user_query) - - model_kwargs = { - "max_completion_tokens": self.max_contextualized_query_len, - # "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, - } - prompt = QUERY_CONTEXTUALIZER_PROMPT.format( - query_language=query_language, - current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S"), - ) - - llm_messages = [ - {"role": "system", "content": prompt}, - {"role": "user", "content": f"Here is the chat history: \n{chat_history}\n"}, - ] - - # Retry once on schema-validation failure; fall back to the raw user query on the second failure. - generator = self.query_generator.bind(**model_kwargs) - for attempt in (1, 2): - try: - return await generator.ainvoke(llm_messages) - except (ValidationError, OutputParserException) as exc: - if attempt == 1: - logger.warning("Query generation schema error — retrying", error=str(exc)) - else: - logger.warning( - "Query generation failed twice — falling back to raw user query", - error=str(exc), - ) - return SearchQueries(query_list=[Query(query=last_user_query)]) - - async def _prepare_for_chat_completion(self, partition: list[str] | None, payload: dict): - messages = payload["messages"] - messages = messages[-self.chat_history_depth :] # limit history depth - - # 1. get the query - queries: SearchQueries = await self.generate_query(messages) - logger.debug("Prepared query for chat completion", queries=str(queries)) - - metadata = payload.get("metadata") or {} - - use_map_reduce = metadata.get("use_map_reduce", False) - spoken_style_answer = metadata.get("spoken_style_answer", False) - use_websearch = metadata.get("websearch", False) - workspace = metadata.get("workspace") - - logger.debug( - "Metadata parameters", - use_map_reduce=use_map_reduce, - spoken_style_answer=spoken_style_answer, - use_websearch=use_websearch, - workspace=workspace, - ) - - # 2. get docs and/or web results concurrently - top_k = config.map_reduce.max_total_documents if use_map_reduce else None - if workspace: - vectordb = ray.get_actor("Vectordb", namespace="openrag") - ws = await call_ray_actor_with_timeout( - vectordb.get_workspace.remote(workspace), - timeout=VECTORDB_TIMEOUT, - task_description=f"get_workspace({workspace})", - ) - if not ws or ("all" not in partition and ws["partition_name"] not in partition): - logger.warning( - "Workspace not found in partition(s) — ignoring workspace filter", - workspace=workspace, - partition=partition, - ) - workspace = None - - filter_params = {"workspace_id": workspace} if workspace else None - - if partition is not None and use_websearch: - # Run one retrieval and one web search per sub-query, all concurrently (Option C). - # Web results from different sub-queries are deduplicated by URL, preserving order. - rag_tasks = [ - self.retriever_pipeline.retrieve_docs( - partition=partition, query=q, top_k=top_k, filter_params=filter_params - ) - for q in queries.query_list - ] - web_tasks = [self.web_search_service.search(q.query) for q in queries.query_list] - all_results = await asyncio.gather(*rag_tasks, *web_tasks) - n = len(queries.query_list) - raw_doc_lists = list(all_results[:n]) - raw_web_lists = list(all_results[n:]) - docs = self.retriever_pipeline.reranker.rrf_reranking(doc_lists=raw_doc_lists) - if top_k is not None: - docs = docs[:top_k] - # Deduplicate web results by URL, preserving first-seen order - seen_urls: set[str] = set() - web_results = [] - for result in (r for web_list in raw_web_lists for r in web_list): - if result.url not in seen_urls: - seen_urls.add(result.url) - web_results.append(result) - elif partition is not None: - docs = await self.retriever_pipeline.get_relevant_docs( - partition=partition, search_queries=queries, top_k=top_k, filter_params=filter_params - ) - web_results = [] - else: - # Web-only mode (partition is None): no RAG retrieval. - # Run one web search per sub-query concurrently and deduplicate by URL. - raw_web_lists = await asyncio.gather(*[self.web_search_service.search(q.query) for q in queries.query_list]) - seen_urls = set() - web_results = [] - for result in (r for web_list in raw_web_lists for r in web_list): - if result.url not in seen_urls: - seen_urls.add(result.url) - web_results.append(result) - docs = [] - - # Web-only with no results: fall back to plain direct LLM mode - if not docs and not web_results and partition is None: - return payload, [], [] - - if use_map_reduce and docs: - docs = await self.map_reduce.map(query=" ".join(q.query for q in queries.query_list), chunks=docs) - - # 3. Format web results first to know actual token usage, then allocate remaining budget to RAG - web_formatted = "" - web_tokens_used = 0 - if web_results: - web_formatted, _, web_tokens_used = format_web_context( - web_results, start_index=1, max_tokens=self.web_search_service.max_tokens - ) - - rag_max_tokens = self.max_context_tokens - web_tokens_used - context, included_indices = format_context(docs, max_context_tokens=rag_max_tokens) - docs = [docs[i] for i in included_indices] - - # Re-number web sources after RAG sources and rebuild if needed - if web_results: - n_rag_sources = len(docs) - if n_rag_sources > 0: - # Re-format with correct start_index now that we know RAG source count - web_formatted, _, _ = format_web_context( - web_results, start_index=n_rag_sources + 1, max_tokens=self.web_search_service.max_tokens - ) - - # Avoid misleading "No document found" when web results provide context - if not docs: - context = "" - - context = f"{context}{SOURCE_SEPARATOR}{web_formatted}" if context else web_formatted - - # 4. prepare the output - messages: list = copy.deepcopy(messages) - - # prepend the messages with the system prompt - prompt = SPOKEN_STYLE_ANSWER_PROMPT if spoken_style_answer else SYS_PROMPT_TMPLT - - messages.insert( - 0, - { - "role": "system", - "content": prompt.format( - context=context, current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S") - ), - }, - ) - payload["messages"] = messages - return payload, docs, web_results - - async def _prepare_for_completions(self, partition: list[str], payload: dict): - prompt = payload["prompt"] - - # 1. get the query - queries: SearchQueries = await self.generate_query(messages=[{"role": "user", "content": prompt}]) - # 2. get docs - docs = await self.retriever_pipeline.get_relevant_docs(partition=partition, search_queries=queries) - - # 3. Format the retrieved docs - context, included_indices = format_context(docs, max_context_tokens=self.max_context_tokens) - docs = [docs[i] for i in included_indices] - - # 4. prepare the output - if docs: - prompt = f"""Given the content - {context} - Complete the following prompt: {prompt} - At the very end of your response, on a new line, list which source numbers you used: [Sources: 1, 3]""" - - payload["prompt"] = prompt - - return payload, docs - - async def completions(self, partition: list[str], payload: dict): - if partition is None: - docs = [] - else: - payload, docs = await self._prepare_for_completions(partition=partition, payload=payload) - llm_output = self.llm_client.completions(request=payload) - return llm_output, docs - - async def chat_completion(self, partition: list[str] | None, payload: dict): - metadata = payload.get("metadata") or {} - use_websearch = metadata.get("websearch", False) - - if partition is None and not use_websearch: - # Direct LLM mode: no RAG, no web search - docs = [] - web_results = [] - else: - payload, docs, web_results = await self._prepare_for_chat_completion(partition=partition, payload=payload) - llm_output = self.llm_client.chat_completion(request=payload) - return llm_output, docs, web_results diff --git a/openrag/components/prompts/prompts.py b/openrag/components/prompts/prompts.py index afb2c3727..9dc113c93 100644 --- a/openrag/components/prompts/prompts.py +++ b/openrag/components/prompts/prompts.py @@ -1,6 +1,20 @@ +"""Backward-compatibility shim — delegates to `openrag.core.prompts.template_loader`. + +The disk-based template loader moved to +`openrag/core/prompts/template_loader.py` in Phase 5C. This module is +retained for legacy imports of `load_prompt(...)` and the eagerly-loaded +SYS_PROMPT_TMPLT / *_PROMPT constants until consumers migrate; +scheduled for removal in Phase 12. + +The new function takes (prompts_dir, mapping, key) explicitly; this +shim's `load_prompt(key)` resolves the first two from the cached +config, matching the legacy call site shape. +""" + from pathlib import Path from config import load_config +from core.prompts.template_loader import load_template_by_key config = load_config() @@ -13,29 +27,18 @@ def load_prompt( prompts_dir: Path = prompts_dir, prompt_mapping=prompt_mapping, ) -> str: - file_name = getattr(prompt_mapping, prompt_name, None) - if not file_name: - raise ValueError(f"No associated file name found for prompt: `{prompt_name}`") - - file_path = prompts_dir / file_name - - if not file_path.exists(): - raise FileNotFoundError(f"Prompt file not found: `{file_path}`") - - with open(file_path) as f: - sys_msg = f.read() - return sys_msg + return load_template_by_key(prompts_dir, prompt_mapping, prompt_name) -# Load prompts +# Eagerly-loaded prompt strings — preserved for legacy callers that +# import these names directly. New code should call `load_template_by_key` +# (or `load_template`) on demand instead. SYS_PROMPT_TMPLT = load_prompt("sys_prompt") QUERY_CONTEXTUALIZER_PROMPT = load_prompt("query_contextualizer") CHUNK_CONTEXTUALIZER_PROMPT = load_prompt("chunk_contextualizer") IMAGE_DESCRIBER = load_prompt("image_describer") -# Retrievers prompts HYDE_PROMPT = load_prompt("hyde") MULTI_QUERY_PROMPT = load_prompt("multi_query") -# Short answer prompt SPOKEN_STYLE_ANSWER_PROMPT = load_prompt("spoken_style_answer") diff --git a/openrag/components/ray_utils.py b/openrag/components/ray_utils.py index baecf358c..83ec534b1 100644 --- a/openrag/components/ray_utils.py +++ b/openrag/components/ray_utils.py @@ -1,90 +1,6 @@ -import asyncio -from collections.abc import Callable -from typing import Any - -import ray -from ray.exceptions import RayTaskError, TaskCancelledError -from utils.logger import get_logger - -logger = get_logger() - - -async def call_ray_actor_with_timeout( - future: ray.ObjectRef, - timeout: float, - task_description: str = "Ray task", -) -> Any: - """ - Wait for a Ray actor call with timeout and proper cancellation handling. - - This utility provides consistent error handling for Ray actor calls: - - Timeout with proper task cancellation - - Propagation of asyncio cancellation to Ray tasks - - Proper handling of Ray-specific exceptions - - Args: - future: The Ray ObjectRef from a remote call - timeout: Timeout in seconds - task_description: Description for error messages - - Returns: - The result of the Ray task - - Raises: - TimeoutError: If the task exceeds the timeout - asyncio.CancelledError: If the calling coroutine is cancelled - TaskCancelledError: If the Ray task was cancelled - RuntimeError: If the Ray task failed with an error - """ - try: - result = await asyncio.wait_for(asyncio.gather(future), timeout=timeout) - return result[0] # gather returns a list - - except TimeoutError: - logger.warning(f"{task_description} timed out, cancelling Ray task") - ray.cancel(future, recursive=True) - raise - - except asyncio.CancelledError: - logger.warning(f"{task_description} cancelled, cancelling Ray task") - ray.cancel(future, recursive=True) - raise - - except TaskCancelledError: - logger.warning(f"{task_description} Ray task was cancelled") - raise - - except RayTaskError as e: - raise RuntimeError(f"{task_description} failed") from e - - -async def retry_with_backoff( - attempt_fn: Callable[[int], Any], - max_retries: int, - base_delay: float, - task_description: str = "task", -) -> Any: - """ - Run `attempt_fn(attempt_index)` (an async callable) with exponential-backoff - retries. The callable owns its own resource acquire/release per attempt. - - Backoff: base_delay * 2**attempt seconds. CancelledError is never retried. - """ - last_exc: Exception | None = None - for attempt in range(max_retries + 1): - try: - return await attempt_fn(attempt) - except asyncio.CancelledError: - raise - except Exception as e: - last_exc = e - if attempt >= max_retries: - logger.error(f"{task_description} failed after {attempt + 1} attempts: {e}") - raise - delay = base_delay * (2**attempt) - logger.warning( - f"{task_description} failed (attempt {attempt + 1}/{max_retries + 1}): {e}. Retrying in {delay:.1f}s..." - ) - await asyncio.sleep(delay) - - raise last_exc # unreachable +# Re-export from canonical location for backward compatibility. +# New code should import from `services.workers.ray_utils` directly. +from services.workers.ray_utils import ( # noqa: F401 + call_ray_actor_with_timeout, + retry_with_backoff, +) diff --git a/openrag/components/reranker/__init__.py b/openrag/components/reranker/__init__.py index aa7e719ac..2b9742509 100644 --- a/openrag/components/reranker/__init__.py +++ b/openrag/components/reranker/__init__.py @@ -1,17 +1,43 @@ +import asyncio + +import services.inference.reranker_clients # noqa: F401 — registers "infinity"/"openai" +from core.config.retrieval import RerankerConfig +from core.rerankers import reranker_registry + from .base import BaseReranker +class _RerankerShim(BaseReranker): + """Wraps a core ``Reranker`` (str-in / (idx, score)-out) behind the + legacy ``BaseReranker`` interface (Document-in / Document-out).""" + + def __init__(self, delegate, semaphore: int = 3): + self._delegate = delegate + self._semaphore = asyncio.Semaphore(semaphore) + + async def rerank(self, query, documents, top_k=None): + async with self._semaphore: + texts = [doc.page_content for doc in documents] + ranked = await self._delegate.rerank(query, texts, top_k=top_k) + output = [] + for index, score in ranked: + if not 0 <= index < len(documents): + continue + doc = documents[index] + doc.metadata["relevance_score"] = score + output.append(doc) + return output + + class RerankerFactory: @staticmethod - def get_reranker(config) -> BaseReranker: - provider = config.reranker.provider - if provider == "infinity": - from .infinity import InfinityReranker - - return InfinityReranker(config) - elif provider == "openai": - from .openai import OpenAIReranker - - return OpenAIReranker(config) - else: - raise ValueError(f"Unsupported reranker provider: {provider}") + def get_reranker(reranker_config: RerankerConfig) -> BaseReranker: + provider = reranker_config.provider + delegate = reranker_registry.create( + provider, + endpoint=reranker_config.base_url, + model_name=reranker_config.model_name, + api_key=reranker_config.api_key, + timeout=reranker_config.timeout, + ) + return _RerankerShim(delegate, semaphore=reranker_config.semaphore) diff --git a/openrag/components/reranker/base.py b/openrag/components/reranker/base.py index 4a0c2367e..19a3defa3 100644 --- a/openrag/components/reranker/base.py +++ b/openrag/components/reranker/base.py @@ -1,5 +1,6 @@ from abc import ABC, abstractmethod +from core.retrieval.rrf import rrf_reranking from langchain_core.documents.base import Document @@ -10,31 +11,8 @@ async def rerank(self, query: str, documents: list[Document], top_k: int | None @staticmethod def rrf_reranking(doc_lists: list[list[Document]], k: int = 60) -> list[Document]: - """Reciprocal_rank_fusion that takes multiple lists of ranked documents - and an optional parameter k used in the RRF formula - RRF formula: \\sum_{i=1}^{n} \frac{1}{k + rank_i} - where rank_i is the rank of the document in the i-th list and n is the number of lists. - - k small: High sensitivity to top ranks - k large: More balanced sensitivity across ranks - k = 60 a common and balanced choice in practice. - """ - - if len(doc_lists) == 1: - return doc_lists[0] - - # Initialize a dictionary to hold fused scores for each unique document - fused_scores = {} - - for doc_list in doc_lists: - doc_list: list[Document] - for rank, doc in enumerate(doc_list, start=1): - doc_id = doc.metadata.get("_id") - doc_key = ("id", doc_id) if doc_id is not None else ("object", id(doc)) - - score, d = fused_scores.get(doc_key, (0, doc)) - fused_scores[doc_key] = (score + 1 / (rank + k), d) - - # sort the docs - reranked_docs = [doc for _, doc in sorted(fused_scores.values(), key=lambda x: x[0], reverse=True)] - return reranked_docs + return rrf_reranking( + doc_lists, + key_fn=lambda doc: doc.metadata.get("_id", id(doc)), + k=k, + ) diff --git a/openrag/components/reranker/infinity.py b/openrag/components/reranker/infinity.py index f55cda05a..32626466d 100644 --- a/openrag/components/reranker/infinity.py +++ b/openrag/components/reranker/infinity.py @@ -1,9 +1,15 @@ +"""Backward-compatibility shim — delegates to services.inference.reranker_clients. + +All new code should import directly from ``services.inference.reranker_clients``. +""" + import asyncio from infinity_client import Client from infinity_client.api.default import rerank from infinity_client.models import RerankInput, ReRankResult from langchain_core.documents.base import Document +from services.inference.reranker_clients import InfinityReranker as InfinityRerankerAdapter # noqa: F401 from utils.logger import get_logger from .base import BaseReranker @@ -12,6 +18,8 @@ class InfinityReranker(BaseReranker): + """Legacy InfinityReranker. New code should use InfinityRerankerAdapter (via DI).""" + def __init__(self, config): self.model_name = config.reranker.model_name self.client = Client( @@ -33,7 +41,7 @@ async def rerank(self, query: str, documents: list[Document], top_k: int | None "documents": [doc.page_content for doc in documents], "top_n": top_k, "return_documents": True, - "raw_scores": True, # Normalized score between 0 and 1 + "raw_scores": True, } ) try: diff --git a/openrag/components/reranker/openai.py b/openrag/components/reranker/openai.py index 43f9d0acf..7bca4fbcf 100644 --- a/openrag/components/reranker/openai.py +++ b/openrag/components/reranker/openai.py @@ -1,7 +1,13 @@ +"""Backward-compatibility shim — delegates to services.inference.reranker_clients. + +All new code should import directly from ``services.inference.reranker_clients``. +""" + import asyncio import httpx from langchain_core.documents.base import Document +from services.inference.reranker_clients import OpenAIReranker as OpenAIRerankerAdapter # noqa: F401 from utils.logger import get_logger from .base import BaseReranker @@ -10,6 +16,8 @@ class OpenAIReranker(BaseReranker): + """Legacy OpenAIReranker. New code should use OpenAIRerankerAdapter (via DI).""" + def __init__(self, config): self.model_name = config.reranker.model_name base_url = config.reranker.base_url.rstrip("/") diff --git a/openrag/components/reranker/test_rrf_reranking.py b/openrag/components/reranker/test_rrf_reranking.py index 13f098baa..40efb456d 100644 --- a/openrag/components/reranker/test_rrf_reranking.py +++ b/openrag/components/reranker/test_rrf_reranking.py @@ -10,10 +10,11 @@ def make_doc(doc_id: str, content: str = "", **metadata) -> Document: class TestRrfRerankingSingleList: - def test_single_list_returned_as_is(self): + def test_single_list_returned_as_list_copy(self): docs = [make_doc("a"), make_doc("b"), make_doc("c")] result = BaseReranker.rrf_reranking([docs]) - assert result is docs + assert result == docs + assert result is not docs class TestRrfRerankingMultipleLists: diff --git a/openrag/components/retriever.py b/openrag/components/retriever.py deleted file mode 100644 index 3f074adaf..000000000 --- a/openrag/components/retriever.py +++ /dev/null @@ -1,333 +0,0 @@ -# Import necessary modules and classes -import asyncio -from abc import ABC, abstractmethod -from itertools import chain -from typing import ClassVar - -from components.prompts import HYDE_PROMPT, MULTI_QUERY_PROMPT -from langchain_core.documents.base import Document -from langchain_core.output_parsers import StrOutputParser -from langchain_core.prompts import ChatPromptTemplate -from langchain_openai import ChatOpenAI -from utils.dependencies import get_vectordb -from utils.logger import get_logger - -logger = get_logger() - - -class ABCRetriever(ABC): - """Abstract class for the base retriever.""" - - @abstractmethod - def __init__( - self, - top_k: int = 6, - similarity_threshold: int = 0.95, - include_related: bool = False, - include_ancestors: bool = False, - related_limit: int = 10, - max_ancestor_depth: int | None = None, - **kwargs, - ) -> None: - pass - - @abstractmethod - async def retrieve( - self, partition: list[str], query: str, filter: str | None = None, filter_params: dict | None = None - ) -> list[Document]: - pass - - async def expand_search_results(self, results: list[Document]) -> list[Document]: - pass - - -# Define the Simple Retriever class -class BaseRetriever(ABCRetriever): - def __init__( - self, - top_k=6, - similarity_threshold=0.95, - with_surrounding_chunks=True, - include_related=False, - include_ancestors=False, - related_limit=10, - max_ancestor_depth: int | None = None, - **kwargs, - ): - super().__init__( - top_k, - similarity_threshold, - include_related=include_related, - include_ancestors=include_ancestors, - related_limit=related_limit, - max_ancestor_depth=max_ancestor_depth, - **kwargs, - ) - self.top_k = top_k - self.similarity_threshold = similarity_threshold - self.with_surrounding_chunks = with_surrounding_chunks - self.include_related = include_related - self.include_ancestors = include_ancestors - self.related_limit = related_limit - self.max_ancestor_depth = max_ancestor_depth - self.expansion_enabled = include_related or include_ancestors - - async def retrieve( - self, - partition: list[str], - query: str, - filter: str | None = None, - filter_params: dict | None = None, - ) -> list[Document]: - db = get_vectordb() - chunks = await db.async_search.remote( - query=query, - partition=partition, - top_k=self.top_k, - filter=filter, - filter_params=filter_params, - similarity_threshold=self.similarity_threshold, - with_surrounding_chunks=self.with_surrounding_chunks, - ) - return chunks - - async def expand_search_results(self, results: list[Document]) -> list[Document]: - """Expand search results with related and ancestor chunks.""" - db = get_vectordb() - return await _expand_with_related_chunks( - db=db, - results=results, - include_related=self.include_related, - include_ancestors=self.include_ancestors, - related_limit=self.related_limit, - max_ancestor_depth=self.max_ancestor_depth, - ) - - -class SingleRetriever(BaseRetriever): - pass - - -class MultiQueryRetriever(BaseRetriever): - def __init__( - self, - top_k=6, - similarity_threshold=0.95, - with_surrounding_chunks=True, - include_related=False, - include_ancestors=False, - related_limit=10, - max_ancestor_depth=None, - k_queries: int = 3, - llm: ChatOpenAI = None, - **kwargs, - ): - super().__init__( - top_k, - similarity_threshold, - with_surrounding_chunks, - include_related, - include_ancestors, - related_limit, - max_ancestor_depth, - **kwargs, - ) - - self.k_queries = k_queries - self.llm = llm - - if llm is None: - raise ValueError("llm must be provided for MultiQueryRetriever") - - prompt: ChatPromptTemplate = ChatPromptTemplate.from_template(MULTI_QUERY_PROMPT) - self.generate_queries = prompt | llm | StrOutputParser() | (lambda x: x.split("[SEP]")) - - async def retrieve( - self, partition: list[str], query: str, filter: str | None = None, filter_params: dict | None = None - ): - db = get_vectordb() - logger.debug("Generating multiple queries", k_queries=self.k_queries) - generated_queries = await self.generate_queries.ainvoke( - { - "query": query, - "k_queries": self.k_queries, - } - ) - chunks = await db.async_multi_query_search.remote( - queries=generated_queries, - partition=partition, - top_k_per_query=self.top_k, - filter=filter, - filter_params=filter_params, - similarity_threshold=self.similarity_threshold, - with_surrounding_chunks=self.with_surrounding_chunks, - ) - return chunks - - -class HyDeRetriever(BaseRetriever): - def __init__( - self, - top_k=6, - similarity_threshold=0.95, - with_surrounding_chunks=True, - include_related=False, - include_ancestors=False, - related_limit=10, - max_ancestor_depth=None, - llm: ChatOpenAI = None, - combine: bool = False, - **kwargs, - ): - super().__init__( - top_k, - similarity_threshold, - with_surrounding_chunks, - include_related, - include_ancestors, - related_limit, - max_ancestor_depth, - **kwargs, - ) - - super().__init__(top_k, similarity_threshold, **kwargs) - if llm is None: - raise ValueError("llm must be provided for HyDeRetriever") - - self.combine = combine - self.llm = llm - - prompt: ChatPromptTemplate = ChatPromptTemplate.from_template(HYDE_PROMPT) - self.hyde_generator = prompt | llm | StrOutputParser() - - async def get_hyde(self, query: str): - logger.debug("Generating HyDe Document") - hyde_document = await self.hyde_generator.ainvoke({"query": query}) - return hyde_document - - async def retrieve( - self, partition: list[str], query: str, filter: str | None = None, filter_params: dict | None = None - ) -> list[Document]: - db = get_vectordb() - hyde = await self.get_hyde(query) - queries = [hyde] - if self.combine: - queries.append(query) - - return await db.async_multi_query_search.remote( - queries=queries, - partition=partition, - top_k_per_query=self.top_k, - filter=filter, - filter_params=filter_params, - similarity_threshold=self.similarity_threshold, - with_surrounding_chunks=self.with_surrounding_chunks, - ) - - -async def _expand_with_related_chunks( - db, - results: list[Document], - include_related: bool, - include_ancestors: bool, - related_limit: int = 10, - max_ancestor_depth: int | None = None, -) -> list[Document]: - """Expand results with related and/or ancestor chunks.""" - if not results or (not include_related and not include_ancestors): - return results - - # Track what we already have to avoid duplicates - seen_ids = {doc.metadata.get("_id") for doc in results} - expanded_results = list(results) - - # Collect unique relationship_ids and file_ids from results - relationship_ids = set() - file_infos = [] # List of (partition, file_id) tuples - - for doc in results: - metadata = doc.metadata - if include_related and metadata.get("relationship_id"): - relationship_ids.add((metadata.get("partition"), metadata.get("relationship_id"))) - if include_ancestors: - file_infos.append((metadata.get("partition"), metadata.get("file_id"))) - - # Create tasks for parallel fetching - async def fetch_related(partition: str, rel_id: str) -> list[Document]: - """Fetch related chunks with error handling.""" - try: - return await db.get_related_chunks.remote( - partition=partition, - relationship_id=rel_id, - limit=related_limit, - ) - except Exception as e: - logger.warning( - "Failed to fetch related chunks", - relationship_id=rel_id, - error=str(e), - ) - return [] - - async def fetch_ancestors(partition: str, file_id: str) -> list[Document]: - """Fetch ancestor chunks with error handling.""" - try: - return await db.get_ancestor_chunks.remote( - partition=partition, - file_id=file_id, - limit=related_limit, - max_ancestor_depth=max_ancestor_depth, - ) - except Exception as e: - logger.warning( - "Failed to fetch ancestor chunks", - file_id=file_id, - error=str(e), - ) - return [] - - # Build list of tasks for parallel execution - tasks = [] - - if include_related: - tasks.extend(fetch_related(partition, rel_id) for partition, rel_id in relationship_ids if partition and rel_id) - - if include_ancestors: - tasks.extend(fetch_ancestors(partition, file_id) for partition, file_id in file_infos if partition and file_id) - - # Execute all tasks in parallel - if tasks: - all_results = await asyncio.gather(*tasks) - for chunk in chain.from_iterable(all_results): - chunk_id = chunk.metadata.get("_id") - if chunk_id and chunk_id not in seen_ids: - seen_ids.add(chunk_id) - expanded_results.append(chunk) - - logger.debug( - "Expanded results with related/ancestor chunks", - original_count=len(results), - expanded_count=len(expanded_results), - ) - return expanded_results - - -class RetrieverFactory: - RETRIEVERS: ClassVar[dict] = { - "single": SingleRetriever, - "multiQuery": MultiQueryRetriever, - "hyde": HyDeRetriever, - } - - @classmethod - def create_retriever(cls, config) -> ABCRetriever: - retrieverConfig = config.retriever.model_dump() - - retriever_type = retrieverConfig.pop("type") - retriever_cls = RetrieverFactory.RETRIEVERS.get(retriever_type, None) - - if retriever_cls is None: - raise ValueError(f"Unknown retriever type: {retriever_type}") - - retrieverConfig["llm"] = ChatOpenAI(**config.llm.model_dump()) - return retriever_cls(**retrieverConfig) diff --git a/openrag/components/utils.py b/openrag/components/utils.py index d3f83931c..f01923277 100644 --- a/openrag/components/utils.py +++ b/openrag/components/utils.py @@ -5,12 +5,13 @@ import threading from typing import ClassVar -import ray -from components.indexer.utils.text_sanitizer import sanitize_text from config import load_config from fast_langdetect import LangDetectConfig, LangDetector from langchain_core.documents.base import Document -from langchain_openai import ChatOpenAI +from services.inference.distributed_semaphore import ( + DistributedSemaphore, # noqa: F401 + DistributedSemaphoreActor, # noqa: F401 +) from utils.logger import get_logger SOURCE_SEPARATOR = "-" * 10 + "\n\n" @@ -33,92 +34,56 @@ def __call__(cls, *args, **kwargs): return cls._instances[cls] -@ray.remote(max_restarts=5, max_concurrency=config.ray.semaphore.concurrency) -class DistributedSemaphoreActor: - def __init__(self, max_concurrent_ops: int): - self.semaphore = asyncio.Semaphore(max_concurrent_ops) - - async def acquire(self): - await self.semaphore.acquire() - - def release(self): - self.semaphore.release() - - -class DistributedSemaphore: - # https://chat.deepseek.com/a/chat/s/890dbcc0-2d3f-4819-af9d-774b892905bc - def __init__( - self, - name: str = "llmSemaphore", - namespace="openrag", - max_concurrent_ops: int = 10, - ): - self._name = name - self._namespace = namespace - self._max_concurrent_ops = max_concurrent_ops - - def _get_or_create_actor(self): - try: - # reuse existing actor if it exists - _actor = ray.get_actor(self._name, namespace=self._namespace) - except ValueError: - # create new actor if it doesn't exist - _actor = DistributedSemaphoreActor.options( - name=self._name, - namespace=self._namespace, - lifetime="detached", - ).remote(self._max_concurrent_ops) - except Exception: - raise - - return _actor - - async def __aenter__(self): - semaphore_actor = self._get_or_create_actor() - await semaphore_actor.acquire.remote() - return self - - async def __aexit__(self, exc_type, exc, tb): - semaphore_actor = self._get_or_create_actor() - await semaphore_actor.release.remote() - - _cached_length_function = None def get_num_tokens(): global _cached_length_function if _cached_length_function is None: - llm = ChatOpenAI(**config.llm.model_dump()) - _cached_length_function = llm.get_num_tokens + try: + from langchain_openai import ChatOpenAI + + llm = ChatOpenAI(**config.llm.model_dump()) + _cached_length_function = llm.get_num_tokens + except Exception as exc: + # ChatOpenAI validates an openai client at construction, which + # requires a non-empty api_key. Token counting itself is local + # (tiktoken) and needs no key or network, so fall back to a + # tiktoken encoder when the client cannot be built (keyless + # deployments / CI mock-vLLM). cl100k_base matches the + # GPT-3.5/4 family OpenRAG targets; counts are equivalent. + import tiktoken + + logger.warning( + "ChatOpenAI unavailable for token counting, falling back to tiktoken cl100k_base", + error=str(exc), + ) + _encoding = tiktoken.get_encoding("cl100k_base") + _cached_length_function = lambda text: len(_encoding.encode(text)) # noqa: E731 return _cached_length_function def format_context( docs: list[Document], max_context_tokens: int = 4096, number_sources: bool = True ) -> tuple[str, list[int]]: - if not docs: - return "No document found from the database", [] + """Backward-compat shim — delegates to `core.prompts.chat_prompt_builder.format_context`. - _length_function = get_num_tokens() - - reduced_docs = [] - included_indices = [] - total_tokens = 0 - - for i, doc in enumerate(docs): - prefix = f"[Source {len(reduced_docs) + 1}]\n" if number_sources else "" - n_tokens = _length_function(doc.page_content) - if prefix: - n_tokens += _length_function(prefix) - if total_tokens + n_tokens > max_context_tokens: - break - reduced_docs.append(f"{prefix}{doc.page_content}") - included_indices.append(i) - total_tokens += n_tokens - - logger.debug("Context formatted", total_tokens=total_tokens, doc_count=len(reduced_docs)) - return SOURCE_SEPARATOR.join(reduced_docs), included_indices + The legacy signature took LangChain Documents and resolved a tokenizer + internally; the core version takes raw strings + an injected + length_function. We adapt by extracting page_content and threading + the cached tokenizer through. + """ + from core.prompts.chat_prompt_builder import format_context as _core_format_context + + texts = [doc.page_content for doc in docs] + text, included = _core_format_context( + texts, + max_context_tokens=max_context_tokens, + length_function=get_num_tokens(), + number_sources=number_sources, + ) + logger.debug("Context formatted", doc_count=len(included)) + return text, included def format_web_context( @@ -126,41 +91,21 @@ def format_web_context( start_index: int = 1, max_tokens: int = 2000, ) -> tuple[str, list[int], int]: - """Format web results as numbered [Source N] blocks within a token budget. + """Backward-compat shim — delegates to `core.prompts.chat_prompt_builder.format_web_context`. - Uses fetched page content when available, falling back to the search snippet. - - Args: - web_results: Results from web search provider (list of WebResult) - start_index: First source number (continues numbering after RAG sources) - max_tokens: Maximum token budget for all web sources combined - - Returns: - (formatted_string, list_of_source_numbers_used, total_tokens_used) + Same adaptation pattern as `format_context`: legacy resolved the + tokenizer internally, core takes it as a parameter. """ - if not web_results: - return "", [], 0 - - _length_function = get_num_tokens() - - parts = [] - source_numbers = [] - total_tokens = 0 - - for i, result in enumerate(web_results): - n = start_index + i - title = sanitize_text(result.title) - body = sanitize_text(result.content) if result.content else sanitize_text(result.snippet) - block = f"[Source {n}]\n{title}\n{body}" - block_tokens = _length_function(block) - if total_tokens + block_tokens > max_tokens and parts: - break - parts.append(block) - source_numbers.append(n) - total_tokens += block_tokens - - logger.debug("Web context formatted", total_tokens=total_tokens, source_count=len(parts)) - return SOURCE_SEPARATOR.join(parts), source_numbers, total_tokens + from core.prompts.chat_prompt_builder import format_web_context as _core_format_web_context + + text, source_numbers, total_tokens = _core_format_web_context( + web_results, + length_function=get_num_tokens(), + start_index=start_index, + max_tokens=max_tokens, + ) + logger.debug("Web context formatted", total_tokens=total_tokens, source_count=len(source_numbers)) + return text, source_numbers, total_tokens # Line-terminal anchor `(?=\n|$)` — matches only when the tag sits flush against diff --git a/openrag/config/__init__.py b/openrag/config/__init__.py index a2e454a45..b7b6b5790 100644 --- a/openrag/config/__init__.py +++ b/openrag/config/__init__.py @@ -1,3 +1,5 @@ +# Re-export from canonical location for backward compatibility. +# New code should import from openrag.core.config directly. """OpenRAG configuration package. Public API: @@ -8,13 +10,13 @@ from functools import lru_cache -from .models import Settings +from openrag.core.config.root import Settings # noqa: F401 @lru_cache def get_settings() -> Settings: """Cached singleton — one Settings instance per process.""" - from .loader import load_config as _load + from openrag.core.config.loader import load_config as _load return _load() @@ -28,7 +30,7 @@ def load_config(config_path=None, overrides=None) -> Settings: The ``overrides`` parameter bypasses the cache (useful for tests). """ if overrides or config_path: - from .loader import load_config as _load + from openrag.core.config.loader import load_config as _load return _load(conf_dir=config_path, overrides=overrides) return get_settings() diff --git a/openrag/config/loader.py b/openrag/config/loader.py index 4a345c3d5..a8c5007d5 100644 --- a/openrag/config/loader.py +++ b/openrag/config/loader.py @@ -9,7 +9,7 @@ import yaml -from .models import Settings +from openrag.core.config.root import Settings logger = logging.getLogger(__name__) diff --git a/openrag/conftest.py b/openrag/conftest.py new file mode 100644 index 000000000..7f0563d99 --- /dev/null +++ b/openrag/conftest.py @@ -0,0 +1,21 @@ +"""Pytest import guards for legacy top-level imports. + +Some modules still import ``utils`` and the third-party ``openai`` package as +top-level names. During collection, pytest can prepend nested test directories +such as ``openrag/routers`` to ``sys.path``, where ``utils.py`` and +``openai.py`` would otherwise shadow those imports. +""" + +from __future__ import annotations + +import importlib +import sys +from pathlib import Path + +OPENRAG_ROOT = Path(__file__).resolve().parent +root = str(OPENRAG_ROOT) +if root not in sys.path: + sys.path.insert(0, root) + +sys.modules.setdefault("utils", importlib.import_module("utils")) +sys.modules.setdefault("openai", importlib.import_module("openai")) diff --git a/openrag/consts.py b/openrag/consts.py index 51bc12904..240bba006 100644 --- a/openrag/consts.py +++ b/openrag/consts.py @@ -1,6 +1,6 @@ -PARTITION_PREFIX = "openrag-" -LEGACY_PARTITION_PREFIX = "ragondin-" - -FILE_READ_CHUNK_SIZE = 1024 * 1024 # Read file in blocks of 1MB to preserve RAM - -IMAGE_PLACEHOLDER = """\n\n[Image Placeholder]\n\n""" +from core.utils.conts import ( # noqa: F401,F403 + FILE_READ_CHUNK_SIZE, + IMAGE_PLACEHOLDER, + LEGACY_PARTITION_PREFIX, + PARTITION_PREFIX, +) diff --git a/openrag/core/__init__.py b/openrag/core/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/core/chunking/__init__.py b/openrag/core/chunking/__init__.py new file mode 100644 index 000000000..a699bc2eb --- /dev/null +++ b/openrag/core/chunking/__init__.py @@ -0,0 +1,24 @@ +"""ChunkingStrategy ABC + registry + concrete strategies.""" + +from .chunking_strategy import ChunkingStrategy +from .markdown_utils import ( + MDElement, + chunk_table, + get_chunk_page_number, + parse_markdown_table, + split_md_elements, +) +from .recursive import BaseChunker, RecursiveSplitter +from .registry import chunking_registry + +__all__ = [ + "ChunkingStrategy", + "chunking_registry", + "BaseChunker", + "RecursiveSplitter", + "MDElement", + "chunk_table", + "get_chunk_page_number", + "parse_markdown_table", + "split_md_elements", +] diff --git a/openrag/core/chunking/chunking_strategy.py b/openrag/core/chunking/chunking_strategy.py new file mode 100644 index 000000000..2f2dbf93d --- /dev/null +++ b/openrag/core/chunking/chunking_strategy.py @@ -0,0 +1,17 @@ +"""Abstract chunking strategy interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from openrag.core.models.chunk import Chunk +from openrag.core.models.document import ProcessedDocument + + +class ChunkingStrategy(ABC): + """Base class for all chunking strategies.""" + + @abstractmethod + def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]: + """Split a processed document into chunks.""" + ... diff --git a/openrag/core/chunking/markdown_utils.py b/openrag/core/chunking/markdown_utils.py new file mode 100644 index 000000000..9112d575d --- /dev/null +++ b/openrag/core/chunking/markdown_utils.py @@ -0,0 +1,224 @@ +"""Markdown parsing primitives used by chunking strategies. + +Pure functions extracted from ``components/indexer/chunker/utils.py``. They +recognize page markers, image-description blocks, and tables; split a +markdown document into typed elements; and split oversize tables along +their semantic groups. + +This module has no IO and no config dependency. +""" + +from __future__ import annotations + +import re +from collections.abc import Callable +from dataclasses import dataclass +from typing import Literal + +from core.utils.text import clean_markdown_table_spacing + +# Header + delimiter + at least one row. +TABLE_RE = re.compile( + r"((?:^|\n)\|.*?\|\r?\n\|\s*[:-]+(?:\s*\|[:-]+)*\|\r?\n(?:\|.*?\|\r?\n)+)", + re.DOTALL | re.MULTILINE, +) + +# `...` block injected by the VLM step. +IMAGE_RE = re.compile(r"((.*?))", re.DOTALL) + +# `[PAGE_N]` page-boundary markers — content BEFORE [PAGE_N] is on page N. +PAGE_RE = re.compile(r"\[PAGE_(\d+)\]") + + +ElementType = Literal["text", "table", "image"] + + +@dataclass +class MDElement: + """A typed segment of markdown content with optional source page number.""" + + type: ElementType + content: str + page_number: int | None = None + + def __repr__(self) -> str: + return f"MDElement(type={self.type}, page_number={self.page_number}, content={self.content[:100]}...)" + + +def span_inside(span: tuple[int, int], container: tuple[int, int]) -> bool: + """Return True if ``span`` is fully contained within ``container``.""" + return container[0] <= span[0] and span[1] <= container[1] + + +def get_page_number(position: int, page_markers: list[tuple[int, int]]) -> int: + """Look up the page number for a position in the source markdown. + + ``page_markers`` is a sorted list of ``(offset, page_n)`` tuples taken + from ``[PAGE_N]`` matches. Content AFTER ``[PAGE_N]`` belongs to page + ``N + 1``; content before any marker is page 1. + """ + current_page = 1 + for marker_pos, page_num in page_markers: + if position >= marker_pos: + current_page = page_num + 1 + else: + break + return current_page + + +def split_md_elements(md_text: str) -> list[MDElement]: + """Split markdown into ``MDElement`` segments of text, table, and image. + + Tables nested inside an ```` block are NOT extracted + as separate elements — they belong to the image. + """ + page_markers: list[tuple[int, int]] = [] + for match in PAGE_RE.finditer(md_text): + page_markers.append((match.start(), int(match.group(1)))) + page_markers.sort() + + all_matches: list[tuple[tuple[int, int], ElementType, str, int | None]] = [] + image_spans: list[tuple[int, int]] = [] + + for match in IMAGE_RE.finditer(md_text): + span = match.span() + page_num = get_page_number(span[0], page_markers) + all_matches.append((span, "image", match.group(1).strip(), page_num)) + image_spans.append(span) + + for match in TABLE_RE.finditer(md_text): + span = match.span() + if not any(span_inside(span, image_span) for image_span in image_spans): + page_num = get_page_number(span[0], page_markers) + all_matches.append((span, "table", match.group(1).strip(), page_num)) + + all_matches.sort(key=lambda x: x[0][0]) + + parts: list[MDElement] = [] + last = 0 + + for (start, end), match_type, content, page_num in all_matches: + if start > last: + text_segment = md_text[last:start] + if text_segment.strip(): + parts.append(MDElement(type="text", content=text_segment.strip())) + parts.append(MDElement(type=match_type, content=content, page_number=page_num)) + last = end + + if last < len(md_text): + remaining = md_text[last:] + if remaining.strip(): + parts.append(MDElement(type="text", content=remaining.strip())) + + return parts + + +def get_chunk_page_number(chunk_str: str, previous_chunk_ending_page: int = 1) -> dict[str, int]: + """Resolve start and end pages for a text chunk containing ``[PAGE_N]`` markers. + + Returns ``{"start_page": int, "end_page": int}``. + """ + matches = list(PAGE_RE.finditer(chunk_str)) + + if not matches: + return { + "start_page": previous_chunk_ending_page, + "end_page": previous_chunk_ending_page, + } + + first_match = matches[0] + last_match = matches[-1] + last_char_idx = len(chunk_str) - 1 + + if first_match.start() == 0: + start_page = int(first_match.group(1)) + 1 + else: + start_page = previous_chunk_ending_page + + if last_match.end() - 1 == last_char_idx: + end_page = int(last_match.group(1)) + else: + end_page = int(last_match.group(1)) + 1 + + return {"start_page": start_page, "end_page": max(start_page, end_page)} + + +def parse_markdown_table(markdown_table: str) -> tuple[list[str], list[list[str]]]: + """Parse a markdown table into header lines + groups of rows. + + Rows are grouped by the first column ("Domain"): a non-empty Domain + starts a new group, an empty Domain continues the current group. This + preserves the document's logical structure when chunking large tables. + """ + lines = markdown_table.strip().split("\n") + header_lines = lines[:2] + data_rows = lines[2:] + + groups: list[list[str]] = [] + current_group: list[str] = [] + + for row in data_rows: + cells = [cell.strip() for cell in row.split("|")[1:-1]] + if not cells: + continue + domain = cells[0] + if domain: + if current_group: + groups.append(current_group) + current_group = [row] + else: + current_group.append(row) + + if current_group: + groups.append(current_group) + + return header_lines, groups + + +def chunk_table( + table_element: MDElement, + chunk_size: int, + length_function: Callable[[str], int], +) -> list[MDElement]: + """Split an oversize markdown table into multiple ``MDElement`` chunks. + + Each chunk repeats the table header. When a new chunk starts, the LAST + row of the previous chunk is replayed as overlap so context is preserved + across the boundary. + """ + txt = clean_markdown_table_spacing(table_element.content) + header_lines, groups = parse_markdown_table(txt) + + header_text = "\n".join(header_lines) + group_texts = ["\n".join(g) for g in groups] + + header_ntoks = length_function(header_text) + groups_ntoks = [length_function(g) for g in group_texts] + + subtables: list[str] = [] + body_rows: list[str] = [] # rows under the current chunk, header excluded + body_size = 0 + prev_last_row: str | None = None + + for group_txt, g_ntoks in zip(group_texts, groups_ntoks, strict=True): + # Only flush when we actually have body content to flush — otherwise an + # oversized first group would emit a header-only chunk. + if body_rows and header_ntoks + body_size + g_ntoks > chunk_size: + subtables.append("\n".join([header_text, *body_rows])) + body_rows = [] + body_size = 0 + # Replay only the last row of the previous chunk as overlap + # (matches the docstring contract; prev_last_row is the trailing + # line of the last admitted group). + if prev_last_row: + body_rows.append(prev_last_row) + body_size += length_function(prev_last_row) + body_rows.append(group_txt) + body_size += g_ntoks + # The "last row" is the trailing line of this group, not the whole group. + prev_last_row = group_txt.rsplit("\n", 1)[-1] + + if body_rows: + subtables.append("\n".join([header_text, *body_rows])) + + return [MDElement(type="table", content=subtable, page_number=table_element.page_number) for subtable in subtables] diff --git a/openrag/core/chunking/recursive.py b/openrag/core/chunking/recursive.py new file mode 100644 index 000000000..7afeade07 --- /dev/null +++ b/openrag/core/chunking/recursive.py @@ -0,0 +1,280 @@ +"""Recursive markdown-aware chunking strategy. + +Pure domain logic — no LLM client, no Ray, no LangChain ``Document``. +The token-counting function is injected (``length_function``); the actual +text splitter is ``langchain.text_splitter.RecursiveCharacterTextSplitter``, +a pure utility kept until a stdlib-only replacement is in place. + +Contextualization (the LLM-driven [CONTEXT] block prepended to each chunk) +lives in ``core/indexing/contextualize.py`` (Phase 5D) and is applied as a +separate stage by the orchestrator — not from inside the chunker. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +from core.chunking.chunking_strategy import ChunkingStrategy +from core.chunking.markdown_utils import ( + MDElement, + chunk_table, + get_chunk_page_number, + split_md_elements, +) +from core.chunking.registry import chunking_registry +from core.models.chunk import Chunk, ChunkType +from core.models.document import ProcessedDocument +from core.utils.text import sanitize_text + +# Substring (case-insensitive) marking a "no useful content" image caption. +# Detection logic mirrors the legacy chunker, which skips these elements so +# they don't pollute the index. +_IMAGE_PLACEHOLDER_MARKER = "[image placeholder]" + +# Tables/images smaller than this token count are inlined with surrounding +# text rather than emitted as standalone chunks. +_INLINE_ELEMENT_TOKEN_THRESHOLD = 100 + + +class BaseChunker(ChunkingStrategy): + """Base markdown-aware chunker. + + Subclasses must set ``self.text_splitter`` to an object with a + ``.split_text(str) -> list[str]`` method (e.g. LangChain's + ``RecursiveCharacterTextSplitter``). + """ + + def __init__( + self, + chunk_size: int = 200, + chunk_overlap_rate: float = 0.2, + length_function: Callable[[str], int] | None = None, + **kwargs: Any, + ) -> None: + if length_function is None: + raise ValueError("length_function is required (e.g. tokenizer.count_tokens)") + self.chunk_size = chunk_size + self.chunk_overlap_rate = chunk_overlap_rate + self.chunk_overlap = int(self.chunk_size * self.chunk_overlap_rate) + self.length_function = length_function + self.text_splitter: Any = None + + # ------------------------------------------------------------------ + # ChunkingStrategy contract + # ------------------------------------------------------------------ + def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]: + """Split a processed document into ``Chunk`` objects.""" + content = self._content_from(document) + if not content.strip(): + return [] + + metadata = self._chunk_metadata_base(document, partition) + md_chunks = self._get_chunks(content=content.strip(), metadata=metadata) + + return [ + Chunk( + document_id=metadata.get("file_id", ""), + text=md_chunks_meta["page_content"], + chunk_index=i, + chunk_type=ChunkType(md_chunks_meta["chunk_type"]), + metadata={k: v for k, v in md_chunks_meta.items() if k not in ("page_content", "chunk_type", "page")}, + partition=partition, + page_number=md_chunks_meta.get("page"), + ) + for i, md_chunks_meta in enumerate(md_chunks) + ] + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + @staticmethod + def _content_from(document: ProcessedDocument) -> str: + """Reconstruct chunkable markdown from a ProcessedDocument. + + Single-block documents on page 1 (or with no page metadata) flow + through unchanged. Anything else gets synthetic ``[PAGE_N]`` markers + injected so downstream chunk-page resolution works correctly. + + Marker semantics: a ``[PAGE_N]`` marker means "everything BEFORE this + marker was on page N" (see ``markdown_utils.get_page_number``). So we + emit the marker for the *outgoing* page just before content from a + new page begins, and we also prepend a marker for the first block if + it doesn't start on page 1. + """ + if not document.text_blocks: + return "" + if len(document.text_blocks) == 1 and document.text_blocks[0].page_number in (None, 1): + return document.text_blocks[0].text + + parts: list[str] = [] + last_page: int | None = None + for index, block in enumerate(document.text_blocks): + if block.page_number is not None: + # Emit `[PAGE_{block.page_number - 1}]` immediately *before* + # this block's text so downstream resolution lands on + # block.page_number. Using `block.page_number - 1` (rather + # than `last_page`) handles non-sequential pages (1 -> 5) + # and a first block already on page > 1. + needs_marker = ( + (index == 0 and block.page_number > 1) + or (last_page is not None and block.page_number != last_page) + or (last_page is None and index > 0) + ) + if needs_marker: + parts.append(f"[PAGE_{block.page_number - 1}]") + parts.append(block.text) + if block.page_number is not None: + last_page = block.page_number + return "\n\n".join(parts) + + @staticmethod + def _chunk_metadata_base(document: ProcessedDocument, partition: str) -> dict[str, Any]: + # Reserved identity fields must win — `chunk()` later reads + # metadata["file_id"] to set Chunk.document_id, so a stray key in + # `document.metadata` would silently reassign chunks to the wrong doc. + return { + **document.metadata, + "file_id": document.document_id, + "partition": partition, + } + + def split_text(self, text: str) -> list[str]: + """Split a text string with the configured text splitter. + + Lazy-initializes a ``RecursiveCharacterTextSplitter`` if a subclass + forgot to set one — preserves legacy behavior. + """ + if self.text_splitter is None: + from langchain.text_splitter import RecursiveCharacterTextSplitter + + self.text_splitter = RecursiveCharacterTextSplitter( + chunk_size=self.chunk_size, + chunk_overlap=self.chunk_overlap, + length_function=self.length_function, + ) + return self.text_splitter.split_text(text) + + def _prepare_md_elements(self, content: str) -> tuple[list[MDElement], list[MDElement]]: + """Separate markdown into (inline-able texts) and (standalone tables/images).""" + md_elements = split_md_elements(content) + tables_and_images: list[MDElement] = [] + texts: list[MDElement] = [] + + for element in md_elements: + if element.type in ("table", "image"): + if element.type == "image" and _IMAGE_PLACEHOLDER_MARKER in element.content.lower(): + continue + if self.length_function(element.content) <= _INLINE_ELEMENT_TOKEN_THRESHOLD: + texts.append(element) + else: + tables_and_images.append(element) + else: + texts.append(element) + + return texts, tables_and_images + + def _get_chunks(self, content: str, metadata: dict[str, Any]) -> list[dict[str, Any]]: + """Produce per-chunk dicts with ``page_content`` + metadata fields. + + The dict shape is intentional — it lets ``chunk()`` build ``Chunk`` + objects without leaking domain types into the lower-level helpers. + """ + texts, tables_and_images = self._prepare_md_elements(content=content) + combined_texts = "\n".join(e.content for e in texts) + + sanitized = sanitize_text( + combined_texts, + normalize_whitespace=True, + remove_control_chars=True, + remove_zero_width_chars=True, + max_consecutive_newlines=2, + normalize_unicode=True, + ) + text_chunks = self.split_text(sanitized) + + chunks: list[dict[str, Any]] = [] + + # Reserved per-chunk keys must win over arbitrary `metadata` values — + # a stray "chunk_type" / "page" / "page_content" in the document's + # metadata would otherwise clobber the resolved value (and crash + # `chunk()` when ChunkType(...) is fed an out-of-enum string). Same + # defensive pattern as `_chunk_metadata_base`. + for element in tables_and_images: + if element.type == "table" and self.length_function(element.content) > self.chunk_size: + subtables = chunk_table( + table_element=element, + chunk_size=self.chunk_size, + length_function=self.length_function, + ) + chunks.extend( + { + **metadata, + "page_content": subtable.content.strip(), + "page": subtable.page_number, + "chunk_type": "table", + } + for subtable in subtables + ) + else: + # MDElement.type is the source-markdown literal ("image"/"table"); + # ChunkType uses "image_caption" for image blocks. + ct = "image_caption" if element.type == "image" else element.type + chunks.append( + { + **metadata, + "page_content": element.content.strip(), + "page": element.page_number, + "chunk_type": ct, + } + ) + + prev_page = 1 + for c in text_chunks: + page_info = get_chunk_page_number(chunk_str=c, previous_chunk_ending_page=prev_page) + prev_page = page_info["end_page"] + chunks.append( + { + **metadata, + "page_content": c.strip(), + "page": page_info["start_page"], + "chunk_type": "text", + } + ) + + if not chunks: + return [] + chunks.sort(key=lambda d: d.get("page") or 0) + return chunks + + +@chunking_registry.register("recursive_splitter") +class RecursiveSplitter(BaseChunker): + """Markdown-aware chunker backed by ``RecursiveCharacterTextSplitter``. + + Splits on paragraph boundaries first, then sentence terminators, then + smaller separators. + """ + + def __init__( + self, + chunk_size: int = 200, + chunk_overlap_rate: float = 0.2, + length_function: Callable[[str], int] | None = None, + **kwargs: Any, + ) -> None: + super().__init__( + chunk_size=chunk_size, + chunk_overlap_rate=chunk_overlap_rate, + length_function=length_function, + **kwargs, + ) + from langchain.text_splitter import RecursiveCharacterTextSplitter + + self.text_splitter = RecursiveCharacterTextSplitter( + chunk_size=self.chunk_size, + chunk_overlap=self.chunk_overlap, + length_function=self.length_function, + is_separator_regex=True, + separators=["\n", r"(?<=[\.\?\!])"], + ) diff --git a/openrag/core/chunking/registry.py b/openrag/core/chunking/registry.py new file mode 100644 index 000000000..9f14bff75 --- /dev/null +++ b/openrag/core/chunking/registry.py @@ -0,0 +1,7 @@ +"""Chunking strategy registry.""" + +from openrag.core.utils.registry import Registry + +from .chunking_strategy import ChunkingStrategy + +chunking_registry: Registry[ChunkingStrategy] = Registry("chunking") diff --git a/openrag/core/chunking/test_markdown_utils.py b/openrag/core/chunking/test_markdown_utils.py new file mode 100644 index 000000000..830bcfa81 --- /dev/null +++ b/openrag/core/chunking/test_markdown_utils.py @@ -0,0 +1,182 @@ +"""Tests for core.chunking.markdown_utils. + +Mirrors components/indexer/chunker/test_chunking.py to verify behavior is +preserved through the move into core/. +""" + +from __future__ import annotations + +from core.chunking.markdown_utils import ( + MDElement, + chunk_table, + get_chunk_page_number, + parse_markdown_table, + span_inside, + split_md_elements, +) + + +def _mock_length(text: str) -> int: + """Estimate token count at ~4 chars per token (matches legacy tests).""" + return len(text) // 4 + + +class TestSplitMdElements: + def test_simple_text_only(self): + md = "This is a simple paragraph.\n\nAnother paragraph here." + elems = split_md_elements(md) + assert len(elems) == 1 + assert elems[0].type == "text" + assert elems[0].content == md + + def test_single_table(self): + md = ( + "Some text before.\n\n| Header 1 | Header 2 |\n|----------|----------|\n" + "| Cell 1 | Cell 2 |\n| Cell 3 | Cell 4 |\n\nSome text after." + ) + elems = split_md_elements(md) + assert [e.type for e in elems] == ["text", "table", "text"] + assert "Header 1" in elems[1].content + + def test_single_image(self): + md = ( + "\nText before image.\n\n\nA beautiful sunset over the ocean.\n" + "\n\nText after image." + ) + elems = split_md_elements(md) + assert [e.type for e in elems] == ["text", "image", "text"] + assert "sunset" in elems[1].content + + def test_table_inside_image_description_is_ignored(self): + md = ( + "\n\nThis image contains a table:\n| Col 1 | Col 2 |\n" + "|-------|-------|\n| A | B |\n\n\n" + "Outside table:\n| Real 1 | Real 2 |\n|--------|--------|\n| X | Y |\n" + ) + elems = split_md_elements(md) + tables = [e for e in elems if e.type == "table"] + assert len(tables) == 1 + assert "Real 1" in tables[0].content + + def test_page_markers_with_table(self): + md = ( + "text on page 1.\n[PAGE_1]\nText on page 2.\n\n" + "| Header 1 | Header 2 |\n|----------|----------|\n| Data 1 | Data 2 |\n\n" + "[PAGE_2]\nMore content.\n" + ) + elems = split_md_elements(md) + tables = [e for e in elems if e.type == "table"] + assert len(tables) == 1 + assert tables[0].page_number == 2 + + def test_page_markers_with_images(self): + md = "\n[PAGE_1]\n[PAGE_2]\n\nImage on page 3.\n\n" + elems = split_md_elements(md) + images = [e for e in elems if e.type == "image"] + assert len(images) == 1 + assert images[0].page_number == 3 + + +class TestGetChunkPageNumber: + def test_no_markers_returns_previous_page(self): + result = get_chunk_page_number("Just some plain text content.", previous_chunk_ending_page=1) + assert result == {"start_page": 1, "end_page": 1} + + def test_chunk_starts_with_marker(self): + result = get_chunk_page_number("[PAGE_2]Content on page 3.", previous_chunk_ending_page=1) + assert result == {"start_page": 3, "end_page": 3} + + def test_chunk_ends_with_marker(self): + result = get_chunk_page_number("Content on page 1.[PAGE_1]", previous_chunk_ending_page=1) + assert result == {"start_page": 1, "end_page": 1} + + def test_marker_in_middle(self): + result = get_chunk_page_number("Start on page 1.[PAGE_1]End on page 2.", previous_chunk_ending_page=1) + assert result == {"start_page": 1, "end_page": 2} + + +class TestChunkTable: + def test_small_table_no_chunking(self): + content = "| Name | Age |\n|------|-----|\n| John | 30 |\n| Jane | 25 |" + elem = MDElement(type="table", content=content, page_number=1) + chunks = chunk_table(elem, chunk_size=1000, length_function=_mock_length) + assert len(chunks) == 1 + assert chunks[0].type == "table" + assert chunks[0].page_number == 1 + assert "John" in chunks[0].content + assert "Jane" in chunks[0].content + + def test_chunking_preserves_groups(self): + header = "| Country | Strategy | Goals |" + g1 = "| USA | Cyber | Goal 1 |\n| | | Goal 2 |\n| | | Goal 3 |" + g2 = "| Mexico | Defense | Goal X |\n| | | Goal Y |\n| | | Goal Z |" + table = f"{header}\n|----|----|----|\n{g1}\n{g2}\n" + elem = MDElement(type="table", content=table, page_number=2) + chunk_size = _mock_length(table) // 2 + chunks = chunk_table(elem, chunk_size=chunk_size, length_function=_mock_length) + assert len(chunks) == 2 + assert all(c.type == "table" for c in chunks) + assert all(header in c.content for c in chunks) + assert "USA" in chunks[0].content + + def test_oversized_first_group_does_not_emit_header_only_chunk(self): + """When the very first group is already larger than chunk_size, the + old code flushed a header-only chunk (CodeRabbit #1).""" + header = "| Country | Strategy | Goals |" + g1 = "| USA | Cyber | Goal 1 |\n| | | Goal 2 |\n| | | Goal 3 |" + g2 = "| Mexico | Defense | Goal X |" + table = f"{header}\n|----|----|----|\n{g1}\n{g2}\n" + # Tight budget: g1 alone already exceeds. + chunks = chunk_table( + MDElement(type="table", content=table, page_number=1), chunk_size=2, length_function=_mock_length + ) + # No chunk may contain only the header. + for c in chunks: + body = c.content.replace(header, "").strip() + assert body, f"header-only chunk emitted: {c.content!r}" + + def test_overlap_replays_only_last_row_not_full_group(self): + """The docstring promises last-row overlap; the old code stored the + whole previous group (CodeRabbit #1).""" + header = "| Country | Strategy | Goal |" + g1 = "| USA | Cyber | first |\n| | | second |\n| | | LAST_ROW_OF_G1 |" + g2 = "| Mexico | Defense | only |" + table = f"{header}\n|----|----|----|\n{g1}\n{g2}\n" + # Force a split between g1 and g2. + chunk_size = _mock_length(g1) + _mock_length(header) + chunks = chunk_table( + MDElement(type="table", content=table, page_number=1), chunk_size=chunk_size, length_function=_mock_length + ) + assert len(chunks) >= 2 + second_chunk = chunks[1].content + assert "LAST_ROW_OF_G1" in second_chunk, "last row should be replayed as overlap" + # The earlier rows of g1 must NOT appear in the second chunk. + assert "first" not in second_chunk + assert "second" not in second_chunk + + +def test_md_element_repr_truncates_content(): + elem = MDElement(type="text", content="x" * 500, page_number=3) + rendered = repr(elem) + assert "type=text" in rendered + assert "page_number=3" in rendered + # Long content is truncated to <=100 chars + ellipsis. + assert "x" * 200 not in rendered + + +def test_span_inside_helper(): + assert span_inside((10, 20), (5, 30)) is True + assert span_inside((5, 30), (10, 20)) is False + assert span_inside((10, 20), (10, 20)) is True + + +def test_parse_markdown_table_skips_blank_data_rows(): + """A pipe-only row (e.g. an extra blank `|`) yields no cells; it must be + skipped without erroring or starting a phantom group.""" + header = "| Country | Goal |" + delim = "|---------|------|" + table = f"{header}\n{delim}\n| USA | A |\n|\n| Mexico | B |" + headers, groups = parse_markdown_table(table) + assert headers == [header, delim] + # Two non-empty rows -> two groups (each row has a non-empty Domain). + assert len(groups) == 2 diff --git a/openrag/core/chunking/test_recursive.py b/openrag/core/chunking/test_recursive.py new file mode 100644 index 000000000..b60f13b67 --- /dev/null +++ b/openrag/core/chunking/test_recursive.py @@ -0,0 +1,251 @@ +"""End-to-end tests for the RecursiveSplitter chunker.""" + +from __future__ import annotations + +from core.chunking.recursive import RecursiveSplitter +from core.chunking.registry import chunking_registry +from core.models.chunk import ChunkType +from core.models.document import ProcessedDocument, TextBlock + + +def _word_tokens(text: str) -> int: + return len(text.split()) + + +def test_recursive_splitter_is_registered(): + assert "recursive_splitter" in chunking_registry + + +def test_recursive_splitter_chunks_simple_document(): + splitter = RecursiveSplitter(chunk_size=10, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[TextBlock(text="alpha beta gamma\ndelta epsilon zeta\neta theta iota.", page_number=1)], + metadata={"source": "test.md"}, + ) + chunks = splitter.chunk(doc, partition="p1") + assert chunks + assert all(c.partition == "p1" for c in chunks) + assert all(c.document_id == "d1" for c in chunks) + assert all(c.chunk_type == ChunkType.TEXT for c in chunks) + + +def test_recursive_splitter_emits_table_chunks(): + table = "| Col | Val |\n|-----|-----|\n" + "\n".join(f"| Group{i} | {' '.join(['x'] * 50)} |" for i in range(6)) + splitter = RecursiveSplitter(chunk_size=20, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[TextBlock(text=f"Some prose here.\n\n{table}\n\nMore prose.", page_number=1)], + ) + chunks = splitter.chunk(doc, partition="p1") + table_chunks = [c for c in chunks if c.chunk_type == ChunkType.TABLE] + assert table_chunks, "expected at least one table-type chunk" + + +def test_recursive_splitter_skips_image_placeholder(): + placeholder_md = ( + "Real text first.\n\n\n\n[Image Placeholder]\n\n\n\nReal text after." + ) + splitter = RecursiveSplitter(chunk_size=200, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[TextBlock(text=placeholder_md, page_number=1)], + ) + chunks = splitter.chunk(doc, partition="p1") + assert all(c.chunk_type != ChunkType.IMAGE_CAPTION for c in chunks) + for c in chunks: + assert "[image placeholder]" not in c.text.lower() + + +def test_recursive_splitter_metadata_passthrough(): + splitter = RecursiveSplitter(chunk_size=10, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[TextBlock(text="alpha beta gamma delta epsilon", page_number=1)], + metadata={"source": "test.md", "filename": "test.md", "tag": "v1"}, + ) + chunks = splitter.chunk(doc, partition="p1") + assert chunks + assert chunks[0].metadata.get("source") == "test.md" + assert chunks[0].metadata.get("tag") == "v1" + + +def test_recursive_splitter_empty_document_returns_empty(): + splitter = RecursiveSplitter(chunk_size=10, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument(document_id="d1", text_blocks=[]) + assert splitter.chunk(doc, partition="p1") == [] + + +def test_recursive_splitter_requires_length_function(): + import pytest + + with pytest.raises(ValueError, match="length_function"): + RecursiveSplitter(chunk_size=10, chunk_overlap_rate=0.0) + + +def test_recursive_splitter_joins_multi_block_document_with_synthetic_page_markers(): + """Multi-block docs need synthetic [PAGE_N] markers so chunks downstream + of a page boundary report the right page. Cover that injection path in + BaseChunker._content_from with a chunk_size small enough to force a split + across pages.""" + block_text = " ".join([f"word{i}" for i in range(20)]) + splitter = RecursiveSplitter(chunk_size=8, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[ + TextBlock(text=block_text, page_number=1), + TextBlock(text=block_text, page_number=2), + TextBlock(text=block_text, page_number=3), + ], + metadata={"source": "multi.md"}, + ) + chunks = splitter.chunk(doc, partition="p1") + assert chunks + pages = {c.page_number for c in chunks} + # First chunk(s) stay on page 1; once a [PAGE_N] marker lands inside a + # chunk's content the next chunk resolves to >=2. + assert 1 in pages + assert any((p or 0) >= 2 for p in pages) + + +def test_recursive_splitter_inlines_small_table(): + """Tables under the inline threshold (<=100 length-function tokens) flow + through the text path rather than emitting a standalone TABLE chunk.""" + table = "| A | B |\n|---|---|\n| 1 | 2 |" + splitter = RecursiveSplitter(chunk_size=200, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[TextBlock(text=f"Lead-in.\n\n{table}\n\nTrailing.", page_number=1)], + ) + chunks = splitter.chunk(doc, partition="p1") + assert chunks + assert all(c.chunk_type != ChunkType.TABLE for c in chunks) + + +def test_recursive_splitter_image_caption_chunk_emitted_when_above_threshold(): + """Image_description blocks above the inline threshold land as their own + chunks (chunk_type=image_caption) — covers the standalone-element path in + _get_chunks's else branch.""" + long_caption = "lorem ipsum dolor sit amet " * 60 # well above inline threshold + md = f"Some text.\n\n\n{long_caption}\n\n\nAfter." + splitter = RecursiveSplitter(chunk_size=400, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[TextBlock(text=md, page_number=1)], + ) + chunks = splitter.chunk(doc, partition="p1") + image_chunks = [c for c in chunks if c.chunk_type == ChunkType.IMAGE_CAPTION] + assert image_chunks, "expected at least one image_caption chunk" + + +def test_recursive_splitter_returns_empty_when_only_image_placeholder(): + """When the only element is a skipped image placeholder, _get_chunks + produces nothing — exercise the `if not chunks: return []` guard.""" + splitter = RecursiveSplitter(chunk_size=200, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[ + TextBlock(text="\n[Image Placeholder]\n", page_number=1), + ], + ) + chunks = splitter.chunk(doc, partition="p1") + assert chunks == [] + + +def test_base_chunker_lazy_initializes_text_splitter(): + """A BaseChunker subclass that forgets to set self.text_splitter still + works — split_text lazy-builds a default RecursiveCharacterTextSplitter.""" + from core.chunking.recursive import BaseChunker + + class BareChunker(BaseChunker): + pass + + bare = BareChunker(chunk_size=12, chunk_overlap_rate=0.0, length_function=_word_tokens) + assert bare.text_splitter is None + pieces = bare.split_text("alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu") + assert pieces + assert bare.text_splitter is not None # cached after first call + + +def test_document_metadata_cannot_override_file_id_or_partition(): + """Reserved identity fields must win over arbitrary metadata keys + (CodeRabbit #3).""" + splitter = RecursiveSplitter(chunk_size=20, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="real-doc-id", + text_blocks=[TextBlock(text="alpha beta gamma delta", page_number=1)], + metadata={ + "file_id": "MALICIOUS_OVERRIDE", + "partition": "MALICIOUS_PARTITION", + "source": "ok.md", + }, + ) + chunks = splitter.chunk(doc, partition="real-partition") + assert chunks + for c in chunks: + assert c.document_id == "real-doc-id" + assert c.partition == "real-partition" + # Other metadata keys still flow through. + assert c.metadata.get("source") == "ok.md" + + +def test_document_metadata_cannot_override_chunk_type_or_page(): + """Per-chunk reserved keys (chunk_type, page, page_content) must win + over `document.metadata`. A poison `chunk_type` value would otherwise + crash `chunk()` when ChunkType(...) is constructed (ultrareview).""" + splitter = RecursiveSplitter(chunk_size=20, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[TextBlock(text="alpha beta gamma delta epsilon", page_number=2)], + metadata={ + "chunk_type": "POISON", + "page": 999, + "page_content": "REPLACED", + "tag": "v1", + }, + ) + # Must not raise ValueError("'POISON' is not a valid ChunkType"). + chunks = splitter.chunk(doc, partition="p1") + assert chunks + for c in chunks: + assert c.chunk_type == ChunkType.TEXT + assert c.page_number != 999 + assert c.text != "REPLACED" + # Other metadata keys still flow through. + assert c.metadata.get("tag") == "v1" + + +def test_recursive_splitter_first_block_on_page_three_resolves_correctly(): + """When the first block already starts on page>1, every chunk used to be + tagged page 1. Now it should land on the actual block page (CodeRabbit #2).""" + block_text = " ".join([f"word{i}" for i in range(20)]) + splitter = RecursiveSplitter(chunk_size=8, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[ + TextBlock(text=block_text, page_number=3), + TextBlock(text=block_text, page_number=4), + ], + ) + chunks = splitter.chunk(doc, partition="p1") + pages = {c.page_number for c in chunks} + assert 1 not in pages, f"chunks tagged page 1 despite first block on page 3: {pages}" + assert any((p or 0) >= 3 for p in pages) + + +def test_recursive_splitter_skips_pages_get_correct_marker(): + """Block pages 1 -> 5 (skipping 2/3/4) — the second block's chunks must + resolve to page 5, not page 2 (= last_page+1).""" + block_text = " ".join([f"word{i}" for i in range(40)]) + splitter = RecursiveSplitter(chunk_size=8, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[ + TextBlock(text=block_text, page_number=1), + TextBlock(text=block_text, page_number=5), + ], + ) + chunks = splitter.chunk(doc, partition="p1") + pages = sorted({c.page_number for c in chunks if c.page_number is not None}) + assert 1 in pages + assert 5 in pages, f"page 5 missing despite second block on page 5: {pages}" diff --git a/openrag/core/config/__init__.py b/openrag/core/config/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/core/config/auth.py b/openrag/core/config/auth.py new file mode 100644 index 000000000..85e21bd8e --- /dev/null +++ b/openrag/core/config/auth.py @@ -0,0 +1,28 @@ +"""Authentication configuration — token mode + OIDC settings.""" + +from __future__ import annotations + +from pydantic import BaseModel + + +class OIDCConfig(BaseModel): + """OIDC configuration for Keycloak / external IdP integration. + + All fields are optional — if OIDC is not enabled, this section is ignored. + Populated from environment variables (OIDC_ENDPOINT, OIDC_CLIENT_ID, etc.). + """ + + enabled: bool = False + issuer_url: str = "" + client_id: str = "" + client_secret: str = "" + redirect_uri: str = "" + scopes: str = "openid email profile offline_access" + token_encryption_key: str = "" + claim_source: str = "id_token" + claim_mapping: str = "" + post_logout_redirect_uri: str = "" + # When True, an unknown ``sub`` at callback time provisions a + # non-admin user on the fly from the ID-token claims. Default keeps + # the strict "admin pre-creates every user" policy. + auto_provision_login: bool = False diff --git a/openrag/core/config/base.py b/openrag/core/config/base.py new file mode 100644 index 000000000..a4a53c4d8 --- /dev/null +++ b/openrag/core/config/base.py @@ -0,0 +1,44 @@ +"""Base config mixin — frozen Pydantic models with dict-like backward compatibility.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel + + +class ConfigMixin(BaseModel): + """Frozen Pydantic model with dict-like access for backward compatibility. + + Existing code using ``config.section.get("key")``, ``config.section["key"]``, + ``dict(config.section)``, and ``**config.section`` keeps working. + """ + + model_config = {"frozen": True} + + def get(self, key: str, default: Any = None) -> Any: + try: + return getattr(self, key) + except AttributeError: + return default + + def __getitem__(self, key: str) -> Any: + try: + return getattr(self, key) + except AttributeError: + raise KeyError(key) + + def keys(self): + return list(type(self).model_fields.keys()) + + def values(self): + return [getattr(self, k) for k in type(self).model_fields] + + def items(self): + return [(k, getattr(self, k)) for k in type(self).model_fields] + + def __iter__(self): + return iter(type(self).model_fields) + + def __contains__(self, key: str) -> bool: + return key in type(self).model_fields diff --git a/openrag/core/config/chunking.py b/openrag/core/config/chunking.py new file mode 100644 index 000000000..53219495f --- /dev/null +++ b/openrag/core/config/chunking.py @@ -0,0 +1,16 @@ +"""Chunking configuration.""" + +from __future__ import annotations + +from .base import ConfigMixin + + +class ChunkerConfig(ConfigMixin): + """Chunking strategy settings.""" + + name: str = "recursive_splitter" + contextual_retrieval: bool = True + contextualization_timeout: int = 120 + max_concurrent_contextualization: int = 10 + chunk_size: int = 512 + chunk_overlap_rate: float = 0.2 diff --git a/openrag/core/config/endpoints.py b/openrag/core/config/endpoints.py new file mode 100644 index 000000000..0bd2ddaf6 --- /dev/null +++ b/openrag/core/config/endpoints.py @@ -0,0 +1,56 @@ +"""Model endpoint configuration — LLM, VLM, embedder, semaphore settings.""" + +from __future__ import annotations + +from pydantic import Field + +from .base import ConfigMixin + + +class LLMParamsConfig(ConfigMixin): + """Shared parameters for LLM/VLM endpoints.""" + + temperature: float = 0.1 + timeout: int = 60 + max_retries: int = 2 + logprobs: bool = True + + +class LLMConfig(LLMParamsConfig): + """LLM endpoint configuration.""" + + base_url: str = "" + model: str = "" + api_key: str = Field(default="", repr=False) + + +class VLMConfig(LLMParamsConfig): + """Vision-Language Model endpoint configuration.""" + + base_url: str = "" + model: str = "" + api_key: str = Field(default="", repr=False) + + +class EmbedderConfig(ConfigMixin): + """Embedding model endpoint configuration.""" + + provider: str = "openai" + model_name: str = "jinaai/jina-embeddings-v3" + base_url: str = "http://vllm:8000/v1" + api_key: str = Field(default="EMPTY", repr=False) + max_model_len: int = 8192 + + +class SemaphoreConfig(ConfigMixin): + """Concurrency limits for LLM and VLM calls.""" + + llm_semaphore: int = 10 + vlm_semaphore: int = 10 + + +class LLMContextConfig(ConfigMixin): + """Token budget settings for LLM context.""" + + max_llm_context_size: int = 8192 + max_output_tokens: int = 1024 diff --git a/openrag/core/config/indexation.py b/openrag/core/config/indexation.py new file mode 100644 index 000000000..9f2d13e87 --- /dev/null +++ b/openrag/core/config/indexation.py @@ -0,0 +1,181 @@ +"""Indexation pipeline configuration — loaders, parsers, transcribers.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import Field, field_validator + +from .base import ConfigMixin + +# --------------------------------------------------------------------------- +# Transcriber (nested under loader) +# --------------------------------------------------------------------------- + +# Audio formats the transcription endpoint accepts as-is — uploads of these +# extensions skip the WAV-conversion pre-step. Configurable via env var +# `TRANSCRIBER_DIRECT_UPLOAD_SUFFIXES` (pipe-delimited string). +_DEFAULT_DIRECT_UPLOAD_SUFFIXES = frozenset( + {".wav", ".flac", ".ogg", ".mp3", ".mp4", ".m4a", ".webm", ".mpeg", ".mpga"} +) + + +def _normalize_suffix(s: str) -> str: + s = s.strip().lower() + if not s: + return "" + return s if s.startswith(".") else f".{s}" + + +class TranscriberConfig(ConfigMixin): + base_url: str = "http://transcriber:8000/v1" + api_key: str = Field(default="EMPTY", repr=False) + model_name: str = "openai/whisper-large-v3-turbo" + timeout: int = 3600 + max_concurrent_chunks: int = 20 + use_whisper_lang_detector: bool = True + direct_upload_suffixes: set[str] = Field(default_factory=lambda: set(_DEFAULT_DIRECT_UPLOAD_SUFFIXES)) + + @field_validator("direct_upload_suffixes", mode="before") + @classmethod + def _split_suffixes(cls, v: Any) -> Any: + if isinstance(v, str): + return {n for raw in v.split("|") if (n := _normalize_suffix(raw))} + return v + + +# --------------------------------------------------------------------------- +# OpenAI Loader (nested under loader) +# --------------------------------------------------------------------------- + + +class OpenAILoaderConfig(ConfigMixin): + base_url: str = "http://openai:8000/v1" + api_key: str = Field(default="EMPTY", repr=False) + model: str = "dotsocr-model" + temperature: float = 0.2 + timeout: int = 180 + max_retries: int = 2 + top_p: float = 0.9 + concurrency_limit: int = 20 + + +# --------------------------------------------------------------------------- +# Local Whisper (nested under loader) +# --------------------------------------------------------------------------- + + +class LocalWhisperConfig(ConfigMixin): + model: str = "base" + whisper_n_workers: int = 3 + whisper_num_gpus: float = 0.01 + whisper_concurrency_per_worker: int = 2 + whisper_timeout: int = 1800 + whisper_max_task_retry: int = 1 + whisper_retry_base_delay: float = 2.0 + + +# --------------------------------------------------------------------------- +# File loaders mapping (nested under loader) +# --------------------------------------------------------------------------- + + +class FileLoadersConfig(ConfigMixin): + txt: str = "TextLoader" + pdf: str = "MarkerLoader" + eml: str = "EmlLoader" + docx: str = "DocxLoader" + pptx: str = "PPTXLoader" + doc: str = "DocLoader" + png: str = "ImageLoader" + jpeg: str = "ImageLoader" + jpg: str = "ImageLoader" + svg: str = "ImageLoader" + wav: str = "LocalWhisperLoader" + mp3: str = "LocalWhisperLoader" + flac: str = "LocalWhisperLoader" + ogg: str = "LocalWhisperLoader" + aac: str = "LocalWhisperLoader" + flv: str = "LocalWhisperLoader" + wma: str = "LocalWhisperLoader" + mp4: str = "LocalWhisperLoader" + md: str = "MarkdownLoader" + + +# --------------------------------------------------------------------------- +# Mimetypes mapping (nested under loader) +# --------------------------------------------------------------------------- + + +class MimetypesConfig(ConfigMixin): + """Maps MIME type strings to file extensions. + + Access via .to_dict() for {mime_type: extension} mapping. + """ + + text_plain: str = Field(default=".txt", alias="text/plain") + text_markdown: str = Field(default=".md", alias="text/markdown") + application_pdf: str = Field(default=".pdf", alias="application/pdf") + message_rfc822: str = Field(default=".eml", alias="message/rfc822") + application_docx: str = Field( + default=".docx", + alias="application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ) + application_pptx: str = Field( + default=".pptx", + alias="application/vnd.openxmlformats-officedocument.presentationml.presentation", + ) + application_msword: str = Field(default=".doc", alias="application/msword") + image_png: str = Field(default=".png", alias="image/png") + image_jpeg: str = Field(default=".jpeg", alias="image/jpeg") + audio_wav: str = Field(default=".wav", alias="audio/wav") + audio_mpeg: str = Field(default=".mp3", alias="audio/mpeg") + audio_flac: str = Field(default=".flac", alias="audio/flac") + audio_ogg: str = Field(default=".ogg", alias="audio/ogg") + audio_aac: str = Field(default=".aac", alias="audio/aac") + video_x_flv: str = Field(default=".flv", alias="video/x-flv") + audio_x_ms_wma: str = Field(default=".wma", alias="audio/x-ms-wma") + video_mp4: str = Field(default=".mp4", alias="video/mp4") + + model_config = {"frozen": True, "extra": "allow", "populate_by_name": True} + + def to_dict(self) -> dict[str, str]: + """Return {mime_type: extension} mapping using aliases as keys.""" + result = {} + for field_name, field_info in type(self).model_fields.items(): + alias = field_info.alias or field_name + result[alias] = getattr(self, field_name) + if self.__pydantic_extra__: + result.update(self.__pydantic_extra__) + return result + + +# --------------------------------------------------------------------------- +# Loader (top-level indexation config) +# --------------------------------------------------------------------------- + + +class LoaderConfig(ConfigMixin): + image_captioning: bool = True + image_captioning_url: bool = True + save_markdown: bool = False + mimetypes: MimetypesConfig = Field(default_factory=MimetypesConfig) + local_whisper: LocalWhisperConfig = Field(default_factory=LocalWhisperConfig) + file_loaders: FileLoadersConfig = Field(default_factory=FileLoadersConfig) + marker_max_tasks_per_child: int = 20 + marker_pool_size: int = 1 + marker_max_processes: int = 2 + marker_num_gpus: float = 0.01 + marker_timeout: int = 3600 + marker_pdftext_workers: int = 2 + marker_chunk_size: int = 10 + marker_max_task_retry: int = 3 + marker_retry_base_delay: float = 2.0 + docling_num_gpus: float = Field(default=0.01, ge=0) + docling_pool_size: int = Field(default=1, ge=1) + docling_max_tasks_per_worker: int = Field(default=2, ge=1) + docling_timeout: int = 3600 + docling_max_task_retry: int = 3 + docling_retry_base_delay: float = 2.0 + transcriber: TranscriberConfig = Field(default_factory=TranscriberConfig) + openai: OpenAILoaderConfig = Field(default_factory=OpenAILoaderConfig) diff --git a/openrag/core/config/infrastructure.py b/openrag/core/config/infrastructure.py new file mode 100644 index 000000000..cad5f996c --- /dev/null +++ b/openrag/core/config/infrastructure.py @@ -0,0 +1,136 @@ +"""Infrastructure configuration — VectorDB, Postgres, Ray, paths, server.""" + +from __future__ import annotations + +from pathlib import Path + +from pydantic import Field + +from .base import ConfigMixin + +# --------------------------------------------------------------------------- +# VectorDB (Milvus) +# --------------------------------------------------------------------------- + + +class VectorDBConfig(ConfigMixin): + host: str = "milvus" + port: int = 19530 + connector_name: str = "milvus" + collection_name: str = "vdb_test" + hybrid_search: bool = True + enable: bool = True + schema_version: int = 1 + + +# --------------------------------------------------------------------------- +# RDB (Postgres) +# --------------------------------------------------------------------------- + + +class RDBConfig(ConfigMixin): + host: str = "rdb" + port: int = 5432 + user: str = "root" + password: str = Field(default="", repr=False) + default_file_quota: int = -1 + # `database` is intentionally optional — historically the database name is + # derived from the Milvus collection name (`partitions_for_collection_{collection}`) + # by the caller wiring the catalog store. The connection manager raises if + # this is still None at initialize() time. + database: str | None = None + pool_min_size: int = 5 + pool_max_size: int = 20 + command_timeout: int = 30 + + +# --------------------------------------------------------------------------- +# Ray — concurrency groups, serve config +# --------------------------------------------------------------------------- + + +class IndexerConcurrencyGroupsConfig(ConfigMixin): + default: int = 1000 + update: int = 100 + search: int = 100 + delete: int = 100 + serialize: int = 50 + chunk: int = 1000 + insert: int = 100 + + +class RayIndexerConfig(ConfigMixin): + max_task_retries: int = 2 + serialize_timeout: int = 3600 + vectordb_timeout: int = 30 + concurrency_groups: IndexerConcurrencyGroupsConfig = Field( + default_factory=IndexerConcurrencyGroupsConfig, + ) + + +class RaySemaphoreConfig(ConfigMixin): + concurrency: int = 100000 + + +class RayServeConfig(ConfigMixin): + enable: bool = False + num_replicas: int = 1 + host: str = "0.0.0.0" + port: int = 8080 + chainlit_port: int = 8090 + + +class RayConfig(ConfigMixin): + num_gpus: float = 0.01 + pool_size: int = 1 + max_tasks_per_worker: int = 8 + indexer: RayIndexerConfig = Field(default_factory=RayIndexerConfig) + semaphore: RaySemaphoreConfig = Field(default_factory=RaySemaphoreConfig) + serve: RayServeConfig = Field(default_factory=RayServeConfig) + + +# --------------------------------------------------------------------------- +# Paths +# --------------------------------------------------------------------------- + + +class PathsConfig(ConfigMixin): + prompts_dir: Path = Path("../prompts/example1") + data_dir: Path = Path("../data") + db_dir: Path = Path("/app/db") + log_dir: Path = Path("/app/logs") + + model_config = {**ConfigMixin.model_config, "arbitrary_types_allowed": True} + + +# --------------------------------------------------------------------------- +# Server +# --------------------------------------------------------------------------- + + +class ServerConfig(ConfigMixin): + preferred_url_scheme: str | None = None + + +# --------------------------------------------------------------------------- +# Verbose / logging +# --------------------------------------------------------------------------- + + +class VerboseConfig(ConfigMixin): + level: str = "DEBUG" + + +# --------------------------------------------------------------------------- +# Prompts (file name mapping) +# --------------------------------------------------------------------------- + + +class PromptsConfig(ConfigMixin): + sys_prompt: str = "sys_prompt_tmpl.txt" + query_contextualizer: str = "query_contextualizer_tmpl.txt" + chunk_contextualizer: str = "chunk_contextualizer_tmpl.txt" + image_describer: str = "image_captioning_tmpl.txt" + spoken_style_answer: str = "spoken_style_answer_tmpl.txt" + hyde: str = "hyde.txt" + multi_query: str = "multi_query_pmpt_tmpl.txt" diff --git a/openrag/core/config/loader.py b/openrag/core/config/loader.py new file mode 100644 index 000000000..030719519 --- /dev/null +++ b/openrag/core/config/loader.py @@ -0,0 +1,288 @@ +"""Configuration loader — reads YAML defaults, merges env var overrides, validates with Pydantic. + +Copied from config/loader.py. The original will be updated to re-export +from here for backward compatibility. +""" + +from __future__ import annotations + +import logging +import os +from pathlib import Path +from typing import Any + +import yaml + +from .root import Settings + +logger = logging.getLogger(__name__) + +_DEFAULT_CONF_DIR = Path(__file__).resolve().parent.parent.parent.parent / "conf" + +# --------------------------------------------------------------------------- +# Env var mappings: {env_var_name: dotted.config.path} +# --------------------------------------------------------------------------- +_ENV_OVERRIDES: list[tuple[str, str, type]] = [ + # LLM + ("BASE_URL", "llm.base_url", str), + ("MODEL", "llm.model", str), + ("API_KEY", "llm.api_key", str), + # VLM + ("VLM_BASE_URL", "vlm.base_url", str), + ("VLM_MODEL", "vlm.model", str), + ("VLM_API_KEY", "vlm.api_key", str), + # Semaphore + ("LLM_SEMAPHORE", "semaphore.llm_semaphore", int), + ("VLM_SEMAPHORE", "semaphore.vlm_semaphore", int), + # Embedder + ("EMBEDDER_MODEL_NAME", "embedder.model_name", str), + ("EMBEDDER_BASE_URL", "embedder.base_url", str), + ("EMBEDDER_API_KEY", "embedder.api_key", str), + ("MAX_MODEL_LEN", "embedder.max_model_len", int), + # VectorDB + ("VDB_HOST", "vectordb.host", str), + ("VDB_iPORT", "vectordb.port", int), + ("VDB_PORT", "vectordb.port", int), + ("VDB_CONNECTOR_NAME", "vectordb.connector_name", str), + ("VDB_COLLECTION_NAME", "vectordb.collection_name", str), + ("VDB_HYBRID_SEARCH", "vectordb.hybrid_search", bool), + ("VDB_ENABLE_INSERTION", "vectordb.enable", bool), + # RDB (Postgres) + ("POSTGRES_HOST", "rdb.host", str), + ("POSTGRES_PORT", "rdb.port", int), + ("POSTGRES_USER", "rdb.user", str), + ("POSTGRES_PASSWORD", "rdb.password", str), + ("DEFAULT_FILE_QUOTA", "rdb.default_file_quota", int), + ("POSTGRES_DATABASE", "rdb.database", str), + ("POSTGRES_POOL_MIN_SIZE", "rdb.pool_min_size", int), + ("POSTGRES_POOL_MAX_SIZE", "rdb.pool_max_size", int), + ("POSTGRES_COMMAND_TIMEOUT", "rdb.command_timeout", int), + # Reranker + ("RERANKER_PROVIDER", "reranker.provider", str), + ("RERANKER_ENABLED", "reranker.enabled", bool), + ("RERANKER_MODEL", "reranker.model_name", str), + ("RERANKER_TOP_K", "reranker.top_k", int), + ("RERANKER_BASE_URL", "reranker.base_url", str), + ("RERANKER_API_KEY", "reranker.api_key", str), + ("RERANKER_TIMEOUT", "reranker.timeout", float), + ("RERANKER_SEMAPHORE", "reranker.semaphore", int), + # Map-Reduce + ("MAP_REDUCE_INITIAL_BATCH_SIZE", "map_reduce.initial_batch_size", int), + ("MAP_REDUCE_EXPANSION_BATCH_SIZE", "map_reduce.expansion_batch_size", int), + ("MAP_REDUCE_MAX_TOTAL_DOCUMENTS", "map_reduce.max_total_documents", int), + ("MAP_REDUCE_DEBUG", "map_reduce.debug", bool), + # Verbose + ("LOG_LEVEL", "verbose.level", str), + # Server + ("PREFERRED_URL_SCHEME", "server.preferred_url_scheme", str), + # LLM Context + ("MAX_LLM_CONTEXT_SIZE", "llm_context.max_llm_context_size", int), + ("MAX_OUTPUT_TOKENS", "llm_context.max_output_tokens", int), + # Paths + ("PROMPTS_DIR", "paths.prompts_dir", str), + ("DATA_DIR", "paths.data_dir", str), + ("DB_DIR", "paths.db_dir", str), + ("LOG_DIR", "paths.log_dir", str), + # Loader + ("IMAGE_CAPTIONING", "loader.image_captioning", bool), + ("IMAGE_CAPTIONING_URL", "loader.image_captioning_url", bool), + ("SAVE_MARKDOWN", "loader.save_markdown", bool), + ("PDFLoader", "loader.file_loaders.pdf", str), + ("AUDIOLOADER", "loader.file_loaders.wav", str), + ("MARKER_MAX_TASKS_PER_CHILD", "loader.marker_max_tasks_per_child", int), + ("MARKER_POOL_SIZE", "loader.marker_pool_size", int), + ("MARKER_MAX_PROCESSES", "loader.marker_max_processes", int), + ("MARKER_NUM_GPUS", "loader.marker_num_gpus", float), + ("MARKER_TIMEOUT", "loader.marker_timeout", int), + ("MARKER_PDFTEXT_WORKERS", "loader.marker_pdftext_workers", int), + ("MARKER_CHUNK_SIZE", "loader.marker_chunk_size", int), + ("DOCLING_NUM_GPUS", "loader.docling_num_gpus", float), + ("DOCLING_POOL_SIZE", "loader.docling_pool_size", int), + ("DOCLING_MAX_TASKS_PER_WORKER", "loader.docling_max_tasks_per_worker", int), + ("WHISPER_MODEL", "loader.local_whisper.model", str), + ("WHISPER_N_WORKERS", "loader.local_whisper.whisper_n_workers", int), + ("WHISPER_NUM_GPUS", "loader.local_whisper.whisper_num_gpus", float), + ("WHISPER_CONCURRENCY_PER_WORKER", "loader.local_whisper.whisper_concurrency_per_worker", int), + ("TRANSCRIBER_BASE_URL", "loader.transcriber.base_url", str), + ("TRANSCRIBER_API_KEY", "loader.transcriber.api_key", str), + ("TRANSCRIBER_MODEL", "loader.transcriber.model_name", str), + ("TRANSCRIBER_TIMEOUT", "loader.transcriber.timeout", int), + ("TRANSCRIBER_MAX_CONCURRENT_CHUNKS", "loader.transcriber.max_concurrent_chunks", int), + ("TRANSCRIBER_DIRECT_UPLOAD_SUFFIXES", "loader.transcriber.direct_upload_suffixes", str), + ("USE_WHISPER_LANG_DETECTOR", "loader.transcriber.use_whisper_lang_detector", bool), + ("OPENAI_LOADER_BASE_URL", "loader.openai.base_url", str), + ("OPENAI_LOADER_API_KEY", "loader.openai.api_key", str), + ("OPENAI_LOADER_MODEL", "loader.openai.model", str), + ("OPENAI_LOADER_TEMPERATURE", "loader.openai.temperature", float), + ("OPENAI_LOADER_TIMEOUT", "loader.openai.timeout", int), + ("OPENAI_LOADER_MAX_RETRIES", "loader.openai.max_retries", int), + ("OPENAI_LOADER_TOP_P", "loader.openai.top_p", float), + ("OPENAI_LOADER_CONCURRENCY_LIMIT", "loader.openai.concurrency_limit", int), + # Ray + ("RAY_NUM_GPUS", "ray.num_gpus", float), + ("RAY_POOL_SIZE", "ray.pool_size", int), + ("RAY_MAX_TASKS_PER_WORKER", "ray.max_tasks_per_worker", int), + ("RAY_MAX_TASK_RETRIES", "ray.indexer.max_task_retries", int), + ("INDEXER_SERIALIZE_TIMEOUT", "ray.indexer.serialize_timeout", int), + ("VECTORDB_TIMEOUT", "ray.indexer.vectordb_timeout", int), + ("INDEXER_DEFAULT_CONCURRENCY", "ray.indexer.concurrency_groups.default", int), + ("INDEXER_UPDATE_CONCURRENCY", "ray.indexer.concurrency_groups.update", int), + ("INDEXER_SEARCH_CONCURRENCY", "ray.indexer.concurrency_groups.search", int), + ("INDEXER_DELETE_CONCURRENCY", "ray.indexer.concurrency_groups.delete", int), + ("INDEXER_SERIALIZE_CONCURRENCY", "ray.indexer.concurrency_groups.serialize", int), + ("INDEXER_CHUNK_CONCURRENCY", "ray.indexer.concurrency_groups.chunk", int), + ("INDEXER_INSERT_CONCURRENCY", "ray.indexer.concurrency_groups.insert", int), + ("RAY_SEMAPHORE_CONCURRENCY", "ray.semaphore.concurrency", int), + ("ENABLE_RAY_SERVE", "ray.serve.enable", bool), + ("RAY_SERVE_NUM_REPLICAS", "ray.serve.num_replicas", int), + ("RAY_SERVE_HOST", "ray.serve.host", str), + ("RAY_SERVE_PORT", "ray.serve.port", int), + ("CHAINLIT_PORT", "ray.serve.chainlit_port", int), + # Chunker + ("CHUNKER", "chunker.name", str), + ("CONTEXTUAL_RETRIEVAL", "chunker.contextual_retrieval", bool), + ("CONTEXTUALIZATION_TIMEOUT", "chunker.contextualization_timeout", int), + ("MAX_CONCURRENT_CONTEXTUALIZATION", "chunker.max_concurrent_contextualization", int), + ("CHUNK_SIZE", "chunker.chunk_size", int), + ("CHUNK_OVERLAP_RATE", "chunker.chunk_overlap_rate", float), + # Retriever + ("RETRIEVER_TYPE", "retriever.type", str), + ("RETRIEVER_TOP_K", "retriever.top_k", int), + ("SIMILARITY_THRESHOLD", "retriever.similarity_threshold", float), + ("WITH_SURROUNDING_CHUNKS", "retriever.with_surrounding_chunks", bool), + ("INCLUDE_RELATED", "retriever.include_related", bool), + ("INCLUDE_ANCESTORS", "retriever.include_ancestors", bool), + ("RELATED_LIMIT", "retriever.related_limit", int), + ("MAX_DEPTH", "retriever.max_ancestor_depth", int), + ("RETRIEVER_ALLOW_FILTERLESS_FALLBACK", "retriever.allow_filterless_fallback", bool), + # RAG + ("RAG_MODE", "rag.mode", str), + # WebSearch + ("WEBSEARCH_PROVIDER", "websearch.provider", str), + ("WEBSEARCH_API_TOKEN", "websearch.api_token", str), + ("WEBSEARCH_BASE_URL", "websearch.base_url", str), + ("WEBSEARCH_TOP_K", "websearch.top_k", int), + ("WEBSEARCH_LANG", "websearch.lang", str), + ("WEBSEARCH_MAX_TOKENS", "websearch.max_tokens", int), + ("WEBSEARCH_FETCH_CONTENT", "websearch.fetch_content", bool), + ("WEBSEARCH_FETCH_MAX_RESULTS", "websearch.fetch_max_results", int), + ("WEBSEARCH_FETCH_TIMEOUT", "websearch.fetch_timeout", float), + ("WEBSEARCH_FETCH_MAX_TOKENS", "websearch.fetch_max_tokens", int), + ("WEBSEARCH_FETCH_VERIFY_SSL", "websearch.fetch_verify_ssl", bool), +] + +_AUDIO_EXTENSIONS = ("mp3", "flac", "ogg", "aac", "flv", "wma", "mp4") + + +def _load_yaml(path: Path) -> dict[str, Any]: + """Load a YAML file, returning empty dict if not found.""" + if not path.exists(): + logger.warning("Config file not found: %s — using defaults", path) + return {} + with open(path) as f: + data = yaml.safe_load(f) + return data or {} + + +def _deep_merge(base: dict, override: dict) -> dict: + """Recursively merge override into base.""" + merged = base.copy() + for key, value in override.items(): + if key in merged and isinstance(merged[key], dict) and isinstance(value, dict): + merged[key] = _deep_merge(merged[key], value) + else: + merged[key] = value + return merged + + +def _set_nested(data: dict, dotted_path: str, value: Any) -> None: + """Set a value in a nested dict using a dotted path.""" + keys = dotted_path.split(".") + current = data + for key in keys[:-1]: + current = current.setdefault(key, {}) + current[keys[-1]] = value + + +def _coerce(value: str, target_type: type, env_var: str = "") -> Any: + """Coerce a string env var value to the target type.""" + if target_type is bool: + lower = value.lower() + if lower in ("true", "1", "yes"): + return True + if lower in ("false", "0", "no"): + return False + raise ValueError(f"Invalid value for {env_var}: expected bool, got {value!r}") + try: + if target_type is int: + return int(value) + if target_type is float: + return float(value) + except ValueError: + raise ValueError(f"Invalid value for {env_var}: expected {target_type.__name__}, got {value!r}") + return value + + +def _apply_env_overrides(data: dict) -> dict: + """Apply environment variable overrides to the config dict.""" + for env_var, dotted_path, target_type in _ENV_OVERRIDES: + value = os.environ.get(env_var) + if value is not None and value != "": + _set_nested(data, dotted_path, _coerce(value, target_type, env_var)) + + semaphore = os.environ.get("SEMAPHORE") + if semaphore: + sem_value = _coerce(semaphore, int, "SEMAPHORE") + sem = data.setdefault("semaphore", {}) + sem.setdefault("llm_semaphore", sem_value) + sem.setdefault("vlm_semaphore", sem_value) + + audio_loader = os.environ.get("AUDIOLOADER") + if audio_loader: + file_loaders = data.setdefault("loader", {}).setdefault("file_loaders", {}) + for ext in _AUDIO_EXTENSIONS: + file_loaders[ext] = audio_loader + + return data + + +def load_config( + conf_dir: Path | str | None = None, + overrides: dict[str, Any] | None = None, +) -> Settings: + """Load configuration: YAML defaults -> env var overrides -> Pydantic validation. + + Args: + conf_dir: Path to the configuration directory. Defaults to ``conf/`` + at the project root, overridable via ``OPENRAG_CONF_DIR``. + overrides: Programmatic overrides (useful for tests). + """ + from dotenv import load_dotenv + + load_dotenv() + + env_conf_dir = os.environ.get("OPENRAG_CONF_DIR") + if conf_dir: + conf_dir = Path(conf_dir) + elif env_conf_dir: + conf_dir = Path(env_conf_dir) + else: + conf_dir = _DEFAULT_CONF_DIR + + data = _load_yaml(conf_dir / "config.yaml") + data = {k: v for k, v in data.items() if not k.startswith("_")} + data = _apply_env_overrides(data) + + reranker = data.get("reranker") + if isinstance(reranker, dict) and not reranker.get("base_url"): + reranker.pop("base_url", None) + + if overrides: + data = _deep_merge(data, overrides) + + paths = data.get("paths", {}) + for key in ("prompts_dir", "data_dir", "db_dir", "log_dir"): + if key in paths and paths[key]: + paths[key] = str(Path(paths[key]).resolve()) + + return Settings(**data) diff --git a/openrag/core/config/retrieval.py b/openrag/core/config/retrieval.py new file mode 100644 index 000000000..29f9f8fe1 --- /dev/null +++ b/openrag/core/config/retrieval.py @@ -0,0 +1,130 @@ +"""Retrieval, reranker, RAG mode, map-reduce, and web search configuration.""" + +from __future__ import annotations + +from typing import Annotated, Literal + +from pydantic import Field + +from .base import ConfigMixin + +# --------------------------------------------------------------------------- +# Reranker +# --------------------------------------------------------------------------- + + +class _BaseRerankerConfig(ConfigMixin): + model_name: str = "Alibaba-NLP/gte-multilingual-reranker-base" + top_k: int = 10 + api_key: str = Field(default="EMPTY", repr=False) + timeout: float = 60.0 + semaphore: int = 5 + enabled: bool = True + + +class InfinityRerankerConfig(_BaseRerankerConfig): + provider: Literal["infinity"] = "infinity" + base_url: str = "http://reranker:7997" + + +class OpenAIRerankerConfig(_BaseRerankerConfig): + provider: Literal["openai"] = "openai" + base_url: str = "http://reranker:8000/v1" + + +RerankerConfig = Annotated[ + InfinityRerankerConfig | OpenAIRerankerConfig, + Field(discriminator="provider"), +] + + +def _default_reranker_config() -> InfinityRerankerConfig: + return InfinityRerankerConfig() + + +# --------------------------------------------------------------------------- +# Retriever +# --------------------------------------------------------------------------- + + +class _BaseRetrieverConfig(ConfigMixin): + top_k: int = 50 + similarity_threshold: float = 0.6 + with_surrounding_chunks: bool = False + include_related: bool = True + include_ancestors: bool = True + related_limit: int = 10 + max_ancestor_depth: int = 10 + allow_filterless_fallback: bool = True + + +class SingleRetrieverConfig(_BaseRetrieverConfig): + type: Literal["single"] = "single" + + +class MultiQueryRetrieverConfig(_BaseRetrieverConfig): + type: Literal["multiQuery"] = "multiQuery" + k_queries: int = 3 + + +class HydeRetrieverConfig(_BaseRetrieverConfig): + type: Literal["hyde"] = "hyde" + combine: bool = False + + +RetrieverConfig = Annotated[ + SingleRetrieverConfig | MultiQueryRetrieverConfig | HydeRetrieverConfig, + Field(discriminator="type"), +] + + +# --------------------------------------------------------------------------- +# RAG +# --------------------------------------------------------------------------- + + +class RAGConfig(ConfigMixin): + mode: str = "ChatBotRag" + chat_history_depth: int = 4 + max_contextualized_query_len: int = 512 + + +# --------------------------------------------------------------------------- +# Map-Reduce +# --------------------------------------------------------------------------- + + +class MapReduceConfig(ConfigMixin): + initial_batch_size: int = 10 + expansion_batch_size: int = 5 + max_total_documents: int = 20 + debug: bool = False + + +# --------------------------------------------------------------------------- +# WebSearch +# --------------------------------------------------------------------------- + + +class _BaseWebSearchConfig(ConfigMixin): + base_url: str + api_token: str = Field(default="", repr=False) + top_k: int = 5 + lang: str = "fr-FR" + max_tokens: int = 2000 + fetch_content: bool = True + fetch_max_results: int = 3 + fetch_timeout: float = 1.0 + fetch_max_tokens: int = 500 + fetch_verify_ssl: bool = False + + +class StaanWebSearchConfig(_BaseWebSearchConfig): + provider: Literal["staan"] = "staan" + base_url: str = "https://api.staan.ai/search/web" + + +WebSearchConfig = Annotated[ + StaanWebSearchConfig, + Field(discriminator="provider"), +] diff --git a/openrag/core/config/root.py b/openrag/core/config/root.py new file mode 100644 index 000000000..cbb219a43 --- /dev/null +++ b/openrag/core/config/root.py @@ -0,0 +1,63 @@ +"""Root configuration — composes all sub-models into a single Settings object.""" + +from __future__ import annotations + +from pydantic import Field + +from .base import ConfigMixin +from .chunking import ChunkerConfig +from .endpoints import ( + EmbedderConfig, + LLMConfig, + LLMContextConfig, + SemaphoreConfig, + VLMConfig, +) +from .indexation import LoaderConfig +from .infrastructure import ( + PathsConfig, + PromptsConfig, + RayConfig, + RDBConfig, + ServerConfig, + VectorDBConfig, + VerboseConfig, +) +from .retrieval import ( + MapReduceConfig, + RAGConfig, + RerankerConfig, + RetrieverConfig, + SingleRetrieverConfig, + StaanWebSearchConfig, + WebSearchConfig, + _default_reranker_config, +) + + +class Settings(ConfigMixin): + """Root configuration. + + Defaults here are fallbacks only. In production, values come from + conf/config.yaml merged with environment variable overrides. + """ + + llm: LLMConfig = Field(default_factory=LLMConfig) + vlm: VLMConfig = Field(default_factory=VLMConfig) + semaphore: SemaphoreConfig = Field(default_factory=SemaphoreConfig) + embedder: EmbedderConfig = Field(default_factory=EmbedderConfig) + vectordb: VectorDBConfig = Field(default_factory=VectorDBConfig) + rdb: RDBConfig = Field(default_factory=RDBConfig) + reranker: RerankerConfig = Field(default_factory=_default_reranker_config) + map_reduce: MapReduceConfig = Field(default_factory=MapReduceConfig) + verbose: VerboseConfig = Field(default_factory=VerboseConfig) + server: ServerConfig = Field(default_factory=ServerConfig) + llm_context: LLMContextConfig = Field(default_factory=LLMContextConfig) + paths: PathsConfig = Field(default_factory=PathsConfig) + prompts: PromptsConfig = Field(default_factory=PromptsConfig) + loader: LoaderConfig = Field(default_factory=LoaderConfig) + ray: RayConfig = Field(default_factory=RayConfig) + chunker: ChunkerConfig = Field(default_factory=ChunkerConfig) + retriever: RetrieverConfig = Field(default_factory=SingleRetrieverConfig) + rag: RAGConfig = Field(default_factory=RAGConfig) + websearch: WebSearchConfig = Field(default_factory=StaanWebSearchConfig) diff --git a/openrag/core/config/test_indexation.py b/openrag/core/config/test_indexation.py new file mode 100644 index 000000000..39edf447d --- /dev/null +++ b/openrag/core/config/test_indexation.py @@ -0,0 +1,31 @@ +"""Tests for indexation config — TranscriberConfig pipe-string parsing.""" + +from __future__ import annotations + +from core.config.indexation import ( + _DEFAULT_DIRECT_UPLOAD_SUFFIXES, + TranscriberConfig, +) + + +def test_transcriber_config_default_direct_upload_suffixes(): + cfg = TranscriberConfig() + assert cfg.direct_upload_suffixes == set(_DEFAULT_DIRECT_UPLOAD_SUFFIXES) + + +def test_transcriber_config_parses_pipe_delimited_string(): + """The YAML default and TRANSCRIBER_DIRECT_UPLOAD_SUFFIXES env var both + arrive as a pipe-delimited string. The validator must split + normalize + into a set of dot-prefixed lowercase suffixes.""" + cfg = TranscriberConfig(direct_upload_suffixes=".wav|FLAC|mp3") + assert cfg.direct_upload_suffixes == {".wav", ".flac", ".mp3"} + + +def test_transcriber_config_drops_empty_components(): + cfg = TranscriberConfig(direct_upload_suffixes="|.wav||.mp3|") + assert cfg.direct_upload_suffixes == {".wav", ".mp3"} + + +def test_transcriber_config_set_input_passes_through(): + cfg = TranscriberConfig(direct_upload_suffixes={".wav", ".m4a"}) + assert cfg.direct_upload_suffixes == {".wav", ".m4a"} diff --git a/openrag/core/embeddings/__init__.py b/openrag/core/embeddings/__init__.py new file mode 100644 index 000000000..237cd0d8a --- /dev/null +++ b/openrag/core/embeddings/__init__.py @@ -0,0 +1,6 @@ +"""Embedder ABC + registry.""" + +from .embedder import Embedder +from .registry import embedder_registry + +__all__ = ["Embedder", "embedder_registry"] diff --git a/openrag/core/embeddings/embedder.py b/openrag/core/embeddings/embedder.py new file mode 100644 index 000000000..1ac34aa05 --- /dev/null +++ b/openrag/core/embeddings/embedder.py @@ -0,0 +1,25 @@ +"""Abstract embedder interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class Embedder(ABC): + """Base class for all embedding providers.""" + + @abstractmethod + async def embed(self, texts: list[str]) -> list[list[float]]: + """Embed a batch of texts, returning vectors.""" + ... + + @abstractmethod + async def embed_single(self, text: str) -> list[float]: + """Embed a single text.""" + ... + + @property + @abstractmethod + def dimension(self) -> int: + """Return the embedding dimension.""" + ... diff --git a/openrag/core/embeddings/registry.py b/openrag/core/embeddings/registry.py new file mode 100644 index 000000000..eddee9387 --- /dev/null +++ b/openrag/core/embeddings/registry.py @@ -0,0 +1,7 @@ +"""Embedder registry.""" + +from openrag.core.utils.registry import Registry + +from .embedder import Embedder + +embedder_registry: Registry[Embedder] = Registry("embedder") diff --git a/openrag/core/indexing/__init__.py b/openrag/core/indexing/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/core/indexing/contextualize.py b/openrag/core/indexing/contextualize.py new file mode 100644 index 000000000..6b60b7c22 --- /dev/null +++ b/openrag/core/indexing/contextualize.py @@ -0,0 +1,138 @@ +"""Chunk contextualization against the ``LLM`` ABC. + +Framework-free implementation of contextual retrieval: for each chunk, +ask an LLM to write a short situating context based on the document's +opening chunks plus the immediate preceding neighbourhood, then prepend +that context to the chunk text so embeddings capture document-level +meaning. + +Inputs and outputs are :class:`core.models.chunk.Chunk` instances. The +caller supplies the LLM, the system prompt, and any concurrency / timeout +limits — core does not reach into Hydra config or the global VLM +semaphore. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Sequence + +from tqdm.asyncio import tqdm + +from ..llm import LLM +from ..models.chunk import Chunk +from ..prompts.contextualization_builder import build_messages, wrap_chunk_with_context + +logger = logging.getLogger(__name__) + +DEFAULT_TIMEOUT_SECONDS = 30.0 +DEFAULT_MAX_CONCURRENT = 4 + + +class ChunkContextualizer: + """Generate a per-chunk context string and prepend it to the chunk text..""" + + def __init__( + self, + llm: LLM, + system_prompt: str, + *, + timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS, + max_concurrent: int = DEFAULT_MAX_CONCURRENT, + semaphore: asyncio.Semaphore | None = None, + ): + self._llm = llm + self._system_prompt = system_prompt + self._timeout = timeout_seconds + self._batch_size = max(1, max_concurrent) + self._semaphore = semaphore or asyncio.Semaphore(self._batch_size) + + async def _generate_context( + self, + first_chunks: Sequence[Chunk], + prev_chunks: Sequence[Chunk], + current_chunk: Chunk, + filename: str, + lang: str, + ) -> str: + messages = build_messages( + system_prompt=self._system_prompt, + filename=filename, + first_chunks_text=[c.text for c in first_chunks], + prev_chunks_text=[c.text for c in prev_chunks], + current_chunk_text=current_chunk.text, + lang=lang, + ) + async with self._semaphore: + try: + return await asyncio.wait_for(self._llm.chat(messages), timeout=self._timeout) + except TimeoutError: + logger.warning("LLM timeout contextualizing chunk (filename=%s)", filename) + return "" + except Exception as exc: + logger.warning("Error contextualizing chunk (filename=%s): %s", filename, exc) + return "" + + async def contextualize( + self, + chunks: Sequence[Chunk], + *, + filename: str = "", + lang: str = "en", + ) -> list[Chunk]: + """Return new chunks with context prepended to ``text``. + + Each returned chunk preserves the input's id, metadata, and other + fields; ``text`` is rewritten to the formatted (context + content) + string used for embedding, ``context`` holds the generated context, + and ``content`` holds the original chunk text. + + Falls back to returning the input chunks unchanged on any + unrecoverable error. + """ + chunks = list(chunks) + if not chunks: + return [] + + try: + first_chunks = chunks[:2] + contexts: list[str] = [] + # Schedule one batch at a time so prompt strings + coroutine + # objects don't all sit in memory upfront on large documents. + for start in range(0, len(chunks), self._batch_size): + end = min(start + self._batch_size, len(chunks)) + batch = [ + self._generate_context( + first_chunks=first_chunks, + prev_chunks=chunks[max(0, i - 2) : i] if i > 0 else [], + current_chunk=chunks[i], + filename=filename, + lang=lang, + ) + for i in range(start, end) + ] + contexts.extend( + await tqdm.gather( + *batch, + desc=f"Contextualizing chunks of *{filename}* [{start + 1}-{end}/{len(chunks)}]", + ) + ) + + return [ + chunk.model_copy( + update={ + "text": wrap_chunk_with_context( + content=chunk.text, + filename=filename, + chunk_context=context, + ), + "context": context, + "content": chunk.text, + } + ) + for chunk, context in zip(chunks, contexts, strict=True) + ] + except (TimeoutError, OSError, RuntimeError, ValueError) as exc: + logger.warning("Error contextualizing chunks from %s: %s", filename, exc) + return chunks diff --git a/openrag/core/indexing/dispatcher.py b/openrag/core/indexing/dispatcher.py new file mode 100644 index 000000000..e8cbb7b17 --- /dev/null +++ b/openrag/core/indexing/dispatcher.py @@ -0,0 +1,87 @@ +"""Transitional port for the indexing dispatch operations. + +``IndexingService`` (Phase 8D.1) owns the business logic around file +ingestion — format/quota/existence checks, metadata assembly, workspace +validation — but the heavy lifting (serialize → chunk → embed → insert) +still runs inside the ``Indexer`` Ray actor, and task bookkeeping lives +in the ``TaskStateManager`` Ray actor. Phase 9 removes that Ray +indirection. + +Defining the operations the service needs on a dedicated port keeps +``IndexingService`` Ray-free (8H: no Ray import / remote call under +``services/orchestrators/``). A small shim in ``services/storage/`` +adapts the two Ray actors to this interface during the shim period; +Phase 9 swaps it for a direct pipeline call and deletes the shim. + +No Ray / pymilvus / LangChain types leak across this boundary — only +plain dicts and strings. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class IndexingDispatcher(ABC): + """Operations the indexing orchestrator needs from the worker layer.""" + + @abstractmethod + async def dispatch_indexing( + self, + *, + path: str, + metadata: dict, + partition: str, + user: dict | None, + workspace_ids: list[str] | None, + replace: bool, + ) -> str: + """Queue an (re)indexing job, register its task state, return its id.""" + ... + + @abstractmethod + async def delete_file(self, file_id: str, partition: str) -> None: + """Delete a file's chunks via the worker layer.""" + ... + + @abstractmethod + async def update_file_metadata( + self, + file_id: str, + metadata: dict, + partition: str, + user: dict | None, + ) -> None: + """Upsert file metadata in place (no re-embedding).""" + ... + + @abstractmethod + async def copy_file( + self, + file_id: str, + metadata: dict, + partition: str, + user: dict | None, + ) -> None: + """Copy a file's chunks into another partition / file id.""" + ... + + @abstractmethod + async def get_task_state(self, task_id: str) -> str | None: + """Current task state, or ``None`` if the task is unknown.""" + ... + + @abstractmethod + async def get_task_error(self, task_id: str) -> str | None: + """Stored traceback for a failed task, or ``None``.""" + ... + + @abstractmethod + async def cancel_task(self, task_id: str) -> bool: + """Cancel a running/queued task. + + Returns ``False`` when no object ref is stored for ``task_id`` + (the caller maps that to a 404), ``True`` once the cancel signal + has been sent. + """ + ... diff --git a/openrag/core/indexing/image_preprocessor.py b/openrag/core/indexing/image_preprocessor.py new file mode 100644 index 000000000..50cead047 --- /dev/null +++ b/openrag/core/indexing/image_preprocessor.py @@ -0,0 +1,122 @@ +"""Image preprocessing helpers for the indexing pipeline. + +Pure helpers — no VLM, no langchain, no infrastructure imports. Used by +parsers (core) and Ray-pool adapters (services) that need to: + +- normalize PIL Image modes for PNG encoding +- encode PIL Images as PNG bytes or base64 data URIs +- detect / decode markdown image references (HTTP / data URI) in extracted text + +Extracted from the legacy ``components/indexer/loaders/base.py``; the +legacy module is kept as a back-compat shim until existing imports are +migrated. +""" + +from __future__ import annotations + +import base64 +import logging +import re +from io import BytesIO +from typing import Any + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Markdown image-reference patterns (compile once; shared regex objects) +# --------------------------------------------------------------------------- + +HTTP_IMAGE_PATTERN = re.compile(r"!\[(.*?)\]\((https?://[^)]+)\)") +DATA_URI_IMAGE_PATTERN = re.compile(r"!\[(.*?)\]\((data:image/[^;]+;base64,[^)]+)\)") + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +# Qwen2.5-VL ``min_pixels`` threshold; images below this break the model. +MIN_IMAGE_PIXELS = 784 + + +# --------------------------------------------------------------------------- +# PIL mode normalization & encoding +# --------------------------------------------------------------------------- + + +def ensure_png_compatible_mode(image: Any) -> Any: + """Convert PIL image modes that PNG can't encode directly. + + CMYK/YCbCr/LAB → RGB; P/LA/PA → RGBA. Others returned unchanged. + """ + if image.mode in ("CMYK", "YCbCr", "LAB"): + return image.convert("RGB") + if image.mode in ("P", "LA", "PA"): + return image.convert("RGBA") + return image + + +def pil_to_png_bytes(image: Any) -> bytes: + """Encode a PIL Image as PNG bytes. ``bytes`` input is passed through.""" + if isinstance(image, bytes): + return image + image = ensure_png_compatible_mode(image) + buf = BytesIO() + image.save(buf, format="PNG") + return buf.getvalue() + + +# --------------------------------------------------------------------------- +# URL / data URI detection +# --------------------------------------------------------------------------- + + +def decode_data_uri(data_uri: str) -> bytes | None: + """Decode a ``data:image/...;base64,...`` URI into raw bytes. ``None`` on failure.""" + try: + _, b64 = data_uri.split(",", 1) + return base64.b64decode(b64) + except Exception as exc: + logger.warning("Failed to decode data URI: %s", exc) + return None + + +def mime_from_data_uri(data_uri: str) -> str: + """Pull the mime type out of a data URI; fall back to ``image/png``. + + Example: ``data:image/jpeg;base64,xxx`` → ``image/jpeg``. + """ + try: + return data_uri.split(",", 1)[0].split(":", 1)[1].split(";", 1)[0] + except Exception: + return "image/png" + + +def extract_data_uri_image_blocks(text: str, *, page_number: int = 1) -> list[Any]: + """Build ``ImageBlock``s for every ``![alt](data:image/...;base64,...)`` ref. + + The original markdown ref is preserved in ``metadata['markdown_ref']`` + so a downstream caption stage can substitute the wrapped caption back + into the corresponding ``TextBlock`` via ``str.replace``. + + Returns ``list[ImageBlock]`` (declared as ``list[Any]`` only because + importing the model would create a cycle in some build orderings — + the caller side is type-correct). + """ + if not text: + return [] + # Local import to avoid a top-level cycle with ``core.models``. + from ..models.document import ImageBlock + + blocks: list[Any] = [] + for alt, data_uri in DATA_URI_IMAGE_PATTERN.findall(text): + payload = decode_data_uri(data_uri) + if payload is None: + continue + blocks.append( + ImageBlock( + image_bytes=payload, + page_number=page_number, + mime_type=mime_from_data_uri(data_uri), + metadata={"markdown_ref": f"![{alt}]({data_uri})", "alt": alt}, + ) + ) + return blocks diff --git a/openrag/core/indexing/parsers/__init__.py b/openrag/core/indexing/parsers/__init__.py new file mode 100644 index 000000000..6025ec27e --- /dev/null +++ b/openrag/core/indexing/parsers/__init__.py @@ -0,0 +1,6 @@ +"""DocumentParser ABC + registry.""" + +from .document_parser import DocumentParser +from .registry import parser_registry + +__all__ = ["DocumentParser", "parser_registry"] diff --git a/openrag/core/indexing/parsers/audio/__init__.py b/openrag/core/indexing/parsers/audio/__init__.py new file mode 100644 index 000000000..7604a3d84 --- /dev/null +++ b/openrag/core/indexing/parsers/audio/__init__.py @@ -0,0 +1,12 @@ +"""Audio parser facades. + +Each backend lives in its own module so its heavy dependencies (Ray +worker pools, cloud SDKs, …) are only pulled in by the concrete impl +in ``services/`` — the core facade just declares the parser type and +delegates ``parse()`` to an injected pool/client. +""" + +from .client_based import ClientAudioParser +from .local_whisper import LocalWhisperParser + +__all__ = ["ClientAudioParser", "LocalWhisperParser"] diff --git a/openrag/core/indexing/parsers/audio/client_based.py b/openrag/core/indexing/parsers/audio/client_based.py new file mode 100644 index 000000000..4718ba5c6 --- /dev/null +++ b/openrag/core/indexing/parsers/audio/client_based.py @@ -0,0 +1,32 @@ +"""Client-backed audio ``DocumentParser`` (thin core facade). + +Holds a ``BaseClientParser`` (the actual HTTP-client / OpenAI-SDK +implementation lives in ``services/`` and is composed in at startup) +and delegates ``parse()`` to it. + +Mirrors the :class:`ClientPdfParser` pattern: core stays free of vendor +SDKs while the facade names "client-backed audio" as a first-class +parser type. +""" + +from __future__ import annotations + +from ....models.document import Document, ProcessedDocument +from ..document_parser import BaseClientParser, DocumentParser +from ..registry import parser_registry + + +@parser_registry.register("audio_client") +class ClientAudioParser(DocumentParser): + """Public audio parser facade backed by an OpenAI-compatible transcription client.""" + + def __init__(self, client: BaseClientParser) -> None: + if not isinstance(client, BaseClientParser): + raise ValueError("ClientAudioParser requires a BaseClientParser instance as client") + self._client = client + + def supported_types(self) -> list[str]: + return self._client.supported_types() + + async def parse(self, document: Document) -> ProcessedDocument: + return await self._client.parse(document) diff --git a/openrag/core/indexing/parsers/audio/local_whisper.py b/openrag/core/indexing/parsers/audio/local_whisper.py new file mode 100644 index 000000000..c880488ee --- /dev/null +++ b/openrag/core/indexing/parsers/audio/local_whisper.py @@ -0,0 +1,34 @@ +"""Local Whisper-backed audio ``DocumentParser`` (thin core facade). + +Holds a reference to a ``BasePooledParser`` (the actual Ray-pool / +GPU-model implementation lives in ``services/`` and is not yet wired +up) and delegates ``parse()`` to it. The split keeps core free of Ray +and GPU lifecycle code while still naming the local-Whisper backend as +a first-class parser type. + +The injected pool is a generic ``BasePooledParser``; if a more specific +``WhisperPoolParser`` ABC emerges in services, this class can tighten +its type without changing call sites. +""" + +from __future__ import annotations + +from ....models.document import Document, DocumentType, ProcessedDocument +from ..document_parser import BasePooledParser, DocumentParser +from ..registry import parser_registry + + +@parser_registry.register("local_whisper") +class LocalWhisperParser(DocumentParser): + """Public audio parser facade backed by a local-Whisper worker pool.""" + + def __init__(self, pool: BasePooledParser) -> None: + if not isinstance(pool, BasePooledParser): + raise ValueError("LocalWhisperParser requires a BasePooledParser instance as pool") + self._pool = pool + + def supported_types(self) -> list[str]: + return [DocumentType.AUDIO.value, DocumentType.VIDEO.value] + + async def parse(self, document: Document) -> ProcessedDocument: + return await self._pool.parse(document) diff --git a/openrag/core/indexing/parsers/doc_parser.py b/openrag/core/indexing/parsers/doc_parser.py new file mode 100644 index 000000000..c7371a593 --- /dev/null +++ b/openrag/core/indexing/parsers/doc_parser.py @@ -0,0 +1,104 @@ +"""Legacy ``.doc`` (binary Word 97-2003) ``DocumentParser``. + +Converts ``.doc`` to ``.docx`` via the ``spire.doc`` library, then +delegates to :class:`DocxParser` for Markdown extraction. Falls back to +plain-text extraction (``Document.GetText()``) if Spire's conversion +fails. + +Spire.Doc requires DOTNET; the constructor sets the invariant-globalization +env var that makes Spire usable without the full ICU data. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import tempfile +from pathlib import Path + +from ...models.document import Document, DocumentType, ProcessedDocument, TextBlock +from .document_parser import DocumentParser +from .docx_parser import DocxParser +from .registry import parser_registry + +logger = logging.getLogger(__name__) + +os.environ.setdefault("DOTNET_SYSTEM_GLOBALIZATION_INVARIANT", "1") + + +@parser_registry.register("doc") +class DocParser(DocumentParser): + """Parse legacy ``.doc`` files via .docx conversion + DocxParser.""" + + def __init__(self, docx_parser: DocxParser | None = None) -> None: + """Pass an explicit ``DocxParser`` to share VLM / semaphore config; + otherwise a captioning-disabled instance is constructed. + """ + self._docx = docx_parser or DocxParser() + + def supported_types(self) -> list[str]: + return [DocumentType.DOC.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + async with document.as_temporary_file() as src_path: + docx_bytes, fallback_text = await asyncio.to_thread(self._convert, str(src_path)) + + if docx_bytes: + docx_doc = document.model_copy(update={"raw_bytes": docx_bytes, "content_type": DocumentType.DOCX}) + return await self._docx.parse(docx_doc) + + text = (fallback_text or "").strip() + text_blocks = [TextBlock(text=text, page_number=1)] if text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + metadata=dict(document.metadata), + page_count=1 if text else 0, + ) + + @staticmethod + def _convert(path: str) -> tuple[bytes | None, str | None]: + """Run blocking Spire.Doc conversion. Returns ``(docx_bytes, fallback_text)``. + + Exactly one of the two will be non-None on success; both ``None`` + means total failure (caller emits an empty ProcessedDocument). + """ + try: + from spire.doc import Document as SpireDocument + from spire.doc import FileFormat + except ImportError: + logger.warning("spire.doc not available; cannot parse legacy .doc files") + return None, None + + spire_doc = SpireDocument() + out_path: str | None = None + try: + spire_doc.LoadFromFile(path) + with tempfile.NamedTemporaryFile(suffix=".docx", delete=False) as out: + out_path = out.name + spire_doc.SaveToFile(out_path, FileFormat.Docx2016) + return Path(out_path).read_bytes(), None + except Exception as exc: + logger.warning("Spire.Doc .doc → .docx conversion failed (%s); falling back to plain text", exc) + try: + return None, spire_doc.GetText() + except Exception as fallback_exc: + logger.warning("Spire.Doc fallback text extraction also failed: %s", fallback_exc) + return None, None + finally: + try: + spire_doc.Close() + except Exception: + pass + if out_path and os.path.exists(out_path): + try: + os.remove(out_path) + except OSError: + pass diff --git a/openrag/core/indexing/parsers/document_parser.py b/openrag/core/indexing/parsers/document_parser.py new file mode 100644 index 000000000..d69d12c8b --- /dev/null +++ b/openrag/core/indexing/parsers/document_parser.py @@ -0,0 +1,51 @@ +"""Abstract document parser interface and category markers. + +``DocumentParser`` is the single port every concrete parser implements. + +Two empty subclasses are exposed alongside it as **type markers** — +they categorize a parser by *how* it gets its work done, without adding +any behaviour: + +- ``BasePooledParser`` — a parser whose ``parse()`` is satisfied by a + pool of workers (Ray actors, ProcessPoolExecutor, asyncio task group, + …). Concrete impls live in ``services/``. +- ``BaseClientParser`` — a parser whose ``parse()`` is satisfied by an + external client (HTTP service, gRPC, vendor SDK, …). Concrete impls + live in ``services/``. + +The markers exist so consumers (e.g. ``MarkerParser(pool: BasePooledParser)``) +can constrain the *kind* of parser they accept tighter than the +generic ``DocumentParser``. Concrete subclasses implement +``parse()`` and ``supported_types()`` directly — no extra hook method. + +If a shared pattern (retry, timeout, semaphore, …) ever materialises +across multiple subclasses, lift it into the base then. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from ...models.document import Document, ProcessedDocument + + +class DocumentParser(ABC): + """Base class for all document parsers (PDF, text, HTML, image, audio, etc.).""" + + @abstractmethod + async def parse(self, document: Document) -> ProcessedDocument: + """Parse a document into text blocks and images.""" + ... + + @abstractmethod + def supported_types(self) -> list[str]: + """Return list of DocumentType values this parser handles.""" + ... + + +class BasePooledParser(DocumentParser, ABC): + """Marker for parsers backed by a worker pool. Concrete impl in services/.""" + + +class BaseClientParser(DocumentParser, ABC): + """Marker for parsers backed by an external client. Concrete impl in services/.""" diff --git a/openrag/core/indexing/parsers/docx_parser.py b/openrag/core/indexing/parsers/docx_parser.py new file mode 100644 index 000000000..a1702eb7c --- /dev/null +++ b/openrag/core/indexing/parsers/docx_parser.py @@ -0,0 +1,200 @@ +"""DOCX ``DocumentParser`` implementation. + +Conversion to Markdown via the ``markitdown`` library; falls back to +plain-text extraction via ``python-docx`` if MarkItDown fails. + +MarkItDown emits a generic ``![](data:image/png;base64...)`` placeholder +(literal, truncated) for every embedded image, with no per-image +identifier. Actual image bytes are pulled from the DOCX zip +(``word/media/``) and matched to placeholders **positionally**, in +document order. Each placeholder is rewritten to a unique synthetic +``![](docx-image-N)`` ref, and the matching :class:`ImageBlock` stores +that ref in ``metadata['markdown_ref']`` for downstream caption +substitution. + +Captioning is not done here — see :class:`ImageBlock` for the +parser→caption contract. + +Output: +- A single ``TextBlock`` containing the rewritten Markdown. +- One ``ImageBlock`` per embedded zip image; ``caption=None``. + +Failures fall back gracefully: missing libraries or malformed zips +degrade to leaving content untouched rather than raising. Images that +can't be decoded by PIL (e.g. EMF, WMF) are skipped — the matching +placeholder in the markdown is left in place for downstream cleanup. +""" + +from __future__ import annotations + +import asyncio +import logging +import re +import zipfile +from io import BytesIO + +from ...models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from ..image_preprocessor import ensure_png_compatible_mode, pil_to_png_bytes +from .document_parser import DocumentParser +from .registry import parser_registry + +logger = logging.getLogger(__name__) + + +# Match MarkItDown's image refs in the rendered markdown. The current +# version emits a truncated placeholder (``![](data:image/png;base64...)``) +# but older / future versions may emit a full data URI or non-empty alt +# text. The pattern below matches both shapes — same regex the legacy +# loader used (``components/indexer/loaders/docx.py``). +_MARKITDOWN_IMAGE_PLACEHOLDER = re.compile(r"!\[.*?\]\(data:image/[^)]*\)") + + +def _image_ref(index: int) -> str: + """Synthetic markdown image ref used as a placeholder for embedded DOCX images.""" + return f"![](docx-image-{index})" + + +@parser_registry.register("docx") +class DocxParser(DocumentParser): + """Parse DOCX into a Markdown TextBlock + one ImageBlock per embedded image.""" + + def supported_types(self) -> list[str]: + return [DocumentType.DOCX.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + async with document.as_temporary_file() as path: + markdown, embedded = await asyncio.to_thread(self._extract, str(path)) + + markdown, images = self._rewrite_placeholders_and_build_blocks(markdown, embedded) + markdown = markdown.strip() + text_blocks = [TextBlock(text=markdown, page_number=1)] if markdown else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=dict(document.metadata), + page_count=1 if markdown else 0, + ) + + # ----- helpers ----- + + @classmethod + def _extract(cls, path: str) -> tuple[str, list[bytes | None]]: + """Run MarkItDown + zip-image extraction in one thread hop.""" + return cls._convert_to_markdown(path), cls._extract_embedded_images(path) + + @staticmethod + def _convert_to_markdown(path: str) -> str: + try: + from markitdown import MarkItDown + except ImportError: + logger.warning("markitdown not available; falling back to python-docx text extraction") + return DocxParser._fallback_extract_text(path) + try: + return MarkItDown().convert(path).text_content + except Exception as exc: + logger.warning("MarkItDown DOCX conversion failed (%s); falling back to plain text", exc) + return DocxParser._fallback_extract_text(path) + + @staticmethod + def _fallback_extract_text(path: str) -> str: + try: + from docx import Document as DocxDocument + except ImportError: + logger.warning("python-docx not available; cannot extract DOCX text") + return "" + try: + doc = DocxDocument(path) + return "\n\n".join(p.text for p in doc.paragraphs if p.text.strip()) + except Exception as exc: + logger.warning("python-docx fallback failed: %s", exc) + return "" + + @staticmethod + def _extract_embedded_images(path: str) -> list[bytes | None]: + """Return PNG bytes for each embedded image in document order. + + ``None`` entries preserve positional alignment with markdown + placeholders for unsupported formats (EMF, WMF, …). + """ + try: + from PIL import Image + except ImportError: + logger.warning("PIL not available; skipping DOCX image extraction") + return [] + try: + with zipfile.ZipFile(path, "r") as zf: + media = [n for n in zf.namelist() if n.startswith("word/media/")] + if not media: + return [] + ordered: dict[int, bytes | None] = {} + for name in media: + raw = zf.read(name) + try: + order_num = int(name.split("media/image")[1].split(".")[0]) + except (IndexError, ValueError): + continue + try: + with Image.open(BytesIO(raw)) as im: + im = ensure_png_compatible_mode(im) + ordered[order_num] = pil_to_png_bytes(im) + except Exception as exc: + logger.warning("Skipping unsupported DOCX media %s: %s", name, exc) + ordered[order_num] = None + if not ordered: + return [] + max_order = max(ordered) + return [ordered.get(i + 1) for i in range(max_order)] + except zipfile.BadZipFile: + logger.warning("DOCX is not a valid zip archive; skipping image extraction") + return [] + except Exception as exc: + logger.warning("DOCX image extraction failed: %s", exc) + return [] + + @staticmethod + def _rewrite_placeholders_and_build_blocks( + markdown: str, embedded: list[bytes | None] + ) -> tuple[str, list[ImageBlock]]: + """Replace each MarkItDown placeholder with a unique synthetic ref + and emit one ``ImageBlock`` per successfully-decoded zip image. + + Positional matching: the i-th placeholder in the markdown maps to + the i-th entry in ``embedded``. ``None`` entries (unsupported + formats) collapse the placeholder to an empty string. + """ + if not markdown or not embedded: + return markdown, [] + + images: list[ImageBlock] = [] + idx = 0 + consumed = 0 # count of `embedded` entries used so far + + def replacer(_match: re.Match[str]) -> str: + nonlocal idx, consumed + if consumed >= len(embedded): + return "" # more placeholders than zip images: drop extras + payload = embedded[consumed] + consumed += 1 + if payload is None: + return "" # zip image was unsupported; drop placeholder + ref = _image_ref(idx) + idx += 1 + images.append( + ImageBlock( + image_bytes=payload, + page_number=1, + mime_type="image/png", + metadata={"markdown_ref": ref}, + ) + ) + return ref + + new_markdown = _MARKITDOWN_IMAGE_PLACEHOLDER.sub(replacer, markdown) + return new_markdown, images diff --git a/openrag/core/indexing/parsers/eml_parser.py b/openrag/core/indexing/parsers/eml_parser.py new file mode 100644 index 000000000..760ff9626 --- /dev/null +++ b/openrag/core/indexing/parsers/eml_parser.py @@ -0,0 +1,216 @@ +"""EML (RFC822 email) ``DocumentParser`` implementation. + +Extracts the message body (``text/plain`` preferred, ``text/html`` +fallback) and dispatches each attachment to a parser supplied via DI. +Email headers (subject, from, to, date, message-id) and an attachment +manifest are merged into the output ``ProcessedDocument.metadata``. + +Attachment dispatch contract: + +- ``attachment_parsers`` maps lowercased extension (``"pdf"``, ``"docx"``, + no leading dot) to a :class:`DocumentParser`. +- Each attachment becomes a synthetic :class:`Document` (raw bytes, the + appropriate ``DocumentType`` if recognised, ``DocumentType.TEXT`` otherwise). +- The dispatched parser's text output is appended after a header block + giving filename, content-type, and size. +- Any ``ImageBlock``s emitted by the dispatched parser are propagated + into the EML's own ``ProcessedDocument.images``. +- Image attachments with no registered parser are emitted directly as + ``ImageBlock``s (no ``markdown_ref`` — there is no in-body placeholder + for them; see :class:`ImageBlock` for the parser→caption contract). +- Unknown non-image attachments include only the manifest header. + +Failures are tolerated: a single attachment that errors does not +propagate; we log and continue. +""" + +from __future__ import annotations + +import email +import logging +from collections.abc import Mapping +from email import policy +from email.utils import parsedate_to_datetime + +from ...models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from .document_parser import DocumentParser +from .registry import parser_registry + +logger = logging.getLogger(__name__) + + +_IMAGE_EXTS = {"png", "jpg", "jpeg", "gif", "webp", "bmp", "svg"} + + +@parser_registry.register("eml") +class EmlParser(DocumentParser): + """Parse ``.eml`` into a single text block plus ImageBlocks; dispatch attachments via DI.""" + + def __init__(self, attachment_parsers: Mapping[str, DocumentParser] | None = None) -> None: + self._attachment_parsers = dict(attachment_parsers or {}) + + def supported_types(self) -> list[str]: + return [DocumentType.EML.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + try: + msg = email.message_from_bytes(document.raw_bytes, policy=policy.default) + except Exception as exc: + logger.warning("Failed to parse EML: %s", exc) + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + headers = self._extract_headers(msg) + body, attachments = self._walk_parts(msg) + + attachments_text, images = await self._render_attachments(attachments) + full_text = (body + attachments_text).strip() + + metadata = dict(document.metadata) + metadata.update( + { + "email_subject": headers["subject"], + "email_from": headers["from"], + "email_to": headers["to"], + "email_date": headers["date"], + "email_message_id": headers["message-id"], + "email_attachment_count": len(attachments), + "email_attachment_filenames": [a["filename"] for a in attachments], + } + ) + if attachments: + metadata["email_attachments"] = [ + {"filename": a["filename"], "content_type": a["content_type"], "size": a["size"]} for a in attachments + ] + + text_blocks = [TextBlock(text=full_text, page_number=1)] if full_text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=metadata, + page_count=1 if full_text else 0, + ) + + # ----- helpers ----- + + @staticmethod + def _extract_headers(msg: email.message.Message) -> dict[str, str]: + # Under policy.default, msg.get(...) returns Header-like objects whose + # str() is the RFC 2047-decoded value; cast eagerly so the metadata is + # plain strings. + headers = { + "subject": str(msg.get("subject", "") or ""), + "from": str(msg.get("from", "") or ""), + "to": str(msg.get("to", "") or ""), + "date": str(msg.get("date", "") or ""), + "message-id": str(msg.get("message-id", "") or ""), + } + if headers["date"]: + try: + headers["date"] = parsedate_to_datetime(headers["date"]).isoformat() + except Exception: + pass + return headers + + @staticmethod + def _walk_parts(msg: email.message.Message) -> tuple[str, list[dict]]: + body = "" + attachments: list[dict] = [] + + for part in msg.walk(): + content_type = part.get_content_type() + disposition = part.get_content_disposition() + + if disposition in ("attachment", "inline"): + filename = part.get_filename() + payload = part.get_payload(decode=True) + if filename and payload: + attachments.append( + { + "filename": filename, + "content_type": content_type, + "size": len(payload), + "raw": payload, + } + ) + continue + + if content_type in ("text/plain", "text/html"): + payload = part.get_payload(decode=True) + if not payload: + continue + try: + text = payload.decode("utf-8") if isinstance(payload, bytes) else str(payload) + except UnicodeDecodeError: + text = payload.decode("latin-1", errors="ignore") if isinstance(payload, bytes) else str(payload) + # text/plain wins; only use text/html if we have nothing yet + if content_type == "text/plain" or not body: + body = text + + return body.strip(), attachments + + async def _render_attachments(self, attachments: list[dict]) -> tuple[str, list[ImageBlock]]: + """Render the attachment-section text and collect any ImageBlocks. + + Returns ``("", [])`` when there are no attachments. + """ + if not attachments: + return "", [] + + rendered: list[str] = ["\n\n--- ATTACHMENTS ---\n"] + images: list[ImageBlock] = [] + for att in attachments: + ext = self._extension(att["filename"]) + header = ( + f"\nAttachment: {att['filename']}\nContent-Type: {att['content_type']}\nSize: {att['size']} bytes\n" + ) + content, att_images = await self._render_one(att, ext) + rendered.append(header + content + "---\n") + images.extend(att_images) + return "".join(rendered), images + + async def _render_one(self, attachment: dict, ext: str) -> tuple[str, list[ImageBlock]]: + """Dispatch one attachment. Returns ``(text_to_inline, image_blocks)``.""" + parser = self._attachment_parsers.get(ext) + if parser is not None: + try: + doc = Document( + filename=attachment["filename"], + raw_bytes=attachment["raw"], + content_type=Document.detect_content_type(attachment["filename"]), + metadata={"source": f"attachment:{attachment['filename']}"}, + ) + processed = await parser.parse(doc) + content = "\n\n".join(b.text for b in processed.text_blocks if b.text) + inline = f"Content:\n{content}\n" if content else "" + return inline, list(processed.images) + except Exception as exc: + logger.warning("Attachment parser failed for %s: %s", attachment["filename"], exc) + + if ext in _IMAGE_EXTS: + # No parser registered — emit the image as an ImageBlock so a + # downstream caption stage can describe it. No ``markdown_ref`` + # because there is no in-body placeholder pointing to it. + return "", [ + ImageBlock( + image_bytes=attachment["raw"], + page_number=1, + mime_type=attachment["content_type"] or "image/png", + metadata={"source": f"attachment:{attachment['filename']}"}, + ) + ] + + return "", [] + + @staticmethod + def _extension(filename: str) -> str: + return filename.rsplit(".", 1)[-1].lower() if "." in filename else "" diff --git a/openrag/core/indexing/parsers/html_parser.py b/openrag/core/indexing/parsers/html_parser.py new file mode 100644 index 000000000..d51d2abb1 --- /dev/null +++ b/openrag/core/indexing/parsers/html_parser.py @@ -0,0 +1,58 @@ +"""HTML ``DocumentParser`` implementation. + +Converts HTML to Markdown via the project's ``html_to_markdown`` dep +(already used by the websearch and pptx pipelines), then emits a single +text block. No file I/O, no JavaScript execution, no image fetching — +purely structural conversion. +""" + +from __future__ import annotations + +import asyncio + +from ...models.document import Document, DocumentType, ProcessedDocument, TextBlock +from ..text_preprocessor import decode_bytes +from .document_parser import DocumentParser +from .registry import parser_registry + + +@parser_registry.register("html") +class HtmlParser(DocumentParser): + """Parse HTML documents into a single Markdown text block.""" + + def __init__(self, *, encoding: str | None = None) -> None: + """``encoding`` forces a specific decode of ``raw_bytes``; ``None`` + auto-detects (UTF-8 first, then chardet). + """ + self._encoding = encoding + + def supported_types(self) -> list[str]: + return [DocumentType.HTML.value] + + async def parse(self, document: Document) -> ProcessedDocument: + markdown = (await asyncio.to_thread(self._html_to_markdown, document)).strip() + text_blocks = [TextBlock(text=markdown, page_number=1)] if markdown else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + metadata=dict(document.metadata), + page_count=1 if markdown else 0, + ) + + def _html_to_markdown(self, document: Document) -> str: + """Decode + HTML→Markdown in one shot. Sync; runs in a thread.""" + html = self._extract_html(document) + return self._to_markdown(html) if html else "" + + def _extract_html(self, document: Document) -> str: + if document.text is not None: + return document.text + if document.raw_bytes: + return decode_bytes(document.raw_bytes, encoding=self._encoding) + return "" + + @staticmethod + def _to_markdown(html: str) -> str: + from html_to_markdown import convert + + return convert(html) diff --git a/openrag/core/indexing/parsers/image_parser.py b/openrag/core/indexing/parsers/image_parser.py new file mode 100644 index 000000000..4e6d5b767 --- /dev/null +++ b/openrag/core/indexing/parsers/image_parser.py @@ -0,0 +1,141 @@ +"""Image ``DocumentParser`` implementation. + +Decodes an image into normalized PNG bytes and emits a single +:class:`ImageBlock`. Supports raster formats (PNG/JPEG/etc. — anything +PIL opens) and SVG (rasterized to PNG via cairosvg). + +Captioning is not done here — see :class:`ImageBlock` for the +parser→caption contract. + +Output: +- A single ``ImageBlock`` with the normalized PNG bytes and no caption. +- No ``TextBlock`` is emitted; downstream stages produce text from the image. + +Failures (decode errors, undersized images) emit an empty +``ProcessedDocument`` rather than raising — RAG pipelines should not die +on a single bad image. +""" + +from __future__ import annotations + +import asyncio +import logging + +from ...models.document import Document, DocumentType, ImageBlock, ProcessedDocument +from ..image_preprocessor import MIN_IMAGE_PIXELS, ensure_png_compatible_mode +from .document_parser import DocumentParser +from .registry import parser_registry + +logger = logging.getLogger(__name__) + + +@parser_registry.register("image") +class ImageParser(DocumentParser): + """Decode an image and emit it as a single ``ImageBlock``.""" + + def __init__(self, *, min_pixels: int = MIN_IMAGE_PIXELS) -> None: + self._min_pixels = max(0, min_pixels) + + def supported_types(self) -> list[str]: + return [DocumentType.IMAGE.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + png_bytes = await asyncio.to_thread(self._normalize_to_png, document) + if png_bytes is None: + logger.warning("ImageParser: failed to decode image (id=%s)", document.id) + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + if await asyncio.to_thread(self._below_min_pixels, png_bytes): + logger.warning("ImageParser: image below min_pixels threshold (id=%s)", document.id) + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + return ProcessedDocument( + document_id=document.id, + images=[ + ImageBlock( + image_bytes=png_bytes, + page_number=1, + mime_type="image/png", + ) + ], + metadata=dict(document.metadata), + page_count=1, + ) + + def _normalize_to_png(self, document: Document) -> bytes | None: + """Return PNG bytes for any supported image input, or None on failure.""" + raw = document.raw_bytes + if not raw: + return None + if self._is_svg(raw, document.filename): + return self._svg_to_png(raw) + return self._raster_to_png(raw) + + @staticmethod + def _is_svg(raw: bytes, filename: str) -> bool: + if filename.lower().endswith(".svg"): + return True + head = raw[:200].lstrip().lower() + return head.startswith((b" bytes | None: + try: + import cairosvg + + return cairosvg.svg2png(bytestring=raw) + except Exception as exc: + logger.warning("Failed to rasterize SVG: %s", exc) + return None + + @staticmethod + def _raster_to_png(raw: bytes) -> bytes | None: + """Decode raw bytes through PIL and re-encode as PNG. + + Re-encoding normalizes the format so downstream consumers + (caption stage, vector-store image fields) only need to handle + one mime type, and validates the image is decodable. + """ + try: + from io import BytesIO + + from PIL import Image + except ImportError: + logger.warning("PIL not available; cannot decode raster image") + return None + try: + with Image.open(BytesIO(raw)) as image: + image = ensure_png_compatible_mode(image) + buf = BytesIO() + image.save(buf, format="PNG") + return buf.getvalue() + except Exception as exc: + logger.warning("Failed to decode image: %s", exc) + return None + + def _below_min_pixels(self, png_bytes: bytes) -> bool: + if self._min_pixels <= 0: + return False + try: + from io import BytesIO + + from PIL import Image + except ImportError: + return False + try: + with Image.open(BytesIO(png_bytes)) as image: + return image.width * image.height < self._min_pixels + except Exception: + return False diff --git a/openrag/core/indexing/parsers/markdown_parser.py b/openrag/core/indexing/parsers/markdown_parser.py new file mode 100644 index 000000000..566ea6de0 --- /dev/null +++ b/openrag/core/indexing/parsers/markdown_parser.py @@ -0,0 +1,69 @@ +"""Markdown ``DocumentParser`` implementation. + +Decodes a Markdown document into a single :class:`TextBlock` and emits +one :class:`ImageBlock` per image reference in the source: + +- Data-URI refs (``![alt](data:image/...;base64,...)``) are decoded and + the bytes stored on the block. +- HTTP/HTTPS refs (``![alt](https://...)``) yield an :class:`ImageBlock` + with empty ``image_bytes`` and ``source_url`` set; a downstream fetch + stage can populate the bytes later. The :attr:`ImageBlock.image_url` + property gives a uniform VLM-friendly URL in either case. + +Captioning is not done here — see :class:`ImageBlock` for the +parser→caption contract. +""" + +from __future__ import annotations + +from ...models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from ..image_preprocessor import HTTP_IMAGE_PATTERN, extract_data_uri_image_blocks +from ..text_preprocessor import decode_bytes +from .document_parser import DocumentParser +from .registry import parser_registry + + +@parser_registry.register("markdown") +class MarkdownParser(DocumentParser): + """Parse Markdown documents and emit ImageBlocks for every image ref.""" + + def __init__(self, *, encoding: str | None = None) -> None: + self._encoding = encoding + + def supported_types(self) -> list[str]: + return [DocumentType.MARKDOWN.value] + + async def parse(self, document: Document) -> ProcessedDocument: + text = self._extract_text(document).strip() + images = self._extract_image_blocks(text) + + text_blocks = [TextBlock(text=text, page_number=1)] if text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=dict(document.metadata), + page_count=1 if text else 0, + ) + + def _extract_text(self, document: Document) -> str: + if document.text is not None: + return document.text + if document.raw_bytes: + return decode_bytes(document.raw_bytes, encoding=self._encoding) + return "" + + @staticmethod + def _extract_image_blocks(text: str) -> list[ImageBlock]: + if not text: + return [] + blocks: list[ImageBlock] = list(extract_data_uri_image_blocks(text, page_number=1)) + for alt, url in HTTP_IMAGE_PATTERN.findall(text): + blocks.append( + ImageBlock( + source_url=url, + page_number=1, + metadata={"markdown_ref": f"![{alt}]({url})", "alt": alt}, + ) + ) + return blocks diff --git a/openrag/core/indexing/parsers/pdf/__init__.py b/openrag/core/indexing/parsers/pdf/__init__.py new file mode 100644 index 000000000..d48cea8b6 --- /dev/null +++ b/openrag/core/indexing/parsers/pdf/__init__.py @@ -0,0 +1,9 @@ +"""PDF parser backends. + +Each backend lives in its own module so its heavy dependencies (Marker, +Docling, DotsOCR, …) are only imported when the backend's submodule is +itself imported. Consumers do +``from core.indexing.parsers.pdf.marker import MarkerParser`` rather +than going through this package, so importing ``pdf`` does not pull in +any backend. +""" diff --git a/openrag/core/indexing/parsers/pdf/client_based.py b/openrag/core/indexing/parsers/pdf/client_based.py new file mode 100644 index 000000000..945aa5088 --- /dev/null +++ b/openrag/core/indexing/parsers/pdf/client_based.py @@ -0,0 +1,35 @@ +"""OpenAI-VLM-backed PDF ``DocumentParser`` (thin core facade). + +Holds a ``BaseClientParser`` (the actual HTTP-client / OpenAI-SDK +implementation lives in ``services/`` and is composed in at startup) and +delegates ``parse()`` to it. + +Mirrors the :class:`MarkerParser` pattern: core stays free of vendor +SDKs while the facade names "OpenAI-VLM PDF" as a first-class parser +type. Concrete subclasses of the services-side base (e.g. DotsOCR) can +be swapped in without changing this facade. +""" + +from __future__ import annotations + +from core.utils.exceptions import ValidationError + +from ....models.document import Document, ProcessedDocument +from ..document_parser import BaseClientParser, DocumentParser +from ..registry import parser_registry + + +@parser_registry.register("pdf_client") +class ClientPdfParser(DocumentParser): + """Public PDF parser facade backed by an OpenAI-compatible VLM client.""" + + def __init__(self, client: BaseClientParser) -> None: + if not isinstance(client, BaseClientParser): + raise ValidationError("ClientPdfParser requires a BaseClientParser instance as client") + self._client = client + + def supported_types(self) -> list[str]: + return self._client.supported_types() + + async def parse(self, document: Document) -> ProcessedDocument: + return await self._client.parse(document) diff --git a/openrag/core/indexing/parsers/pdf/docling.py b/openrag/core/indexing/parsers/pdf/docling.py new file mode 100644 index 000000000..323f2ac52 --- /dev/null +++ b/openrag/core/indexing/parsers/pdf/docling.py @@ -0,0 +1,29 @@ +"""Docling-backed PDF ``DocumentParser`` (thin core facade). + +Holds a reference to a ``BasePooledParser`` (the Ray pool implementation +lives in ``services/workers/parsers/docling_workers.py``) and delegates +``parse()`` to it. Keeping this facade in ``core/`` lets the indexing +pipeline reference the Docling backend without importing Ray. +""" + +from __future__ import annotations + +from ....models.document import Document, DocumentType, ProcessedDocument +from ..document_parser import BasePooledParser, DocumentParser +from ..registry import parser_registry + + +@parser_registry.register("docling") +class DoclingParser(DocumentParser): + """Public PDF parser facade backed by a Docling worker pool.""" + + def __init__(self, pool: BasePooledParser) -> None: + if not isinstance(pool, BasePooledParser): + raise ValueError("DoclingParser requires a BasePooledParser instance as pool") + self._pool = pool + + def supported_types(self) -> list[str]: + return [DocumentType.PDF.value] + + async def parse(self, document: Document) -> ProcessedDocument: + return await self._pool.parse(document) diff --git a/openrag/core/indexing/parsers/pdf/marker.py b/openrag/core/indexing/parsers/pdf/marker.py new file mode 100644 index 000000000..db17f0ce8 --- /dev/null +++ b/openrag/core/indexing/parsers/pdf/marker.py @@ -0,0 +1,36 @@ +"""Marker-backed PDF ``DocumentParser`` (thin core facade). + +Holds a reference to a ``BasePooledParser`` (the actual Ray-pool / +GPU-model / process-pool implementation lives in ``services/`` and is +not yet wired up) and delegates ``parse()`` to it. The split keeps core +free of Ray and GPU lifecycle code while still naming the Marker +backend as a first-class parser type. + +The injected pool is a generic ``BasePooledParser``; if a more specific +``MarkerPoolParser`` ABC emerges in services, this class can tighten +its type without changing call sites. +""" + +from __future__ import annotations + +from ....models.document import Document, ProcessedDocument +from ..document_parser import BasePooledParser, DocumentParser +from ..registry import parser_registry + + +@parser_registry.register("marker") +class MarkerParser(DocumentParser): + """Public PDF parser facade backed by a Marker worker pool.""" + + def __init__(self, pool: BasePooledParser) -> None: + # check pool is a BasePooledParser? and not empty + if not isinstance(pool, BasePooledParser) or pool is None: + raise ValueError("MarkerParser requires a BasePooledParser instance as pool") + + self._pool = pool + + def supported_types(self) -> list[str]: + return self._pool.supported_types() + + async def parse(self, document: Document) -> ProcessedDocument: + return await self._pool.parse(document) diff --git a/openrag/core/indexing/parsers/pdf/pymupdf.py b/openrag/core/indexing/parsers/pdf/pymupdf.py new file mode 100644 index 000000000..f42b33755 --- /dev/null +++ b/openrag/core/indexing/parsers/pdf/pymupdf.py @@ -0,0 +1,114 @@ +"""PyMuPDF-backed PDF ``DocumentParser``. + +The lightweight, no-VLM, no-GPU PDF backend. Uses ``pymupdf`` (a.k.a. +``fitz``) for plain-text extraction and ``pymupdf4llm`` for Markdown +extraction. Operates on ``Document.raw_bytes`` — file I/O is upstream. + +In ``mode="markdown"``, embedded images are surfaced as ``ImageBlock``s +via ``pymupdf4llm``'s ``embed_images=True`` (each image becomes a +``data:image/png;base64,…`` ref in the markdown, which we decode into +an :class:`ImageBlock` with ``markdown_ref`` set so a downstream caption +stage can substitute a description back in). ``mode="text"`` does not +extract images. + +Threading note: PyMuPDF is **not** thread-safe — concurrent calls to +``page.get_text`` / ``pymupdf4llm.to_markdown`` from different threads +can raise ``ValueError: not a textpage of this page`` (upstream +maintainer position: documented limitation, won't fix). We therefore +serialize all pymupdf work onto a single dedicated worker thread via +``_PYMUPDF_EXECUTOR``. The async ``parse`` method stays concurrent — +multiple callers will queue on the executor, but only one pymupdf +operation runs at a time. +""" + +from __future__ import annotations + +import asyncio +from concurrent.futures import ThreadPoolExecutor +from typing import Literal + +import pymupdf +import pymupdf4llm + +from ....models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from ...image_preprocessor import extract_data_uri_image_blocks +from ..document_parser import DocumentParser +from ..registry import parser_registry + +ParseMode = Literal["markdown", "text"] + +# Single dedicated worker for pymupdf — see "Threading note" in module docstring. +_PYMUPDF_EXECUTOR = ThreadPoolExecutor(max_workers=1, thread_name_prefix="pymupdf") + + +def _extract_text(raw: bytes) -> tuple[list[str], list[ImageBlock]]: + """Return one stripped plain-text string per page; no images.""" + with pymupdf.open(stream=raw, filetype="pdf") as doc: + return [page.get_text().strip() for page in doc], [] + + +def _extract_markdown(raw: bytes) -> tuple[list[str], list[ImageBlock]]: + """Return Markdown per page + ``ImageBlock``s built from embedded data URIs. + + ``embed_images=True`` makes ``pymupdf4llm`` write images as base64 + data URIs in-line. We decode each ref into an ``ImageBlock`` and + leave the ref in the page text untouched so the caption stage can + substitute later via ``ImageBlock.metadata['markdown_ref']``. + """ + with pymupdf.open(stream=raw, filetype="pdf") as doc: + chunks = pymupdf4llm.to_markdown( + doc, + page_chunks=True, + embed_images=True, + write_images=False, + dpi=300, + ) + pages: list[str] = [] + images: list[ImageBlock] = [] + for i, chunk in enumerate(chunks, start=1): + text = (chunk.get("text") or "").strip() + pages.append(text) + if text: + images.extend(extract_data_uri_image_blocks(text, page_number=i)) + return pages, images + + +@parser_registry.register("pymupdf") +class PyMuPDFParser(DocumentParser): + """Extract text from a PDF as one ``TextBlock`` per page (+ ImageBlocks in markdown mode). + + ``mode="markdown"`` (default) uses ``pymupdf4llm`` for layout-preserving + Markdown — better for downstream embedding and chunking, and surfaces + embedded images. ``mode="text"`` uses raw ``pymupdf`` for plain text — + slightly faster, no formatting, no images. + """ + + def __init__(self, *, mode: ParseMode = "markdown") -> None: + if mode not in ("markdown", "text"): + raise ValueError(f"PyMuPDFParser: unsupported mode {mode!r}") + self._mode = mode + self._extract = _extract_text if mode == "text" else _extract_markdown + + def supported_types(self) -> list[str]: + return [DocumentType.PDF.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + pages, images = await asyncio.get_running_loop().run_in_executor( + _PYMUPDF_EXECUTOR, self._extract, document.raw_bytes + ) + # Keep one TextBlock per source page (including empties) so callers + # can preserve a 1-to-1 mapping with the original PDF's pagination. + text_blocks = [TextBlock(text=text, page_number=i) for i, text in enumerate(pages, start=1)] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=dict(document.metadata), + page_count=len(pages), + ) diff --git a/openrag/core/indexing/parsers/pdf/test_docling.py b/openrag/core/indexing/parsers/pdf/test_docling.py new file mode 100644 index 000000000..2d02f7bed --- /dev/null +++ b/openrag/core/indexing/parsers/pdf/test_docling.py @@ -0,0 +1,60 @@ +"""Unit tests for :class:`DoclingParser`.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from core.indexing.parsers.document_parser import BasePooledParser +from core.indexing.parsers.pdf.docling import DoclingParser +from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock + + +def _make_doc() -> Document: + return Document(filename="doc.pdf", content_type=DocumentType.PDF, raw_bytes=b"%PDF-1.4") + + +class _FakePool(BasePooledParser): + def __init__(self, result: ProcessedDocument) -> None: + self._result = result + + def supported_types(self) -> list[str]: + return [DocumentType.PDF.value] + + async def parse(self, document: Document) -> ProcessedDocument: + return self._result + + +class TestDoclingParser: + def test_rejects_non_pool(self): + with pytest.raises(ValueError, match="BasePooledParser"): + DoclingParser(pool=object()) # type: ignore[arg-type] + + def test_supported_types_delegates_to_pool(self): + pool = _FakePool(ProcessedDocument(document_id="x")) + assert DoclingParser(pool=pool).supported_types() == [DocumentType.PDF.value] + + @pytest.mark.asyncio + async def test_parse_delegates_to_pool(self): + expected = ProcessedDocument( + document_id="d1", + text_blocks=[TextBlock(text="hello", page_number=1)], + ) + pool = _FakePool(expected) + parser = DoclingParser(pool=pool) + result = await parser.parse(_make_doc()) + assert result is expected + + @pytest.mark.asyncio + async def test_parse_propagates_pool_exception(self): + pool = MagicMock(spec=BasePooledParser) + pool.supported_types.return_value = [DocumentType.PDF.value] + pool.parse = AsyncMock(side_effect=RuntimeError("docling crashed")) + + with pytest.raises(RuntimeError, match="docling crashed"): + await DoclingParser(pool=pool).parse(_make_doc()) + + def test_registered_as_docling(self): + from core.indexing.parsers.registry import parser_registry + + assert "docling" in parser_registry diff --git a/openrag/core/indexing/parsers/pptx_parser.py b/openrag/core/indexing/parsers/pptx_parser.py new file mode 100644 index 000000000..0ccd1b8d6 --- /dev/null +++ b/openrag/core/indexing/parsers/pptx_parser.py @@ -0,0 +1,195 @@ +"""PPTX ``DocumentParser`` implementation. + +Walks slides via ``python-pptx``, converting each slide to Markdown: +title → ``#`` heading, text frames → paragraphs, tables → HTML→Markdown, +charts → Markdown tables, pictures → ``![](pptx-image-N)`` synthetic +markdown image refs. Speaker notes are appended as ``### Notes:``. + +Captioning is not done here — see :class:`ImageBlock` for the +parser→caption contract. + +Output is one :class:`TextBlock` per slide (1-indexed ``page_number``) +plus one :class:`ImageBlock` per slide picture. ``page_number`` on each +``ImageBlock`` is the slide number it came from. + +Implementation derived from the legacy ``PPTXConverter`` (which itself +mirrored the MarkItDown PPTX converter). +""" + +from __future__ import annotations + +import asyncio +import html +import logging +from io import BytesIO +from typing import Any + +from ...models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from ..image_preprocessor import ensure_png_compatible_mode, pil_to_png_bytes +from .document_parser import DocumentParser +from .registry import parser_registry + +logger = logging.getLogger(__name__) + + +def _image_ref(index: int) -> str: + """Synthetic markdown image ref used as a placeholder for slide pictures.""" + return f"![](pptx-image-{index})" + + +@parser_registry.register("pptx") +class PptxParser(DocumentParser): + """Parse PPTX into one TextBlock per slide plus one ImageBlock per picture.""" + + def supported_types(self) -> list[str]: + return [DocumentType.PPTX.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + async with document.as_temporary_file() as path: + slide_count, slides, images = await asyncio.to_thread(self._convert, str(path)) + + text_blocks = [TextBlock(text=text, page_number=slide_num) for slide_num, text in slides] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=dict(document.metadata), + page_count=slide_count, + ) + + # ----- conversion ----- + + def _convert(self, path: str) -> tuple[int, list[tuple[int, str]], list[ImageBlock]]: + try: + import pptx + from PIL import Image + except ImportError: + logger.warning("python-pptx or PIL not available; cannot parse PPTX") + return 0, [], [] + + try: + presentation = pptx.Presentation(path) + except Exception as exc: + logger.warning("Failed to open PPTX: %s", exc) + return 0, [], [] + + slides: list[tuple[int, str]] = [] + images: list[ImageBlock] = [] + + for slide_num, slide in enumerate(presentation.slides, start=1): + md = "" + title = slide.shapes.title + + for shape in slide.shapes: + if self._is_picture(shape): + try: + with Image.open(BytesIO(shape.image.blob)) as im: + im = ensure_png_compatible_mode(im) + png_bytes = pil_to_png_bytes(im) + ref = _image_ref(len(images)) + images.append( + ImageBlock( + image_bytes=png_bytes, + page_number=slide_num, + mime_type="image/png", + metadata={"markdown_ref": ref}, + ) + ) + md += ref + except Exception as exc: + logger.warning("Skipping unreadable PPTX picture: %s", exc) + elif self._is_table(shape): + md += "\n" + self._table_to_markdown(shape.table) + "\n" + elif getattr(shape, "has_chart", False): + md += self._chart_to_markdown(shape.chart) + elif getattr(shape, "has_text_frame", False): + if shape == title: + md += "# " + shape.text.lstrip() + "\n" + else: + md += shape.text + "\n" + + md = md.strip() + if slide.has_notes_slide: + notes_frame = slide.notes_slide.notes_text_frame + if notes_frame is not None: + md += "\n\n### Notes:\n" + notes_frame.text + md = md.strip() + + if md: + slides.append((slide_num, md)) + + return len(presentation.slides), slides, images + + @staticmethod + def _is_picture(shape: Any) -> bool: + try: + from pptx.enum.shapes import MSO_SHAPE_TYPE + + if shape.shape_type == MSO_SHAPE_TYPE.PICTURE: + return True + if shape.shape_type == MSO_SHAPE_TYPE.PLACEHOLDER and hasattr(shape, "image"): + return True + except NotImplementedError: + logger.debug("Encountered an unimplemented shape type") + except Exception: + return False + return False + + @staticmethod + def _is_table(shape: Any) -> bool: + try: + from pptx.enum.shapes import MSO_SHAPE_TYPE + + return shape.shape_type == MSO_SHAPE_TYPE.TABLE + except NotImplementedError: + logger.debug("Encountered an unimplemented shape type") + return False + except Exception: + return False + + @staticmethod + def _table_to_markdown(table: Any) -> str: + from html_to_markdown import convert + + html_rows = [""] + first_row = True + for row in table.rows: + html_rows.append("") + for cell in row.cells: + tag = "th" if first_row else "td" + html_rows.append(f"<{tag}>{html.escape(cell.text)}") + html_rows.append("") + first_row = False + html_rows.append("
") + return convert("".join(html_rows)).strip() + + @staticmethod + def _chart_to_markdown(chart: Any) -> str: + try: + md = "\n\n### Chart" + if chart.has_title: + md += f": {chart.chart_title.text_frame.text}" + md += "\n\n" + category_names = [c.label for c in chart.plots[0].categories] + series_names = [s.name for s in chart.series] + data: list[list[str]] = [["Category"] + series_names] + for idx, category in enumerate(category_names): + row = [category] + for series in chart.series: + row.append(series.values[idx]) + data.append(row) + rows = ["| " + " | ".join(map(str, r)) + " |" for r in data] + separator = "|" + "|".join(["---"] * len(data[0])) + "|" + return md + "\n".join([rows[0], separator] + rows[1:]) + except ValueError as exc: + if "unsupported plot type" in str(exc): + return "\n\n[unsupported chart]\n\n" + return "\n\n[unsupported chart]\n\n" + except Exception: + return "\n\n[unsupported chart]\n\n" diff --git a/openrag/core/indexing/parsers/registry.py b/openrag/core/indexing/parsers/registry.py new file mode 100644 index 000000000..c7b8fed37 --- /dev/null +++ b/openrag/core/indexing/parsers/registry.py @@ -0,0 +1,7 @@ +"""Document parser registry.""" + +from openrag.core.utils.registry import Registry + +from .document_parser import DocumentParser + +parser_registry: Registry[DocumentParser] = Registry("parser") diff --git a/openrag/core/indexing/parsers/test_doc_parser.py b/openrag/core/indexing/parsers/test_doc_parser.py new file mode 100644 index 000000000..454de181f --- /dev/null +++ b/openrag/core/indexing/parsers/test_doc_parser.py @@ -0,0 +1,118 @@ +"""Unit tests for :class:`DocParser` (.doc → DocxParser delegation + fallback).""" + +from __future__ import annotations + +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from ...models.document import Document, DocumentType, ProcessedDocument, TextBlock +from .doc_parser import DocParser + + +@pytest.fixture +def fake_spire(): + """Inject a fake ``spire.doc`` into ``sys.modules`` for the duration of a test. + + Returns the ``Document`` mock class so tests can configure the + instance returned by ``Document()``. + """ + spire = MagicMock() + spire_doc = MagicMock() + spire.doc = spire_doc + saved = {k: sys.modules.get(k) for k in ("spire", "spire.doc")} + sys.modules["spire"] = spire + sys.modules["spire.doc"] = spire_doc + try: + yield spire_doc.Document + finally: + for k, v in saved.items(): + if v is None: + sys.modules.pop(k, None) + else: + sys.modules[k] = v + + +def _doc_document(raw: bytes = b"\xd0\xcf\x11\xe0fake-doc") -> Document: + return Document(filename="x.doc", content_type=DocumentType.DOC, raw_bytes=raw) + + +class TestParse: + @pytest.mark.asyncio + async def test_empty_raw_bytes_returns_empty(self): + doc = _doc_document(raw=b"") + result = await DocParser().parse(doc) + assert result.text_blocks == [] and result.images == [] and result.page_count == 0 + + @pytest.mark.asyncio + async def test_successful_conversion_delegates_to_docx(self, fake_spire, tmp_path): + # Spire writes a real .docx file at the path given to SaveToFile. + dummy_docx = b"DOCX-CONTENT" + + instance = MagicMock() + + def save_to_file(path: str, _fmt) -> None: + with open(path, "wb") as fh: + fh.write(dummy_docx) + + instance.SaveToFile.side_effect = save_to_file + fake_spire.return_value = instance + + docx_parser = MagicMock() + expected = ProcessedDocument( + document_id="test", + text_blocks=[TextBlock(text="from-docx", page_number=1)], + page_count=1, + ) + docx_parser.parse = AsyncMock(return_value=expected) + + parser = DocParser(docx_parser=docx_parser) + result = await parser.parse(_doc_document()) + + assert result is expected + docx_parser.parse.assert_awaited_once() + forwarded = docx_parser.parse.await_args.args[0] + assert forwarded.raw_bytes == dummy_docx + assert forwarded.content_type is DocumentType.DOCX + instance.LoadFromFile.assert_called_once() + instance.Close.assert_called_once() + + @pytest.mark.asyncio + async def test_save_failure_falls_back_to_get_text(self, fake_spire): + instance = MagicMock() + instance.SaveToFile.side_effect = RuntimeError("Spire crashed") + instance.GetText.return_value = " plain text content " + fake_spire.return_value = instance + + docx_parser = MagicMock() + docx_parser.parse = AsyncMock() + result = await DocParser(docx_parser=docx_parser).parse(_doc_document()) + + assert result.text_blocks == [TextBlock(text="plain text content", page_number=1)] + assert result.page_count == 1 + instance.GetText.assert_called_once() + instance.Close.assert_called_once() + docx_parser.parse.assert_not_awaited() # never delegated + + @pytest.mark.asyncio + async def test_total_failure_returns_empty(self, fake_spire): + instance = MagicMock() + instance.SaveToFile.side_effect = RuntimeError("Spire crashed") + instance.GetText.side_effect = RuntimeError("GetText crashed") + fake_spire.return_value = instance + + result = await DocParser().parse(_doc_document()) + assert result.text_blocks == [] and result.page_count == 0 + instance.Close.assert_called_once() + + @pytest.mark.asyncio + async def test_missing_spire_returns_empty(self, monkeypatch): + # spire-doc is in the runtime deps, so just omitting fake_spire would + # actually drive a real Spire instance against malformed bytes. + # Pin the import to None so ``_convert``'s ``import spire.doc`` raises + # ImportError deterministically. + monkeypatch.setitem(sys.modules, "spire", None) + monkeypatch.setitem(sys.modules, "spire.doc", None) + result = await DocParser().parse(_doc_document()) + assert result.text_blocks == [] and result.page_count == 0 diff --git a/openrag/core/indexing/parsers/test_docx_parser.py b/openrag/core/indexing/parsers/test_docx_parser.py new file mode 100644 index 000000000..cf612b069 --- /dev/null +++ b/openrag/core/indexing/parsers/test_docx_parser.py @@ -0,0 +1,141 @@ +"""Unit tests for :class:`DocxParser`.""" + +from __future__ import annotations + +import tempfile +import zipfile +from io import BytesIO +from pathlib import Path + +import pytest +from PIL import Image + +from ...models.document import Document, DocumentType +from .docx_parser import DocxParser, _image_ref + + +def _png_bytes(color: str = "red") -> bytes: + img = Image.new("RGBA", (10, 10), color) + buf = BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + + +def _fake_docx(media_files: dict[str, bytes]) -> Path: + """Build a minimal .docx zip with given ``word/media/`` entries.""" + tmp = tempfile.NamedTemporaryFile(suffix=".docx", delete=False) + with zipfile.ZipFile(tmp, "w") as zf: + for name, data in media_files.items(): + zf.writestr(f"word/media/{name}", data) + return Path(tmp.name) + + +class TestExtractEmbeddedImages: + """Mirrors legacy ``TestGetImagesFromZip`` against the new staticmethod.""" + + def test_valid_images_kept_in_order(self): + docx = _fake_docx({"image2.png": _png_bytes("blue"), "image1.png": _png_bytes("red")}) + result = DocxParser._extract_embedded_images(str(docx)) + assert len(result) == 2 + assert result[0] is not None and result[1] is not None + + def test_unsupported_format_collapses_to_none_at_position(self): + docx = _fake_docx( + { + "image1.png": _png_bytes(), + "image2.emf": b"\x01\x00\x00\x00garbage", + "image3.png": _png_bytes(), + } + ) + result = DocxParser._extract_embedded_images(str(docx)) + assert len(result) == 3 + assert result[0] is not None + assert result[1] is None + assert result[2] is not None + + def test_non_image_media_skipped(self): + docx = _fake_docx( + { + "image1.png": _png_bytes(), + "oleObject1.bin": b"OLE", + "hdphoto1.wdp": b"WDP", + } + ) + result = DocxParser._extract_embedded_images(str(docx)) + assert sum(1 for x in result if x is not None) == 1 + + def test_all_unsupported_returns_empty(self): + docx = _fake_docx({"image1.emf": b"EMF", "image2.wmf": b"WMF"}) + # Two unsupported entries: positional list still has length 2 with None slots. + result = DocxParser._extract_embedded_images(str(docx)) + assert all(x is None for x in result) + + def test_no_media_returns_empty(self): + tmp = tempfile.NamedTemporaryFile(suffix=".docx", delete=False) + with zipfile.ZipFile(tmp, "w") as zf: + zf.writestr("word/document.xml", "") + result = DocxParser._extract_embedded_images(tmp.name) + assert result == [] + + def test_invalid_zip_returns_empty(self): + tmp = tempfile.NamedTemporaryFile(suffix=".docx", delete=False) + tmp.write(b"not a zip") + tmp.flush() + assert DocxParser._extract_embedded_images(tmp.name) == [] + + +class TestRewritePlaceholdersAndBuildBlocks: + """The parser→caption contract: synthetic refs + ImageBlock metadata.""" + + def test_assigns_unique_refs_in_order(self): + md = "before ![](data:image/png;base64,trunc...) middle ![](data:image/png;base64,trunc...) end" + embedded = [_png_bytes("red"), _png_bytes("blue")] + new_md, blocks = DocxParser._rewrite_placeholders_and_build_blocks(md, embedded) + + assert _image_ref(0) in new_md + assert _image_ref(1) in new_md + assert len(blocks) == 2 + assert blocks[0].metadata["markdown_ref"] == _image_ref(0) + assert blocks[1].metadata["markdown_ref"] == _image_ref(1) + assert blocks[0].image_bytes == embedded[0] + assert blocks[1].image_bytes == embedded[1] + assert all(b.page_number == 1 and b.mime_type == "image/png" for b in blocks) + + def test_none_entry_drops_placeholder(self): + md = "![](data:image/png;base64,a) ![](data:image/png;base64,b)" + embedded = [None, _png_bytes()] + new_md, blocks = DocxParser._rewrite_placeholders_and_build_blocks(md, embedded) + + # First placeholder collapses to empty; second becomes ref-0 (only one block emitted). + assert _image_ref(0) in new_md + assert _image_ref(1) not in new_md + assert len(blocks) == 1 + + def test_more_placeholders_than_zip_entries_drops_extras(self): + md = "![](data:image/png;base64,a) ![](data:image/png;base64,b)" + embedded = [_png_bytes()] + new_md, blocks = DocxParser._rewrite_placeholders_and_build_blocks(md, embedded) + + assert _image_ref(0) in new_md + # The extra placeholder is dropped — the regex match collapses to "". + assert "data:image" not in new_md + assert len(blocks) == 1 + + def test_no_placeholders_passthrough(self): + new_md, blocks = DocxParser._rewrite_placeholders_and_build_blocks("plain text", [_png_bytes()]) + assert new_md == "plain text" + assert blocks == [] + + def test_empty_inputs(self): + assert DocxParser._rewrite_placeholders_and_build_blocks("", []) == ("", []) + assert DocxParser._rewrite_placeholders_and_build_blocks("text", []) == ("text", []) + + +class TestParse: + @pytest.mark.asyncio + async def test_empty_raw_bytes_returns_empty(self): + doc = Document(filename="x.docx", content_type=DocumentType.DOCX, raw_bytes=b"") + result = await DocxParser().parse(doc) + assert result.text_blocks == [] + assert result.images == [] + assert result.page_count == 0 diff --git a/openrag/core/indexing/parsers/text_parser.py b/openrag/core/indexing/parsers/text_parser.py new file mode 100644 index 000000000..b23b1a2b4 --- /dev/null +++ b/openrag/core/indexing/parsers/text_parser.py @@ -0,0 +1,49 @@ +"""Plain-text ``DocumentParser`` implementation. + +Decodes ``Document.raw_bytes`` (or uses ``Document.text`` if already +populated) into a single :class:`TextBlock`. Performs no image captioning, +no markdown-image extraction, and no file I/O — those are upstream +concerns. Handles the ``TEXT`` content type; Markdown (with image +captioning) lives in :class:`MarkdownParser`; HTML, PDF, and richer +formats live in their own parsers. +""" + +from __future__ import annotations + +import asyncio + +from ...models.document import Document, DocumentType, ProcessedDocument, TextBlock +from ..text_preprocessor import decode_bytes +from .document_parser import DocumentParser +from .registry import parser_registry + + +@parser_registry.register("text") +class TextParser(DocumentParser): + """Parse plain-text documents into a single text block.""" + + def __init__(self, *, encoding: str | None = None) -> None: + """If ``encoding`` is ``None``, raw bytes are auto-detected via + :func:`core.indexing.text_preprocessor.decode_bytes`. + """ + self._encoding = encoding + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + async def parse(self, document: Document) -> ProcessedDocument: + text = (await asyncio.to_thread(self._extract_text, document)).strip() + text_blocks = [TextBlock(text=text, page_number=1)] if text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + metadata=dict(document.metadata), + page_count=1 if text else 0, + ) + + def _extract_text(self, document: Document) -> str: + if document.text is not None: + return document.text + if document.raw_bytes: + return decode_bytes(document.raw_bytes, encoding=self._encoding) + return "" diff --git a/openrag/core/indexing/serializer.py b/openrag/core/indexing/serializer.py new file mode 100644 index 000000000..a2b8da941 --- /dev/null +++ b/openrag/core/indexing/serializer.py @@ -0,0 +1,26 @@ +"""Transitional port for the document-serialization operation. + +``ConversionService`` (Phase 8E) exposes the ``extractText`` tool — +serialize an uploaded file to raw text. The work runs in the +``DocSerializer`` Ray actor; defining it on a dedicated port keeps the +orchestrator Ray-free (8H: no Ray import / remote call under +``services/orchestrators/``). A small shim in ``services/storage/`` +adapts the actor to this interface during the shim period; Phase 9 +swaps it for a direct serializer call and deletes the shim. + +No Ray / LangChain types leak across this boundary — the serialized +document is returned as its plain text content. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class FileSerializer(ABC): + """The single serialize operation the conversion orchestrator needs.""" + + @abstractmethod + async def serialize(self, path: str, metadata: dict) -> str: + """Serialize the file at ``path`` and return its raw text content.""" + ... diff --git a/openrag/core/indexing/test_image_preprocessor.py b/openrag/core/indexing/test_image_preprocessor.py new file mode 100644 index 000000000..29ae9ec5d --- /dev/null +++ b/openrag/core/indexing/test_image_preprocessor.py @@ -0,0 +1,98 @@ +"""Unit tests for ``core.indexing.image_preprocessor``.""" + +from __future__ import annotations + +import base64 + +from PIL import Image + +from .image_preprocessor import ( + MIN_IMAGE_PIXELS, + decode_data_uri, + ensure_png_compatible_mode, + extract_data_uri_image_blocks, + mime_from_data_uri, + pil_to_png_bytes, +) + + +class TestEnsurePngCompatibleMode: + def test_cmyk_to_rgb(self): + assert ensure_png_compatible_mode(Image.new("CMYK", (10, 10))).mode == "RGB" + + def test_palette_to_rgba(self): + assert ensure_png_compatible_mode(Image.new("P", (10, 10))).mode == "RGBA" + + def test_la_to_rgba(self): + assert ensure_png_compatible_mode(Image.new("LA", (10, 10))).mode == "RGBA" + + def test_rgb_unchanged(self): + assert ensure_png_compatible_mode(Image.new("RGB", (10, 10))).mode == "RGB" + + def test_rgba_unchanged(self): + assert ensure_png_compatible_mode(Image.new("RGBA", (10, 10))).mode == "RGBA" + + +class TestPilToPngBytes: + def test_rgb_round_trip(self): + img = Image.new("RGB", (32, 32), "red") + png = pil_to_png_bytes(img) + assert png[:8] == b"\x89PNG\r\n\x1a\n" + + def test_cmyk_normalised_then_encoded(self): + png = pil_to_png_bytes(Image.new("CMYK", (16, 16))) + assert png[:8] == b"\x89PNG\r\n\x1a\n" + + def test_bytes_passthrough(self): + raw = b"already-bytes" + assert pil_to_png_bytes(raw) is raw + + +class TestDecodeDataUri: + def test_round_trip(self): + payload = b"hello" + uri = f"data:image/png;base64,{base64.b64encode(payload).decode()}" + assert decode_data_uri(uri) == payload + + def test_malformed_returns_none(self): + assert decode_data_uri("not-a-data-uri") is None + assert decode_data_uri("data:image/png;base64,!!!not-base64") is None + + +class TestMimeFromDataUri: + def test_jpeg(self): + assert mime_from_data_uri("data:image/jpeg;base64,xxx") == "image/jpeg" + + def test_png(self): + assert mime_from_data_uri("data:image/png;base64,xxx") == "image/png" + + def test_malformed_falls_back_to_png(self): + assert mime_from_data_uri("garbage") == "image/png" + + +class TestExtractDataUriImageBlocks: + def _data_uri(self, payload: bytes = b"x", mime: str = "image/png") -> str: + return f"data:{mime};base64,{base64.b64encode(payload).decode()}" + + def test_emits_one_block_per_match(self): + uri = self._data_uri(b"hello") + text = f"intro ![alt-1]({uri}) middle ![alt-2]({uri}) end" + blocks = extract_data_uri_image_blocks(text, page_number=3) + + assert len(blocks) == 2 + assert all(b.image_bytes == b"hello" for b in blocks) + assert all(b.page_number == 3 for b in blocks) + assert blocks[0].metadata["alt"] == "alt-1" + assert blocks[0].metadata["markdown_ref"] == f"![alt-1]({uri})" + + def test_no_matches_returns_empty(self): + assert extract_data_uri_image_blocks("plain text") == [] + assert extract_data_uri_image_blocks("") == [] + + def test_skips_undecodable(self): + text = "![](data:image/png;base64,!!!not-base64)" + assert extract_data_uri_image_blocks(text) == [] + + +def test_min_image_pixels_constant(): + assert MIN_IMAGE_PIXELS == 784 diff --git a/openrag/core/indexing/test_validators.py b/openrag/core/indexing/test_validators.py new file mode 100644 index 000000000..0ffae81ac --- /dev/null +++ b/openrag/core/indexing/test_validators.py @@ -0,0 +1,69 @@ +"""Unit tests for ``core.indexing.validators``.""" + +from __future__ import annotations + +import pytest + +from ..utils.exceptions import ValidationError +from .validators import parse_metadata, validate_file_format, validate_file_id + + +class TestParseMetadata: + def test_none_returns_empty(self): + assert parse_metadata(None) == {} + + def test_empty_string_returns_empty(self): + assert parse_metadata("") == {} + + def test_dict_passthrough(self): + d = {"a": 1, "b": [1, 2]} + assert parse_metadata(d) is d + + def test_valid_json_string(self): + assert parse_metadata('{"k": "v"}') == {"k": "v"} + + def test_invalid_json_raises_400(self): + with pytest.raises(ValidationError) as exc: + parse_metadata("{not-json") + assert exc.value.status_code == 400 + + def test_non_object_json_raises_400(self): + with pytest.raises(ValidationError) as exc: + parse_metadata('["a", "b"]') + assert exc.value.status_code == 400 + + +class TestValidateFileId: + def test_valid(self): + assert validate_file_id("abc-123") == "abc-123" + + def test_default_forbidden_slash(self): + with pytest.raises(ValidationError) as exc: + validate_file_id("a/b") + assert exc.value.status_code == 400 + + def test_empty_or_whitespace_raises(self): + for bad in ("", " "): + with pytest.raises(ValidationError): + validate_file_id(bad) + + def test_custom_forbidden_chars(self): + with pytest.raises(ValidationError): + validate_file_id("hello?world", forbidden_chars="?") + assert validate_file_id("hello/world", forbidden_chars="?") == "hello/world" + + +class TestValidateFileFormat: + formats = ("pdf", "docx") + mimetypes = ("application/pdf", "application/vnd.openxmlformats-officedocument.wordprocessingml.document") + + def test_extension_match(self): + assert validate_file_format("doc.PDF", self.formats, self.mimetypes) == "pdf" + + def test_mimetype_match_when_no_extension(self): + assert validate_file_format("noext", self.formats, self.mimetypes, mimetype="application/pdf") == "" + + def test_unsupported_raises_415(self): + with pytest.raises(ValidationError) as exc: + validate_file_format("img.exe", self.formats, self.mimetypes, mimetype="application/x-msdownload") + assert exc.value.status_code == 415 diff --git a/openrag/core/indexing/text_preprocessor.py b/openrag/core/indexing/text_preprocessor.py new file mode 100644 index 000000000..8f31137ef --- /dev/null +++ b/openrag/core/indexing/text_preprocessor.py @@ -0,0 +1,16 @@ +"""Text preprocessing utilities for the indexing pipeline. + +Re-exports the canonical implementations from `core.utils.text`. Kept as a +named entry point under `core.indexing` so callers can import preprocessing +helpers alongside parsers, validators, and contextualization without +reaching into the generic utils package. +""" + +from ..utils.text import clean_markdown_table_spacing, decode_bytes, sanitize_extracted_text, sanitize_text + +__all__ = [ + "clean_markdown_table_spacing", + "sanitize_extracted_text", + "sanitize_text", + "decode_bytes", +] diff --git a/openrag/core/indexing/validators.py b/openrag/core/indexing/validators.py new file mode 100644 index 000000000..25f459562 --- /dev/null +++ b/openrag/core/indexing/validators.py @@ -0,0 +1,78 @@ +"""Framework-free validators for indexing inputs. + +Pure functions on plain types — no FastAPI, no Hydra. Routers translate +incoming HTTP requests into these inputs and let the global ``OpenRAGError`` +handler convert raised ``ValidationError`` instances into HTTP responses. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable +from typing import Any + +from ..utils.exceptions import ValidationError + +DEFAULT_FORBIDDEN_CHARS_IN_FILE_ID: frozenset[str] = frozenset("/") + + +def parse_metadata(raw: Any | None) -> dict: + """Parse JSON-encoded metadata into a dict. + + Accepts ``None`` / empty string (returns ``{}``), an existing dict + (returned as-is), or a JSON string that decodes to a dict. + """ + if raw is None or raw == "": + return {} + if isinstance(raw, dict): + return raw + try: + decoded = json.loads(raw) + except (json.JSONDecodeError, TypeError) as exc: + raise ValidationError("Invalid JSON in metadata", status_code=400) from exc + if not isinstance(decoded, dict): + raise ValidationError("Metadata must be a JSON object", status_code=400) + return decoded + + +def validate_file_id( + file_id: str, + forbidden_chars: Iterable[str] = DEFAULT_FORBIDDEN_CHARS_IN_FILE_ID, +) -> str: + """Return normalized ``file_id`` if valid, else raise ``ValidationError`` (HTTP 400).""" + if not isinstance(file_id, str): + raise ValidationError("File ID must be a string.", status_code=400) + file_id = file_id.strip() + if not file_id: + raise ValidationError("File ID cannot be empty.", status_code=400) + forbidden = frozenset(forbidden_chars) + if any(c in file_id for c in forbidden): + raise ValidationError( + f"File ID contains forbidden characters: {', '.join(sorted(forbidden))}", + status_code=400, + ) + return file_id + + +def validate_file_format( + filename: str, + accepted_formats: Iterable[str], + accepted_mimetypes: Iterable[str], + mimetype: str | None = None, +) -> str: + """Validate the file by extension or mimetype. + + Returns the lowercased file extension (without the leading dot, possibly + empty). Raises ``ValidationError`` (HTTP 415) on rejection. + """ + file_extension = filename.rsplit(".", 1)[-1].lower() if "." in filename else "" + formats = set(accepted_formats) + mimetypes = set(accepted_mimetypes) + if file_extension not in formats and mimetype not in mimetypes: + details = ( + f"Unsupported file format: {file_extension} or file mimetype.\n" + f"Supported formats: {', '.join(sorted(formats))}\n" + f"Supported mimetypes: {', '.join(sorted(mimetypes))}" + ) + raise ValidationError(details, status_code=415) + return file_extension diff --git a/openrag/core/llm/__init__.py b/openrag/core/llm/__init__.py new file mode 100644 index 000000000..bb7e15818 --- /dev/null +++ b/openrag/core/llm/__init__.py @@ -0,0 +1,6 @@ +"""LLM ABC + registry.""" + +from .llm import LLM +from .registry import llm_registry + +__all__ = ["LLM", "llm_registry"] diff --git a/openrag/core/llm/llm.py b/openrag/core/llm/llm.py new file mode 100644 index 000000000..54c0bc275 --- /dev/null +++ b/openrag/core/llm/llm.py @@ -0,0 +1,30 @@ +"""Abstract LLM interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import AsyncIterator + + +class LLM(ABC): + """Base class for all LLM providers.""" + + @abstractmethod + async def generate(self, prompt: str, **kwargs) -> dict: + """Generate a text completion for a prompt.""" + ... + + @abstractmethod + async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: + """Chat completion with message list.""" + ... + + @abstractmethod + def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIterator[str]: + """Stream chat completion as raw SSE lines. + + Implementations must be ``async def`` generators yielding ``str`` chunks. + Declared without ``async def`` here so the abstract signature matches the + ``AsyncIterator[str]`` return type without forcing an empty ``yield``. + """ + ... diff --git a/openrag/core/llm/registry.py b/openrag/core/llm/registry.py new file mode 100644 index 000000000..2a7af4769 --- /dev/null +++ b/openrag/core/llm/registry.py @@ -0,0 +1,7 @@ +"""LLM registry.""" + +from openrag.core.utils.registry import Registry + +from .llm import LLM + +llm_registry: Registry[LLM] = Registry("llm") diff --git a/openrag/core/models/__init__.py b/openrag/core/models/__init__.py new file mode 100644 index 000000000..1ddb25e67 --- /dev/null +++ b/openrag/core/models/__init__.py @@ -0,0 +1,46 @@ +"""Domain models — pure Pydantic, no infrastructure imports.""" + +from .catalog import DocumentRecord, DocumentStatus, IndexationJob, JobStatus +from .chunk import Chunk, ChunkType +from .contextualization import ContextualizedQuery +from .conversation import Conversation, Message +from .document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from .prompt import Prompt, PromptType +from .query import Query, RetrievalQuery, SearchQueries, TemporalPredicate +from .retrieval_response import RetrievalResponse +from .retrieval_result import RetrievalResult, ScoredChunk +from .user import ApiKey, OIDCSession, PartitionRole, TokenPayload, User, UserPartition +from .workspace import Workspace + +__all__ = [ + "ApiKey", + "Chunk", + "ChunkType", + "ContextualizedQuery", + "Conversation", + "Document", + "DocumentRecord", + "DocumentStatus", + "DocumentType", + "ImageBlock", + "IndexationJob", + "JobStatus", + "Message", + "OIDCSession", + "PartitionRole", + "ProcessedDocument", + "Prompt", + "PromptType", + "Query", + "RetrievalQuery", + "RetrievalResponse", + "RetrievalResult", + "ScoredChunk", + "SearchQueries", + "TemporalPredicate", + "TextBlock", + "TokenPayload", + "User", + "UserPartition", + "Workspace", +] diff --git a/openrag/core/models/catalog.py b/openrag/core/models/catalog.py new file mode 100644 index 000000000..014ed94ed --- /dev/null +++ b/openrag/core/models/catalog.py @@ -0,0 +1,56 @@ +"""Catalog domain models — document records, indexation jobs, status tracking.""" + +from __future__ import annotations + +import uuid +from datetime import UTC, datetime +from enum import Enum +from typing import Any + +from pydantic import BaseModel, Field + + +class DocumentStatus(str, Enum): + QUEUED = "QUEUED" + SERIALIZING = "SERIALIZING" + CHUNKING = "CHUNKING" + INSERTING = "INSERTING" + COMPLETED = "COMPLETED" + FAILED = "FAILED" + CANCELLED = "CANCELLED" + + +class JobStatus(str, Enum): + QUEUED = "QUEUED" + RUNNING = "RUNNING" + SUCCESS = "SUCCESS" + FAILED = "FAILED" + PARTIAL = "PARTIAL" + + +class DocumentRecord(BaseModel): + """A document entry in the catalog (PostgreSQL).""" + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + file_id: str = "" + filename: str = "" + partition: str = "default" + metadata: dict[str, Any] = Field(default_factory=dict) + status: DocumentStatus = DocumentStatus.QUEUED + error_message: str | None = None + created_by: int | None = None + relationship_id: str | None = None + parent_id: str | None = None + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + + +class IndexationJob(BaseModel): + """An indexation job tracking batch document processing.""" + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + status: JobStatus = JobStatus.QUEUED + total_documents: int = 0 + partition: str = "default" + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + started_at: datetime | None = None + completed_at: datetime | None = None diff --git a/openrag/core/models/chunk.py b/openrag/core/models/chunk.py new file mode 100644 index 000000000..8489c0fb9 --- /dev/null +++ b/openrag/core/models/chunk.py @@ -0,0 +1,104 @@ +"""Chunk — the unit of indexable and retrievable text.""" + +from __future__ import annotations + +import uuid +from enum import Enum +from typing import Any + +from pydantic import BaseModel, Field + + +class ChunkType(str, Enum): + TEXT = "text" + IMAGE_CAPTION = "image_caption" + TABLE = "table" + CONTEXTUALIZED = "contextualized" + + +# Pre-Phase-5 chunkers stamped Document metadata with the raw MDElement +# literal (`"image"`) for image elements. Deployments upgraded without +# re-indexing have those values in Milvus; map them to the current enum +# at read time so retrieval doesn't crash on legacy data. +_CHUNK_TYPE_LEGACY_ALIASES = {"image": ChunkType.IMAGE_CAPTION} + + +def _coerce_chunk_type(value: Any) -> ChunkType: + if isinstance(value, ChunkType): + return value + if value in _CHUNK_TYPE_LEGACY_ALIASES: + return _CHUNK_TYPE_LEGACY_ALIASES[value] + try: + return ChunkType(value) + except (ValueError, TypeError): + # Unknown value from upstream/legacy data — fall back to TEXT rather + # than crash the retrieval call. + return ChunkType.TEXT + + +class Chunk(BaseModel): + """A chunk of text extracted from a document, optionally embedded.""" + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + document_id: str = "" + text: str = "" + chunk_index: int = 0 + chunk_type: ChunkType = ChunkType.TEXT + embedding: list[float] | None = None + metadata: dict[str, Any] = Field(default_factory=dict) + partition: str = "default" + page_number: int | None = None + token_count: int | None = None + header: str | None = None + context: str | None = None + content: str | None = None + + def with_embedding(self, embedding: list[float]) -> Chunk: + """Return a copy with the embedding set.""" + return self.model_copy(update={"embedding": embedding}) + + @classmethod + def from_langchain(cls, doc: Any) -> Chunk: + """Convert a LangChain Document to a Chunk. + + Import is deferred to method body so core/ stays pure at import time. + """ + metadata = dict(doc.metadata) if doc.metadata else {} + # Milvus assigns the primary key `_id` as INT64 (auto_id), so the value + # comes back as a Python int. Chunk.id is typed `str`, so coerce here + # rather than loosen the model — keeps the domain type strict while the + # store-specific shape is contained in the conversion boundary. + raw_id = metadata.pop("_id", None) + chunk_id = str(raw_id) if raw_id is not None else str(uuid.uuid4()) + return cls( + id=chunk_id, + document_id=metadata.pop("file_id", ""), + text=doc.page_content, + partition=metadata.pop("partition", "default"), + page_number=metadata.pop("page", None), + chunk_type=_coerce_chunk_type(metadata.pop("chunk_type", "text")), + metadata=metadata, + ) + + def to_langchain(self, *, with_id: bool = True) -> Any: + """Convert back to a LangChain Document. + + Import is deferred to method body so core/ stays pure at import time. + + Args: + with_id: When True (default), stamp ``_id`` into metadata. + Pass False on the write path — ``_id`` is auto-assigned by + Milvus (auto_id=True) and must not appear in the insert payload. + """ + from langchain_core.documents.base import Document + + metadata = { + **self.metadata, + "file_id": self.document_id, + "partition": self.partition, + "page": self.page_number, + "chunk_type": self.chunk_type.value, + } + if with_id: + metadata["_id"] = self.id + return Document(page_content=self.text, metadata=metadata) diff --git a/openrag/core/models/contextualization.py b/openrag/core/models/contextualization.py new file mode 100644 index 000000000..e30c2313a --- /dev/null +++ b/openrag/core/models/contextualization.py @@ -0,0 +1,19 @@ +"""Contextualized query models — query rewriting, HyDE, reasoning.""" + +from __future__ import annotations + +from pydantic import BaseModel, Field + + +class ContextualizedQuery(BaseModel): + """A user query after contextualization by the LLM. + + May include rewritten variants, a hypothetical document (HyDE), + and sub-queries for multi-step retrieval. + """ + + original: str = "" + query_list: list[str] = Field(default_factory=list) + intent: str = "qa" + hypothetical_doc: str | None = None + sub_queries: list[str] = Field(default_factory=list) diff --git a/openrag/core/models/conversation.py b/openrag/core/models/conversation.py new file mode 100644 index 000000000..61b6977cf --- /dev/null +++ b/openrag/core/models/conversation.py @@ -0,0 +1,32 @@ +"""Conversation and message domain models.""" + +from __future__ import annotations + +import uuid +from datetime import UTC, datetime +from typing import Any + +from pydantic import BaseModel, Field + + +class Message(BaseModel): + """A single message within a conversation.""" + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + conversation_id: str = "" + role: str = "user" + content: str = "" + sources_json: list[dict[str, Any]] | None = None + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + + +class Conversation(BaseModel): + """A persistent conversation between a user and the RAG system.""" + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + user_id: int = 0 + partition_scope: str = "default" + title: str = "" + messages: list[Message] = Field(default_factory=list) + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) diff --git a/openrag/core/models/document.py b/openrag/core/models/document.py new file mode 100644 index 000000000..459a8e637 --- /dev/null +++ b/openrag/core/models/document.py @@ -0,0 +1,226 @@ +"""Document — the input to the indexing pipeline.""" + +from __future__ import annotations + +import asyncio +import base64 +import os +import tempfile +import uuid +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from datetime import UTC, datetime +from enum import Enum +from pathlib import Path +from typing import Any + +from pydantic import BaseModel, Field + + +class DocumentType(str, Enum): + PDF = "pdf" + TEXT = "text" + HTML = "html" + MARKDOWN = "markdown" + IMAGE = "image" + AUDIO = "audio" + VIDEO = "video" + DOCX = "docx" + PPTX = "pptx" + DOC = "doc" + EML = "eml" + + +class TextBlock(BaseModel): + """A block of text extracted from a document.""" + + text: str + page_number: int | None = None + block_type: str = "paragraph" + metadata: dict[str, Any] = Field(default_factory=dict) + + +class ImageBlock(BaseModel): + """An image extracted from a document. + + Parser→caption contract: + - Parsers emit ``ImageBlock`` with ``caption=None``. A downstream + caption stage fills it in via a VLM. + - When the source text contains a placeholder for the image + (``![alt](data:image/...)``, ``![](pptx-image-3)``, ``![](marker-key-7)``, + …), the parser stores the exact placeholder string in + ``metadata["markdown_ref"]``. The caption stage substitutes the + wrapped caption back into the corresponding ``TextBlock`` via + ``str.replace`` on that ref. + - When there is no in-text placeholder (standalone image uploads, + EML image attachments), ``metadata["markdown_ref"]`` is omitted; + the caption stage produces a free-standing captioned ``TextBlock`` + instead. + + Bytes vs. URL: + - Locally-extracted images set ``image_bytes`` (raw PNG / JPEG bytes) + and leave ``source_url`` as ``None``. + - Remote images parsed from a markdown ``![](http://…)`` ref leave + ``image_bytes`` empty and set ``source_url`` to the URL. A + downstream fetch stage may populate ``image_bytes`` later. + - The :attr:`image_url` property is the unified VLM-friendly form: + a ``data:`` URI built from the bytes when present, otherwise the + ``source_url`` as-is. + """ + + image_bytes: bytes = Field(default=b"", exclude=True, repr=False) + source_url: str | None = None + page_number: int | None = None + caption: str | None = None + mime_type: str = "image/png" + metadata: dict[str, Any] = Field(default_factory=dict) + + @property + def image_url(self) -> str: + """A VLM-friendly URL for this image. + + - Bytes present → ``data:{mime_type};base64,{...}`` URI. + - Otherwise → ``source_url`` if set, else empty string. + - On any encoding failure → falls back to ``source_url`` (or ""). + """ + if self.image_bytes: + try: + b64 = base64.b64encode(self.image_bytes).decode() + return f"data:{self.mime_type};base64,{b64}" + except Exception: + pass + return self.source_url or "" + + +class Document(BaseModel): + """A document before or during indexing.""" + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + filename: str = "" + content_type: DocumentType = DocumentType.TEXT + text: str | None = None + raw_bytes: bytes | None = Field(None, exclude=True) + partition: str = "default" + tags: list[str] = Field(default_factory=list) + metadata: dict[str, Any] = Field(default_factory=dict) + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + + @staticmethod + def detect_content_type(filename: str) -> DocumentType: + """Detect DocumentType from filename extension.""" + ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else "" + mapping = { + "pdf": DocumentType.PDF, + "txt": DocumentType.TEXT, + "md": DocumentType.MARKDOWN, + "html": DocumentType.HTML, + "htm": DocumentType.HTML, + "png": DocumentType.IMAGE, + "jpg": DocumentType.IMAGE, + "jpeg": DocumentType.IMAGE, + "mp3": DocumentType.AUDIO, + "wav": DocumentType.AUDIO, + "flac": DocumentType.AUDIO, + "ogg": DocumentType.AUDIO, + "aac": DocumentType.AUDIO, + "wma": DocumentType.AUDIO, + "mp4": DocumentType.VIDEO, + "flv": DocumentType.VIDEO, + "docx": DocumentType.DOCX, + "pptx": DocumentType.PPTX, + "doc": DocumentType.DOC, + "eml": DocumentType.EML, + } + return mapping.get(ext, DocumentType.TEXT) + + @classmethod + def from_langchain(cls, doc: Any) -> Document: + """Convert a LangChain Document to a domain Document.""" + metadata = dict(doc.metadata) if doc.metadata else {} + return cls( + filename=metadata.pop("source", ""), + text=doc.page_content, + partition=metadata.pop("partition", "default"), + metadata=metadata, + ) + + def to_langchain(self) -> Any: + """Convert back to a LangChain Document.""" + from langchain_core.documents.base import Document as LCDocument + + metadata = { + **self.metadata, + "source": self.filename, + "partition": self.partition, + } + return LCDocument(page_content=self.text or "", metadata=metadata) + + @asynccontextmanager + async def as_temporary_file(self, *, suffix: str | None = None) -> AsyncIterator[Path]: + """Materialize ``raw_bytes`` to a temporary file and yield its ``Path``. + + Parsers wrapping a sync library that requires a path on disk + (Marker, Whisper, MarkItDown, python-pptx, Spire.Doc, …) use this + helper instead of rolling their own ``NamedTemporaryFile`` dance. + The file is removed on context exit even if the body raises. + + ``suffix`` defaults to ``filename``'s extension, falling back to + a content-type-appropriate default. + """ + if self.raw_bytes is None: + raise ValueError("Document.as_temporary_file requires raw_bytes") + + if suffix is None: + suffix = Path(self.filename).suffix or _DEFAULT_TEMPFILE_SUFFIX.get(self.content_type, "") + + raw = self.raw_bytes + + def _write_temp() -> str: + # Close before yielding so sync callers (Marker/Whisper/MarkItDown/ + # python-pptx/Spire.Doc) can reopen the path on Windows, where + # NamedTemporaryFile(delete=True) holds an exclusive handle. + tf = tempfile.NamedTemporaryFile(suffix=suffix, delete=False) + try: + tf.write(raw) + finally: + tf.close() + return tf.name + + path = await asyncio.to_thread(_write_temp) + try: + yield Path(path) + finally: + await asyncio.to_thread(_safe_unlink, path) + + +def _safe_unlink(path: str) -> None: + """``os.unlink`` that swallows missing-file errors (sync callers may have already removed it).""" + try: + os.unlink(path) + except FileNotFoundError: + pass + + +_DEFAULT_TEMPFILE_SUFFIX: dict[DocumentType, str] = { + DocumentType.PDF: ".pdf", + DocumentType.DOCX: ".docx", + DocumentType.PPTX: ".pptx", + DocumentType.DOC: ".doc", + DocumentType.AUDIO: ".wav", + DocumentType.VIDEO: ".mp4", + DocumentType.EML: ".eml", + DocumentType.IMAGE: ".png", + DocumentType.HTML: ".html", + DocumentType.MARKDOWN: ".md", + DocumentType.TEXT: ".txt", +} + + +class ProcessedDocument(BaseModel): + """Document after parsing/extraction — contains text blocks and images.""" + + document_id: str = "" + text_blocks: list[TextBlock] = Field(default_factory=list) + images: list[ImageBlock] = Field(default_factory=list) + metadata: dict[str, Any] = Field(default_factory=dict) + page_count: int = 0 diff --git a/openrag/core/models/prompt.py b/openrag/core/models/prompt.py new file mode 100644 index 000000000..23f2b38a0 --- /dev/null +++ b/openrag/core/models/prompt.py @@ -0,0 +1,31 @@ +"""Prompt domain models.""" + +from __future__ import annotations + +import uuid +from datetime import UTC, datetime +from enum import Enum + +from pydantic import BaseModel, Field + + +class PromptType(str, Enum): + SYS_PROMPT = "sys_prompt" + QUERY_CONTEXTUALIZER = "query_contextualizer" + CHUNK_CONTEXTUALIZER = "chunk_contextualizer" + IMAGE_CAPTIONING = "image_captioning" + HYDE = "hyde" + MULTI_QUERY = "multi_query" + SPOKEN_STYLE_ANSWER = "spoken_style_answer" + + +class Prompt(BaseModel): + """A prompt template stored in the library.""" + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + prompt_type: str = "" + name: str = "" + content: str = "" + is_default: bool = False + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) diff --git a/openrag/core/models/query.py b/openrag/core/models/query.py new file mode 100644 index 000000000..b5dad680d --- /dev/null +++ b/openrag/core/models/query.py @@ -0,0 +1,108 @@ +"""Retrieval query domain model.""" + +from __future__ import annotations + +import logging +from datetime import datetime +from typing import Any, Literal + +from pydantic import BaseModel, Field + +logger = logging.getLogger(__name__) + + +class RetrievalQuery(BaseModel): + """A user query with retrieval parameters.""" + + text: str + partition: str = "default" + top_k: int = 10 + similarity_threshold: float = 0.6 + filters: dict[str, Any] = Field(default_factory=dict) + include_related: bool = False + include_ancestors: bool = False + related_limit: int = 10 + max_ancestor_depth: int | None = None + with_surrounding_chunks: bool = True + rerank: bool = True + + +class TemporalPredicate(BaseModel): + """A single date constraint on a document's creation date. + + Multiple predicates on the same ``Query`` are AND-combined. Closed + ranges (e.g. "last month") are encoded as two predicates, one per side. + """ + + field: Literal["created_at"] = Field( + default="created_at", + description="Document metadata field to filter on. Always `created_at` for now.", + ) + operator: Literal[">", "<", ">=", "<="] = Field( + description="Comparison operator applied to the date field.", + ) + value: str = Field( + description='ISO 8601 datetime with timezone, e.g. "2026-03-15T00:00:00+00:00".', + ) + + +class Query(BaseModel): + """A single vector-database search query plus optional temporal filters. + + Two predicates yield an AND-range; an exclusion range (e.g. "last year + except March") is expressed as two separate ``Query`` objects. + """ + + query: str = Field( + description="A semantically enriched, descriptive query for vector similarity search.", + ) + temporal_filters: list[TemporalPredicate] | None = Field( + default=None, + description="Date predicates on `created_at`, AND-combined.", + ) + + def to_milvus_filter(self) -> str | None: + """Render the AND-combined predicates as a Milvus filter expression. + + Pydantic validates the field/operator types up front. The ``value`` + field is parsed as ISO 8601 here defensively — predicates with an + unparseable value are dropped rather than crashing the search. + """ + if not self.temporal_filters: + return None + parts: list[str] = [] + for p in self.temporal_filters: + try: + parsed = datetime.fromisoformat(p.value) + except (TypeError, ValueError): + logger.warning( + "Dropping temporal predicate with non-ISO value: field=%s operator=%s value=%r", + p.field, + p.operator, + p.value, + ) + continue + if parsed.tzinfo is None: + logger.warning( + "Dropping temporal predicate without timezone: field=%s operator=%s value=%r", + p.field, + p.operator, + p.value, + ) + continue + parts.append(f'{p.field} {p.operator} ISO "{p.value}"') + if not parts: + return None + return " and ".join(parts) + + def __str__(self) -> str: + return f"Query: {self.query}, Filter: {self.to_milvus_filter()}" + + +class SearchQueries(BaseModel): + """Collection of sub-queries produced by query decomposition.""" + + query_list: list[Query] = Field(..., description="Search sub-queries to retrieve relevant documents.") + + def __str__(self) -> str: + return " --- ".join(str(q) for q in self.query_list) diff --git a/openrag/core/models/retrieval_response.py b/openrag/core/models/retrieval_response.py new file mode 100644 index 000000000..d4072015b --- /dev/null +++ b/openrag/core/models/retrieval_response.py @@ -0,0 +1,18 @@ +"""Retrieval response — end-to-end retrieval output.""" + +from __future__ import annotations + +from pydantic import BaseModel, Field + +from .retrieval_result import RetrievalResult + + +class RetrievalResponse(BaseModel): + """Complete response from a retrieval pipeline execution.""" + + query: str = "" + results: list[RetrievalResult] = Field(default_factory=list) + pipeline_used: str | None = None + partition: str = "default" + total_candidates: int = 0 + latency_ms: float | None = None diff --git a/openrag/core/models/retrieval_result.py b/openrag/core/models/retrieval_result.py new file mode 100644 index 000000000..9e544c35d --- /dev/null +++ b/openrag/core/models/retrieval_result.py @@ -0,0 +1,32 @@ +"""Retrieval result domain models — per-chunk scored results.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + + +class RetrievalResult(BaseModel): + """A single scored chunk from retrieval.""" + + chunk_id: str = "" + document_id: str = "" + text: str = "" + score: float = 0.0 + metadata: dict[str, Any] = Field(default_factory=dict) + rerank_score: float | None = None + page_number: int | None = None + + +class ScoredChunk(BaseModel): + """A chunk with both vector and rerank scores.""" + + chunk_id: str = "" + document_id: str = "" + text: str = "" + vector_score: float = 0.0 + rerank_score: float | None = None + combined_score: float = 0.0 + metadata: dict[str, Any] = Field(default_factory=dict) + page_number: int | None = None diff --git a/openrag/core/models/test_chunk.py b/openrag/core/models/test_chunk.py new file mode 100644 index 000000000..4195a4e56 --- /dev/null +++ b/openrag/core/models/test_chunk.py @@ -0,0 +1,49 @@ +"""Tests for Chunk model — backward-compat coercions on from_langchain.""" + +from __future__ import annotations + +from core.models.chunk import Chunk, ChunkType, _coerce_chunk_type +from langchain_core.documents.base import Document + + +def test_from_langchain_maps_legacy_image_chunk_type(): + """Pre-Phase-5 chunkers stamped chunk_type='image' (raw MDElement + literal). Upgraded deployments still have those values in Milvus — + Chunk.from_langchain must not crash on them (ultrareview).""" + doc = Document(page_content="caption", metadata={"chunk_type": "image", "_id": "x", "file_id": "f1"}) + chunk = Chunk.from_langchain(doc) + assert chunk.chunk_type == ChunkType.IMAGE_CAPTION + + +def test_from_langchain_unknown_chunk_type_falls_back_to_text(): + """Defensive: any historical value that isn't in the enum and isn't + in the legacy alias map should land on TEXT, not crash retrieval.""" + doc = Document(page_content="x", metadata={"chunk_type": "unknown_legacy_value"}) + chunk = Chunk.from_langchain(doc) + assert chunk.chunk_type == ChunkType.TEXT + + +def test_from_langchain_accepts_canonical_values(): + for value, expected in [ + ("text", ChunkType.TEXT), + ("table", ChunkType.TABLE), + ("image_caption", ChunkType.IMAGE_CAPTION), + ("contextualized", ChunkType.CONTEXTUALIZED), + ]: + doc = Document(page_content="x", metadata={"chunk_type": value}) + assert Chunk.from_langchain(doc).chunk_type == expected + + +def test_coerce_chunk_type_passthrough_for_enum_input(): + assert _coerce_chunk_type(ChunkType.TABLE) == ChunkType.TABLE + + +def test_from_langchain_coerces_int_milvus_id_to_string(): + """Milvus' `_id` primary key is INT64 (auto_id), so the value comes back + from the Ray actor as a Python int. Chunk.id is typed `str`; the + conversion boundary must coerce to avoid a ValidationError on every + retrieval call (CI api-tests regression).""" + doc = Document(page_content="hello", metadata={"_id": 466085833598567840, "file_id": "f1"}) + chunk = Chunk.from_langchain(doc) + assert chunk.id == "466085833598567840" + assert isinstance(chunk.id, str) diff --git a/openrag/core/models/user.py b/openrag/core/models/user.py new file mode 100644 index 000000000..5811b6032 --- /dev/null +++ b/openrag/core/models/user.py @@ -0,0 +1,101 @@ +"""User, role, partition assignment, API key, and OIDC session domain models.""" + +from __future__ import annotations + +import uuid +from datetime import UTC, datetime +from enum import Enum + +from pydantic import BaseModel, Field + + +class PartitionRole(str, Enum): + VIEWER = "viewer" + EDITOR = "editor" + OWNER = "owner" + + +class User(BaseModel): + """An OpenRAG user account. + + Supports three auth modes: + - OIDC/SSO: user matched by external_user_id (Keycloak sub claim), + session managed via OIDCSession. For browser users. + - API token: opaque or- prefixed token, SHA-256 hashed in DB. + For scripts, CI/CD, programmatic access. Legacy mode. + - Password + JWT: user logs in with email/password, receives + JWT access + refresh tokens. For programmatic access and + users without Keycloak. + """ + + id: int = 0 + display_name: str | None = None + external_user_id: str | None = None + email: str | None = None + password_hash: str | None = Field(None, exclude=True, repr=False) + is_admin: bool = False + is_active: bool = True + file_quota: int | None = None + file_count: int = 0 + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + partitions: list[UserPartition] = Field(default_factory=list) + + +class UserPartition(BaseModel): + """A user's membership in a partition with a role.""" + + user_id: int + partition: str + role: PartitionRole = PartitionRole.VIEWER + added_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + + +class ApiKey(BaseModel): + """An API key for programmatic access. + + The raw key is shown once on creation (key_prefix + random hex). + Only the hash is stored in DB. + """ + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + user_id: int = 0 + key_hash: str = "" + key_prefix: str = "" + name: str = "" + is_active: bool = True + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + expires_at: datetime | None = None + + +class TokenPayload(BaseModel): + """JWT token claims payload.""" + + sub: str = "" + type: str = "access" + role: str = "user" + exp: int = 0 + + +class OIDCSession(BaseModel): + """An active OIDC session linking a user to IdP tokens. + + Session token is opaque (stored hashed in DB). + IdP tokens (access, refresh, id) are Fernet-encrypted in DB — the auth + service encrypts before passing them in and decrypts after reading them + back; the repository just stores the bytes verbatim. + """ + + id: int = 0 + session_token_hash: str = "" + user_id: int = 0 + sid: str | None = None + sub: str = "" + id_token_encrypted: bytes | None = None + access_token_encrypted: bytes | None = None + refresh_token_encrypted: bytes | None = None + access_token_expires_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + session_expires_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + last_refresh_at: datetime | None = None + revoked_at: datetime | None = None diff --git a/openrag/core/models/workspace.py b/openrag/core/models/workspace.py new file mode 100644 index 000000000..7840013a7 --- /dev/null +++ b/openrag/core/models/workspace.py @@ -0,0 +1,25 @@ +"""Workspace domain model. + +A workspace is a named subset of files within a partition, used to scope +search and chat to a curated document set. Workspaces share a partition's +files (no copy) and a single file may belong to many workspaces. +""" + +from __future__ import annotations + +from datetime import UTC, datetime + +from pydantic import BaseModel, Field + + +class Workspace(BaseModel): + """A named subset of files within a partition.""" + + workspace_id: str + partition: str + display_name: str | None = None + created_by: int | None = None + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + + +__all__ = ["Workspace"] diff --git a/openrag/core/observability/__init__.py b/openrag/core/observability/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/core/ports/__init__.py b/openrag/core/ports/__init__.py new file mode 100644 index 000000000..db9c0cb5c --- /dev/null +++ b/openrag/core/ports/__init__.py @@ -0,0 +1,39 @@ +"""Port interfaces — repository ABCs + CatalogStore aggregate root.""" + +from .audit_log_repo import AuditLogRepository +from .catalog_store import CatalogStore +from .chunk_repo import ChunkRepository +from .conversation_repo import ConversationRepository +from .document_repo import DocumentRepository +from .entity_repo import EntityRepository +from .idempotency_repo import IdempotencyRepository +from .job_repo import JobRepository +from .model_endpoint_repo import ModelEndpointRepository +from .oidc_session_repo import OIDCSessionRepository +from .partition_membership_repo import PartitionMembershipRepository +from .partition_repo import PartitionRepository +from .preset_repo import PresetRepository +from .prompt_repo import PromptRepository +from .topic_tag_repo import TopicTagRepository +from .user_repo import UserRepository +from .workspace_repo import WorkspaceRepository + +__all__ = [ + "AuditLogRepository", + "CatalogStore", + "ChunkRepository", + "ConversationRepository", + "DocumentRepository", + "EntityRepository", + "IdempotencyRepository", + "JobRepository", + "ModelEndpointRepository", + "OIDCSessionRepository", + "PartitionMembershipRepository", + "PartitionRepository", + "PresetRepository", + "PromptRepository", + "TopicTagRepository", + "UserRepository", + "WorkspaceRepository", +] diff --git a/openrag/core/ports/audit_log_repo.py b/openrag/core/ports/audit_log_repo.py new file mode 100644 index 000000000..2318bc4d1 --- /dev/null +++ b/openrag/core/ports/audit_log_repo.py @@ -0,0 +1,24 @@ +"""Audit log repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +class AuditLogRepository(ABC): + """Append-only audit trail.""" + + @abstractmethod + async def insert( + self, + user_id: int | None, + action: str, + resource_type: str, + resource_id: str | None = None, + details_json: dict | None = None, + request_id: str | None = None, + ) -> None: ... + + @abstractmethod + async def query(self, filters: dict[str, Any], offset: int = 0, limit: int = 50) -> list[dict]: ... diff --git a/openrag/core/ports/catalog_store.py b/openrag/core/ports/catalog_store.py new file mode 100644 index 000000000..cf0ff6f26 --- /dev/null +++ b/openrag/core/ports/catalog_store.py @@ -0,0 +1,103 @@ +"""CatalogStore — aggregate root composing all repository ports. + +Concrete implementations (e.g. PostgresStore) own the connection pool +and compose per-entity repository instances. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from .audit_log_repo import AuditLogRepository +from .chunk_repo import ChunkRepository +from .conversation_repo import ConversationRepository +from .document_repo import DocumentRepository +from .entity_repo import EntityRepository +from .idempotency_repo import IdempotencyRepository +from .job_repo import JobRepository +from .model_endpoint_repo import ModelEndpointRepository +from .oidc_session_repo import OIDCSessionRepository +from .partition_membership_repo import PartitionMembershipRepository +from .partition_repo import PartitionRepository +from .preset_repo import PresetRepository +from .prompt_repo import PromptRepository +from .topic_tag_repo import TopicTagRepository +from .user_repo import UserRepository +from .workspace_repo import WorkspaceRepository + + +class CatalogStore(ABC): + """Abstract interface for the relational catalog backing store. + + Concrete implementations (e.g. PostgresStore) live in the services layer. + """ + + @abstractmethod + async def initialize(self) -> None: ... + + @abstractmethod + async def shutdown(self) -> None: ... + + @property + @abstractmethod + def document_repo(self) -> DocumentRepository: ... + + @property + @abstractmethod + def job_repo(self) -> JobRepository: ... + + @property + @abstractmethod + def user_repo(self) -> UserRepository: ... + + @property + @abstractmethod + def prompt_repo(self) -> PromptRepository: ... + + @property + @abstractmethod + def partition_repo(self) -> PartitionRepository: ... + + @property + @abstractmethod + def membership_repo(self) -> PartitionMembershipRepository: ... + + @property + @abstractmethod + def model_endpoint_repo(self) -> ModelEndpointRepository: ... + + @property + @abstractmethod + def preset_repo(self) -> PresetRepository: ... + + @property + @abstractmethod + def chunk_repo(self) -> ChunkRepository: ... + + @property + @abstractmethod + def entity_repo(self) -> EntityRepository: ... + + @property + @abstractmethod + def topic_tag_repo(self) -> TopicTagRepository: ... + + @property + @abstractmethod + def conversation_repo(self) -> ConversationRepository: ... + + @property + @abstractmethod + def audit_log_repo(self) -> AuditLogRepository: ... + + @property + @abstractmethod + def idempotency_repo(self) -> IdempotencyRepository: ... + + @property + @abstractmethod + def oidc_session_repo(self) -> OIDCSessionRepository: ... + + @property + @abstractmethod + def workspace_repo(self) -> WorkspaceRepository: ... diff --git a/openrag/core/ports/chunk_repo.py b/openrag/core/ports/chunk_repo.py new file mode 100644 index 000000000..e4b986d8e --- /dev/null +++ b/openrag/core/ports/chunk_repo.py @@ -0,0 +1,39 @@ +"""Chunk repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class ChunkRepository(ABC): + """Bulk CRUD operations for text chunks.""" + + @abstractmethod + async def bulk_insert(self, chunks: list[dict]) -> int: + """Insert multiple chunks. Returns count of inserted rows.""" + ... + + @abstractmethod + async def get_by_ids(self, chunk_ids: list[str]) -> list[dict]: + """Batch fetch chunks by IDs.""" + ... + + @abstractmethod + async def get_by_document_id(self, document_id: str) -> list[dict]: + """Fetch all chunks for a document, ordered by chunk_index.""" + ... + + @abstractmethod + async def delete_by_document_id(self, document_id: str) -> int: + """Delete all chunks belonging to a document. Returns count.""" + ... + + @abstractmethod + async def delete_by_partition(self, partition: str) -> int: + """Delete all chunks in a partition. Returns count.""" + ... + + @abstractmethod + async def bm25_search(self, query_text: str, partition: str, top_k: int = 20) -> list[dict]: + """Full-text search using tsvector column with ts_rank scoring.""" + ... diff --git a/openrag/core/ports/conversation_repo.py b/openrag/core/ports/conversation_repo.py new file mode 100644 index 000000000..6b1938f11 --- /dev/null +++ b/openrag/core/ports/conversation_repo.py @@ -0,0 +1,29 @@ +"""Conversation repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from openrag.core.models.conversation import Conversation, Message + + +class ConversationRepository(ABC): + """CRUD operations for conversations and messages.""" + + @abstractmethod + async def create_conversation(self, conversation: Conversation) -> Conversation: ... + + @abstractmethod + async def get_conversation(self, conversation_id: str) -> Conversation | None: ... + + @abstractmethod + async def list_conversations(self, user_id: int, partition: str | None = None) -> list[Conversation]: ... + + @abstractmethod + async def delete_conversation(self, conversation_id: str) -> bool: ... + + @abstractmethod + async def add_message(self, message: Message) -> Message: ... + + @abstractmethod + async def list_messages(self, conversation_id: str) -> list[Message]: ... diff --git a/openrag/core/ports/document_repo.py b/openrag/core/ports/document_repo.py new file mode 100644 index 000000000..b2962bfa0 --- /dev/null +++ b/openrag/core/ports/document_repo.py @@ -0,0 +1,50 @@ +"""Document repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from openrag.core.models.catalog import DocumentRecord + + +class DocumentRepository(ABC): + """CRUD operations for documents.""" + + @abstractmethod + async def create_document(self, doc: DocumentRecord) -> DocumentRecord: ... + + @abstractmethod + async def get_document(self, document_id: str) -> DocumentRecord | None: ... + + @abstractmethod + async def list_documents( + self, + partition: str | list[str] | None = None, + status: str | None = None, + offset: int = 0, + limit: int = 50, + ) -> list[DocumentRecord]: ... + + @abstractmethod + async def update_document(self, document_id: str, **fields: Any) -> DocumentRecord | None: ... + + @abstractmethod + async def delete_document(self, document_id: str) -> bool: ... + + @abstractmethod + async def delete_documents_by_partition(self, partition: str) -> int: ... + + @abstractmethod + async def count_documents(self, partition: str | list[str] | None = None, status: str | None = None) -> int: ... + + @abstractmethod + async def file_exists_in_partition(self, file_id: str, partition: str) -> bool: ... + + @abstractmethod + async def get_file_ids_by_relationship(self, partition: str, relationship_id: str) -> list[str]: ... + + @abstractmethod + async def get_ancestor_file_ids( + self, partition: str, file_id: str, max_ancestor_depth: int | None = None + ) -> list[str]: ... diff --git a/openrag/core/ports/entity_repo.py b/openrag/core/ports/entity_repo.py new file mode 100644 index 000000000..a8520484f --- /dev/null +++ b/openrag/core/ports/entity_repo.py @@ -0,0 +1,21 @@ +"""Entity repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class EntityRepository(ABC): + """CRUD operations for extracted entities.""" + + @abstractmethod + async def upsert(self, partition: str, entity_type: str, canonical_name: str, aliases: list[str]) -> str: ... + + @abstractmethod + async def search(self, partition: str, query: str, top_k: int = 10) -> list[dict]: ... + + @abstractmethod + async def get_by_document(self, document_id: str) -> list[dict]: ... + + @abstractmethod + async def delete_by_document(self, document_id: str) -> int: ... diff --git a/openrag/core/ports/idempotency_repo.py b/openrag/core/ports/idempotency_repo.py new file mode 100644 index 000000000..8a5f65dcd --- /dev/null +++ b/openrag/core/ports/idempotency_repo.py @@ -0,0 +1,15 @@ +"""Idempotency key repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class IdempotencyRepository(ABC): + """Cache for request idempotency keys.""" + + @abstractmethod + async def get_by_hash(self, key_hash: str) -> dict | None: ... + + @abstractmethod + async def store(self, key_hash: str, http_method: str, status_code: int, response_body: bytes) -> None: ... diff --git a/openrag/core/ports/job_repo.py b/openrag/core/ports/job_repo.py new file mode 100644 index 000000000..5e20b6830 --- /dev/null +++ b/openrag/core/ports/job_repo.py @@ -0,0 +1,24 @@ +"""Job repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from openrag.core.models.catalog import IndexationJob + + +class JobRepository(ABC): + """CRUD operations for indexation jobs.""" + + @abstractmethod + async def create_job(self, job: IndexationJob) -> IndexationJob: ... + + @abstractmethod + async def get_job(self, job_id: str) -> IndexationJob | None: ... + + @abstractmethod + async def list_jobs(self, status: str | None = None, offset: int = 0, limit: int = 50) -> list[IndexationJob]: ... + + @abstractmethod + async def update_job(self, job_id: str, **fields: Any) -> IndexationJob | None: ... diff --git a/openrag/core/ports/model_endpoint_repo.py b/openrag/core/ports/model_endpoint_repo.py new file mode 100644 index 000000000..e907a0ce0 --- /dev/null +++ b/openrag/core/ports/model_endpoint_repo.py @@ -0,0 +1,21 @@ +"""Model endpoint repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class ModelEndpointRepository(ABC): + """CRUD operations for model endpoint configurations.""" + + @abstractmethod + async def get(self, name: str, model_type: str) -> dict | None: ... + + @abstractmethod + async def list_all(self, model_type: str | None = None) -> list[dict]: ... + + @abstractmethod + async def upsert(self, name: str, model_type: str, config: dict) -> dict: ... + + @abstractmethod + async def delete(self, name: str, model_type: str) -> bool: ... diff --git a/openrag/core/ports/oidc_session_repo.py b/openrag/core/ports/oidc_session_repo.py new file mode 100644 index 000000000..e42b9cfa6 --- /dev/null +++ b/openrag/core/ports/oidc_session_repo.py @@ -0,0 +1,44 @@ +"""OIDC session repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from openrag.core.models.user import OIDCSession + + +class OIDCSessionRepository(ABC): + """CRUD operations for OIDC sessions.""" + + @abstractmethod + async def create_session(self, session: OIDCSession) -> OIDCSession: ... + + @abstractmethod + async def get_by_token_hash(self, token_hash: str) -> OIDCSession | None: ... + + @abstractmethod + async def get_by_id(self, session_id: int) -> OIDCSession | None: ... + + @abstractmethod + async def get_by_sid(self, sid: str) -> list[OIDCSession]: ... + + @abstractmethod + async def update_session(self, session_id: int, **fields) -> OIDCSession | None: ... + + @abstractmethod + async def revoke_session(self, session_id: int) -> bool: ... + + @abstractmethod + async def revoke_by_sid(self, sid: str) -> int: + """Revoke all sessions with a given OIDC session ID (back-channel logout).""" + ... + + @abstractmethod + async def revoke_by_user(self, user_id: int) -> int: + """Revoke all sessions for a user.""" + ... + + @abstractmethod + async def delete_expired(self) -> int: + """Delete sessions past their session_expires_at. Returns count.""" + ... diff --git a/openrag/core/ports/partition_membership_repo.py b/openrag/core/ports/partition_membership_repo.py new file mode 100644 index 000000000..786d02e47 --- /dev/null +++ b/openrag/core/ports/partition_membership_repo.py @@ -0,0 +1,36 @@ +"""Partition membership repository interface. + +Split out of :class:`~openrag.core.ports.user_repo.UserRepository` so the +catalog matches the 7A.2 one-repo-per-entity layout — ``partition_memberships`` +is its own table and gets its own port. ``UserRepository`` still reads +memberships internally to hydrate the ``User`` aggregate's ``partitions`` +field, but all membership *management* lives here. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from openrag.core.models.user import PartitionRole, UserPartition + + +class PartitionMembershipRepository(ABC): + """CRUD operations for partition memberships (owner/editor/viewer).""" + + @abstractmethod + async def assign_partition(self, assignment: UserPartition) -> UserPartition: ... + + @abstractmethod + async def remove_partition(self, user_id: int, partition: str) -> bool: ... + + @abstractmethod + async def list_user_partitions(self, user_id: int) -> list[UserPartition]: ... + + @abstractmethod + async def list_partition_users(self, partition: str) -> list[UserPartition]: ... + + @abstractmethod + async def update_partition_role(self, user_id: int, partition: str, role: PartitionRole) -> bool: ... + + @abstractmethod + async def count_partition_users(self, partition: str) -> int: ... diff --git a/openrag/core/ports/partition_repo.py b/openrag/core/ports/partition_repo.py new file mode 100644 index 000000000..00898ce76 --- /dev/null +++ b/openrag/core/ports/partition_repo.py @@ -0,0 +1,24 @@ +"""Partition repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class PartitionRepository(ABC): + """CRUD operations for partitions.""" + + @abstractmethod + async def create_partition(self, name: str, user_id: int | None = None) -> dict: ... + + @abstractmethod + async def get_partition(self, name: str) -> dict | None: ... + + @abstractmethod + async def list_partitions(self) -> list[dict]: ... + + @abstractmethod + async def delete_partition(self, name: str) -> bool: ... + + @abstractmethod + async def partition_exists(self, name: str) -> bool: ... diff --git a/openrag/core/ports/preset_repo.py b/openrag/core/ports/preset_repo.py new file mode 100644 index 000000000..363b654fc --- /dev/null +++ b/openrag/core/ports/preset_repo.py @@ -0,0 +1,21 @@ +"""Preset repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class PresetRepository(ABC): + """CRUD operations for pipeline presets.""" + + @abstractmethod + async def get(self, name: str, preset_type: str) -> dict | None: ... + + @abstractmethod + async def list_all(self, preset_type: str | None = None) -> list[dict]: ... + + @abstractmethod + async def upsert(self, name: str, preset_type: str, config: dict) -> dict: ... + + @abstractmethod + async def delete(self, name: str, preset_type: str) -> bool: ... diff --git a/openrag/core/ports/prompt_repo.py b/openrag/core/ports/prompt_repo.py new file mode 100644 index 000000000..43336f2f5 --- /dev/null +++ b/openrag/core/ports/prompt_repo.py @@ -0,0 +1,32 @@ +"""Prompt repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from openrag.core.models.prompt import Prompt + + +class PromptRepository(ABC): + """CRUD operations for prompt templates.""" + + @abstractmethod + async def create_prompt(self, prompt: Prompt) -> Prompt: ... + + @abstractmethod + async def get_prompt(self, prompt_id: str) -> Prompt | None: ... + + @abstractmethod + async def get_by_type(self, prompt_type: str) -> list[Prompt]: ... + + @abstractmethod + async def get_active(self, prompt_type: str) -> Prompt | None: ... + + @abstractmethod + async def list_prompts(self) -> list[Prompt]: ... + + @abstractmethod + async def update_prompt(self, prompt_id: str, content: str) -> Prompt | None: ... + + @abstractmethod + async def delete_prompt(self, prompt_id: str) -> bool: ... diff --git a/openrag/core/ports/topic_tag_repo.py b/openrag/core/ports/topic_tag_repo.py new file mode 100644 index 000000000..99287f46c --- /dev/null +++ b/openrag/core/ports/topic_tag_repo.py @@ -0,0 +1,21 @@ +"""Topic tag repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class TopicTagRepository(ABC): + """CRUD operations for document topic tags.""" + + @abstractmethod + async def bulk_insert(self, tags: list[dict]) -> int: ... + + @abstractmethod + async def get_by_document(self, document_id: str) -> list[dict]: ... + + @abstractmethod + async def delete_by_document(self, document_id: str) -> int: ... + + @abstractmethod + async def search(self, partition: str, tag: str, top_k: int = 10) -> list[dict]: ... diff --git a/openrag/core/ports/user_repo.py b/openrag/core/ports/user_repo.py new file mode 100644 index 000000000..747eadf0c --- /dev/null +++ b/openrag/core/ports/user_repo.py @@ -0,0 +1,66 @@ +"""User repository interface — users and API keys. + +Partition memberships moved to +:class:`~openrag.core.ports.partition_membership_repo.PartitionMembershipRepository` +(7A.2 one-repo-per-entity layout). +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from openrag.core.models.user import ApiKey, User + + +class UserRepository(ABC): + """CRUD operations for users and API keys. + + Supports three auth modes: + - OIDC/SSO: lookup by external_user_id + - API token: lookup by token hash (legacy or- tokens) + - Password + JWT: lookup by email, verify password hash + """ + + # ── User CRUD ───────────────────────────────────────────────────── + + @abstractmethod + async def create_user(self, user: User) -> User: ... + + @abstractmethod + async def get_user(self, user_id: int) -> User | None: ... + + @abstractmethod + async def get_user_by_email(self, email: str) -> User | None: ... + + @abstractmethod + async def get_user_by_token(self, token_hash: str) -> User | None: ... + + @abstractmethod + async def get_user_by_external_id(self, external_id: str) -> User | None: ... + + @abstractmethod + async def list_users(self, offset: int = 0, limit: int = 50) -> list[User]: ... + + @abstractmethod + async def update_user(self, user_id: int, **fields: Any) -> User | None: ... + + @abstractmethod + async def delete_user(self, user_id: int) -> bool: ... + + @abstractmethod + async def count_users(self) -> int: ... + + # ── API keys ────────────────────────────────────────────────────── + + @abstractmethod + async def create_api_key(self, key: ApiKey) -> ApiKey: ... + + @abstractmethod + async def get_api_keys_by_prefix(self, prefix: str) -> list[ApiKey]: ... + + @abstractmethod + async def list_api_keys_for_user(self, user_id: int) -> list[ApiKey]: ... + + @abstractmethod + async def delete_api_key(self, key_id: str) -> bool: ... diff --git a/openrag/core/ports/workspace_repo.py b/openrag/core/ports/workspace_repo.py new file mode 100644 index 000000000..7d3361029 --- /dev/null +++ b/openrag/core/ports/workspace_repo.py @@ -0,0 +1,59 @@ +"""Workspace repository interface — workspaces + workspace_files join.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from openrag.core.models.workspace import Workspace + + +class WorkspaceRepository(ABC): + """CRUD operations for workspaces and their file membership. + + A workspace is a named subset of files within a partition. The actual + file contents are not copied — the join table ``workspace_files`` + references the canonical ``files`` row by integer PK so file deletion + cascades correctly. + """ + + # ── Workspace lifecycle ─────────────────────────────────────────── + + @abstractmethod + async def create_workspace(self, workspace: Workspace) -> Workspace: ... + + @abstractmethod + async def get_workspace(self, workspace_id: str) -> Workspace | None: ... + + @abstractmethod + async def list_workspaces(self, partition: str) -> list[Workspace]: ... + + @abstractmethod + async def delete_workspace(self, workspace_id: str) -> list[str]: + """Delete a workspace, return file_ids that no longer belong to any workspace.""" + ... + + # ── Workspace ↔ file membership ─────────────────────────────────── + + @abstractmethod + async def add_files_to_workspace(self, workspace_id: str, file_ids: list[str]) -> list[str]: + """Attach files to a workspace. Returns the file_ids that could not be resolved.""" + ... + + @abstractmethod + async def remove_file_from_workspace(self, workspace_id: str, file_id: str) -> bool: ... + + @abstractmethod + async def list_workspace_files(self, workspace_id: str) -> list[str]: ... + + @abstractmethod + async def get_file_workspaces(self, file_id: str, partition: str) -> list[str]: ... + + @abstractmethod + async def get_existing_file_ids(self, partition: str, file_ids: list[str]) -> set[str]: + """Return the subset of ``file_ids`` that actually exist in ``partition``.""" + ... + + @abstractmethod + async def remove_file_from_all_workspaces(self, file_id: str, partition: str) -> None: + """Detach ``file_id`` from every workspace in ``partition``.""" + ... diff --git a/openrag/core/prompts/__init__.py b/openrag/core/prompts/__init__.py new file mode 100644 index 000000000..275f496f4 --- /dev/null +++ b/openrag/core/prompts/__init__.py @@ -0,0 +1,68 @@ +"""Prompt assembly helpers — pure string-formatting builders + disk loader.""" + +from .chat_prompt_builder import ( + EMPTY_CONTEXT_MESSAGE, + SOURCE_SEPARATOR, + WebSourceLike, + format_context, + format_web_context, + prepend_system_prompt, +) +from .contextualization_builder import ( + BASE_CHUNK_FORMAT, + CHUNK_FORMAT, + wrap_chunk_with_context, +) +from .contextualization_builder import ( + build_messages as build_contextualization_messages, +) +from .contextualization_builder import ( + build_user_message as build_contextualization_user_message, +) +from .map_reduce_builder import ( + SYSTEM_PROMPT_MAP, + USER_PROMPT_TEMPLATE, + build_map_messages, +) +from .query_rewriter import ( + MULTI_QUERY_SEPARATOR, + build_hyde_prompt, + build_multi_query_prompt, + split_multi_query_response, +) +from .template_loader import load_template, load_template_by_key +from .vlm_prompt_builder import ( + build_caption_messages, + wrap_caption, +) + +__all__ = [ + # template loader + "load_template", + "load_template_by_key", + # chat + "format_context", + "format_web_context", + "prepend_system_prompt", + "SOURCE_SEPARATOR", + "EMPTY_CONTEXT_MESSAGE", + "WebSourceLike", + # contextualization + "BASE_CHUNK_FORMAT", + "CHUNK_FORMAT", + "build_contextualization_messages", + "build_contextualization_user_message", + "wrap_chunk_with_context", + # query rewriter + "MULTI_QUERY_SEPARATOR", + "build_hyde_prompt", + "build_multi_query_prompt", + "split_multi_query_response", + # map-reduce + "build_map_messages", + "SYSTEM_PROMPT_MAP", + "USER_PROMPT_TEMPLATE", + # VLM + "build_caption_messages", + "wrap_caption", +] diff --git a/openrag/core/prompts/chat_prompt_builder.py b/openrag/core/prompts/chat_prompt_builder.py new file mode 100644 index 000000000..bf63093aa --- /dev/null +++ b/openrag/core/prompts/chat_prompt_builder.py @@ -0,0 +1,138 @@ +"""Chat-completion prompt builder. + +Pure helpers extracted from ``components/utils.py`` and ``components/pipeline.py``: + +* ``format_context`` — fit document snippets into a token budget, + numbering each as ``[Source N]``. +* ``format_web_context`` — same, for web-search results, with continuous + numbering across RAG and web sources. +* ``prepend_system_prompt`` — clone a message list and prepend a system + prompt rendered against ``context`` and + ``current_date``. +* ``SOURCE_SEPARATOR`` — separator emitted between consecutive sources. + +Tokenizers are injected as ``Callable[[str], int]`` so this module stays pure +(no LLM client, no LangChain). +""" + +from __future__ import annotations + +import copy +from collections.abc import Callable +from typing import Protocol + +from core.utils.text import sanitize_text + +SOURCE_SEPARATOR = "-" * 10 + "\n\n" +EMPTY_CONTEXT_MESSAGE = "No document found from the database" + + +class WebSourceLike(Protocol): + """Minimal shape needed from a web-search result.""" + + title: str + url: str + snippet: str + content: str | None + + +def format_context( + texts: list[str], + max_context_tokens: int, + length_function: Callable[[str], int], + *, + number_sources: bool = True, +) -> tuple[str, list[int]]: + """Render ``texts`` as numbered ``[Source N]`` blocks within a token budget. + + Args: + texts: Document texts (e.g. ``[d.page_content for d in docs]``). + max_context_tokens: Maximum total tokens for the context. + length_function: Token counter, e.g. ``llm.get_num_tokens``. + number_sources: If ``True``, prefix each block with ``[Source N]\\n``. + + Returns: + ``(formatted_text, included_indices)`` — ``included_indices`` is the + positions in ``texts`` that fit within the budget; callers use it to + filter associated metadata down to the same set. + """ + if not texts: + return EMPTY_CONTEXT_MESSAGE, [] + + reduced: list[str] = [] + included: list[int] = [] + total_tokens = 0 + + for i, text in enumerate(texts): + prefix = f"[Source {len(reduced) + 1}]\n" if number_sources else "" + n_tokens = length_function(text) + if prefix: + n_tokens += length_function(prefix) + if total_tokens + n_tokens > max_context_tokens: + break + reduced.append(f"{prefix}{text}") + included.append(i) + total_tokens += n_tokens + + return SOURCE_SEPARATOR.join(reduced), included + + +def format_web_context( + web_results: list[WebSourceLike], + length_function: Callable[[str], int], + *, + start_index: int = 1, + max_tokens: int = 2000, +) -> tuple[str, list[int], int]: + """Render web results as numbered ``[Source N]`` blocks within a token budget. + + Uses ``result.content`` when present, falling back to ``result.snippet``. + + Args: + web_results: Web-search result objects (matching ``WebSourceLike``). + length_function: Token counter. + start_index: First source number — set to ``len(rag_sources) + 1`` so + web sources continue numbering after RAG sources. + max_tokens: Maximum total tokens for the web context. + + Returns: + ``(formatted_text, source_numbers_used, total_tokens_used)``. + """ + if not web_results: + return "", [], 0 + + parts: list[str] = [] + source_numbers: list[int] = [] + total_tokens = 0 + + for i, result in enumerate(web_results): + n = start_index + i + title = sanitize_text(result.title) + body_raw = result.content if result.content else result.snippet + body = sanitize_text(body_raw) if body_raw else "" + block = f"[Source {n}]\n{title}\n{body}" + block_tokens = length_function(block) + if total_tokens + block_tokens > max_tokens: + break + parts.append(block) + source_numbers.append(n) + total_tokens += block_tokens + + return SOURCE_SEPARATOR.join(parts), source_numbers, total_tokens + + +def prepend_system_prompt( + messages: list[dict], + system_template: str, + *, + context: str, + current_date: str, +) -> list[dict]: + """Return a deep-copied message list with a rendered system prompt prepended. + + ``system_template`` must contain ``{context}`` and ``{current_date}``. + """ + out = copy.deepcopy(messages) + rendered = system_template.format(context=context, current_date=current_date) + out.insert(0, {"role": "system", "content": rendered}) + return out diff --git a/openrag/core/prompts/contextualization_builder.py b/openrag/core/prompts/contextualization_builder.py new file mode 100644 index 000000000..c5c45dccd --- /dev/null +++ b/openrag/core/prompts/contextualization_builder.py @@ -0,0 +1,77 @@ +"""Chunk-contextualization prompt builder. + +Pure helpers extracted from ``components/indexer/chunker/chunker.py``. They +produce the system+user message pair sent to the LLM when generating +chunk-level context, and assemble the final wrapped chunk text used downstream. + +Format strings: + BASE_CHUNK_FORMAT — chunk wrapping when no LLM context is generated + CHUNK_FORMAT — chunk wrapping with leading [CONTEXT] block +""" + +from __future__ import annotations + +BASE_CHUNK_FORMAT = "* filename: {filename}\n\n[CHUNK_START]\n\n{content}\n\n[CHUNK_END]" +CHUNK_FORMAT = "[CONTEXT]\n\n{chunk_context}\n\n" + BASE_CHUNK_FORMAT + + +def build_user_message( + filename: str, + first_chunks_text: list[str], + prev_chunks_text: list[str], + current_chunk_text: str, + lang: str = "en", +) -> str: + """Render the user-message body for a single chunk-contextualization call. + + The system prompt is loaded from disk (``CHUNK_CONTEXTUALIZER_PROMPT``) and + paired with this user message by the caller. + """ + first = "\n--\n".join(first_chunks_text) + previous = "\n--\n".join(prev_chunks_text) + return ( + "\n" + " Here is the context to consider for generating the context:\n" + f" - Filename: {filename}\n" + " - First chunks:\n" + f" {first}\n\n" + " - Previous chunks:\n" + f" {previous}\n\n" + f" Here is the current chunk to contextualize strictly in this {lang} language:\n" + " - Current chunk:\n\n" + f" {current_chunk_text}\n " + ) + + +def build_messages( + system_prompt: str, + filename: str, + first_chunks_text: list[str], + prev_chunks_text: list[str], + current_chunk_text: str, + lang: str = "en", +) -> list[dict[str, str]]: + """Build the system+user message list for a chunk-contextualization call.""" + user = build_user_message( + filename=filename, + first_chunks_text=first_chunks_text, + prev_chunks_text=prev_chunks_text, + current_chunk_text=current_chunk_text, + lang=lang, + ) + return [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user}, + ] + + +def wrap_chunk_with_context(content: str, filename: str, chunk_context: str = "") -> str: + """Wrap a chunk in the ``[CONTEXT] ... [CHUNK_START] ... [CHUNK_END]`` envelope. + + If ``chunk_context`` is empty or whitespace-only, only the BASE_CHUNK_FORMAT + (no [CONTEXT] block) is used — preserves the legacy behavior for chunkers + that don't run contextualization. + """ + if chunk_context and chunk_context.strip(): + return CHUNK_FORMAT.format(content=content, chunk_context=chunk_context, filename=filename) + return BASE_CHUNK_FORMAT.format(content=content, filename=filename) diff --git a/openrag/core/prompts/map_reduce_builder.py b/openrag/core/prompts/map_reduce_builder.py new file mode 100644 index 000000000..632db5f8c --- /dev/null +++ b/openrag/core/prompts/map_reduce_builder.py @@ -0,0 +1,38 @@ +"""Map-reduce prompt builder. + +The orchestrator (Phase 8) loops over chunks, calls the LLM with these +messages, and reduces the structured outputs. The system + user-template +strings live here so they're testable in isolation and can be evolved +without touching pipeline code. +""" + +from __future__ import annotations + +SYSTEM_PROMPT_MAP = """You are an AI assistant specialized in extracting and synthesizing relevant information from text. + +Your task: +1. Analyze the provided text in relation to the user's question +2. Extract only the essential information that directly addresses the query +3. Preserve necessary context (Key words, project names or initiatives, dates, etc.) to maintain accuracy and clarity of the summary for it to be self-understandable + +Guidelines: +- Present information clearly and concisely without unnecessary rephrasing or commentary +- Focus on precision: include what matters, exclude what doesn't. +- If a document does not have any relevant content with respect to the query, classify it as irrelevant without providing a `synthesis`. +""" + +USER_PROMPT_TEMPLATE = """ +Here is a text: +{content} + +From this document, identify and comprehensively summarize the information useful for answering the following question: +{query} +""" + + +def build_map_messages(query: str, content: str) -> list[dict[str, str]]: + """Build the system+user message list for one map-step LLM call.""" + return [ + {"role": "system", "content": SYSTEM_PROMPT_MAP}, + {"role": "user", "content": USER_PROMPT_TEMPLATE.format(query=query, content=content)}, + ] diff --git a/openrag/core/prompts/query_rewriter.py b/openrag/core/prompts/query_rewriter.py new file mode 100644 index 000000000..a0b449ad7 --- /dev/null +++ b/openrag/core/prompts/query_rewriter.py @@ -0,0 +1,35 @@ +"""Query-rewriting prompt builders for HyDe and Multi-Query retrieval. + +Templates live on disk under ``prompts//`` and are loaded via +``template_loader``. These functions are pure: they take a template string + +substitution variables and return the formatted prompt. + +Template variables expected: + HyDe template: ``{question}`` + Multi-query template: ``{query}``, ``{k_queries}`` + +The multi-query helper also exposes the ``[SEP]`` separator used to split the +LLM response into individual queries. +""" + +from __future__ import annotations + +MULTI_QUERY_SEPARATOR = "[SEP]" + + +def build_hyde_prompt(template: str, query: str) -> str: + """Format a HyDe prompt. ``template`` must contain ``{question}``.""" + return template.format(question=query) + + +def build_multi_query_prompt(template: str, query: str, k_queries: int) -> str: + """Format a multi-query prompt. ``template`` must contain ``{query}`` and ``{k_queries}``.""" + return template.format(query=query, k_queries=k_queries) + + +def split_multi_query_response(response: str, separator: str = MULTI_QUERY_SEPARATOR) -> list[str]: + """Split an LLM multi-query response into individual queries. + + Drops empty entries and trims surrounding whitespace. + """ + return [q.strip() for q in response.split(separator) if q.strip()] diff --git a/openrag/core/prompts/template_loader.py b/openrag/core/prompts/template_loader.py new file mode 100644 index 000000000..3bd31822d --- /dev/null +++ b/openrag/core/prompts/template_loader.py @@ -0,0 +1,62 @@ +"""Disk-based prompt template loader. + +Pure I/O helper: given a directory and a filename, read and return the file +contents as a string. Callers (typically the DI/composition layer) resolve +the directory and the filename mapping from config; this function has no +config dependency of its own. +""" + +from __future__ import annotations + +from pathlib import Path + + +def load_template(prompts_dir: str | Path, file_name: str) -> str: + """Read a prompt template from disk. + + Args: + prompts_dir: Directory containing prompt template files. + file_name: Template filename (relative to ``prompts_dir``). + + Returns: + The template contents as a string. + + Raises: + FileNotFoundError: if the resolved path does not exist. + """ + base = Path(prompts_dir).resolve() + file_path = (base / file_name).resolve() + if not file_path.is_relative_to(base): + raise ValueError(f"Prompt path escapes base directory: `{file_name}`") + if not file_path.exists(): + raise FileNotFoundError(f"Prompt file not found: `{file_path}`") + return file_path.read_text(encoding="utf-8") + + +def load_template_by_key( + prompts_dir: str | Path, + prompt_mapping: object, + prompt_key: str, +) -> str: + """Read a prompt by logical key, looking up the filename on a mapping object. + + The mapping object is typically the ``PromptsConfig`` Pydantic model with + attributes like ``sys_prompt``, ``hyde``, ``multi_query`` whose values are + template filenames. + + Args: + prompts_dir: Directory containing prompt template files. + prompt_mapping: Object exposing prompt keys as attributes. + prompt_key: Attribute name on ``prompt_mapping`` (e.g. ``"hyde"``). + + Returns: + The template contents as a string. + + Raises: + ValueError: if ``prompt_key`` is not defined on ``prompt_mapping``. + FileNotFoundError: if the resolved path does not exist. + """ + file_name = getattr(prompt_mapping, prompt_key, None) + if not file_name: + raise ValueError(f"No associated file name found for prompt: `{prompt_key}`") + return load_template(prompts_dir, file_name) diff --git a/openrag/core/prompts/test_chat_prompt_builder.py b/openrag/core/prompts/test_chat_prompt_builder.py new file mode 100644 index 000000000..0cd7be904 --- /dev/null +++ b/openrag/core/prompts/test_chat_prompt_builder.py @@ -0,0 +1,110 @@ +"""Tests for chat_prompt_builder — these lock in the exact wire format.""" + +from __future__ import annotations + +from openrag.core.prompts.chat_prompt_builder import ( + EMPTY_CONTEXT_MESSAGE, + SOURCE_SEPARATOR, + format_context, + format_web_context, + prepend_system_prompt, +) + + +def _word_tokens(text: str) -> int: + """Simple deterministic token counter: 1 token per whitespace-delimited word.""" + return len(text.split()) + + +def test_format_context_empty_returns_placeholder(): + text, included = format_context([], max_context_tokens=100, length_function=_word_tokens) + assert text == EMPTY_CONTEXT_MESSAGE + assert included == [] + + +def test_format_context_numbers_sources_and_separates(): + docs = ["alpha beta", "gamma delta epsilon"] + text, included = format_context(docs, max_context_tokens=100, length_function=_word_tokens) + assert "[Source 1]\nalpha beta" in text + assert "[Source 2]\ngamma delta epsilon" in text + assert SOURCE_SEPARATOR in text + assert included == [0, 1] + + +def test_format_context_drops_to_fit_budget(): + docs = ["one two", "three four five", "six"] + # _word_tokens("[Source 1]\n") = 2, doc1 = 2, sep = 1, prefix2 = 2, doc2 = 3 -> total 10 + text, included = format_context(docs, max_context_tokens=4, length_function=_word_tokens) + assert "[Source 1]" in text + assert "[Source 2]" not in text + assert included == [0] + + +def test_format_context_no_numbering(): + docs = ["a", "b"] + text, included = format_context(docs, max_context_tokens=100, length_function=_word_tokens, number_sources=False) + assert "[Source" not in text + assert text == f"a{SOURCE_SEPARATOR}b" + assert included == [0, 1] + + +class _FakeWeb: + def __init__(self, title: str, url: str, snippet: str, content: str | None = None): + self.title = title + self.url = url + self.snippet = snippet + self.content = content + + +def test_format_web_context_uses_content_when_present(): + results = [_FakeWeb("T1", "u1", "snip1", content="full body")] + text, nums, _ = format_web_context(results, length_function=_word_tokens, max_tokens=100) + assert "full body" in text + assert "snip1" not in text + assert nums == [1] + + +def test_format_web_context_falls_back_to_snippet(): + results = [_FakeWeb("T1", "u1", "snip1", content=None)] + text, _, _ = format_web_context(results, length_function=_word_tokens, max_tokens=100) + assert "snip1" in text + + +def test_format_web_context_continues_numbering_with_start_index(): + results = [_FakeWeb("T1", "u1", "snip1")] + text, nums, _ = format_web_context(results, length_function=_word_tokens, start_index=4, max_tokens=100) + assert "[Source 4]" in text + assert nums == [4] + + +def test_prepend_system_prompt_does_not_mutate_input(): + msgs = [{"role": "user", "content": "hi"}] + out = prepend_system_prompt( + msgs, + system_template="ctx={context} date={current_date}", + context="C", + current_date="2026-04-29", + ) + assert msgs == [{"role": "user", "content": "hi"}] + assert out[0] == {"role": "system", "content": "ctx=C date=2026-04-29"} + assert out[1] == {"role": "user", "content": "hi"} + + +def test_format_web_context_empty_returns_empty_tuple(): + text, nums, total = format_web_context([], length_function=_word_tokens) + assert text == "" + assert nums == [] + assert total == 0 + + +def test_format_web_context_drops_overflow_block_after_first_fits(): + """If a later block would push past max_tokens we break — but only after + at least one block has been admitted (parts truthy guard).""" + results = [ + _FakeWeb("T1", "u1", "short body"), + _FakeWeb("T2", "u2", "this snippet has many many many many many words that will overflow"), + ] + text, nums, _ = format_web_context(results, length_function=_word_tokens, max_tokens=10) + assert "[Source 1]" in text + assert "[Source 2]" not in text + assert nums == [1] diff --git a/openrag/core/prompts/test_contextualization_builder.py b/openrag/core/prompts/test_contextualization_builder.py new file mode 100644 index 000000000..556ee107d --- /dev/null +++ b/openrag/core/prompts/test_contextualization_builder.py @@ -0,0 +1,82 @@ +"""Tests for the chunk-contextualization prompt builder.""" + +from __future__ import annotations + +from core.prompts.contextualization_builder import ( + BASE_CHUNK_FORMAT, + CHUNK_FORMAT, + build_messages, + build_user_message, + wrap_chunk_with_context, +) + + +def test_build_user_message_includes_filename_and_lang(): + out = build_user_message( + filename="doc.pdf", + first_chunks_text=["intro chunk A", "intro chunk B"], + prev_chunks_text=["prev chunk"], + current_chunk_text="here is the current chunk", + lang="fr", + ) + assert "doc.pdf" in out + assert "intro chunk A" in out + assert "intro chunk B" in out + assert "prev chunk" in out + assert "here is the current chunk" in out + assert "fr language" in out + + +def test_build_user_message_handles_empty_history(): + out = build_user_message( + filename="doc.pdf", + first_chunks_text=[], + prev_chunks_text=[], + current_chunk_text="solo chunk", + ) + assert "solo chunk" in out + assert "en language" in out # default lang + + +def test_build_messages_returns_system_then_user(): + msgs = build_messages( + system_prompt="SYS", + filename="doc.pdf", + first_chunks_text=["a"], + prev_chunks_text=["b"], + current_chunk_text="c", + ) + assert len(msgs) == 2 + assert msgs[0] == {"role": "system", "content": "SYS"} + assert msgs[1]["role"] == "user" + assert "doc.pdf" in msgs[1]["content"] + + +def test_wrap_chunk_with_context_uses_full_format_when_context_given(): + out = wrap_chunk_with_context(content="body", filename="f.pdf", chunk_context="ctx") + assert "[CONTEXT]" in out + assert "ctx" in out + assert "[CHUNK_START]" in out + assert "body" in out + assert "[CHUNK_END]" in out + # Sanity: result uses the documented CHUNK_FORMAT template. + assert out == CHUNK_FORMAT.format(content="body", chunk_context="ctx", filename="f.pdf") + + +def test_wrap_chunk_with_context_uses_base_format_when_context_empty(): + out = wrap_chunk_with_context(content="body", filename="f.pdf", chunk_context="") + assert "[CONTEXT]" not in out + assert "[CHUNK_START]" in out + assert "body" in out + assert out == BASE_CHUNK_FORMAT.format(content="body", filename="f.pdf") + + +def test_wrap_chunk_with_context_defaults_chunk_context_to_empty(): + out = wrap_chunk_with_context(content="body", filename="f.pdf") + assert "[CONTEXT]" not in out + + +def test_wrap_chunk_with_context_treats_whitespace_only_as_empty(): + out = wrap_chunk_with_context(content="body", filename="f.pdf", chunk_context=" \n\t") + assert "[CONTEXT]" not in out + assert out == BASE_CHUNK_FORMAT.format(content="body", filename="f.pdf") diff --git a/openrag/core/prompts/test_map_reduce_builder.py b/openrag/core/prompts/test_map_reduce_builder.py new file mode 100644 index 000000000..26f702d93 --- /dev/null +++ b/openrag/core/prompts/test_map_reduce_builder.py @@ -0,0 +1,27 @@ +"""Tests for the map-reduce prompt builder.""" + +from __future__ import annotations + +from core.prompts.map_reduce_builder import ( + SYSTEM_PROMPT_MAP, + USER_PROMPT_TEMPLATE, + build_map_messages, +) + + +def test_build_map_messages_returns_system_then_user_with_substitution(): + msgs = build_map_messages(query="what is rag?", content="some doc body") + assert len(msgs) == 2 + assert msgs[0] == {"role": "system", "content": SYSTEM_PROMPT_MAP} + assert msgs[1]["role"] == "user" + user_text = msgs[1]["content"] + assert "what is rag?" in user_text + assert "some doc body" in user_text + + +def test_user_prompt_template_uses_named_placeholders(): + """Lock the template's named placeholders so accidental positional refactors + don't silently change the wire format.""" + rendered = USER_PROMPT_TEMPLATE.format(query="Q", content="C") + assert "Q" in rendered + assert "C" in rendered diff --git a/openrag/core/prompts/test_query_rewriter.py b/openrag/core/prompts/test_query_rewriter.py new file mode 100644 index 000000000..79bdb56a0 --- /dev/null +++ b/openrag/core/prompts/test_query_rewriter.py @@ -0,0 +1,40 @@ +"""Tests for the HyDe / multi-query prompt builders.""" + +from __future__ import annotations + +from core.prompts.query_rewriter import ( + MULTI_QUERY_SEPARATOR, + build_hyde_prompt, + build_multi_query_prompt, + split_multi_query_response, +) + + +def test_build_hyde_prompt_substitutes_question(): + out = build_hyde_prompt("Q: {question}", "what is rag?") + assert out == "Q: what is rag?" + + +def test_build_multi_query_prompt_substitutes_query_and_k(): + out = build_multi_query_prompt("Generate {k_queries} variants of: {query}", "what is rag?", 5) + assert "5" in out + assert "what is rag?" in out + + +def test_split_multi_query_response_splits_on_separator_and_trims(): + raw = f" first query {MULTI_QUERY_SEPARATOR} second query {MULTI_QUERY_SEPARATOR}third" + assert split_multi_query_response(raw) == ["first query", "second query", "third"] + + +def test_split_multi_query_response_drops_empty_entries(): + raw = f" {MULTI_QUERY_SEPARATOR}only{MULTI_QUERY_SEPARATOR}{MULTI_QUERY_SEPARATOR} " + assert split_multi_query_response(raw) == ["only"] + + +def test_split_multi_query_response_empty_string_returns_empty_list(): + assert split_multi_query_response("") == [] + + +def test_split_multi_query_response_accepts_custom_separator(): + raw = "a||b||c" + assert split_multi_query_response(raw, separator="||") == ["a", "b", "c"] diff --git a/openrag/core/prompts/test_template_loader.py b/openrag/core/prompts/test_template_loader.py new file mode 100644 index 000000000..c20aa837a --- /dev/null +++ b/openrag/core/prompts/test_template_loader.py @@ -0,0 +1,58 @@ +"""Tests for the disk-based prompt template loader.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from core.prompts.template_loader import load_template, load_template_by_key + + +def test_load_template_reads_file_contents(tmp_path: Path): + target = tmp_path / "sys.txt" + target.write_text("hello {name}", encoding="utf-8") + assert load_template(tmp_path, "sys.txt") == "hello {name}" + + +def test_load_template_accepts_str_path(tmp_path: Path): + target = tmp_path / "sys.txt" + target.write_text("hi", encoding="utf-8") + assert load_template(str(tmp_path), "sys.txt") == "hi" + + +def test_load_template_raises_on_missing_file(tmp_path: Path): + with pytest.raises(FileNotFoundError, match="Prompt file not found"): + load_template(tmp_path, "nope.txt") + + +class _Mapping: + """Stand-in for the PromptsConfig pydantic model.""" + + def __init__(self, **kwargs: str) -> None: + for k, v in kwargs.items(): + setattr(self, k, v) + + +def test_load_template_by_key_resolves_filename(tmp_path: Path): + (tmp_path / "hyde.txt").write_text("hyde body", encoding="utf-8") + mapping = _Mapping(hyde="hyde.txt") + assert load_template_by_key(tmp_path, mapping, "hyde") == "hyde body" + + +def test_load_template_by_key_raises_when_attr_missing(tmp_path: Path): + mapping = _Mapping(hyde="hyde.txt") + with pytest.raises(ValueError, match="No associated file name"): + load_template_by_key(tmp_path, mapping, "multi_query") + + +def test_load_template_by_key_raises_when_attr_falsy(tmp_path: Path): + """A mapping value of empty string / None should be treated as "not set".""" + mapping = _Mapping(multi_query="") + with pytest.raises(ValueError, match="No associated file name"): + load_template_by_key(tmp_path, mapping, "multi_query") + + +def test_load_template_by_key_propagates_file_not_found(tmp_path: Path): + mapping = _Mapping(hyde="missing.txt") + with pytest.raises(FileNotFoundError): + load_template_by_key(tmp_path, mapping, "hyde") diff --git a/openrag/core/prompts/test_vlm_prompt_builder.py b/openrag/core/prompts/test_vlm_prompt_builder.py new file mode 100644 index 000000000..0b6bac2fa --- /dev/null +++ b/openrag/core/prompts/test_vlm_prompt_builder.py @@ -0,0 +1,43 @@ +"""Tests for the VLM (image-captioning) prompt builder.""" + +from __future__ import annotations + +from core.prompts.vlm_prompt_builder import ( + build_caption_messages, + wrap_caption, +) +from core.utils.conts import IMG_WRAPPER_CLOSE, IMG_WRAPPER_OPEN + + +def test_build_caption_messages_shapes_for_openai_multimodal(): + msgs = build_caption_messages(template="Describe this image.", image_url="https://example.com/x.png") + assert len(msgs) == 1 + msg = msgs[0] + assert msg["role"] == "user" + parts = msg["content"] + assert {p["type"] for p in parts} == {"image_url", "text"} + image_part = next(p for p in parts if p["type"] == "image_url") + text_part = next(p for p in parts if p["type"] == "text") + assert image_part["image_url"] == {"url": "https://example.com/x.png"} + assert text_part["text"] == "Describe this image." + + +def test_build_caption_messages_supports_data_uri(): + msgs = build_caption_messages(template="caption", image_url="data:image/png;base64,abc") + image_part = next(p for p in msgs[0]["content"] if p["type"] == "image_url") + assert image_part["image_url"]["url"].startswith("data:image/png;base64,") + + +def test_wrap_caption_uses_image_description_markers(): + out = wrap_caption("a sunset over the sea") + assert out.startswith(IMG_WRAPPER_OPEN) + assert out.endswith(IMG_WRAPPER_CLOSE) + assert "a sunset over the sea" in out + + +def test_wrap_caption_format_is_load_bearing_for_chunker(): + """The chunker's image-element regex matches `...`, + so this exact pairing is part of the contract.""" + out = wrap_caption("x") + assert IMG_WRAPPER_OPEN in out + assert IMG_WRAPPER_CLOSE in out diff --git a/openrag/core/prompts/vlm_prompt_builder.py b/openrag/core/prompts/vlm_prompt_builder.py new file mode 100644 index 000000000..a0bebcaab --- /dev/null +++ b/openrag/core/prompts/vlm_prompt_builder.py @@ -0,0 +1,43 @@ +"""VLM (vision-language model) prompt builder. + +Pure helpers: given a captioning template and an image reference, produce a +multimodal chat-message payload (OpenAI-style) and wrap captions in the +```` markers downstream pipelines expect. +""" + +from __future__ import annotations + +from typing import Any + +from core.utils.conts import IMG_WRAPPER_CLOSE, IMG_WRAPPER_OPEN + + +def build_caption_messages(template: str, image_url: str) -> list[dict[str, Any]]: + """Build a multimodal chat-message list for image captioning. + + Args: + template: Image-captioning prompt text (no substitution required). + image_url: ``https://...`` URL or ``data:image/...;base64,...`` data URI. + + Returns: + A single-message list shaped for OpenAI / vLLM chat completions: + ``[{"role": "user", "content": [{"type": "image_url", ...}, {"type": "text", ...}]}]`` + """ + return [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": image_url}}, + {"type": "text", "text": template}, + ], + } + ] + + +def wrap_caption(caption: str) -> str: + """Wrap a raw caption in ```` markers. + + Pipelines downstream (markdown image replacement, chunk parsing) look for + this exact marker, so the wrapping format is part of the contract. + """ + return f"{IMG_WRAPPER_OPEN}\n\n{caption}\n\n{IMG_WRAPPER_CLOSE}" diff --git a/openrag/core/rerankers/__init__.py b/openrag/core/rerankers/__init__.py new file mode 100644 index 000000000..31ed9f68d --- /dev/null +++ b/openrag/core/rerankers/__init__.py @@ -0,0 +1,6 @@ +"""Reranker ABC + registry.""" + +from .registry import reranker_registry +from .reranker import Reranker + +__all__ = ["Reranker", "reranker_registry"] diff --git a/openrag/core/rerankers/registry.py b/openrag/core/rerankers/registry.py new file mode 100644 index 000000000..c637e1c8c --- /dev/null +++ b/openrag/core/rerankers/registry.py @@ -0,0 +1,7 @@ +"""Reranker registry.""" + +from openrag.core.utils.registry import Registry + +from .reranker import Reranker + +reranker_registry: Registry[Reranker] = Registry("reranker") diff --git a/openrag/core/rerankers/reranker.py b/openrag/core/rerankers/reranker.py new file mode 100644 index 000000000..9f2ea60da --- /dev/null +++ b/openrag/core/rerankers/reranker.py @@ -0,0 +1,17 @@ +"""Abstract reranker interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class Reranker(ABC): + """Base class for all reranking providers.""" + + @abstractmethod + async def rerank(self, query: str, documents: list[str], top_k: int | None = None) -> list[tuple[int, float]]: + """Rerank documents for a query. + + Returns list of (original_index, score) sorted by relevance. + """ + ... diff --git a/openrag/core/retrieval/__init__.py b/openrag/core/retrieval/__init__.py new file mode 100644 index 000000000..de83e5492 --- /dev/null +++ b/openrag/core/retrieval/__init__.py @@ -0,0 +1,25 @@ +"""Retrieval domain logic: retriever strategies, RRF, and pipeline.""" + +from .pipeline import RetrieverPipeline +from .retriever import ( + BaseRetriever, + HyDeRetriever, + MultiQueryRetriever, + Retriever, + SingleRetriever, + retriever_registry, +) +from .rrf import rrf_reranking +from .searcher import RetrievalSearcher + +__all__ = [ + "Retriever", + "BaseRetriever", + "SingleRetriever", + "MultiQueryRetriever", + "HyDeRetriever", + "retriever_registry", + "RetrievalSearcher", + "RetrieverPipeline", + "rrf_reranking", +] diff --git a/openrag/core/retrieval/pipeline.py b/openrag/core/retrieval/pipeline.py new file mode 100644 index 000000000..a4cd17408 --- /dev/null +++ b/openrag/core/retrieval/pipeline.py @@ -0,0 +1,152 @@ +"""Retrieval pipeline: per-query retrieval, optional temporal-filter fallback, +optional reranking, optional related/ancestor expansion, and RRF fusion across +sub-queries. + +Extracted from ``components/pipeline.py:RetrieverPipeline``. The legacy +``RagPipeline`` (LLM-driven query generation, system-prompt assembly, +streaming) lives in the orchestrator layer and is rebuilt in Phase 8. + +This pipeline depends only on core ABCs: + + * ``Retriever`` — strategy that produces candidate chunks + * ``Reranker`` — optional cross-encoder reranker (per Phase 4 ABC: + returns ``[(idx, score), ...]`` over a list of texts) + * ``RetrievalSearcher`` is consumed by the retriever, not directly here. + +Config knobs are constructor arguments; there is no module-level config load. +""" + +from __future__ import annotations + +import asyncio +import copy +from typing import Any + +from core.models.chunk import Chunk +from core.models.query import Query, SearchQueries +from core.rerankers.reranker import Reranker +from core.retrieval.retriever import Retriever +from core.retrieval.rrf import rrf_reranking + + +def _chunk_key(c: Chunk) -> Any: + """Identity key for fusion / dedup. Falls back to object id when missing.""" + return c.id or id(c) + + +async def _rerank_chunks(reranker: Reranker, query: str, chunks: list[Chunk]) -> list[Chunk]: + """Reorder chunks via the Reranker ABC. + + The ABC scores text+query pairs and returns ``[(orig_index, score), ...]``; + we look up the original chunk for each ranked index. Items the reranker + drops are excluded. + """ + if not chunks: + return chunks + ranking = await reranker.rerank(query=query, documents=[c.text for c in chunks], top_k=None) + return [chunks[idx] for idx, _ in ranking] + + +class RetrieverPipeline: + """Orchestrates retrieval + reranking + expansion for a list of sub-queries. + + Args: + retriever: Concrete retrieval strategy (Single / MultiQuery / HyDe). + reranker: Reranker implementation, or ``None`` to skip reranking. + reranker_top_k: When expansion is enabled, the top-K size used to + decide which results to expand. + allow_filterless_fallback: If a temporal filter wipes out all + candidates, retry once without it. When ``False``, + return zero docs rather than ones outside the + temporal range. + """ + + def __init__( + self, + retriever: Retriever, + reranker: Reranker | None = None, + reranker_top_k: int = 5, + allow_filterless_fallback: bool = True, + ) -> None: + self.retriever = retriever + self.reranker = reranker + self.reranker_top_k = reranker_top_k + self.allow_filterless_fallback = allow_filterless_fallback + + @property + def reranker_enabled(self) -> bool: + return self.reranker is not None + + @property + def expansion_enabled(self) -> bool: + # The retriever's BaseRetriever sets this; non-Base implementations + # may not. Treat absent attribute as no expansion. + return getattr(self.retriever, "expansion_enabled", False) + + async def retrieve_docs( + self, + partition: list[str], + query: Query, + top_k: int | None = None, + filter_params: dict | None = None, + ) -> list[Chunk]: + """Run a single ``Query`` through retrieval, expansion, and reranking.""" + milvus_filter = query.to_milvus_filter() + chunks = await self.retriever.retrieve( + partition=partition, + query=query.query, + filter=milvus_filter, + filter_params=filter_params, + ) + + if not chunks and milvus_filter and self.allow_filterless_fallback: + # Temporal filter killed every candidate — retry without it so + # the user gets some results rather than none. + chunks = await self.retriever.retrieve( + partition=partition, + query=query.query, + filter=None, + filter_params=filter_params, + ) + + if not chunks: + return chunks + + if self.reranker_enabled: + chunks = await _rerank_chunks(self.reranker, query.query, chunks) + + if self.expansion_enabled: + limit = self.reranker_top_k if top_k is None else max(self.reranker_top_k, top_k) + head = copy.deepcopy(chunks[:limit]) + expanded = await self.retriever.expand_search_results(results=head) + if len(expanded) > len(head): + chunks = expanded + if self.reranker_enabled: + chunks = await _rerank_chunks(self.reranker, query.query, chunks) + + if top_k is not None: + chunks = chunks[:top_k] + return chunks + + async def get_relevant_docs( + self, + partition: list[str], + search_queries: SearchQueries, + top_k: int | None = None, + filter_params: dict | None = None, + ) -> list[Chunk]: + """Run every sub-query in parallel and fuse the per-query rankings via RRF.""" + tasks = [ + self.retrieve_docs( + partition=partition, + query=q, + top_k=top_k, + filter_params=filter_params, + ) + for q in search_queries.query_list + ] + ranked_lists = await asyncio.gather(*tasks) + fused = rrf_reranking(ranked_lists, key_fn=_chunk_key) + if top_k is not None: + fused = fused[:top_k] + return fused diff --git a/openrag/core/retrieval/retriever.py b/openrag/core/retrieval/retriever.py new file mode 100644 index 000000000..f5f783968 --- /dev/null +++ b/openrag/core/retrieval/retriever.py @@ -0,0 +1,283 @@ +"""Retriever strategies: Single, MultiQuery, HyDe. + +Rewritten from ``components/retriever.py``. Differences from the legacy: + + * ``RetrievalSearcher`` (clean ABC) replaces ``get_vectordb()`` / Ray actor + direct access. The retriever has no Ray imports. + * ``LLM`` (clean ABC) replaces ``ChatOpenAI`` + LangChain chain assembly. + * Prompt templates are passed in as strings. The DI layer loads them + from disk via ``core/prompts/template_loader``. + * Returns are domain ``Chunk`` objects, not LangChain ``Document``. + +A ``retriever_registry`` is exposed so the composition root can pick a +strategy by name (``single`` / ``multiQuery`` / ``hyde``) per the +strategy doc's "every factory becomes a Registry" rule. +""" + +from __future__ import annotations + +import asyncio +import logging +from abc import ABC, abstractmethod +from itertools import chain as ichain +from typing import Any + +from core.llm.llm import LLM +from core.models.chunk import Chunk +from core.prompts.query_rewriter import ( + build_hyde_prompt, + build_multi_query_prompt, + split_multi_query_response, +) +from core.retrieval.searcher import RetrievalSearcher +from core.utils.registry import Registry + +logger = logging.getLogger(__name__) + + +class Retriever(ABC): + """Common surface for all retrieval strategies.""" + + @abstractmethod + async def retrieve( + self, + partition: list[str], + query: str, + filter: str | None = None, + filter_params: dict | None = None, + ) -> list[Chunk]: + """Run the strategy and return scored chunks.""" + ... + + @abstractmethod + async def expand_search_results(self, results: list[Chunk]) -> list[Chunk]: + """Optionally enrich a result set with related/ancestor chunks.""" + ... + + +class BaseRetriever(Retriever): + """Single-query retriever — the building block for the others.""" + + def __init__( + self, + searcher: RetrievalSearcher, + top_k: int = 6, + similarity_threshold: float = 0.95, + with_surrounding_chunks: bool = True, + include_related: bool = False, + include_ancestors: bool = False, + related_limit: int = 10, + max_ancestor_depth: int | None = None, + **_: Any, + ) -> None: + self.searcher = searcher + self.top_k = top_k + self.similarity_threshold = similarity_threshold + self.with_surrounding_chunks = with_surrounding_chunks + self.include_related = include_related + self.include_ancestors = include_ancestors + self.related_limit = related_limit + self.max_ancestor_depth = max_ancestor_depth + self.expansion_enabled = include_related or include_ancestors + + async def retrieve( + self, + partition: list[str], + query: str, + filter: str | None = None, + filter_params: dict | None = None, + ) -> list[Chunk]: + return await self.searcher.search( + query=query, + partition=partition, + top_k=self.top_k, + filter=filter, + filter_params=filter_params, + similarity_threshold=self.similarity_threshold, + with_surrounding_chunks=self.with_surrounding_chunks, + ) + + async def expand_search_results(self, results: list[Chunk]) -> list[Chunk]: + return await _expand_with_related_chunks( + searcher=self.searcher, + results=results, + include_related=self.include_related, + include_ancestors=self.include_ancestors, + related_limit=self.related_limit, + max_ancestor_depth=self.max_ancestor_depth, + ) + + +class SingleRetriever(BaseRetriever): + """Default strategy — issues exactly one similarity search per query.""" + + +class MultiQueryRetriever(BaseRetriever): + """Generates K query variants via the LLM and unions their results.""" + + def __init__( + self, + searcher: RetrievalSearcher, + llm: LLM, + multi_query_template: str, + k_queries: int = 3, + **kwargs: Any, + ) -> None: + super().__init__(searcher=searcher, **kwargs) + if llm is None: + raise ValueError("llm must be provided for MultiQueryRetriever") + self.llm = llm + self.multi_query_template = multi_query_template + self.k_queries = k_queries + + async def _generate_queries(self, query: str) -> list[str]: + prompt = build_multi_query_prompt(self.multi_query_template, query, self.k_queries) + response = await self.llm.chat([{"role": "user", "content": prompt}]) + # Cap to k_queries — a non-compliant LLM response can otherwise fan + # out far more searches than configured. + queries = split_multi_query_response(response)[: self.k_queries] + return queries or [query] + + async def retrieve( + self, + partition: list[str], + query: str, + filter: str | None = None, + filter_params: dict | None = None, + ) -> list[Chunk]: + queries = await self._generate_queries(query) + return await self.searcher.multi_query_search( + queries=queries, + partition=partition, + top_k_per_query=self.top_k, + filter=filter, + filter_params=filter_params, + similarity_threshold=self.similarity_threshold, + with_surrounding_chunks=self.with_surrounding_chunks, + ) + + +class HyDeRetriever(BaseRetriever): + """Generates a hypothetical answer document and searches with it. + + If ``combine`` is set, the original query is also issued and results + are unioned via the searcher's multi-query path. + """ + + def __init__( + self, + searcher: RetrievalSearcher, + llm: LLM, + hyde_template: str, + combine: bool = False, + **kwargs: Any, + ) -> None: + super().__init__(searcher=searcher, **kwargs) + if llm is None: + raise ValueError("llm must be provided for HyDeRetriever") + self.llm = llm + self.hyde_template = hyde_template + self.combine = combine + + async def get_hyde(self, query: str) -> str: + prompt = build_hyde_prompt(self.hyde_template, query) + return await self.llm.chat([{"role": "user", "content": prompt}]) + + async def retrieve( + self, + partition: list[str], + query: str, + filter: str | None = None, + filter_params: dict | None = None, + ) -> list[Chunk]: + hyde = (await self.get_hyde(query)).strip() + if not hyde: + queries = [query] + else: + queries = [hyde, query] if self.combine else [hyde] + return await self.searcher.multi_query_search( + queries=queries, + partition=partition, + top_k_per_query=self.top_k, + filter=filter, + filter_params=filter_params, + similarity_threshold=self.similarity_threshold, + with_surrounding_chunks=self.with_surrounding_chunks, + ) + + +async def _expand_with_related_chunks( + searcher: RetrievalSearcher, + results: list[Chunk], + include_related: bool, + include_ancestors: bool, + related_limit: int = 10, + max_ancestor_depth: int | None = None, +) -> list[Chunk]: + """Append related and/or ancestor chunks to a result set, deduplicated by id. + + Failures on individual related/ancestor lookups are logged and treated + as empty results, matching legacy behavior so retrieval remains + resilient to per-document errors. + """ + if not results or (not include_related and not include_ancestors): + return results + + seen_ids = {c.id for c in results if c.id} + expanded: list[Chunk] = list(results) + + relationship_ids: set[tuple[str, str]] = set() + file_infos: set[tuple[str, str]] = set() + + for c in results: + if include_related: + rel_id = c.metadata.get("relationship_id") + if rel_id and c.partition: + relationship_ids.add((c.partition, rel_id)) + if include_ancestors and c.partition and c.document_id: + file_infos.add((c.partition, c.document_id)) + + async def _safe_related(part: str, rel_id: str) -> list[Chunk]: + try: + return await searcher.get_related_chunks(partition=part, relationship_id=rel_id, limit=related_limit) + except Exception: + logger.warning("get_related_chunks failed (partition=%s, relationship_id=%s)", part, rel_id, exc_info=True) + return [] + + async def _safe_ancestors(part: str, file_id: str) -> list[Chunk]: + try: + return await searcher.get_ancestor_chunks( + partition=part, + file_id=file_id, + limit=related_limit, + max_ancestor_depth=max_ancestor_depth, + ) + except Exception: + logger.warning("get_ancestor_chunks failed (partition=%s, file_id=%s)", part, file_id, exc_info=True) + return [] + + tasks: list[asyncio.Future] = [] + if include_related: + tasks.extend(_safe_related(part, rid) for part, rid in relationship_ids) + if include_ancestors: + tasks.extend(_safe_ancestors(part, fid) for part, fid in file_infos if part and fid) + + if tasks: + all_results = await asyncio.gather(*tasks) + for chunk in ichain.from_iterable(all_results): + if chunk.id and chunk.id in seen_ids: + continue + if chunk.id: + seen_ids.add(chunk.id) + expanded.append(chunk) + + return expanded + + +# --------------------------------------------------------------------------- +# Registry — config-driven factory replacement +# --------------------------------------------------------------------------- +retriever_registry: Registry[Retriever] = Registry("retriever") +retriever_registry.register("single")(SingleRetriever) +retriever_registry.register("multiQuery")(MultiQueryRetriever) +retriever_registry.register("hyde")(HyDeRetriever) diff --git a/openrag/core/retrieval/rrf.py b/openrag/core/retrieval/rrf.py new file mode 100644 index 000000000..13e7f5baf --- /dev/null +++ b/openrag/core/retrieval/rrf.py @@ -0,0 +1,66 @@ +"""Reciprocal Rank Fusion — pure math, no domain coupling. + +Combines multiple ranked lists into a single ranking by summing reciprocal +ranks across lists. Items present in more lists, or higher-ranked in any +list, sort to the top of the fused result. + +Formula: + score(item) = Σ_i 1 / (k + rank_i) + +with ``rank_i`` the 1-based rank of the item in list ``i``. Smaller ``k`` +amplifies the top of each list; ``k=60`` is the canonical default and +balances rank sensitivity across lists. + +Identification of "the same item" is delegated to the caller via +``key_fn`` — typically returning the chunk id, document id, or URL. +""" + +from __future__ import annotations + +from collections.abc import Callable, Hashable, Sequence +from typing import TypeVar + +T = TypeVar("T") + + +def rrf_reranking( + ranked_lists: Sequence[Sequence[T]], + key_fn: Callable[[T], Hashable] | None = None, + k: int = 60, +) -> list[T]: + """Fuse multiple ranked lists into one via Reciprocal Rank Fusion. + + Args: + ranked_lists: Each inner sequence is a ranked list (best first). + key_fn: Returns the identity key for an item; items sharing a key + across lists have their RRF scores summed. Defaults to + ``id(item)`` (object identity), which prevents fusion across + lists for items lacking a logical id. + k: RRF dampening constant. ``60`` is canonical. + + Returns: + A single ranked list, best first. Empty input -> empty list. + Single input list is shallow-copied so callers always get a ``list``. + + Raises: + ValueError: if ``k < 0`` (would produce a zero or negative + denominator at rank 1 or below and crash with ZeroDivisionError). + """ + if k < 0: + raise ValueError(f"RRF k must be non-negative, got {k}") + if not ranked_lists: + return [] + if len(ranked_lists) == 1: + return list(ranked_lists[0]) + + if key_fn is None: + key_fn = id # type: ignore[assignment] + + fused: dict[Hashable, tuple[float, T]] = {} + for ranked in ranked_lists: + for rank, item in enumerate(ranked, start=1): + key = key_fn(item) + score, kept = fused.get(key, (0.0, item)) + fused[key] = (score + 1.0 / (rank + k), kept) + + return [item for _, item in sorted(fused.values(), key=lambda x: x[0], reverse=True)] diff --git a/openrag/core/retrieval/searcher.py b/openrag/core/retrieval/searcher.py new file mode 100644 index 000000000..f981f6585 --- /dev/null +++ b/openrag/core/retrieval/searcher.py @@ -0,0 +1,81 @@ +"""Transitional port for chunk-level retrieval operations. + +The strict ``VectorStore`` ABC in ``core/vector_stores`` is intentionally +narrow — ``search(embedding, top_k, ...)``. Phase 5 retrievers, however, +still go through the legacy Milvus Ray actor which exposes higher-level +operations: + + * search by query string (embedding done internally, plus BM25) + * multi-query search (one round per query, dedup at the bottom) + * related-chunk lookup by ``relationship_id`` + * ancestor lookup by ``file_id`` with depth bound + +Defining these on a dedicated port lets the retriever depend on a clean +interface from day one, while a small shim in ``services/storage/`` +adapts the Ray actor to it. When the god object is decomposed in Phase 7 +this port either retires (operations move to ``VectorStore`` + +``ChunkRepository``) or evolves into the shape MilvusVectorStore exposes +directly. + +Returns are domain ``Chunk`` objects throughout — no LangChain types +leak across this boundary. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from core.models.chunk import Chunk + + +class RetrievalSearcher(ABC): + """Operations a retriever needs from the chunk store.""" + + @abstractmethod + async def search( + self, + query: str, + partition: list[str], + top_k: int, + filter: str | None = None, + filter_params: dict | None = None, + similarity_threshold: float = 0.0, + with_surrounding_chunks: bool = True, + ) -> list[Chunk]: + """Single-query similarity search.""" + ... + + @abstractmethod + async def multi_query_search( + self, + queries: list[str], + partition: list[str], + top_k_per_query: int, + filter: str | None = None, + filter_params: dict | None = None, + similarity_threshold: float = 0.0, + with_surrounding_chunks: bool = True, + ) -> list[Chunk]: + """Run one similarity search per query, return the merged result.""" + ... + + @abstractmethod + async def get_related_chunks( + self, + partition: str, + relationship_id: str, + limit: int, + ) -> list[Chunk]: + """Fetch other chunks belonging to the same relationship group.""" + ... + + @abstractmethod + async def get_ancestor_chunks( + self, + partition: str, + file_id: str, + limit: int, + max_ancestor_depth: int | None = None, + ) -> list[Chunk]: + """Walk parent links up the document tree from a file.""" + ... diff --git a/openrag/core/retrieval/test_pipeline.py b/openrag/core/retrieval/test_pipeline.py new file mode 100644 index 000000000..a7702d0f2 --- /dev/null +++ b/openrag/core/retrieval/test_pipeline.py @@ -0,0 +1,205 @@ +"""Tests for RetrieverPipeline using fake Retriever / Reranker.""" + +from __future__ import annotations + +import pytest +from core.models.chunk import Chunk +from core.models.query import Query, SearchQueries, TemporalPredicate +from core.retrieval.pipeline import RetrieverPipeline +from core.retrieval.retriever import Retriever + + +class FakeRetriever(Retriever): + """Plays back canned per-call results; records call kwargs.""" + + def __init__(self, expansion_enabled: bool = False) -> None: + self.calls: list[dict] = [] + self.results_queue: list[list[Chunk]] = [] + self.expand_input: list[Chunk] | None = None + self.expand_result: list[Chunk] | None = None + self.expansion_enabled = expansion_enabled + + async def retrieve(self, partition, query, filter=None, filter_params=None): + self.calls.append({"partition": partition, "query": query, "filter": filter, "filter_params": filter_params}) + if self.results_queue: + return self.results_queue.pop(0) + return [] + + async def expand_search_results(self, results): + self.expand_input = list(results) + return list(self.expand_result) if self.expand_result is not None else list(results) + + +class FakeReranker: + """Reverses input ordering — easy to detect in assertions.""" + + def __init__(self) -> None: + self.calls: list[dict] = [] + + async def rerank(self, query, documents, top_k=None): + self.calls.append({"query": query, "documents": list(documents), "top_k": top_k}) + # Reverse ranking, perfect score for the (now-)first item + return [(i, float(len(documents) - i)) for i in range(len(documents) - 1, -1, -1)] + + +def _chunks(*ids: str) -> list[Chunk]: + return [Chunk(id=i, text=f"text-{i}", partition="p1") for i in ids] + + +@pytest.mark.asyncio +async def test_retrieve_docs_no_filter_no_rerank_no_expand(): + r = FakeRetriever() + r.results_queue = [_chunks("a", "b", "c")] + p = RetrieverPipeline(retriever=r) + out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi")) + assert [c.id for c in out] == ["a", "b", "c"] + assert r.calls[0]["filter"] is None + + +@pytest.mark.asyncio +async def test_retrieve_docs_temporal_filter_passed_through(): + r = FakeRetriever() + r.results_queue = [_chunks("a")] + p = RetrieverPipeline(retriever=r) + q = Query( + query="hi", + temporal_filters=[TemporalPredicate(operator=">=", value="2026-01-01T00:00:00+00:00")], + ) + await p.retrieve_docs(partition=["p1"], query=q) + assert "created_at" in r.calls[0]["filter"] + + +@pytest.mark.asyncio +async def test_retrieve_docs_filterless_fallback_when_filter_returns_zero(): + r = FakeRetriever() + r.results_queue = [[], _chunks("a")] + p = RetrieverPipeline(retriever=r, allow_filterless_fallback=True) + q = Query( + query="hi", + temporal_filters=[TemporalPredicate(operator=">=", value="2026-01-01T00:00:00+00:00")], + ) + out = await p.retrieve_docs(partition=["p1"], query=q) + assert [c.id for c in out] == ["a"] + assert r.calls[0]["filter"] is not None + assert r.calls[1]["filter"] is None + + +@pytest.mark.asyncio +async def test_retrieve_docs_no_fallback_when_disabled(): + r = FakeRetriever() + r.results_queue = [[]] + p = RetrieverPipeline(retriever=r, allow_filterless_fallback=False) + q = Query( + query="hi", + temporal_filters=[TemporalPredicate(operator=">=", value="2026-01-01T00:00:00+00:00")], + ) + out = await p.retrieve_docs(partition=["p1"], query=q) + assert out == [] + assert len(r.calls) == 1 + + +@pytest.mark.asyncio +async def test_retrieve_docs_runs_reranker_when_enabled(): + r = FakeRetriever() + r.results_queue = [_chunks("a", "b", "c")] + rer = FakeReranker() + p = RetrieverPipeline(retriever=r, reranker=rer) + out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi")) + assert [c.id for c in out] == ["c", "b", "a"] + assert rer.calls[0]["query"] == "hi" + + +@pytest.mark.asyncio +async def test_retrieve_docs_expansion_path_re_reranks(): + r = FakeRetriever(expansion_enabled=True) + r.results_queue = [_chunks("a", "b")] + r.expand_result = _chunks("a", "b", "c") + rer = FakeReranker() + p = RetrieverPipeline(retriever=r, reranker=rer, reranker_top_k=2) + out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi")) + # Two reranker invocations: pre-expansion (2 chunks), post-expansion (3 chunks) + assert len(rer.calls) == 2 + assert len(rer.calls[0]["documents"]) == 2 + assert len(rer.calls[1]["documents"]) == 3 + assert {c.id for c in out} == {"a", "b", "c"} + + +@pytest.mark.asyncio +async def test_get_relevant_docs_runs_one_call_per_subquery_and_fuses(): + r = FakeRetriever() + r.results_queue = [_chunks("a", "b"), _chunks("b", "c")] + p = RetrieverPipeline(retriever=r) + sq = SearchQueries(query_list=[Query(query="q1"), Query(query="q2")]) + out = await p.get_relevant_docs(partition=["p1"], search_queries=sq) + assert len(r.calls) == 2 + assert {c.id for c in out} == {"a", "b", "c"} + # 'b' appears in both lists -> highest fused score + assert out[0].id == "b" + + +@pytest.mark.asyncio +async def test_get_relevant_docs_applies_top_k_cap(): + r = FakeRetriever() + r.results_queue = [_chunks("a", "b", "c")] + p = RetrieverPipeline(retriever=r) + sq = SearchQueries(query_list=[Query(query="q1")]) + out = await p.get_relevant_docs(partition=["p1"], search_queries=sq, top_k=2) + assert len(out) == 2 + + +@pytest.mark.asyncio +async def test_retrieve_docs_expansion_no_new_chunks_skips_second_rerank(): + r = FakeRetriever(expansion_enabled=True) + r.results_queue = [_chunks("a", "b")] + r.expand_result = _chunks("a", "b") # expansion returns same set + rer = FakeReranker() + p = RetrieverPipeline(retriever=r, reranker=rer, reranker_top_k=2) + out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi")) + # Only the pre-expansion rerank fired. + assert len(rer.calls) == 1 + assert {c.id for c in out} == {"a", "b"} + + +@pytest.mark.asyncio +async def test_rerank_chunks_short_circuits_on_empty_input(): + """Direct cover of the early-return guard inside _rerank_chunks.""" + from core.retrieval.pipeline import _rerank_chunks + + rer = FakeReranker() + out = await _rerank_chunks(rer, "q", []) + assert out == [] + assert rer.calls == [] + + +def test_pipeline_expansion_enabled_false_for_non_base_retriever(): + """Retrievers without an expansion_enabled attr (e.g. custom impls) are + treated as non-expanding via getattr default.""" + + class MinimalRetriever(Retriever): + async def retrieve(self, partition, query, filter=None, filter_params=None): + return [] + + async def expand_search_results(self, results): + return results + + p = RetrieverPipeline(retriever=MinimalRetriever()) + assert p.expansion_enabled is False + + +@pytest.mark.asyncio +async def test_retrieve_docs_caps_to_top_k(): + r = FakeRetriever() + r.results_queue = [_chunks("a", "b", "c", "d")] + p = RetrieverPipeline(retriever=r) + out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi"), top_k=2) + assert [c.id for c in out] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_retrieve_docs_top_k_zero_returns_empty(): + """top_k=0 must mean "zero results", not "treated as None" (the legacy bug).""" + r = FakeRetriever() + r.results_queue = [_chunks("a", "b")] + p = RetrieverPipeline(retriever=r) + out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi"), top_k=0) + assert out == [] diff --git a/openrag/core/retrieval/test_retriever.py b/openrag/core/retrieval/test_retriever.py new file mode 100644 index 000000000..688eaa24e --- /dev/null +++ b/openrag/core/retrieval/test_retriever.py @@ -0,0 +1,274 @@ +"""Retriever strategy tests with fake searcher + LLM. + +These exercise the strategy logic without Ray, OpenAI, or LangChain — proving +the new core/ retriever has clean dependencies. +""" + +from __future__ import annotations + +import pytest +from core.models.chunk import Chunk +from core.retrieval.retriever import ( + HyDeRetriever, + MultiQueryRetriever, + SingleRetriever, + retriever_registry, +) +from core.retrieval.searcher import RetrievalSearcher + + +class FakeSearcher(RetrievalSearcher): + """Records calls; returns canned chunks.""" + + def __init__(self) -> None: + self.search_calls: list[dict] = [] + self.multi_calls: list[dict] = [] + self.related_calls: list[dict] = [] + self.ancestor_calls: list[dict] = [] + self.search_result: list[Chunk] = [] + self.multi_result: list[Chunk] = [] + self.related_result: list[Chunk] = [] + self.ancestor_result: list[Chunk] = [] + + async def search(self, **kwargs): + self.search_calls.append(kwargs) + return list(self.search_result) + + async def multi_query_search(self, **kwargs): + self.multi_calls.append(kwargs) + return list(self.multi_result) + + async def get_related_chunks(self, **kwargs): + self.related_calls.append(kwargs) + return list(self.related_result) + + async def get_ancestor_chunks(self, **kwargs): + self.ancestor_calls.append(kwargs) + return list(self.ancestor_result) + + +class FakeLLM: + def __init__(self, response: str) -> None: + self.response = response + self.chat_calls: list[list[dict]] = [] + + async def generate(self, prompt: str, **kwargs) -> str: + return self.response + + async def chat(self, messages: list[dict], **kwargs) -> str: + self.chat_calls.append(messages) + return self.response + + +def _chunk(idv: str, text: str = "x", document_id: str = "", partition: str = "p1") -> Chunk: + return Chunk(id=idv, text=text, document_id=document_id, partition=partition) + + +def test_registry_has_three_strategies(): + assert set(retriever_registry.list_registered()) == {"single", "multiQuery", "hyde"} + + +@pytest.mark.asyncio +async def test_single_retriever_passes_through_to_searcher(): + s = FakeSearcher() + s.search_result = [_chunk("1"), _chunk("2")] + r = SingleRetriever(searcher=s, top_k=4, similarity_threshold=0.3, with_surrounding_chunks=False) + out = await r.retrieve(partition=["p1"], query="hello", filter="x>0", filter_params={"a": 1}) + assert [c.id for c in out] == ["1", "2"] + assert s.search_calls == [ + { + "query": "hello", + "partition": ["p1"], + "top_k": 4, + "filter": "x>0", + "filter_params": {"a": 1}, + "similarity_threshold": 0.3, + "with_surrounding_chunks": False, + } + ] + + +@pytest.mark.asyncio +async def test_multi_query_retriever_splits_llm_response(): + s = FakeSearcher() + s.multi_result = [_chunk("a")] + llm = FakeLLM(response="Q one[SEP]Q two[SEP]Q three") + r = MultiQueryRetriever( + searcher=s, + llm=llm, + multi_query_template="generate {k_queries} variants of: {query}", + k_queries=3, + top_k=5, + ) + await r.retrieve(partition=["p1"], query="seed") + assert s.multi_calls[0]["queries"] == ["Q one", "Q two", "Q three"] + assert s.multi_calls[0]["top_k_per_query"] == 5 + + +@pytest.mark.asyncio +async def test_multi_query_falls_back_to_seed_on_empty_response(): + s = FakeSearcher() + llm = FakeLLM(response="") + r = MultiQueryRetriever( + searcher=s, + llm=llm, + multi_query_template="{query} {k_queries}", + k_queries=3, + ) + await r.retrieve(partition=["p1"], query="seed") + assert s.multi_calls[0]["queries"] == ["seed"] + + +@pytest.mark.asyncio +async def test_hyde_retriever_uses_hyde_only_by_default(): + s = FakeSearcher() + llm = FakeLLM(response="A hypothetical answer paragraph.") + r = HyDeRetriever(searcher=s, llm=llm, hyde_template="Answer: {question}") + await r.retrieve(partition=["p1"], query="real question") + assert s.multi_calls[0]["queries"] == ["A hypothetical answer paragraph."] + + +@pytest.mark.asyncio +async def test_hyde_retriever_combine_appends_original_query(): + s = FakeSearcher() + llm = FakeLLM(response="hypothetical") + r = HyDeRetriever(searcher=s, llm=llm, hyde_template="Answer: {question}", combine=True) + await r.retrieve(partition=["p1"], query="real") + assert s.multi_calls[0]["queries"] == ["hypothetical", "real"] + + +@pytest.mark.asyncio +async def test_expansion_disabled_returns_unchanged(): + s = FakeSearcher() + r = SingleRetriever(searcher=s) + initial = [_chunk("1")] + out = await r.expand_search_results(initial) + assert out is initial + assert not s.related_calls + assert not s.ancestor_calls + + +@pytest.mark.asyncio +async def test_expansion_with_related_dedupes_by_id(): + s = FakeSearcher() + s.related_result = [_chunk("1"), _chunk("3")] # "1" already in results + r = SingleRetriever(searcher=s, include_related=True) + initial = [ + Chunk(id="1", text="x", partition="p1", metadata={"relationship_id": "r1"}), + ] + out = await r.expand_search_results(initial) + assert [c.id for c in out] == ["1", "3"] + + +@pytest.mark.asyncio +async def test_expansion_with_ancestors_calls_searcher(): + s = FakeSearcher() + s.ancestor_result = [_chunk("99", document_id="f1")] + r = SingleRetriever(searcher=s, include_ancestors=True, related_limit=20, max_ancestor_depth=2) + initial = [Chunk(id="1", text="x", partition="p1", document_id="f1")] + out = await r.expand_search_results(initial) + assert [c.id for c in out] == ["1", "99"] + assert s.ancestor_calls[0]["partition"] == "p1" + assert s.ancestor_calls[0]["file_id"] == "f1" + assert s.ancestor_calls[0]["limit"] == 20 + assert s.ancestor_calls[0]["max_ancestor_depth"] == 2 + + +@pytest.mark.asyncio +async def test_expansion_swallows_per_call_errors(): + class BoomSearcher(FakeSearcher): + async def get_related_chunks(self, **kwargs): + raise RuntimeError("kaboom") + + s = BoomSearcher() + r = SingleRetriever(searcher=s, include_related=True) + initial = [Chunk(id="1", text="x", partition="p1", metadata={"relationship_id": "r1"})] + out = await r.expand_search_results(initial) + assert [c.id for c in out] == ["1"] + + +@pytest.mark.asyncio +async def test_expansion_swallows_ancestor_errors(): + class BoomSearcher(FakeSearcher): + async def get_ancestor_chunks(self, **kwargs): + raise RuntimeError("ancestor exploded") + + s = BoomSearcher() + r = SingleRetriever(searcher=s, include_ancestors=True) + initial = [Chunk(id="1", text="x", partition="p1", document_id="f1")] + out = await r.expand_search_results(initial) + assert [c.id for c in out] == ["1"] + + +def test_multi_query_retriever_rejects_missing_llm(): + s = FakeSearcher() + with pytest.raises(ValueError, match="llm must be provided"): + MultiQueryRetriever(searcher=s, llm=None, multi_query_template="{query} {k_queries}") + + +def test_hyde_retriever_rejects_missing_llm(): + s = FakeSearcher() + with pytest.raises(ValueError, match="llm must be provided"): + HyDeRetriever(searcher=s, llm=None, hyde_template="{question}") + + +@pytest.mark.asyncio +async def test_multi_query_retriever_caps_response_to_k_queries(): + """A non-compliant LLM that returns more variants than requested must + not fan out additional searches.""" + s = FakeSearcher() + llm = FakeLLM(response="Q1[SEP]Q2[SEP]Q3[SEP]Q4[SEP]Q5") + r = MultiQueryRetriever( + searcher=s, + llm=llm, + multi_query_template="{query} {k_queries}", + k_queries=2, + ) + await r.retrieve(partition=["p1"], query="seed") + assert s.multi_calls[0]["queries"] == ["Q1", "Q2"] + + +@pytest.mark.asyncio +async def test_hyde_retriever_falls_back_to_seed_on_blank_generation(): + s = FakeSearcher() + llm = FakeLLM(response=" \n\t ") + r = HyDeRetriever(searcher=s, llm=llm, hyde_template="{question}") + await r.retrieve(partition=["p1"], query="seed") + assert s.multi_calls[0]["queries"] == ["seed"] + + +@pytest.mark.asyncio +async def test_hyde_retriever_falls_back_to_seed_when_combine_and_blank(): + """Combine mode should also fall back to just [seed] on blank generation, + not [blank, seed].""" + s = FakeSearcher() + llm = FakeLLM(response="") + r = HyDeRetriever(searcher=s, llm=llm, hyde_template="{question}", combine=True) + await r.retrieve(partition=["p1"], query="seed") + assert s.multi_calls[0]["queries"] == ["seed"] + + +@pytest.mark.asyncio +async def test_expansion_dedupes_ancestor_fetches_per_file(): + """Two chunks from the same (partition, document_id) must enqueue a + single ancestor fetch, not two.""" + + class CountingSearcher(FakeSearcher): + def __init__(self) -> None: + super().__init__() + self.ancestor_call_count = 0 + + async def get_ancestor_chunks(self, **kwargs): + self.ancestor_call_count += 1 + return list(self.ancestor_result) + + s = CountingSearcher() + s.ancestor_result = [_chunk("99", document_id="f1")] + r = SingleRetriever(searcher=s, include_ancestors=True) + initial = [ + Chunk(id="1", text="x", partition="p1", document_id="f1"), + Chunk(id="2", text="y", partition="p1", document_id="f1"), + Chunk(id="3", text="z", partition="p1", document_id="f1"), + ] + await r.expand_search_results(initial) + assert s.ancestor_call_count == 1 diff --git a/openrag/core/retrieval/test_rrf.py b/openrag/core/retrieval/test_rrf.py new file mode 100644 index 000000000..9ad9bc617 --- /dev/null +++ b/openrag/core/retrieval/test_rrf.py @@ -0,0 +1,49 @@ +"""RRF unit tests — fusion semantics and edge cases.""" + +from __future__ import annotations + +import pytest +from core.retrieval.rrf import rrf_reranking + + +def test_rrf_empty_returns_empty(): + assert rrf_reranking([]) == [] + + +def test_rrf_single_list_returned_as_is(): + items = [{"id": "a"}, {"id": "b"}, {"id": "c"}] + assert rrf_reranking([items]) == items + + +def test_rrf_fuses_overlapping_results(): + # 'a' is rank 1 in list1 and rank 2 in list2 → top + # 'c' is rank 1 in list2 → second + # 'b' is rank 2 in list1 only + list1 = [{"id": "a"}, {"id": "b"}] + list2 = [{"id": "c"}, {"id": "a"}] + fused = rrf_reranking([list1, list2], key_fn=lambda x: x["id"]) + ids = [item["id"] for item in fused] + assert ids[0] == "a" + assert set(ids) == {"a", "b", "c"} + + +def test_rrf_without_key_fn_does_not_fuse(): + list1 = [{"id": "a"}] + list2 = [{"id": "a"}] # different object, same logical id + fused = rrf_reranking([list1, list2]) + # Object identity → two separate items in fused result + assert len(fused) == 2 + + +def test_rrf_smaller_k_emphasizes_top_ranks(): + list1 = [{"id": "a"}, {"id": "b"}] + list2 = [{"id": "b"}, {"id": "a"}] + fused = rrf_reranking([list1, list2], key_fn=lambda x: x["id"], k=1) + # k=1: top-rank in any list dominates; with two top-1s for different items, + # both score the same — order is stable across implementations though + assert {item["id"] for item in fused} == {"a", "b"} + + +def test_rrf_rejects_negative_k(): + with pytest.raises(ValueError, match="non-negative"): + rrf_reranking([[{"id": "a"}], [{"id": "b"}]], key_fn=lambda x: x["id"], k=-1) diff --git a/openrag/core/utils/__init__.py b/openrag/core/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/core/utils/conts.py b/openrag/core/utils/conts.py new file mode 100644 index 000000000..9f457374d --- /dev/null +++ b/openrag/core/utils/conts.py @@ -0,0 +1,10 @@ +PARTITION_PREFIX = "openrag-" +LEGACY_PARTITION_PREFIX = "ragondin-" + +FILE_READ_CHUNK_SIZE = 1024 * 1024 # Read file in blocks of 1MB to preserve RAM + + +IMG_WRAPPER_OPEN = "\n\n" +IMG_WRAPPER_CLOSE = "\n\n" + +IMAGE_PLACEHOLDER = f"""{IMG_WRAPPER_OPEN}[Image Placeholder]{IMG_WRAPPER_CLOSE}""" diff --git a/openrag/core/utils/exceptions.py b/openrag/core/utils/exceptions.py new file mode 100644 index 000000000..db2410fde --- /dev/null +++ b/openrag/core/utils/exceptions.py @@ -0,0 +1,374 @@ +"""Unified exception hierarchy for OpenRAG. + +All exceptions inherit from OpenRAGError and carry a machine-readable +``code``, an HTTP ``status_code``, and an optional ``extra`` dict. + +The hierarchy is organised by concern: + + OpenRAGError + +-- ConfigError + +-- RegistryError + +-- PipelineError + +-- AuthError + | +-- AuthenticationError (401) + +-- ValidationError (422) + +-- NotFoundError (404) + | +-- DocumentNotFoundError + | +-- PartitionNotFoundError + | +-- UserNotFoundError + +-- QuotaExceededError (429) + +-- ServiceUnavailableError (503) + | +-- CircuitBreakerOpenError + +-- InferenceError (503) + | +-- LLMParsingError (502) + | +-- InferenceTimeoutError (504) + | +-- InferenceConnectionError (503) + +-- StorageError (500) + | +-- MilvusError + | +-- PostgresError + +-- EmbeddingError (500) + | +-- EmbeddingAPIError + | +-- EmbeddingResponseError (422) + | +-- UnexpectedEmbeddingError + +-- VDBError (500) + +-- VDBConnectionError (503) + +-- VDBInsertError (422) + +-- VDBDeleteError (422) + +-- VDBSearchError (422) + +-- VDBFileIDAlreadyExistsError (409) + +-- VDBPartitionNotFound (404) + +-- VDBFileNotFoundError (404) + +-- VDBUserNotFound (404) + +-- VDBMembershipNotFound (404) + +-- VDBSchemaMigrationRequiredError (503) + +-- VDBCreateOrLoadCollectionError (422) + +-- UnexpectedVDBError (500) +""" + +from __future__ import annotations + +# --------------------------------------------------------------------------- +# Root +# --------------------------------------------------------------------------- + + +class OpenRAGError(Exception): + """Base class for all OpenRAG exceptions. + + Preserves the existing API: message, code, status_code, to_dict(). + """ + + def __init__( + self, + message: str, + code: str = "OPENRAG_ERROR", + status_code: int = 500, + **kwargs, + ): + self.message = message + self.code = code + self.status_code = status_code + self.extra = kwargs or {} + super().__init__(f"{self.code}: {self.message}") + + def to_dict(self) -> dict: + return { + "detail": f"[{self.code}]: {self.message}", + "extra": self.extra, + } + + +# --------------------------------------------------------------------------- +# Config & registry +# --------------------------------------------------------------------------- + + +class ConfigError(OpenRAGError): + """Configuration-related errors.""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="CONFIG_ERROR", status_code=500, **kwargs) + + +class RegistryError(OpenRAGError): + """Registry lookup errors (unknown component name).""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="REGISTRY_ERROR", status_code=500, **kwargs) + + +class PipelineError(OpenRAGError): + """Pipeline execution errors.""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="PIPELINE_ERROR", status_code=500, **kwargs) + + +# --------------------------------------------------------------------------- +# Auth +# --------------------------------------------------------------------------- + + +class AuthError(OpenRAGError): + """Authentication / authorization errors.""" + + def __init__(self, message: str, *, code: str = "AUTH_ERROR", status_code: int = 403, **kwargs): + super().__init__(message, code=code, status_code=status_code, **kwargs) + + +class AuthenticationError(AuthError): + """Missing or invalid credentials. Maps to HTTP 401.""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="AUTHENTICATION_ERROR", status_code=401, **kwargs) + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + + +class ValidationError(OpenRAGError): + """Input validation or business rule violation. Maps to HTTP 422 by default. + + Accepts a custom ``status_code`` so callers can preserve more specific + semantics (e.g. 400 Bad Request for malformed input, 415 Unsupported + Media Type for rejected file formats). + """ + + def __init__(self, message: str, *, status_code: int = 422, code: str = "VALIDATION_ERROR", **kwargs): + super().__init__(message, code=code, status_code=status_code, **kwargs) + + +# --------------------------------------------------------------------------- +# Not found +# --------------------------------------------------------------------------- + + +class NotFoundError(OpenRAGError): + """Requested resource not found. Maps to HTTP 404.""" + + def __init__(self, message: str, code: str = "NOT_FOUND", **kwargs): + super().__init__(message, code=code, status_code=404, **kwargs) + + +class DocumentNotFoundError(NotFoundError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="DOCUMENT_NOT_FOUND", **kwargs) + + +class PartitionNotFoundError(NotFoundError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="PARTITION_NOT_FOUND", **kwargs) + + +class UserNotFoundError(NotFoundError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="USER_NOT_FOUND", **kwargs) + + +# --------------------------------------------------------------------------- +# Quota +# --------------------------------------------------------------------------- + + +class QuotaExceededError(OpenRAGError): + """File quota exceeded. Maps to HTTP 429.""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="QUOTA_EXCEEDED", status_code=429, **kwargs) + + +# --------------------------------------------------------------------------- +# Infrastructure — service availability +# --------------------------------------------------------------------------- + + +class ServiceUnavailableError(OpenRAGError): + """External service unavailable after retry exhaustion. Maps to HTTP 503.""" + + def __init__(self, message: str, *, code: str = "SERVICE_UNAVAILABLE", status_code: int = 503, **kwargs): + super().__init__(message, code=code, status_code=status_code, **kwargs) + + +class CircuitBreakerOpenError(ServiceUnavailableError): + """Circuit breaker is open. Maps to HTTP 503.""" + + def __init__(self, service_type: str, **kwargs): + self.service_type = service_type + super().__init__( + f"Circuit breaker open for {service_type} — service unavailable", + code="CIRCUIT_BREAKER_OPEN", + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# Inference +# --------------------------------------------------------------------------- + + +class InferenceError(OpenRAGError): + """Base for all inference service failures. Maps to HTTP 503.""" + + def __init__(self, message: str, *, code: str = "INFERENCE_ERROR", status_code: int = 503, **kwargs): + super().__init__(message, code=code, status_code=status_code, **kwargs) + + +class LLMParsingError(InferenceError): + """LLM returned invalid JSON. Maps to HTTP 502.""" + + def __init__(self, raw_response: str, parse_error: str | None = None, **kwargs): + self.raw_response = raw_response[:500] + self.parse_error = parse_error + super().__init__( + f"LLM returned invalid JSON: {self.raw_response[:100]}...", + code="LLM_PARSING_ERROR", + status_code=502, + **kwargs, + ) + + +class InferenceTimeoutError(InferenceError): + """Inference request timed out. Maps to HTTP 504.""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="INFERENCE_TIMEOUT", status_code=504, **kwargs) + + +class InferenceConnectionError(InferenceError): + """Cannot reach inference service. Maps to HTTP 503.""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="INFERENCE_CONNECTION_ERROR", **kwargs) + + +# --------------------------------------------------------------------------- +# Storage +# --------------------------------------------------------------------------- + + +class StorageError(OpenRAGError): + """Base for storage failures. Maps to HTTP 500.""" + + def __init__(self, message: str, *, code: str = "STORAGE_ERROR", status_code: int = 500, **kwargs): + super().__init__(message, code=code, status_code=status_code, **kwargs) + + +class MilvusError(StorageError): + """Milvus-specific failures.""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="MILVUS_ERROR", **kwargs) + + +class PostgresError(StorageError): + """Postgres-specific failures.""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="POSTGRES_ERROR", **kwargs) + + +# --------------------------------------------------------------------------- +# Embedding (preserves existing OpenRAG exception classes) +# --------------------------------------------------------------------------- + + +class EmbeddingError(OpenRAGError): + """Base exception for all embedding-related errors.""" + + def __init__(self, message: str, code: str = "EMBEDDING_ERROR", status_code: int = 500, **kwargs): + super().__init__(message, code=code, status_code=status_code, **kwargs) + + +class EmbeddingAPIError(EmbeddingError): + """API error with the embedding provider.""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="EMBEDDING_API_ERROR", status_code=500, **kwargs) + + +class EmbeddingResponseError(EmbeddingError): + """Invalid or unexpected response from embedding provider.""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="EMBEDDING_RESPONSE_ERROR", status_code=422, **kwargs) + + +class UnexpectedEmbeddingError(EmbeddingError): + """Unexpected error in embedding operations.""" + + def __init__(self, message: str, **kwargs): + super().__init__(message, code="EMBEDDING_UNEXPECTED_ERROR", status_code=500, **kwargs) + + +# --------------------------------------------------------------------------- +# Vector database (preserves existing OpenRAG exception classes) +# --------------------------------------------------------------------------- + + +class VDBError(OpenRAGError): + """Base exception for all vector database-related errors.""" + + def __init__(self, message: str, code: str = "VDB_ERROR", status_code: int = 500, **kwargs): + super().__init__(message, code=code, status_code=status_code, **kwargs) + + +class VDBConnectionError(VDBError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="VDB_CONNECTION_ERROR", status_code=503, **kwargs) + + +class VDBCreateOrLoadCollectionError(VDBError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="VDB_COLLECTION_ERROR", status_code=422, **kwargs) + + +class VDBInsertError(VDBError): + def __init__(self, message: str, status_code: int = 422, **kwargs): + super().__init__(message, code="VDB_INSERT_ERROR", status_code=status_code, **kwargs) + + +class VDBFileIDAlreadyExistsError(VDBError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="VDB_FILE_ALREADY_EXISTS", status_code=409, **kwargs) + + +class VDBDeleteError(VDBError): + def __init__(self, message: str, status_code: int = 422, **kwargs): + super().__init__(message, code="VDB_DELETE_ERROR", status_code=status_code, **kwargs) + + +class VDBSearchError(VDBError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="VDB_SEARCH_ERROR", status_code=422, **kwargs) + + +class VDBPartitionNotFound(VDBError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="VDB_PARTITION_NOT_FOUND", status_code=404, **kwargs) + + +class VDBFileNotFoundError(VDBError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="VDB_FILE_NOT_FOUND", status_code=404, **kwargs) + + +class VDBUserNotFound(VDBError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="VDB_USER_NOT_FOUND", status_code=404, **kwargs) + + +class VDBMembershipNotFound(VDBError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="VDB_MEMBERSHIP_NOT_FOUND", status_code=404, **kwargs) + + +class VDBSchemaMigrationRequiredError(VDBError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="VDB_SCHEMA_MIGRATION_REQUIRED", status_code=503, **kwargs) + + +class UnexpectedVDBError(VDBError): + def __init__(self, message: str, **kwargs): + super().__init__(message, code="VDB_UNEXPECTED_ERROR", status_code=500, **kwargs) diff --git a/openrag/core/utils/external_errors.py b/openrag/core/utils/external_errors.py new file mode 100644 index 000000000..3cd24460b --- /dev/null +++ b/openrag/core/utils/external_errors.py @@ -0,0 +1,62 @@ +"""Utilities for detecting external resource access errors. + +When VLM models fetch external image URLs, HTTP errors (403, 404, etc.) +from remote servers get wrapped in InternalServerError, which is +misleading. This module detects such errors for better logging. + +Pure functions — no infrastructure imports. + +Moved from: utils/external_resource_errors.py +""" + +import re + +EXTERNAL_ERROR_CODES = frozenset( + { + # 4xx client errors + "400", + "401", + "403", + "404", + "405", + "408", + "410", + "429", + "451", + # 5xx gateway errors + "502", + "503", + "504", + } +) + +EXTERNAL_ERROR_INDICATORS = ( + "ClientResponseError", + "HTTPError", + "ConnectionError", + "TimeoutError", + "SSLError", +) + + +def is_external_resource_error(error: Exception) -> tuple[bool, str, str]: + """Check if an error is caused by an external resource access issue. + + Returns: + (is_external_error, status_code, url) — status_code and url are + empty strings if not detected. + """ + error_str = str(error) + + status_code = "" + for match in re.finditer(r"\b([45]\d{2})\b", error_str): + if match.group(1) in EXTERNAL_ERROR_CODES: + status_code = match.group(1) + break + + url_match = re.search(r"https?://[^\s'\"\)>]+", error_str) + url = url_match.group(0) if url_match else "" + + has_indicator = any(ind in error_str for ind in EXTERNAL_ERROR_INDICATORS) + + return bool(status_code) or has_indicator, status_code, url diff --git a/openrag/core/utils/filename.py b/openrag/core/utils/filename.py new file mode 100644 index 000000000..7fe852970 --- /dev/null +++ b/openrag/core/utils/filename.py @@ -0,0 +1,49 @@ +"""Filename sanitization and generation utilities. + +Pure functions — no infrastructure imports. + +Extracted from: components/indexer/utils/files.py (pure parts only). +""" + +import re +import secrets +import time +from pathlib import Path + + +def sanitize_filename(filename: str) -> str: + """Sanitize a filename by removing special characters. + + Keeps only word characters and underscores. Hyphens are converted + to underscores. Multiple underscores are collapsed. + + Args: + filename: Original filename (with extension) + + Returns: + Sanitized filename with extension preserved + """ + path = Path(filename) + name = path.stem + ext = path.suffix + + name = re.sub(r"[^\w\-]", "_", name) + name = name.replace("-", "_") + name = re.sub(r"_+", "_", name) + name = name.strip("_") + + return name + ext + + +def make_unique_filename(filename: str) -> str: + """Generate a unique filename by prepending timestamp + random hex. + + Args: + filename: Original filename + + Returns: + Unique filename like "1713700000000_a1b2_original.pdf" + """ + ts = int(time.time() * 1000) + rand = secrets.token_hex(2) + return f"{ts}_{rand}_{filename}" diff --git a/openrag/core/utils/registry.py b/openrag/core/utils/registry.py new file mode 100644 index 000000000..cfae4f73c --- /dev/null +++ b/openrag/core/utils/registry.py @@ -0,0 +1,66 @@ +"""Generic registry pattern for pluggable components. + +Usage: + from .registry import Registry + + embedder_registry: Registry[Embedder] = Registry("embedder") + + @embedder_registry.register("vllm") + class VLLMEmbedder(Embedder): + ... + + instance = embedder_registry.create("vllm", endpoint="http://...") +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +from core.utils.exceptions import RegistryError as RegistryError # noqa: F401 + + +class Registry[T]: + """Generic type-safe registry mapping string names to component classes. + + Each component domain (embedder, reranker, llm, vlm, chunking, parser) + has its own Registry instance. Implementations register via the + ``@registry.register("name")`` decorator and are instantiated via + ``registry.create("name", **kwargs)``. + """ + + def __init__(self, kind: str) -> None: + self._kind = kind + self._registry: dict[str, type[T]] = {} + + def register(self, name: str) -> Callable[[type[T]], type[T]]: + """Decorator to register a class under *name*.""" + + def decorator(cls: type[T]) -> type[T]: + self._registry[name] = cls + return cls + + return decorator + + def create(self, name: str, **kwargs: Any) -> T: + """Instantiate a registered class by name.""" + cls = self._registry.get(name) + if cls is None: + available = ", ".join(sorted(self._registry)) + raise RegistryError(f"{self._kind} '{name}' not found. Available: [{available}]") + return cls(**kwargs) + + def get_class(self, name: str) -> type[T]: + """Return the registered class without instantiating.""" + cls = self._registry.get(name) + if cls is None: + available = ", ".join(sorted(self._registry)) + raise RegistryError(f"{self._kind} '{name}' not found. Available: [{available}]") + return cls + + def list_registered(self) -> list[str]: + """Return sorted list of registered names.""" + return sorted(self._registry) + + def __contains__(self, name: str) -> bool: + return name in self._registry diff --git a/openrag/core/utils/test_external_errors.py b/openrag/core/utils/test_external_errors.py new file mode 100644 index 000000000..f0ee3b290 --- /dev/null +++ b/openrag/core/utils/test_external_errors.py @@ -0,0 +1,129 @@ +""" +Tests for external resource error detection utilities. +Related to: https://github.com/linagora/openrag/issues/182 +""" + +import pytest +from core.utils.external_errors import is_external_resource_error + + +class TestIsExternalResourceError: + """Test suite for is_external_resource_error function.""" + + @pytest.mark.parametrize( + "error_msg,expected_code,url_contains", + [ + # Issue #182 + ( + "aiohttp.client_exceptions.ClientResponseError: 403, message='Forbidden', " + "url='https://upload.wikimedia.org/wikipedia/commons/thumb/d/d5/Logo.png'", + "403", + "upload.wikimedia.org", + ), + # Other HTTP status codes + ( + "ClientResponseError: 404, url='https://example.com/missing.png'", + "404", + "example.com", + ), + ( + "HTTPError: 401 Unauthorized for url: https://api.example.com/image.jpg", + "401", + "api.example.com", + ), + ( + "ClientResponseError: 429 Too Many Requests - https://cdn.example.com/img.png", + "429", + "cdn.example.com", + ), + # 5xx gateway errors + ( + "502 Bad Gateway: https://api.example.com/image.png", + "502", + "api.example.com", + ), + ( + "ClientResponseError: 503 Service Unavailable - https://cdn.example.com/img.png", + "503", + "cdn.example.com", + ), + # vLLM wrapped error (the real-world scenario) + ( + "openai.InternalServerError: Error code: 500 - {'error': {'message': " + "'litellm.InternalServerError: aiohttp.client_exceptions.ClientResponseError: " + "403, message=Forbidden, url=https://example.com/path/to/image.png'}}", + "403", + "example.com/path/to/image.png", + ), + ], + ) + def test_detects_http_errors_with_urls(self, error_msg, expected_code, url_contains): + """Test detection of HTTP errors with URL extraction.""" + is_external, status_code, url = is_external_resource_error(Exception(error_msg)) + + assert is_external is True + assert status_code == expected_code + assert url_contains in url + + @pytest.mark.parametrize( + "error_msg", + [ + "TimeoutError: Connection timed out while fetching resource", + "SSLError: Certificate verification failed", + "ConnectionError: Failed to connect to server", + "aiohttp.client_exceptions.ClientResponseError: some error", + "requests.exceptions.HTTPError: 500 Server Error", + ], + ) + def test_detects_error_indicators(self, error_msg): + """Test detection via error type indicators.""" + is_external, _, _ = is_external_resource_error(Exception(error_msg)) + assert is_external is True + + @pytest.mark.parametrize( + "error", + [ + Exception("ValueError: Invalid input parameter"), + Exception("Something went wrong during processing"), + TypeError("'NoneType' object is not subscriptable"), + AttributeError("'dict' object has no attribute 'content'"), + Exception(""), + # vLLM error without external cause details + Exception( + "openai.InternalServerError: Error code: 500 - {'error': {'message': " + "'litellm.InternalServerError: InternalServerError: OpenAIException'}}" + ), + ], + ) + def test_does_not_flag_internal_errors(self, error): + """Test that internal/generic errors are not flagged as external.""" + is_external, status_code, url = is_external_resource_error(error) + + assert is_external is False + assert status_code == "" + assert url == "" + + def test_extracts_url_with_query_params(self): + """Test URL extraction with query parameters.""" + error = Exception("403 Forbidden: https://api.example.com/image?id=123&size=large") + _, _, url = is_external_resource_error(error) + + assert "api.example.com/image?id=123" in url + + def test_indicator_substring_causes_false_positive(self): + """Document known limitation: indicator substrings cause false positives. + + This test documents that internal errors mentioning HTTP error class names + will be incorrectly classified as external. This is accepted because: + 1. Real error messages use these as exception class names, not prose + 2. Stricter matching (word boundaries) would break legitimate matches + like 'aiohttp.client_exceptions.ClientResponseError' + 3. This scenario is unlikely in practice + """ + error = Exception("InternalServerError: Failed to handle ClientResponseError in retry logic") + is_external, status_code, url = is_external_resource_error(error) + + # This IS classified as external (false positive) due to substring match + assert is_external is True + assert status_code == "" # No HTTP status code + assert url == "" # No URL diff --git a/openrag/core/utils/text.py b/openrag/core/utils/text.py new file mode 100644 index 000000000..b992191f6 --- /dev/null +++ b/openrag/core/utils/text.py @@ -0,0 +1,129 @@ +"""Text sanitization utilities for cleaning extracted text. + +Pure functions — no infrastructure imports. Used by chunking, indexing, +and document processing pipelines. + +Moved from: components/indexer/utils/text_sanitizer.py +""" + +import re +import unicodedata + +DEFAULT_FALLBACK_ENCODING = "utf-8" + + +def decode_bytes(raw: bytes, encoding: str | None = None) -> str: + """Decode ``raw`` to ``str`` with a UTF-8-first detection strategy. + + chardet alone misclassifies short ASCII-heavy UTF-8 as Latin-1, which + produces mojibake on common short inputs. Trying strict UTF-8 first + catches the common case; chardet handles genuinely non-UTF-8 inputs. + Falls back to UTF-8 with ``errors="replace"`` so this never raises. + """ + if encoding: + try: + return raw.decode(encoding, errors="replace") + except LookupError: + # Invalid codec name — fall through to detection. + pass + try: + return raw.decode("utf-8") + except UnicodeDecodeError: + pass + try: + import chardet + except ImportError: + return raw.decode(DEFAULT_FALLBACK_ENCODING, errors="replace") + guess = chardet.detect(raw) + detected = guess.get("encoding") or DEFAULT_FALLBACK_ENCODING + try: + return raw.decode(detected, errors="replace") + except LookupError: + return raw.decode(DEFAULT_FALLBACK_ENCODING, errors="replace") + + +def sanitize_text( + text: str, + normalize_whitespace: bool = True, + remove_control_chars: bool = True, + remove_zero_width_chars: bool = True, + max_consecutive_newlines: int = 2, + normalize_unicode: bool = True, +) -> str: + """Sanitize text by removing useless characters and normalizing whitespace. + + Performs comprehensive text cleaning including: + - Removing or normalizing control characters + - Removing zero-width spaces and invisible characters + - Normalizing excessive whitespace (spaces, tabs) + - Limiting consecutive newlines + - Unicode normalization + + Args: + text: The input text to sanitize + normalize_whitespace: If True, normalize spaces and tabs to single spaces + remove_control_chars: If True, remove control characters (except \\n, \\r, \\t) + remove_zero_width_chars: If True, remove zero-width spaces and similar chars + max_consecutive_newlines: Maximum number of consecutive newlines to keep (0 = unlimited) + normalize_unicode: If True, normalize unicode to NFC form + + Returns: + Sanitized text string + """ + if not text: + return text + + if normalize_unicode: + text = unicodedata.normalize("NFC", text) + + if remove_zero_width_chars: + text = re.sub(r"[\u200B-\u200D\u2060\uFEFF]", "", text) + + if remove_control_chars: + text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F-\x9F]", "", text) + + if normalize_whitespace: + text = re.sub(r" {2,}", " ", text) + text = re.sub(r"\t+", " ", text) + text = re.sub(r"(?m)^ +", "", text) + text = re.sub(r"(?m) +$", "", text) + + text = re.sub(r"\r\n", "\n", text) + text = re.sub(r"\r", "\n", text) + + if max_consecutive_newlines > 0: + pattern = r"\n{" + str(max_consecutive_newlines + 1) + r",}" + replacement = "\n" * max_consecutive_newlines + text = re.sub(pattern, replacement, text) + + text = text.strip() + return text + + +def clean_markdown_table_spacing(markdown_table: str) -> str: + """Normalize spacing inside a markdown table. + + Trims each cell while keeping table shape intact. + """ + cleaned_lines = [] + + for line in markdown_table.strip().split("\n"): + if "|" not in line: + cleaned_lines.append(line.strip()) + continue + + parts = line.split("|") + cleaned_cells = [cell.strip() for cell in parts] + new_line = "| " + " | ".join(cleaned_cells[1:-1]) + " |" + cleaned_lines.append(new_line) + + return "\n".join(cleaned_lines) + + +def sanitize_extracted_text(text: str) -> str: + """Convenience function for sanitizing text extracted from documents. + + Applies default sanitization settings suitable for text extraction + endpoints and general document processing. + """ + return sanitize_text(text) diff --git a/openrag/core/vector_stores/__init__.py b/openrag/core/vector_stores/__init__.py new file mode 100644 index 000000000..7655349fe --- /dev/null +++ b/openrag/core/vector_stores/__init__.py @@ -0,0 +1,5 @@ +"""VectorStore ABC.""" + +from .vector_store import VectorStore + +__all__ = ["VectorStore"] diff --git a/openrag/core/vector_stores/vector_store.py b/openrag/core/vector_stores/vector_store.py new file mode 100644 index 000000000..78761790c --- /dev/null +++ b/openrag/core/vector_stores/vector_store.py @@ -0,0 +1,76 @@ +"""Abstract vector store interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from openrag.core.models.chunk import Chunk + + +class VectorStore(ABC): + """Base class for vector database backends.""" + + @abstractmethod + async def upsert(self, chunks: list[Chunk], collection: str = "default") -> int: + """Insert or update chunks. Returns count of upserted items.""" + ... + + @abstractmethod + async def search( + self, + embedding: list[float], + query_text: str | None = None, + top_k: int = 10, + collection: str = "default", + filters: dict[str, Any] | None = None, + similarity_threshold: float | None = None, + ) -> list[dict[str, Any]]: + """Similarity search returning raw result dicts. + + Hybrid (dense + lexical) retrieval is a backend configuration + concern, not a separate entry point: when a backend has it enabled + it fuses a dense vector match with a lexical match, and ``query_text`` + carries the raw query such backends compute the sparse vector from + server-side. Dense-only backends ignore ``query_text``. + + ``similarity_threshold`` (when set) lower-bounds the dense leg's + similarity; backends supporting range search drop anything scoring at + or below it. ``None`` disables the bound. + """ + ... + + @abstractmethod + async def delete(self, ids: list[str], collection: str = "default") -> int: + """Delete chunks by ID. Returns count of deleted items.""" + ... + + @abstractmethod + async def ensure_collection(self, name: str, dimension: int, **kwargs: Any) -> None: + """Create collection if it doesn't exist.""" + ... + + @abstractmethod + async def drop_collection(self, name: str) -> None: + """Drop a collection entirely.""" + ... + + @abstractmethod + async def collection_exists(self, name: str) -> bool: + """Check if collection exists.""" + ... + + @abstractmethod + async def query_ids_by_filter(self, collection: str, filters: dict[str, Any]) -> list[str]: + """Return chunk IDs matching the given filter expression.""" + ... + + @abstractmethod + async def query_chunks_by_filter( + self, + collection: str, + filters: dict[str, Any], + output_fields: list[str] | None = None, + ) -> list[dict[str, Any]]: + """Return full chunk data matching the given filter expression.""" + ... diff --git a/openrag/core/vlm/__init__.py b/openrag/core/vlm/__init__.py new file mode 100644 index 000000000..0ae67a20b --- /dev/null +++ b/openrag/core/vlm/__init__.py @@ -0,0 +1,6 @@ +"""VLM ABC + registry.""" + +from .registry import vlm_registry +from .vlm import VLM + +__all__ = ["VLM", "vlm_registry"] diff --git a/openrag/core/vlm/registry.py b/openrag/core/vlm/registry.py new file mode 100644 index 000000000..253a67b37 --- /dev/null +++ b/openrag/core/vlm/registry.py @@ -0,0 +1,7 @@ +"""VLM registry.""" + +from openrag.core.utils.registry import Registry + +from .vlm import VLM + +vlm_registry: Registry[VLM] = Registry("vlm") diff --git a/openrag/core/vlm/vlm.py b/openrag/core/vlm/vlm.py new file mode 100644 index 000000000..2da69b6d7 --- /dev/null +++ b/openrag/core/vlm/vlm.py @@ -0,0 +1,19 @@ +"""Abstract VLM (Vision-Language Model) interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class VLM(ABC): + """Base class for all vision-language model providers.""" + + @abstractmethod + async def caption_image(self, image_bytes: bytes, prompt: str | None = None) -> str: + """Generate a caption/description for an image.""" + ... + + @abstractmethod + async def caption_images_batch(self, images: list[bytes], prompt: str | None = None) -> list[str]: + """Batch caption multiple images.""" + ... diff --git a/openrag/di/__init__.py b/openrag/di/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/di/container.py b/openrag/di/container.py new file mode 100644 index 000000000..8ddf53c85 --- /dev/null +++ b/openrag/di/container.py @@ -0,0 +1,474 @@ +"""Service container — wires registries and exposes component factories. + +The container is the composition root for the refactored stack. It does +three things: + +1. Populates the inference registries (Phase 6) so factory helpers can spin + up embedders, LLMs, rerankers and VLMs by name. +2. Builds the storage adapters (Phase 7E) when a :class:`Settings` instance + is supplied — a :class:`~core.ports.catalog_store.CatalogStore` and a + :class:`~core.vector_stores.VectorStore`. +3. Owns the async :meth:`initialize` / :meth:`shutdown` lifecycle that opens + and closes the asyncpg pool. + +The ``settings`` argument is optional so the legacy test paths that only +care about registry side effects (``ServiceContainer()`` with no config) +keep working. Code that wants storage adapters must pass a +:class:`Settings` and ``await container.initialize()`` before issuing +queries. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +from core.embeddings import embedder_registry +from core.llm import llm_registry +from core.rerankers import reranker_registry +from core.vlm import vlm_registry +from di.embedders import register_embedders +from di.llms import register_llms +from di.repositories import create_catalog_store +from di.rerankers import register_rerankers +from di.vector_stores import create_vector_store +from di.vlms import register_vlms + +if TYPE_CHECKING: + from core.config.root import Settings + from core.ports.audit_log_repo import AuditLogRepository + from core.ports.catalog_store import CatalogStore + from core.ports.chunk_repo import ChunkRepository + from core.ports.conversation_repo import ConversationRepository + from core.ports.document_repo import DocumentRepository + from core.ports.entity_repo import EntityRepository + from core.ports.idempotency_repo import IdempotencyRepository + from core.ports.job_repo import JobRepository + from core.ports.model_endpoint_repo import ModelEndpointRepository + from core.ports.oidc_session_repo import OIDCSessionRepository + from core.ports.partition_membership_repo import PartitionMembershipRepository + from core.ports.partition_repo import PartitionRepository + from core.ports.preset_repo import PresetRepository + from core.ports.prompt_repo import PromptRepository + from core.ports.topic_tag_repo import TopicTagRepository + from core.ports.user_repo import UserRepository + from core.ports.workspace_repo import WorkspaceRepository + from core.vector_stores import VectorStore + from services.orchestrators.auth_service import AuthService + from services.orchestrators.conversion_service import ConversionService + from services.orchestrators.indexing_service import IndexingService + from services.orchestrators.job_service import JobService + from services.orchestrators.partition_service import PartitionService + from services.orchestrators.query_service import QueryService + from services.orchestrators.retrieval_service import RetrievalService + from services.orchestrators.user_service import UserService + from services.orchestrators.workspace_service import WorkspaceService + + +_NO_SETTINGS_MESSAGE = ( + "ServiceContainer was constructed without a Settings instance — " + "pass Settings to ServiceContainer(...) to wire storage adapters." +) + + +def _oidc_config_from_env(): + """Build :class:`OIDCConfig` from the same env vars ``main.py`` validates. + + Phase 8A.1 keeps OIDC config env-sourced (it is not yet wired into the + root :class:`Settings`); ``enabled`` mirrors ``AUTH_MODE=oidc``. + """ + from core.config.auth import OIDCConfig + + return OIDCConfig( + enabled=os.getenv("AUTH_MODE", "token").strip().lower() == "oidc", + issuer_url=os.getenv("OIDC_ENDPOINT", "") or "", + client_id=os.getenv("OIDC_CLIENT_ID", "") or "", + client_secret=os.getenv("OIDC_CLIENT_SECRET", "") or "", + redirect_uri=os.getenv("OIDC_REDIRECT_URI", "") or "", + scopes=os.getenv("OIDC_SCOPES", "openid email profile offline_access"), + token_encryption_key=os.getenv("OIDC_TOKEN_ENCRYPTION_KEY", "") or "", + claim_source=os.getenv("OIDC_CLAIM_SOURCE", "id_token").strip().lower(), + claim_mapping=os.getenv("OIDC_CLAIM_MAPPING", "").strip(), + post_logout_redirect_uri=os.getenv("OIDC_POST_LOGOUT_REDIRECT_URI", "") or "", + auto_provision_login=os.getenv("OIDC_AUTO_PROVISION_LOGIN", "false").strip().lower() == "true", + ) + + +class ServiceContainer: + """Populates registries and provides typed factory access.""" + + def __init__(self, settings: Settings | None = None) -> None: + register_embedders() + register_llms() + register_rerankers() + register_vlms() + + self._settings = settings + self._catalog_store: CatalogStore | None = create_catalog_store(settings) if settings is not None else None + self._vector_store: VectorStore | None = create_vector_store(settings) if settings is not None else None + self._auth_service: AuthService | None = None + self._user_service: UserService | None = None + self._partition_service: PartitionService | None = None + self._workspace_service: WorkspaceService | None = None + self._retrieval_service: RetrievalService | None = None + self._query_service: QueryService | None = None + self._indexing_service: IndexingService | None = None + self._job_service: JobService | None = None + self._conversion_service: ConversionService | None = None + + def _require_settings(self) -> Settings: + """Settings guard for the settings-dependent service properties. + + Without this, ``ServiceContainer()`` (no-settings legacy path) + fails with a bare ``AttributeError`` on ``self._settings.x`` — + inconsistent with the ``catalog_store`` / ``vector_store`` + contract, which raises ``RuntimeError(_NO_SETTINGS_MESSAGE)``. + """ + if self._settings is None: + raise RuntimeError(_NO_SETTINGS_MESSAGE) + return self._settings + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def initialize(self) -> None: + """Open the storage adapters (asyncpg pool + Alembic migrations).""" + if self._catalog_store is not None: + await self._catalog_store.initialize() + await self.user_repo.ensure_admin_user(os.getenv("AUTH_TOKEN")) + + async def shutdown(self) -> None: + """Close the storage adapters cleanly.""" + if self._catalog_store is not None: + await self._catalog_store.shutdown() + + # ------------------------------------------------------------------ + # Storage adapters + # ------------------------------------------------------------------ + + @property + def catalog_store(self) -> CatalogStore: + if self._catalog_store is None: + raise RuntimeError(_NO_SETTINGS_MESSAGE) + return self._catalog_store + + @property + def vector_store(self) -> VectorStore: + """The Phase 7B :class:`MilvusVectorStore` built from settings. + + Cached at construction so repeated property reads return the same + instance — every fresh build would open a new pymilvus gRPC + channel. + """ + if self._vector_store is None: + raise RuntimeError(_NO_SETTINGS_MESSAGE) + return self._vector_store + + # ------------------------------------------------------------------ + # Per-repo accessors (Phase 8 orchestrators take one repo, not the + # whole store). All fifteen repos are exposed for symmetry and + # grep-findability: shortcuts for the five real repos plus the ten + # post-refactoring stubs. + # ------------------------------------------------------------------ + + @property + def document_repo(self) -> DocumentRepository: + return self.catalog_store.document_repo + + @property + def user_repo(self) -> UserRepository: + return self.catalog_store.user_repo + + @property + def partition_repo(self) -> PartitionRepository: + return self.catalog_store.partition_repo + + @property + def membership_repo(self) -> PartitionMembershipRepository: + return self.catalog_store.membership_repo + + @property + def oidc_session_repo(self) -> OIDCSessionRepository: + return self.catalog_store.oidc_session_repo + + @property + def workspace_repo(self) -> WorkspaceRepository: + return self.catalog_store.workspace_repo + + @property + def job_repo(self) -> JobRepository: + return self.catalog_store.job_repo + + @property + def chunk_repo(self) -> ChunkRepository: + return self.catalog_store.chunk_repo + + @property + def prompt_repo(self) -> PromptRepository: + return self.catalog_store.prompt_repo + + @property + def conversation_repo(self) -> ConversationRepository: + return self.catalog_store.conversation_repo + + @property + def audit_log_repo(self) -> AuditLogRepository: + return self.catalog_store.audit_log_repo + + @property + def idempotency_repo(self) -> IdempotencyRepository: + return self.catalog_store.idempotency_repo + + @property + def entity_repo(self) -> EntityRepository: + return self.catalog_store.entity_repo + + @property + def topic_tag_repo(self) -> TopicTagRepository: + return self.catalog_store.topic_tag_repo + + @property + def model_endpoint_repo(self) -> ModelEndpointRepository: + return self.catalog_store.model_endpoint_repo + + @property + def preset_repo(self) -> PresetRepository: + return self.catalog_store.preset_repo + + # ------------------------------------------------------------------ + # Orchestrators (Phase 8) + # ------------------------------------------------------------------ + + @property + def auth_service(self) -> AuthService: + """AuthService — lazily built, cached for the container's lifetime. + + The OIDC client is only constructed in ``AUTH_MODE=oidc`` (it reads + required env vars and would raise otherwise); in token mode it is + ``None`` and the OIDC flow methods refuse cleanly. + """ + if self._auth_service is None: + from services.orchestrators.auth_service import AuthService + + cfg = _oidc_config_from_env() + client = None + if cfg.enabled: + from components.auth import get_oidc_client + + client = get_oidc_client() + self._auth_service = AuthService( + user_repo=self.user_repo, + oidc_session_repo=self.oidc_session_repo, + membership_repo=self.membership_repo, + oidc_client=client, + config=cfg, + ) + return self._auth_service + + @property + def user_service(self) -> UserService: + """UserService — lazily built, cached for the container's lifetime.""" + if self._user_service is None: + from services.orchestrators.user_service import UserService + + settings = self._require_settings() + self._user_service = UserService( + user_repo=self.user_repo, + auth_service=self.auth_service, + default_file_quota=settings.rdb.default_file_quota, + partition_service=self.partition_service, + membership_repo=self.membership_repo, + job_service=self.job_service, + ) + return self._user_service + + @property + def partition_service(self) -> PartitionService: + """PartitionService — lazily built, cached for the container's lifetime.""" + if self._partition_service is None: + from services.orchestrators.partition_service import PartitionService + + settings = self._require_settings() + self._partition_service = PartitionService( + partition_repo=self.partition_repo, + membership_repo=self.membership_repo, + document_repo=self.document_repo, + vector_store=self.vector_store, + user_repo=self.user_repo, + collection=settings.vectordb.collection_name, + ) + return self._partition_service + + @property + def workspace_service(self) -> WorkspaceService: + """WorkspaceService — lazily built, cached for the container's lifetime.""" + if self._workspace_service is None: + from services.orchestrators.workspace_service import WorkspaceService + + settings = self._require_settings() + self._workspace_service = WorkspaceService( + workspace_repo=self.workspace_repo, + document_repo=self.document_repo, + vector_store=self.vector_store, + collection=settings.vectordb.collection_name, + ) + return self._workspace_service + + @property + def retrieval_service(self) -> RetrievalService: + """RetrievalService — lazily built, cached for the container's lifetime.""" + if self._retrieval_service is None: + from services.orchestrators.retrieval_service import RetrievalService + from services.storage.vector_store_searcher import VectorStoreSearcher + + settings = self._require_settings() + embed_cfg = settings.embedder + embedder = self.create_embedder( + "vllm", + endpoint=embed_cfg.base_url, + model_name=embed_cfg.model_name, + api_key=embed_cfg.api_key, + max_model_len=embed_cfg.max_model_len, + ) + searcher = VectorStoreSearcher( + vector_store=self.vector_store, + embedder=embedder, + document_repo=self.document_repo, + collection=settings.vectordb.collection_name, + ) + llm_cfg = settings.llm.model_dump() + llm = self.create_llm( + "vllm", + endpoint=llm_cfg["base_url"], + model_name=llm_cfg["model"], + api_key=llm_cfg.get("api_key", ""), + **{k: v for k, v in llm_cfg.items() if k not in ("base_url", "model", "api_key")}, + ) + reranker = None + rcfg = settings.reranker + if rcfg.enabled: + reranker = self.create_reranker( + rcfg.provider, + endpoint=rcfg.base_url, + model_name=rcfg.model_name, + api_key=rcfg.api_key, + timeout=rcfg.timeout, + ) + self._retrieval_service = RetrievalService( + searcher=searcher, + reranker=reranker, + llm=llm, + config=settings, + ) + return self._retrieval_service + + @property + def query_service(self) -> QueryService: + """QueryService — lazily built, cached for the container's lifetime. + + Shares the same core LLM construction as ``retrieval_service`` + (built from ``settings.llm``); the web-search service comes from + the legacy ``WebSearchFactory`` (provider is ``None`` when + ``WEBSEARCH_API_TOKEN`` is unset — web search silently disabled). + """ + if self._query_service is None: + from components.websearch import WebSearchFactory + from services.orchestrators.query_service import QueryService + + settings = self._require_settings() + llm_cfg = settings.llm.model_dump() + llm = self.create_llm( + "vllm", + endpoint=llm_cfg["base_url"], + model_name=llm_cfg["model"], + api_key=llm_cfg.get("api_key", ""), + **{k: v for k, v in llm_cfg.items() if k not in ("base_url", "model", "api_key")}, + ) + self._query_service = QueryService( + retrieval_service=self.retrieval_service, + llm=llm, + config=settings, + web_search_service=WebSearchFactory.create_service(settings), + workspace_service=self.workspace_service, + ) + return self._query_service + + @property + def indexing_service(self) -> IndexingService: + """IndexingService — lazily built, cached for the container's lifetime. + + Phase 9B routes new indexing jobs through the thin ``IndexerPool`` + actor while delete/update/copy remain on the legacy actor path. + """ + if self._indexing_service is None: + from services.orchestrators.indexing_service import IndexingService + from services.workers.dispatcher import from_ray_namespace + + settings = self._require_settings() + self._indexing_service = IndexingService( + document_repo=self.document_repo, + workspace_repo=self.workspace_repo, + dispatcher=from_ray_namespace( + vector_store=self.vector_store, + document_repo=self.document_repo, + workspace_repo=self.workspace_repo, + collection=settings.vectordb.collection_name, + ), + ) + return self._indexing_service + + @property + def job_service(self) -> JobService: + """JobService — lazily built, cached for the container's lifetime. + + Wraps the ``TaskStateManager`` Ray actor directly (8H excepts + JobService); resolved lazily so the actor only needs to exist at + first request. + """ + if self._job_service is None: + from services.orchestrators.job_service import JobService + from services.workers.bootstrap import get_task_state_manager + + self._job_service = JobService(task_state_manager=get_task_state_manager()) + return self._job_service + + @property + def conversion_service(self) -> ConversionService: + """ConversionService — lazily built, cached for the container's lifetime. + + The serializer is the Ray-backed ``SerializerRayShim`` during the + Phase-8 shim period (Ray cleanup is Phase 9); the DocSerializer + actor is resolved lazily per call inside the shim. + """ + if self._conversion_service is None: + from services.orchestrators.conversion_service import ConversionService + from services.workers.parsers.doc_serializer_adapter import from_ray_namespace + + settings = self._require_settings() + self._conversion_service = ConversionService( + serializer=from_ray_namespace(), + vector_store=self.vector_store, + collection=settings.vectordb.collection_name, + ) + return self._conversion_service + + # ------------------------------------------------------------------ + # Registry-based inference factories (Phase 6) + # ------------------------------------------------------------------ + + @staticmethod + def create_embedder(name: str = "vllm", **kwargs): + return embedder_registry.create(name, **kwargs) + + @staticmethod + def create_llm(name: str = "vllm", **kwargs): + return llm_registry.create(name, **kwargs) + + @staticmethod + def create_reranker(name: str = "infinity", **kwargs): + return reranker_registry.create(name, **kwargs) + + @staticmethod + def create_vlm(name: str = "vllm", **kwargs): + return vlm_registry.create(name, **kwargs) diff --git a/openrag/di/embedders.py b/openrag/di/embedders.py new file mode 100644 index 000000000..5f09ecda2 --- /dev/null +++ b/openrag/di/embedders.py @@ -0,0 +1,5 @@ +"""Register embedder implementations with the core registry.""" + + +def register_embedders() -> None: + import services.inference.vllm_client # noqa: F401 diff --git a/openrag/di/inference.py b/openrag/di/inference.py new file mode 100644 index 000000000..8447f4e42 --- /dev/null +++ b/openrag/di/inference.py @@ -0,0 +1,16 @@ +"""Convenience wrapper — registers all inference adapters at once. + +Delegates to the per-domain registration modules. +""" + +from di.embedders import register_embedders +from di.llms import register_llms +from di.rerankers import register_rerankers +from di.vlms import register_vlms + + +def register_inference() -> None: + register_embedders() + register_llms() + register_rerankers() + register_vlms() diff --git a/openrag/di/llms.py b/openrag/di/llms.py new file mode 100644 index 000000000..8620f7159 --- /dev/null +++ b/openrag/di/llms.py @@ -0,0 +1,5 @@ +"""Register LLM implementations with the core registry.""" + + +def register_llms() -> None: + import services.inference.vllm_client # noqa: F401 diff --git a/openrag/di/providers.py b/openrag/di/providers.py new file mode 100644 index 000000000..fecd16e2c --- /dev/null +++ b/openrag/di/providers.py @@ -0,0 +1,74 @@ +"""FastAPI dependency providers. + +Thin accessors over the request-scoped :class:`ServiceContainer` that +``main.py`` attaches at ``app.state.container``. Phase 8 keeps these as +one-liners — the container (``di/container.py``) is the composition +root. Phase 11 moves the attachment into a proper FastAPI lifespan and +wires ``container.initialize()``; until then the OIDC flow that needs +the asyncpg pool is dormant (token-mode auth routes already short-circuit +before reaching a service). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from fastapi import HTTPException, Request, status + +if TYPE_CHECKING: + from di.container import ServiceContainer + from services.orchestrators.auth_service import AuthService + from services.orchestrators.conversion_service import ConversionService + from services.orchestrators.indexing_service import IndexingService + from services.orchestrators.job_service import JobService + from services.orchestrators.partition_service import PartitionService + from services.orchestrators.query_service import QueryService + from services.orchestrators.retrieval_service import RetrievalService + from services.orchestrators.user_service import UserService + from services.orchestrators.workspace_service import WorkspaceService + + +def get_container(request: Request) -> ServiceContainer: + container = getattr(request.app.state, "container", None) + if container is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Service container is not available.", + ) + return container + + +def get_auth_service(request: Request) -> AuthService: + return get_container(request).auth_service + + +def get_user_service(request: Request) -> UserService: + return get_container(request).user_service + + +def get_partition_service(request: Request) -> PartitionService: + return get_container(request).partition_service + + +def get_workspace_service(request: Request) -> WorkspaceService: + return get_container(request).workspace_service + + +def get_retrieval_service(request: Request) -> RetrievalService: + return get_container(request).retrieval_service + + +def get_query_service(request: Request) -> QueryService: + return get_container(request).query_service + + +def get_indexing_service(request: Request) -> IndexingService: + return get_container(request).indexing_service + + +def get_job_service(request: Request) -> JobService: + return get_container(request).job_service + + +def get_conversion_service(request: Request) -> ConversionService: + return get_container(request).conversion_service diff --git a/openrag/di/repositories.py b/openrag/di/repositories.py new file mode 100644 index 000000000..4987727c0 --- /dev/null +++ b/openrag/di/repositories.py @@ -0,0 +1,49 @@ +"""Factory for the :class:`CatalogStore` adapter. + +The container calls :func:`create_catalog_store` once at startup to build the +concrete :class:`~services.storage.postgres_store.PostgresStore`. Centralising +construction here keeps two pieces of knowledge out of the container itself: + +* **Database-name fallback.** The legacy ``MilvusDB`` actor derives the + Postgres database name from the Milvus collection name + (``partitions_for_collection_``) at ``vectordb.py:238``. The + factory keeps that contract so wiring code never has to mention the + ``partitions_for_collection_`` prefix. +* **Migration trigger.** The factory always builds a store that will run + Alembic at :meth:`PostgresStore.initialize` unless the caller opts out via + ``run_migrations=False`` (useful in tests against a pre-migrated database). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from services.storage.postgres_store import PostgresStore + +if TYPE_CHECKING: + from core.config.root import Settings + from core.ports.catalog_store import CatalogStore + + +def create_catalog_store( + settings: Settings, + *, + run_migrations: bool = True, +) -> CatalogStore: + """Build the relational catalog store from the root settings. + + When ``settings.rdb.database`` is unset the database name is derived from + ``settings.vectordb.collection_name`` so the new adapter targets the same + Postgres database the legacy actor has always used. + """ + rdb = settings.rdb + if rdb.database is None: + rdb = rdb.model_copy( + update={ + "database": f"partitions_for_collection_{settings.vectordb.collection_name}", + }, + ) + return PostgresStore(rdb, run_migrations=run_migrations) + + +__all__ = ["create_catalog_store"] diff --git a/openrag/di/rerankers.py b/openrag/di/rerankers.py new file mode 100644 index 000000000..d16aab1bb --- /dev/null +++ b/openrag/di/rerankers.py @@ -0,0 +1,5 @@ +"""Register reranker implementations with the core registry.""" + + +def register_rerankers() -> None: + import services.inference.reranker_clients # noqa: F401 diff --git a/openrag/di/test_container.py b/openrag/di/test_container.py new file mode 100644 index 000000000..d9a9e7275 --- /dev/null +++ b/openrag/di/test_container.py @@ -0,0 +1,246 @@ +"""Phase 7E — ServiceContainer storage wiring.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from core.config.infrastructure import RDBConfig, VectorDBConfig +from core.config.root import Settings +from core.ports.audit_log_repo import AuditLogRepository +from core.ports.catalog_store import CatalogStore +from core.ports.chunk_repo import ChunkRepository +from core.ports.conversation_repo import ConversationRepository +from core.ports.document_repo import DocumentRepository +from core.ports.entity_repo import EntityRepository +from core.ports.idempotency_repo import IdempotencyRepository +from core.ports.job_repo import JobRepository +from core.ports.model_endpoint_repo import ModelEndpointRepository +from core.ports.oidc_session_repo import OIDCSessionRepository +from core.ports.partition_repo import PartitionRepository +from core.ports.preset_repo import PresetRepository +from core.ports.prompt_repo import PromptRepository +from core.ports.topic_tag_repo import TopicTagRepository +from core.ports.user_repo import UserRepository +from core.ports.workspace_repo import WorkspaceRepository +from di.container import ServiceContainer +from di.repositories import create_catalog_store +from di.vector_stores import create_vector_store + + +def _settings(database: str | None = None, collection: str = "vdb_test") -> Settings: + return Settings( + rdb=RDBConfig(password="x", database=database), + vectordb=VectorDBConfig(collection_name=collection), + ) + + +@pytest.fixture(autouse=True) +def _stub_milvus_clients(monkeypatch): + """Replace the pymilvus gRPC clients with mocks for the whole module. + + ``ServiceContainer(settings)`` eagerly builds a :class:`MilvusVectorStore` + in its constructor; the real pymilvus client tries to open a channel + immediately and hangs if Milvus is unreachable. These tests only need + to verify the *wiring*, so we stub the clients out. Real Milvus + coverage lives in ``tests/integration/test_milvus_store_integration.py``. + """ + from unittest.mock import MagicMock + + import services.storage.milvus_store as ms + + monkeypatch.setattr(ms, "MilvusClient", MagicMock()) + monkeypatch.setattr(ms, "AsyncMilvusClient", MagicMock()) + + +class TestLegacyContainerStillWorks: + """The pre-Phase-7E callers do ``ServiceContainer()`` with no settings.""" + + def test_constructs_without_settings(self): + ServiceContainer() # must not raise + + def test_catalog_store_raises_when_unconfigured(self): + c = ServiceContainer() + with pytest.raises(RuntimeError, match="without a Settings instance"): + _ = c.catalog_store + + def test_vector_store_raises_when_unconfigured(self): + c = ServiceContainer() + with pytest.raises(RuntimeError, match="without a Settings instance"): + _ = c.vector_store + + @pytest.mark.parametrize( + "name", + [ + "document_repo", + "user_repo", + "partition_repo", + "oidc_session_repo", + "workspace_repo", + "job_repo", + "chunk_repo", + "prompt_repo", + "conversation_repo", + "audit_log_repo", + "idempotency_repo", + "entity_repo", + "topic_tag_repo", + "model_endpoint_repo", + "preset_repo", + ], + ) + def test_repo_properties_raise_when_unconfigured(self, name): + c = ServiceContainer() + with pytest.raises(RuntimeError, match="without a Settings instance"): + getattr(c, name) + + +class TestCatalogStoreWiring: + def test_catalog_store_satisfies_port(self): + c = ServiceContainer(_settings()) + assert isinstance(c.catalog_store, CatalogStore) + + def test_database_name_derived_from_collection(self): + c = ServiceContainer(_settings(database=None, collection="my_collection")) + assert c.catalog_store._conn._conn_kwargs["database"] == "partitions_for_collection_my_collection" + + def test_explicit_database_overrides_fallback(self): + c = ServiceContainer(_settings(database="custom_db", collection="my_collection")) + assert c.catalog_store._conn._conn_kwargs["database"] == "custom_db" + + @pytest.mark.parametrize( + ("name", "port"), + [ + ("document_repo", DocumentRepository), + ("user_repo", UserRepository), + ("partition_repo", PartitionRepository), + ("oidc_session_repo", OIDCSessionRepository), + ("workspace_repo", WorkspaceRepository), + ("job_repo", JobRepository), + ("chunk_repo", ChunkRepository), + ("prompt_repo", PromptRepository), + ("conversation_repo", ConversationRepository), + ("audit_log_repo", AuditLogRepository), + ("idempotency_repo", IdempotencyRepository), + ("entity_repo", EntityRepository), + ("topic_tag_repo", TopicTagRepository), + ("model_endpoint_repo", ModelEndpointRepository), + ("preset_repo", PresetRepository), + ], + ) + def test_repo_property_returns_port_typed_instance(self, name, port): + c = ServiceContainer(_settings()) + repo = getattr(c, name) + assert isinstance(repo, port) + # Container shortcuts must be the same object exposed by the store — + # otherwise orchestrator injection drifts from store state. + assert repo is getattr(c.catalog_store, name) + + @pytest.mark.asyncio + async def test_initialize_seeds_admin_token(self, monkeypatch): + calls = [] + + async def ensure_admin_user(token): + calls.append(token) + + class FakeCatalogStore: + user_repo = SimpleNamespace(ensure_admin_user=ensure_admin_user) + + async def initialize(self): + calls.append("initialize") + + monkeypatch.setenv("AUTH_TOKEN", "admin-token") + c = ServiceContainer(_settings()) + c._catalog_store = FakeCatalogStore() + + await c.initialize() + + assert calls == ["initialize", "admin-token"] + + +class TestVectorStoreWiring: + """The factory returns a real :class:`MilvusVectorStore`; pymilvus gRPC + clients are patched out so these unit tests don't need Milvus reachable. + The full integration coverage lives in + ``tests/integration/test_milvus_store_integration.py``.""" + + def test_factory_returns_milvus_vector_store(self): + from services.storage.milvus_store import MilvusVectorStore + + store = create_vector_store(_settings()) + assert isinstance(store, MilvusVectorStore) + + def test_container_property_returns_milvus_vector_store(self): + from services.storage.milvus_store import MilvusVectorStore + + c = ServiceContainer(_settings()) + assert isinstance(c.vector_store, MilvusVectorStore) + + def test_container_caches_vector_store(self): + c = ServiceContainer(_settings()) + # Repeated property reads must return the same instance — every + # construction opens a fresh pymilvus gRPC channel. + assert c.vector_store is c.vector_store + + +class TestRepositoriesFactory: + def test_returns_a_catalog_store(self): + assert isinstance(create_catalog_store(_settings()), CatalogStore) + + def test_run_migrations_flag_propagates(self): + store = create_catalog_store(_settings(), run_migrations=False) + assert store._run_migrations is False + + def test_does_not_mutate_input_settings(self): + s = _settings(database=None, collection="abc") + original_database = s.rdb.database + create_catalog_store(s) + # The factory uses model_copy with an update — the source must stay None. + assert s.rdb.database == original_database + + +# Phase 8F — every orchestrator is a lazy cached property on the +# container and is reachable through a one-liner provider. Keyed by the +# container property name; value is the matching provider function. +_ORCHESTRATORS = [ + ("auth_service", "get_auth_service"), + ("user_service", "get_user_service"), + ("partition_service", "get_partition_service"), + ("workspace_service", "get_workspace_service"), + ("retrieval_service", "get_retrieval_service"), + ("query_service", "get_query_service"), + ("indexing_service", "get_indexing_service"), + ("job_service", "get_job_service"), + ("conversion_service", "get_conversion_service"), +] + + +class TestPhase8OrchestratorWiring: + """8F: all nine orchestrators wired consistently (container + providers).""" + + @pytest.mark.parametrize("prop,_provider", _ORCHESTRATORS) + def test_property_is_lazy_and_cache_slot_starts_none(self, prop, _provider): + # The public accessor is a property (lazy), not an eager attribute. + assert isinstance(getattr(ServiceContainer, prop), property) + # The cache slot exists and is None before first access (no + # settings needed — the legacy no-arg path must keep working). + assert getattr(ServiceContainer(), f"_{prop}") is None + + @pytest.mark.parametrize("prop,provider", _ORCHESTRATORS) + def test_provider_delegates_to_container_property(self, prop, provider): + from types import SimpleNamespace + + from di import providers + + sentinel = object() + fake_container = SimpleNamespace(**{prop: sentinel}) + request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(container=fake_container))) + + resolved = getattr(providers, provider)(request) + assert resolved is sentinel + + def test_no_orchestrator_is_missing_a_provider(self): + from di import providers + + wired = {name for name in vars(providers) if name.startswith("get_") and name.endswith("_service")} + assert wired == {p for _, p in _ORCHESTRATORS} diff --git a/openrag/di/test_inference.py b/openrag/di/test_inference.py new file mode 100644 index 000000000..f38de06b2 --- /dev/null +++ b/openrag/di/test_inference.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from core.embeddings import embedder_registry +from core.llm import llm_registry +from core.rerankers import reranker_registry +from core.vlm import vlm_registry +from di.container import ServiceContainer +from di.inference import register_inference + + +class TestRegisterInference: + def test_registries_populated(self): + register_inference() + + assert "vllm" in llm_registry + assert "vllm" in embedder_registry + assert "vllm" in vlm_registry + assert "infinity" in reranker_registry + assert "openai" in reranker_registry + + def test_idempotent(self): + register_inference() + register_inference() + + +class TestServiceContainer: + def test_container_populates_all_registries(self): + ServiceContainer() + + assert "vllm" in llm_registry + assert "vllm" in embedder_registry + assert "vllm" in vlm_registry + assert "infinity" in reranker_registry + assert "openai" in reranker_registry + + def test_create_llm(self): + container = ServiceContainer() + client = container.create_llm(endpoint="http://vllm:8000/v1", model_name="m") + assert client is not None + + def test_create_embedder(self): + container = ServiceContainer() + client = container.create_embedder(endpoint="http://vllm:8000/v1", model_name="m") + assert client is not None + + def test_create_reranker(self): + container = ServiceContainer() + client = container.create_reranker(endpoint="http://reranker:7997", model_name="m") + assert client is not None + + def test_create_vlm(self): + container = ServiceContainer() + client = container.create_vlm(endpoint="http://vllm:8000/v1", model_name="m") + assert client is not None diff --git a/openrag/di/vector_stores.py b/openrag/di/vector_stores.py new file mode 100644 index 000000000..ee70755b4 --- /dev/null +++ b/openrag/di/vector_stores.py @@ -0,0 +1,26 @@ +"""Factory for the :class:`VectorStore` adapter. + +Returns the Phase 7B :class:`services.storage.milvus_store.MilvusVectorStore` +built from ``settings.vectordb``. Construction is I/O-free; the embedder +dependency is materialised later via ``await store.initialize(dim)`` in the +composition root, mirroring the :class:`PostgresStore` lifecycle (cheap +construct, async materialise — see Phase 7B decision #2). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from services.storage.milvus_store import MilvusVectorStore + +if TYPE_CHECKING: + from core.config.root import Settings + from core.vector_stores import VectorStore + + +def create_vector_store(settings: Settings) -> VectorStore: + """Build a :class:`MilvusVectorStore` from the root settings.""" + return MilvusVectorStore(settings.vectordb) + + +__all__ = ["create_vector_store"] diff --git a/openrag/di/vlms.py b/openrag/di/vlms.py new file mode 100644 index 000000000..873c78b81 --- /dev/null +++ b/openrag/di/vlms.py @@ -0,0 +1,5 @@ +"""Register VLM implementations with the core registry.""" + + +def register_vlms() -> None: + import services.inference.vllm_client # noqa: F401 diff --git a/openrag/api.py b/openrag/main.py similarity index 81% rename from openrag/api.py rename to openrag/main.py index d0cd7bb8e..510f433d5 100644 --- a/openrag/api.py +++ b/openrag/main.py @@ -23,6 +23,10 @@ # flake8: noqa: E402 +# Bootstrap the long-lived worker actors (TaskStateManager, DocSerializer, +# parser pools, semaphores). Imported for side effects; the routes look the +# actors up by name via ray.get_actor. +import services.workers.bootstrap # noqa: F401, E402 from components.auth.middleware import AuthMiddleware from routers.actors import router as actors_router from routers.auth import router as auth_router @@ -39,7 +43,6 @@ from routers.utils import require_admin from routers.workspaces import router as workspaces_router from starlette.middleware.base import BaseHTTPMiddleware -from utils.dependencies import get_vectordb from utils.exceptions import OpenRAGError from utils.logger import get_logger @@ -214,7 +217,10 @@ async def dispatch(self, request: Request, call_next): # Register middlewares (order matters - last added runs first) -app.add_middleware(AuthMiddleware, get_vectordb=get_vectordb) +app.add_middleware( + AuthMiddleware, + get_auth_service=lambda request: request.app.state.container.auth_service, +) app.add_middleware(TokenRedactingMiddleware) app.add_middleware(MonitoringMiddleware) @@ -256,6 +262,61 @@ async def unhandled_exception_handler(request: Request, exc: Exception): ) app.state.app_state = AppState(config) + +# Phase 8 composition root. Attached here as minimal wiring so the thinned +# routers can resolve services via di.providers; Phase 11 moves this into a +# proper FastAPI lifespan. Construction is best-effort: a Milvus/PG hiccup at +# import must not stop the app from booting (the providers raise 503 if the +# container is absent). +try: + from di.container import ServiceContainer + + app.state.container = ServiceContainer(config) +except Exception: # pragma: no cover - defensive boot guard + # Best-effort by design (a Milvus/PG hiccup must not block boot, and + # di/providers.py serves a 503 while the container is absent), but log + # at exception level with a full traceback so an unexpected failure + # (import/syntax/misconfig) is loud rather than a one-line warning. + logger.exception("ServiceContainer wiring skipped") + app.state.container = None + + +@app.on_event("startup") +async def _initialize_container() -> None: + """Open the container's asyncpg pool + run idempotent migrations. + + The thinned Phase-8 routers resolve repositories through the container's + own :class:`PostgresStore`, which is a *separate* instance from the one + the legacy Ray ``Vectordb`` actor owns. Without this the catalog-backed + routes raise (uninitialised pool). The asyncpg layer (Phase 7) is + idempotent, so initialising alongside the actor's store is safe. Phase 11 + folds this into a lifespan; the deferral originally logged in the Phase-8 + decision log (§1) is corrected here because it broke the live app. + """ + container = getattr(app.state, "container", None) + if container is None: + return + try: + await container.initialize() + except Exception: # pragma: no cover - defensive boot guard + # A half-initialised container (e.g. asyncpg pool never opened) + # would route requests into broken repos and 500. Drop it so + # di/providers.py serves the intended degraded 503 instead. + logger.exception("ServiceContainer.initialize failed; serving degraded (503)") + app.state.container = None + + +@app.on_event("shutdown") +async def _shutdown_container() -> None: + container = getattr(app.state, "container", None) + if container is None: + return + try: + await container.shutdown() + except Exception: # pragma: no cover - defensive shutdown guard + logger.exception("ServiceContainer.shutdown skipped") + + app.mount("/static", StaticFiles(directory=DATA_DIR.resolve(), check_dir=True), name="static") @@ -346,4 +407,4 @@ class OpenRagAPI: serve.run(OpenRagAPI.bind(), route_prefix="/", blocking=True) else: - uvicorn.run("api:app", host="0.0.0.0", port=8080, reload=True, proxy_headers=True) + uvicorn.run("main:app", host="0.0.0.0", port=8080, reload=True, proxy_headers=True) diff --git a/openrag/routers/actors.py b/openrag/routers/actors.py index 6191199c1..18b551069 100644 --- a/openrag/routers/actors.py +++ b/openrag/routers/actors.py @@ -2,7 +2,7 @@ from fastapi import APIRouter, Depends, HTTPException, status from fastapi.responses import JSONResponse from ray.util.state import list_actors -from utils.dependencies import ( +from services.workers.bootstrap import ( actor_creation_map, ) from utils.logger import get_logger @@ -64,8 +64,6 @@ async def list_ray_actors(): - `TaskStateManager`: Manages task states - `MarkerPool`: PDF processing actor pool - `SerializerQueue`: Document serialization queue -- `Indexer`: Document indexing coordinator -- `Vectordb`: Vector database interface - `llmSemaphore`: LLM request semaphore - `vlmSemaphore`: Vision LM request semaphore diff --git a/openrag/routers/auth.py b/openrag/routers/auth.py index f23fb8434..169b0e398 100644 --- a/openrag/routers/auth.py +++ b/openrag/routers/auth.py @@ -1,4 +1,4 @@ -"""OIDC authentication routes — phase 4 of the OIDC integration. +"""OIDC authentication routes — thin HTTP layer over :class:`AuthService`. Routes exposed (all bypassed by ``AuthMiddleware``): - ``GET /auth/login`` — start Authorization Code + PKCE flow @@ -11,43 +11,35 @@ All routes return ``400`` when ``AUTH_MODE != "oidc"`` — the feature is dormant in ``token`` mode. + +Phase 8A.1: every business decision (PKCE/state generation, code exchange, +user lookup / provisioning, session creation, logout-URL construction) now +lives in :class:`services.orchestrators.auth_service.AuthService`. This +module only does HTTP transport: the ``AUTH_MODE`` gate, cookie set/clear, +the Secure-flag heuristic, and mapping :class:`OIDCFlowError` to responses. """ from __future__ import annotations import os -from datetime import datetime, timedelta -from typing import Any -from urllib.parse import urlencode, urlparse - -from components.auth import ( - OIDCClient, - StateCookiePayload, - StateCookieSerializer, - decrypt_token, - encrypt_token, - get_oidc_client, - issue_session_token, -) -from fastapi import APIRouter, Form, HTTPException, Request, Response, status -from fastapi.responses import JSONResponse, RedirectResponse -from models.user import UserCreate -from utils.dependencies import get_vectordb -from utils.logger import get_logger, mask_email -# Whitelist mirrors ``api._OIDC_CLAIM_MAPPING_ALLOWED_FIELDS`` — kept in sync -# at the DB layer too (``PartitionFileManager.update_user_fields``). -_OIDC_CLAIM_MAPPING_ALLOWED_FIELDS = {"display_name", "email"} +from components.auth import StateCookieSerializer +from di.providers import get_auth_service +from fastapi import APIRouter, Depends, Form, HTTPException, Request, Response, status +from fastapi.responses import JSONResponse, RedirectResponse +from services.orchestrators.auth_service import ( + SESSION_COOKIE_NAME, + AuthService, + OIDCFlowError, +) +from utils.logger import get_logger logger = get_logger() router = APIRouter() -SESSION_COOKIE_NAME = "openrag_session" - - # --------------------------------------------------------------------------- -# Env helpers — read lazily so tests can monkeypatch os.environ +# HTTP-transport helpers (kept in the router by design) # --------------------------------------------------------------------------- @@ -55,94 +47,7 @@ def _auth_mode() -> str: return os.getenv("AUTH_MODE", "token").strip().lower() -def _token_encryption_key() -> str: - key = os.getenv("OIDC_TOKEN_ENCRYPTION_KEY") - if not key: - raise RuntimeError("OIDC_TOKEN_ENCRYPTION_KEY is not set") - return key - - -def _claim_source() -> str: - return os.getenv("OIDC_CLAIM_SOURCE", "id_token").strip().lower() - - -def _auto_provision_login() -> bool: - """Whether to auto-provision a non-admin user on first OIDC login. - - Defaults to ``False`` — keeping the historical "admin pre-creates every - user" model. Set ``OIDC_AUTO_PROVISION_LOGIN=true`` to enable: when the - callback receives a ``sub`` that isn't yet mapped to an OpenRAG user, - a row is created on the fly using the ID-token claims (``name`` / - ``preferred_username`` for the display name, ``email`` if present). - - Auto-provisioned users are **never** admin and inherit the default file - quota — operators can promote / adjust afterwards via ``/users/``. - """ - return os.getenv("OIDC_AUTO_PROVISION_LOGIN", "false").strip().lower() == "true" - - -def _display_name_from_claims(claims: dict[str, Any], sub: str) -> str: - """Pick a sensible display name from the standard OIDC claims. - - Falls back to a short ``sub`` prefix when nothing readable is available - so the user row always has something printable for the UI. - """ - for key in ("name", "preferred_username"): - value = claims.get(key) - if isinstance(value, str) and value.strip(): - return value.strip() - given = claims.get("given_name") or "" - family = claims.get("family_name") or "" - composed = f"{given} {family}".strip() - if composed: - return composed - # Last-resort fallback — keep enough of the sub to be unique-ish in the UI. - return f"oidc-{sub[:8]}" - - -def _claim_mapping() -> dict[str, str]: - """Parse ``OIDC_CLAIM_MAPPING`` at request time so tests can monkeypatch it. - - Shares the validation rules with ``api._parse_oidc_claim_mapping``: entries - whose ``db_field`` is not whitelisted are silently dropped here because the - hard-failure path belongs to the startup validator in ``api.py`` — at login - time we prefer to log and continue rather than break the flow on a - misconfiguration the operator has already been warned about. - """ - raw = os.getenv("OIDC_CLAIM_MAPPING", "").strip() - if not raw: - return {} - mapping: dict[str, str] = {} - for pair in raw.split(","): - pair = pair.strip() - if not pair or ":" not in pair: - continue - db_field, claim = pair.split(":", 1) - db_field = db_field.strip() - claim = claim.strip() - if db_field not in _OIDC_CLAIM_MAPPING_ALLOWED_FIELDS or not claim: - continue - mapping[db_field] = claim - return mapping - - -def _post_logout_redirect_uri() -> str | None: - """Return the configured post-logout redirect URI, or None if unset. - - No default is provided: a default of "/" would land the user back on - OpenRag's root which immediately re-triggers OIDC login (silent re-auth - if the IdP session is still alive, or a loop on the IdP form if not). - Operators deliberately choose a URL outside OpenRag (corporate intranet, - a static 'you are logged out' page, the IdP's own post-logout page). - """ - return os.getenv("OIDC_POST_LOGOUT_REDIRECT_URI") - - -def _oidc_client_id() -> str: - return os.environ["OIDC_CLIENT_ID"] - - -def _require_oidc_mode(): +def _require_oidc_mode() -> None: if _auth_mode() != "oidc": raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -153,70 +58,19 @@ def _require_oidc_mode(): def _is_request_secure(request: Request) -> bool: """True if the client-observed scheme is HTTPS. - Checks multiple indicators: - 1. ``PREFERRED_URL_SCHEME`` env var (set when behind a TLS-terminating proxy) - 2. ``X-Forwarded-Proto`` header (set by reverse proxies like Traefik/Nginx) - 3. ``request.url.scheme`` (accounts for proxy_headers=True in uvicorn) + Checks ``PREFERRED_URL_SCHEME``, the ``X-Forwarded-Proto`` header + (client-most hop when comma-separated), then ``request.url.scheme``. """ if os.environ.get("PREFERRED_URL_SCHEME", "").lower() == "https": return True - # X-Forwarded-Proto can be comma-separated when chained through multiple - # proxies (e.g. "https, http"); the client-most hop is the first entry. xfp = request.headers.get("x-forwarded-proto", "") if xfp.split(",", 1)[0].strip().lower() == "https": return True return request.url.scheme == "https" -def _state_serializer() -> StateCookieSerializer: - return StateCookieSerializer(secret_key=_token_encryption_key()) - - -def _allowed_next_origins() -> set[str]: - """Origins (scheme://host[:port]) accepted as redirect targets after login. - - Mirrors the CORS allow_origins from ``api.py``: localhost dev ports plus - ``INDEXERUI_URL`` so that the indexer-ui (served on a different port) can - receive the user back after the OIDC flow completes. - """ - origins = {"http://localhost:3042", "http://localhost:5173"} - indexer_ui = os.getenv("INDEXERUI_URL") - if indexer_ui: - origins.add(indexer_ui.rstrip("/")) - return origins - - -def _sanitize_next_url(next_url: str | None) -> str: - """Accept either a same-origin relative path (``/...`` but not ``//...``) - or an absolute URL whose origin is explicitly whitelisted (indexer-ui, - dev-only localhost). Fall back to ``/`` on any mismatch — protects against - open-redirect attacks. - """ - if not next_url: - return "/" - if next_url.startswith("/") and not next_url.startswith("//"): - return next_url - # Absolute URL: only allow whitelisted origins. - parsed = urlparse(next_url) - if parsed.scheme in ("http", "https") and parsed.netloc: - origin = f"{parsed.scheme}://{parsed.netloc}" - if origin in _allowed_next_origins(): - return next_url - return "/" - - -def _utcnow() -> datetime: - # DB-side timestamps are naive local time (models' default is ``datetime.now``), - # and every read site compares against ``datetime.now()``. Using ``datetime.now()`` - # here keeps newly-issued sessions from appearing pre-expired on non-UTC hosts. - return datetime.now() - - def _delete_state_cookie(response: Response) -> None: - response.delete_cookie( - key=StateCookieSerializer.COOKIE_NAME, - path="/", - ) + response.delete_cookie(key=StateCookieSerializer.COOKIE_NAME, path="/") def _json_error(status_code: int, detail: str, *, delete_state_cookie: bool = False) -> JSONResponse: @@ -232,35 +86,22 @@ def _json_error(status_code: int, detail: str, *, delete_state_cookie: bool = Fa @router.get("/auth/login", include_in_schema=False) -async def login(request: Request, next: str | None = None): - _require_oidc_mode() - client: OIDCClient = get_oidc_client() - - state, nonce = OIDCClient.generate_state_and_nonce() - code_verifier, code_challenge = OIDCClient.generate_pkce_pair() - +async def login( + request: Request, + next: str | None = None, + _oidc: None = Depends(_require_oidc_mode), + service: AuthService = Depends(get_auth_service), +): try: - auth_url = await client.build_authorization_url(state=state, nonce=nonce, code_challenge=code_challenge) - except Exception as e: - logger.error(f"Failed to build OIDC authorization URL: {e}") - raise HTTPException( - status_code=status.HTTP_502_BAD_GATEWAY, - detail="OIDC discovery failed — see server logs.", - ) from e - - payload = StateCookiePayload( - state=state, - nonce=nonce, - code_verifier=code_verifier, - next_url=_sanitize_next_url(next), - ) - cookie_value = _state_serializer().dumps(payload) + result = await service.start_oidc_login(next) + except OIDCFlowError as e: + raise HTTPException(status_code=e.status_code, detail=e.message) from e - response = RedirectResponse(url=auth_url, status_code=302) + response = RedirectResponse(url=result.authorization_url, status_code=302) response.set_cookie( - key=StateCookieSerializer.COOKIE_NAME, - value=cookie_value, - max_age=StateCookieSerializer.DEFAULT_TTL_SECONDS, + key=result.state_cookie_name, + value=result.state_cookie_value, + max_age=result.state_cookie_max_age, httponly=True, secure=_is_request_secure(request), samesite="lax", @@ -275,243 +116,33 @@ async def login(request: Request, next: str | None = None): @router.get("/auth/callback", include_in_schema=False) -async def callback(request: Request, code: str | None = None, state: str | None = None): - _require_oidc_mode() - - if not code or not state: - return _json_error( - status.HTTP_400_BAD_REQUEST, - "Missing 'code' or 'state' query parameter.", - delete_state_cookie=True, - ) - - # --- 1. Parse state cookie ------------------------------------------------- - cookie_raw = request.cookies.get(StateCookieSerializer.COOKIE_NAME) - if not cookie_raw: - return _json_error( - status.HTTP_400_BAD_REQUEST, - "OIDC state cookie missing.", - delete_state_cookie=True, - ) - +async def callback( + request: Request, + code: str | None = None, + state: str | None = None, + _oidc: None = Depends(_require_oidc_mode), + service: AuthService = Depends(get_auth_service), +): try: - payload = _state_serializer().loads(cookie_raw) - except ValueError as e: - logger.warning(f"Invalid OIDC state cookie: {e}") - return _json_error( - status.HTTP_400_BAD_REQUEST, - "Invalid or expired OIDC state cookie.", - delete_state_cookie=True, - ) - - # --- 2. CSRF check -------------------------------------------------------- - if state != payload.state: - logger.warning("OIDC state mismatch between query and cookie") - return _json_error( - status.HTTP_400_BAD_REQUEST, - "OIDC state mismatch.", - delete_state_cookie=True, - ) - - # --- 3. Exchange code ------------------------------------------------------ - client: OIDCClient = get_oidc_client() - try: - bundle = await client.exchange_code( + result = await service.handle_oidc_callback( code=code, - code_verifier=payload.code_verifier, - expected_nonce=payload.nonce, - ) - except Exception: - # Log full exception for operators; return a generic message so IdP - # URLs / stack-adjacent internals don't leak via the HTTP response. - logger.exception("OIDC code exchange failed") - return _json_error( - status.HTTP_400_BAD_REQUEST, - "OIDC code exchange failed", - delete_state_cookie=True, + state=state, + state_cookie_raw=request.cookies.get(StateCookieSerializer.COOKIE_NAME), ) + except OIDCFlowError as e: + return _json_error(e.status_code, e.message, delete_state_cookie=True) - # --- 4. Extract sub and match user ---------------------------------------- - sub = bundle.claims.get("sub") - if not sub: - return _json_error( - status.HTTP_400_BAD_REQUEST, - "ID token missing 'sub' claim.", - delete_state_cookie=True, - ) - - vdb = get_vectordb() - user: dict[str, Any] | None = await vdb.get_user_by_external_id.remote(sub) - if user is None: - if not _auto_provision_login(): - logger.warning(f"OIDC login rejected — user not registered (sub={sub!r})") - return _json_error( - status.HTTP_403_FORBIDDEN, - "User not registered", - delete_state_cookie=True, - ) - - # Auto-provision: create a non-admin user from the ID-token claims. - # Email is best-effort — populated when the IdP exposes it on the - # ``email`` claim (typically via the ``email`` scope, which is in the - # default ``OIDC_SCOPES``). Display name falls back to the sub when - # the IdP exposes nothing readable. - display_name = _display_name_from_claims(bundle.claims, sub) - email = bundle.claims.get("email") - try: - user = await vdb.create_user.remote( - UserCreate( - display_name=display_name, - external_user_id=sub, - email=email if isinstance(email, str) and email.strip() else None, - is_admin=False, - ) - ) - except Exception as e: - # create_user failed. Separate the two causes: - # 1. Concurrent first-login on the same sub — another request - # already inserted the row; re-read by external_id and proceed. - # 2. The unique email index rejected the insert because a row with - # this email already exists under a different identity. Matching - # is external_id-only, so it wasn't found above, and only an - # admin can reconcile it — surface an actionable 409 rather than - # an opaque 500. - logger.exception(f"OIDC auto-provisioning failed for sub={sub!r}: {e}") - user = await vdb.get_user_by_external_id.remote(sub) - if user is None: - if isinstance(email, str) and email.strip() and await vdb.get_user_by_email.remote(email): - logger.error( - f"OIDC auto-provisioning blocked for sub={sub!r}: an account with email " - f"{mask_email(email)} already exists under a different identity. Set that " - f"user's external_user_id to this sub to allow login." - ) - return _json_error( - status.HTTP_409_CONFLICT, - "An account with this email already exists. Ask your administrator to " - "link it to your identity provider login.", - delete_state_cookie=True, - ) - return _json_error( - status.HTTP_500_INTERNAL_SERVER_ERROR, - "Failed to provision user", - delete_state_cookie=True, - ) - else: - # display_name (the user's real name) is intentionally not logged — id + sub - # identify the row without writing PII to the logs. - logger.info(f"OIDC user auto-provisioned (id={user['id']}, sub={sub!r})") - - # --- 4b. Auto-provision: keep display_name + email in sync with claims ---- - # When OIDC_AUTO_PROVISION_LOGIN is on, the IdP is treated as the source of - # truth for these two fields on every login (not just at creation), so a - # user renamed in the IdP doesn't drift out of sync. No-op for users whose - # row already matches the claims (including the user we just created). - if _auto_provision_login(): - derived_display = _display_name_from_claims(bundle.claims, sub) - derived_email_raw = bundle.claims.get("email") - derived_email = ( - derived_email_raw.strip() if isinstance(derived_email_raw, str) and derived_email_raw.strip() else None - ) - - sync_updates: dict[str, Any] = {} - if derived_display and user.get("display_name") != derived_display: - sync_updates["display_name"] = derived_display - if derived_email is not None and user.get("email") != derived_email: - sync_updates["email"] = derived_email - - if sync_updates: - try: - await vdb.update_user_fields.remote(user["id"], sync_updates) - except Exception as e: - logger.warning(f"OIDC auto-provision sync failed for user_id={user['id']}: {e}") - else: - refreshed = await vdb.get_user_by_external_id.remote(sub) - if refreshed is not None: - user = refreshed - - # --- 5. Optional claim-mapping update -------------------------------------- - mapping = _claim_mapping() - if mapping: - if _claim_source() == "userinfo": - try: - claims_for_mapping: dict[str, Any] = await client.fetch_userinfo(bundle.access_token) - except Exception as e: - logger.warning(f"OIDC userinfo fetch failed: {e}") - return _json_error( - status.HTTP_400_BAD_REQUEST, - "Failed to fetch userinfo from IdP.", - delete_state_cookie=True, - ) - else: - claims_for_mapping = bundle.claims - - updates: dict[str, Any] = {} - for db_field, claim in mapping.items(): - value = claims_for_mapping.get(claim) - if value is None: - continue - # No-op filter: skip fields already matching, so we don't churn the DB. - if user.get(db_field) == value: - continue - updates[db_field] = value - - if updates: - try: - await vdb.update_user_fields.remote(user["id"], updates) - except Exception as e: - logger.warning(f"update_user_fields failed for user_id={user['id']}: {e}") - else: - # Refresh the user dict so anything downstream sees the new values. - refreshed = await vdb.get_user_by_external_id.remote(sub) - if refreshed is not None: - user = refreshed - - # --- 6. Timestamps --------------------------------------------------------- - now = _utcnow() - expires_in = max(int(bundle.expires_in or 0), 60) - access_token_expires_at = now + timedelta(seconds=expires_in) - if bundle.refresh_token: - session_expires_at = now + timedelta(days=7) - else: - session_expires_at = access_token_expires_at - - # --- 7. Issue session & encrypt ------------------------------------------ - plain, _hashed = issue_session_token() - key = _token_encryption_key() - id_token_encrypted = encrypt_token(bundle.id_token, key=key) - access_token_encrypted = encrypt_token(bundle.access_token, key=key) - refresh_token_encrypted = encrypt_token(bundle.refresh_token, key=key) - sid = bundle.claims.get("sid") - - await vdb.create_oidc_session.remote( - user_id=user["id"], - sub=sub, - sid=sid, - session_token_plain=plain, - id_token_encrypted=id_token_encrypted, - access_token_encrypted=access_token_encrypted, - refresh_token_encrypted=refresh_token_encrypted, - access_token_expires_at=access_token_expires_at, - session_expires_at=session_expires_at, - ) - - # --- 8. Build redirect: clear state cookie, set session cookie ----------- - next_url = _sanitize_next_url(payload.next_url) - redirect = RedirectResponse(url=next_url, status_code=302) + redirect = RedirectResponse(url=result.next_url, status_code=302) _delete_state_cookie(redirect) - - max_age = max(int((session_expires_at - now).total_seconds()), 1) redirect.set_cookie( - key=SESSION_COOKIE_NAME, - value=plain, - max_age=max_age, + key=result.session_cookie_name, + value=result.session_cookie_value, + max_age=result.session_cookie_max_age, httponly=True, secure=_is_request_secure(request), samesite="lax", path="/", ) - - logger.info(f"OIDC login success — user_id={user['id']}, sid={sid!r}, next={next_url!r}") return redirect @@ -521,44 +152,27 @@ async def callback(request: Request, code: str | None = None, state: str | None @router.post("/auth/backchannel-logout", include_in_schema=False) -async def backchannel_logout(logout_token: str = Form(...)): +async def backchannel_logout( + logout_token: str = Form(...), + _oidc: None = Depends(_require_oidc_mode), + service: AuthService = Depends(get_auth_service), +): """IdP-initiated logout per OIDC Back-Channel Logout spec. Content-Type: ``application/x-www-form-urlencoded`` with field ``logout_token``. """ - _require_oidc_mode() - - client: OIDCClient = get_oidc_client() - try: - claims = await client.verify_logout_token(logout_token) - except ValueError as e: - logger.warning(f"Invalid back-channel logout token: {e}") + await service.handle_backchannel_logout(logout_token) + except OIDCFlowError as e: + content: dict[str, str] = {"error": "invalid_request"} + if e.error_description: + content["error_description"] = e.error_description return JSONResponse( status_code=status.HTTP_400_BAD_REQUEST, - content={"error": "invalid_request", "error_description": str(e)}, - headers={"Cache-Control": "no-store"}, - ) - except Exception as e: - logger.warning(f"Back-channel logout token verification failed: {e}") - return JSONResponse( - status_code=status.HTTP_400_BAD_REQUEST, - content={"error": "invalid_request"}, + content=content, headers={"Cache-Control": "no-store"}, ) - if claims.sid: - vdb = get_vectordb() - count = await vdb.revoke_oidc_sessions_by_sid.remote(claims.sid) - logger.info(f"Back-channel logout revoked sessions — sid={claims.sid!r}, count={count}") - else: - # Plan §2 #10 limits back-channel logout scope to sid only. - # Still return 200 to keep the IdP happy. - logger.warning( - f"Received sid-less back-channel logout token — not supported; " - f"ignoring per implementation policy (sub={claims.sub!r})" - ) - return Response( status_code=status.HTTP_200_OK, headers={"Cache-Control": "no-store"}, @@ -571,52 +185,17 @@ async def backchannel_logout(logout_token: str = Form(...)): @router.get("/auth/logout", include_in_schema=False) -async def logout(request: Request): - _require_oidc_mode() - - vdb = get_vectordb() - client: OIDCClient = get_oidc_client() - - # Look up & revoke the session; keep the id_token to forward as id_token_hint. - id_token_hint: str | None = None - cookie_value = request.cookies.get(SESSION_COOKIE_NAME) - if cookie_value: - session = await vdb.get_oidc_session_by_token.remote(cookie_value) - if session: - enc = session.get("id_token_encrypted") - if enc: - try: - id_token_hint = decrypt_token(enc, key=_token_encryption_key()) - except ValueError as e: - logger.warning(f"Failed to decrypt id_token for logout: {e}") - try: - await vdb.revoke_oidc_session_by_id.remote(session["id"]) - except Exception as e: - logger.warning(f"Failed to revoke oidc_session during logout: {e}") - - # Build redirect target: IdP end_session if discovery provides one, - # otherwise the configured post-logout URL. If neither is available - # we return a plain 200 with the cookie deleted — better than a 302 - # loop through the root. - local_target = _post_logout_redirect_uri() - redirect_target: str | None = local_target - try: - meta = await client.discover() - end_session = meta.get("end_session_endpoint") - if end_session: - params: dict[str, str] = {"client_id": _oidc_client_id()} - if local_target: - params["post_logout_redirect_uri"] = local_target - if id_token_hint: - params["id_token_hint"] = id_token_hint - redirect_target = f"{end_session}?{urlencode(params)}" - except Exception as e: - logger.warning(f"OIDC discovery failed during logout, skipping IdP redirect: {e}") +async def logout( + request: Request, + _oidc: None = Depends(_require_oidc_mode), + service: AuthService = Depends(get_auth_service), +): + redirect_target = await service.logout(request.cookies.get(SESSION_COOKIE_NAME)) if redirect_target: - response = RedirectResponse(url=redirect_target, status_code=302) + response: Response = RedirectResponse(url=redirect_target, status_code=302) else: - # No IdP end_session and no local post-logout URL → just confirm the + # No IdP end_session and no local post-logout URL → confirm the # logout in-place. The cookie deletion below still takes effect. response = JSONResponse(status_code=200, content={"detail": "Logged out"}) response.delete_cookie(key=SESSION_COOKIE_NAME, path="/") @@ -642,7 +221,6 @@ async def me(request: Request): if oidc_session and oidc_session.get("session_expires_at"): exp = oidc_session["session_expires_at"] try: - # naive datetime → iso str session_expires_at = exp.isoformat() except AttributeError: session_expires_at = str(exp) diff --git a/openrag/routers/extract.py b/openrag/routers/extract.py index c47d39121..86f3f371f 100644 --- a/openrag/routers/extract.py +++ b/openrag/routers/extract.py @@ -1,13 +1,23 @@ -from fastapi import APIRouter, Depends, HTTPException, Request, status +"""Extract route — thin HTTP layer over :class:`ConversionService`. + +Phase 8E: the chunk-by-id lookup moved to +``services.orchestrators.conversion_service.ConversionService`` (clean +``VectorStore`` port, no Ray). This module keeps HTTP transport only: +the request-scoped partition authorization and the not-found / forbidden +guards whose exact ``{"detail": ...}`` body the legacy endpoint +returned via ``HTTPException``. +""" + +from di.providers import get_conversion_service +from fastapi import APIRouter, Depends, HTTPException, status from fastapi.responses import JSONResponse -from utils.dependencies import get_vectordb +from services.orchestrators.conversion_service import ConversionService from utils.logger import get_logger from .utils import current_user_or_admin_partitions_list logger = get_logger() -# Create an APIRouter instance router = APIRouter() @@ -39,21 +49,26 @@ """, ) async def get_extract( - request: Request, extract_id: str, - vectordb=Depends(get_vectordb), user_partitions=Depends(current_user_or_admin_partitions_list), + service: ConversionService = Depends(get_conversion_service), ): log = logger.bind(extract_id=extract_id) - chunk = await vectordb.get_chunk_by_id.remote(extract_id) + chunk = await service.get_chunk(extract_id) if chunk is None: log.warning("Extract not found.") raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Extract '{extract_id}' not found.", ) - chunk_partition = chunk.metadata["partition"] + chunk_partition = chunk.get("metadata", {}).get("partition") + if not chunk_partition: + log.warning("Extract metadata missing partition.") + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Extract '{extract_id}' not found.", + ) log.info(f"User partitions: {user_partitions}, Chunk partition: {chunk_partition}") if chunk_partition not in user_partitions and user_partitions != ["all"]: log.warning("User does not have access to this extract.") @@ -65,5 +80,5 @@ async def get_extract( return JSONResponse( status_code=status.HTTP_200_OK, - content={"page_content": chunk.page_content, "metadata": chunk.metadata}, + content={"page_content": chunk["page_content"], "metadata": chunk["metadata"]}, ) diff --git a/openrag/routers/indexer.py b/openrag/routers/indexer.py index 9a8568198..8abdc9423 100644 --- a/openrag/routers/indexer.py +++ b/openrag/routers/indexer.py @@ -1,11 +1,23 @@ +"""Indexing routes — thin HTTP layer over :class:`IndexingService`. + +Phase 8D.1: metadata assembly, existence/workspace checks and task +dispatch moved to +``services.orchestrators.indexing_service.IndexingService`` (the Ray +``Indexer`` / ``TaskStateManager`` actors now sit behind the +``IndexingDispatcher`` port). This module keeps HTTP transport only: +the saved-file IO, ``request.url_for`` link building, the shared +``Depends`` auth wrappers, and the conflict / not-found / bad-input +guards whose exact non-bracketed ``{"detail": ...}`` body the legacy +endpoints returned via ``HTTPException``. +""" + import json from pathlib import Path from typing import Any -import ray -from components.indexer.utils.files import extract_temporal_fields, sanitize_filename, save_file_to_disk -from components.ray_utils import call_ray_actor_with_timeout +from components.indexer.utils.files import sanitize_filename, save_file_to_disk from config import load_config +from di.providers import get_auth_service, get_indexing_service, get_partition_service from fastapi import ( APIRouter, Depends, @@ -17,14 +29,15 @@ status, ) from fastapi.responses import JSONResponse -from utils.dependencies import get_indexer, get_task_state_manager, get_vectordb +from services.orchestrators.auth_service import AuthService +from services.orchestrators.indexing_service import IndexingService +from services.orchestrators.partition_service import PartitionService from utils.logger import get_logger from .utils import ( check_user_file_quota, current_user_partitions, ensure_partition_role, - human_readable_size, require_partition_editor, require_task_owner, validate_file_format, @@ -32,27 +45,18 @@ validate_metadata, ) -# load logger logger = get_logger() -# load config config = load_config() DATA_DIR = config.paths.data_dir -VECTORDB_TIMEOUT = config.ray.indexer.vectordb_timeout - -FORBIDDEN_CHARS_IN_FILE_ID = set("/") # set('"<>#%{}|\\^`[]') LOG_FILE = Path(config.paths.log_dir or "logs") / "app.json" # supported file formats or mimetypes ACCEPTED_FILE_FORMATS = config.loader.file_loaders.model_dump().keys() DICT_MIMETYPES = config.loader.mimetypes.to_dict() -# URL scheme configuration PREFERRED_URL_SCHEME = config.server.preferred_url_scheme -# DATETIME FIELDS: Fields provided by the client -TEMPORAL_FIELDS = ["created_at"] - def build_url(request: Request, route_name: str, **path_params) -> str: """Build a URL using the preferred scheme if configured.""" @@ -62,7 +66,6 @@ def build_url(request: Request, route_name: str, **path_params) -> str: return str(url) -# Create an APIRouter instance router = APIRouter() @@ -83,9 +86,7 @@ async def get_supported_types(): - `extensions`: List of supported file extensions. - `mimetypes`: List of supported MIME types. """ - list_extensions = list(ACCEPTED_FILE_FORMATS) - list_mimetypes = list(DICT_MIMETYPES) - resp = {"extensions": list_extensions, "mimetypes": list_mimetypes} + resp = {"extensions": list(ACCEPTED_FILE_FORMATS), "mimetypes": list(DICT_MIMETYPES)} return JSONResponse(content=resp) @@ -129,23 +130,20 @@ async def add_file( file: UploadFile = Depends(validate_file_format), metadata: dict = Depends(validate_metadata), workspace_ids: str | None = Form(None, description="JSON array of workspace IDs to add the file to"), - indexer=Depends(get_indexer), - task_state_manager=Depends(get_task_state_manager), - vectordb=Depends(get_vectordb), user=Depends(require_partition_editor), _quota_check=Depends(check_user_file_quota), + service: IndexingService = Depends(get_indexing_service), ): - if await vectordb.file_exists.remote(file_id, partition): + if await service.file_exists(file_id, partition): raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail=f"File '{file_id}' already exists in partition {partition}", ) - save_dir = Path(DATA_DIR) original_filename = file.filename file.filename = sanitize_filename(file.filename) try: - file_path = await save_file_to_disk(file, save_dir, with_random_prefix=True) + file_path = await save_file_to_disk(file, Path(DATA_DIR), with_random_prefix=True) except Exception as e: logger.exception("Failed to save file to disk.", error=str(e)) raise HTTPException( @@ -153,24 +151,6 @@ async def add_file( detail=str(e), ) - metadata.update( - { - "source": str(file_path), - "filename": file.filename, - "original_filename": original_filename, - } - ) - file_stat = Path(file_path).stat() - - # Append extra metadata - metadata["file_size"] = human_readable_size(file_stat.st_size) - metadata["file_id"] = file_id - - ## Add temporal fields to metadata, using provided values if available, otherwise extracting from file system - temporal_fields = extract_temporal_fields(metadata, temporal_fields=TEMPORAL_FIELDS) - metadata.update(temporal_fields) - - # Validate and parse workspace_ids parsed_workspace_ids = None if workspace_ids: try: @@ -183,27 +163,27 @@ async def add_file( detail="workspace_ids must be a JSON array of strings", ) for ws_id in parsed_workspace_ids: - ws = await call_ray_actor_with_timeout( - vectordb.get_workspace.remote(ws_id), - timeout=VECTORDB_TIMEOUT, - task_description=f"get_workspace({ws_id})", - ) + ws = await service.get_workspace(ws_id) if not ws or ws["partition_name"] != partition: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Workspace '{ws_id}' not found in partition '{partition}'", ) - # Indexing the file (workspace association happens inside add_file after successful indexing) - task = indexer.add_file.remote( - path=file_path, metadata=metadata, partition=partition, user=user, workspace_ids=parsed_workspace_ids + task_id = await service.add_file( + file_path=str(file_path), + file_id=file_id, + partition=partition, + metadata=metadata, + sanitized_filename=file.filename, + original_filename=original_filename, + user=user, + workspace_ids=parsed_workspace_ids, ) - await task_state_manager.set_state.remote(task.task_id().hex(), "QUEUED") - await task_state_manager.set_object_ref.remote(task.task_id().hex(), {"ref": task}) return JSONResponse( status_code=status.HTTP_201_CREATED, - content={"task_status_url": build_url(request, "get_task_status", task_id=task.task_id().hex())}, + content={"task_status_url": build_url(request, "get_task_status", task_id=task_id)}, ) @@ -222,16 +202,15 @@ async def add_file( async def delete_file( partition: str, file_id: str, - indexer=Depends(get_indexer), - vectordb=Depends(get_vectordb), user=Depends(require_partition_editor), + service: IndexingService = Depends(get_indexing_service), ): - if not await vectordb.file_exists.remote(file_id, partition): + if not await service.file_exists(file_id, partition): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"'{file_id}' not found in partition '{partition}'", ) - await indexer.delete_file.remote(file_id, partition) + await service.delete_file(file_id, partition) return Response(status_code=status.HTTP_204_NO_CONTENT) @@ -275,12 +254,10 @@ async def put_file( file_id: str = Depends(validate_file_id), file: UploadFile = Depends(validate_file_format), metadata: dict = Depends(validate_metadata), - indexer=Depends(get_indexer), - task_state_manager=Depends(get_task_state_manager), - vectordb=Depends(get_vectordb), user=Depends(require_partition_editor), + service: IndexingService = Depends(get_indexing_service), ): - if not await vectordb.file_exists.remote(file_id, partition): + if not await service.file_exists(file_id, partition): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"'{file_id}' not found in partition '{partition}'", @@ -289,45 +266,24 @@ async def put_file( # No Milvus deletion here. The Indexer's add_file(replace=True) flow uses # insert-before-delete: it snapshots old chunk IDs, inserts new chunks, # then deletes old ones — so the file is never left in a half-replaced state. - - save_dir = Path(DATA_DIR) original_filename = file.filename file.filename = sanitize_filename(file.filename) - file_path = await save_file_to_disk(file, save_dir, with_random_prefix=True) - - metadata.update( - { - "source": str(file_path), - "filename": file.filename, - "original_filename": original_filename, - } - ) - - file_stat = Path(file_path).stat() - - # Append extra metadata - metadata["file_size"] = human_readable_size(file_stat.st_size) - metadata["file_id"] = file_id + file_path = await save_file_to_disk(file, Path(DATA_DIR), with_random_prefix=True) - ## Add temporal fields to metadata, using provided values if available, otherwise extracting from file system - temporal_fields = extract_temporal_fields(metadata, temporal_fields=TEMPORAL_FIELDS) - metadata.update(temporal_fields) - - # Re-index: serialize → chunk → embed → insert into Milvus + update PG row in-place. - # replace=True tells add_file to update the existing PG File row rather than creating a new one. - task = indexer.add_file.remote( - path=file_path, - metadata=metadata, + task_id = await service.add_file( + file_path=str(file_path), + file_id=file_id, partition=partition, + metadata=metadata, + sanitized_filename=file.filename, + original_filename=original_filename, user=user, replace=True, ) - await task_state_manager.set_state.remote(task.task_id().hex(), "QUEUED") - await task_state_manager.set_object_ref.remote(task.task_id().hex(), {"ref": task}) return JSONResponse( status_code=status.HTTP_202_ACCEPTED, - content={"task_status_url": build_url(request, "get_task_status", task_id=task.task_id().hex())}, + content={"task_status_url": build_url(request, "get_task_status", task_id=task_id)}, ) @@ -353,12 +309,12 @@ async def patch_file( partition: str, file_id: str = Depends(validate_file_id), metadata: Any | None = Depends(validate_metadata), - indexer=Depends(get_indexer), user=Depends(require_partition_editor), user_partitions=Depends(current_user_partitions), + service: IndexingService = Depends(get_indexing_service), + auth_service: AuthService = Depends(get_auth_service), + partition_service: PartitionService = Depends(get_partition_service), ): - metadata["file_id"] = file_id - # Make sure partition role is valid if partition is being changed if "partition" in metadata: await ensure_partition_role( @@ -366,9 +322,11 @@ async def patch_file( user=user, user_partitions=user_partitions, required_role="editor", + auth_service=auth_service, + partition_service=partition_service, ) - await indexer.update_file_metadata.remote(file_id, metadata, partition, user=user) + await service.update_metadata(file_id, metadata, partition, user) return JSONResponse( status_code=status.HTTP_200_OK, content={"message": f"Metadata for file '{file_id}' successfully updated."}, @@ -400,22 +358,31 @@ async def copy_file_between_partitions( metadata: Any | None = Depends(validate_metadata), source_partition: str = Form(...), source_file_id: str = Form(...), - indexer=Depends(get_indexer), user=Depends(require_partition_editor), user_partitions=Depends(current_user_partitions), _quota_check=Depends(check_user_file_quota), + service: IndexingService = Depends(get_indexing_service), + auth_service: AuthService = Depends(get_auth_service), + partition_service: PartitionService = Depends(get_partition_service), ): - # Make sure user has access to destination partition + # Make sure user has access to the source partition await ensure_partition_role( partition=source_partition, user=user, user_partitions=user_partitions, required_role="viewer", + auth_service=auth_service, + partition_service=partition_service, ) - metadata["file_id"] = file_id - metadata["partition"] = partition - await indexer.copy_file.remote(file_id=source_file_id, metadata=metadata, partition=source_partition, user=user) + await service.copy_file( + source_file_id=source_file_id, + source_partition=source_partition, + target_file_id=file_id, + target_partition=partition, + metadata=metadata, + user=user, + ) return JSONResponse( status_code=status.HTTP_201_CREATED, content={"message": "File copied successfully."}, @@ -440,18 +407,16 @@ async def copy_file_between_partitions( async def get_task_status( request: Request, task_id: str, - task_state_manager=Depends(get_task_state_manager), task_details=Depends(require_task_owner), + service: IndexingService = Depends(get_indexing_service), ): - # fetch task state - state = await task_state_manager.get_state.remote(task_id) + state = await service.get_task_state(task_id) if state is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Task '{task_id}' not found.", ) - # format the response content: dict[str, Any] = { "task_id": task_id, "task_state": state, @@ -481,10 +446,10 @@ async def get_task_status( ) async def get_task_error( task_id: str, - task_state_manager=Depends(get_task_state_manager), task_details=Depends(require_task_owner), + service: IndexingService = Depends(get_indexing_service), ): - error = await task_state_manager.get_error.remote(task_id) + error = await service.get_task_error(task_id) if error is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -552,15 +517,10 @@ async def get_task_logs(task_id: str, max_lines: int = 100, task_details=Depends ) async def cancel_task( task_id: str, - task_state_manager=Depends(get_task_state_manager), task_details=Depends(require_task_owner), + service: IndexingService = Depends(get_indexing_service), ): - obj_ref = await task_state_manager.get_object_ref.remote(task_id) - if obj_ref is None: + cancelled = await service.cancel_task(task_id) + if not cancelled: raise HTTPException(404, f"No ObjectRef stored for task {task_id}") - - ray.cancel(obj_ref["ref"], recursive=True) - current_state = await task_state_manager.get_state.remote(task_id) - if current_state not in {"COMPLETED", "FAILED"}: - await task_state_manager.set_state.remote(task_id, "CANCELLED") return {"message": f"Cancellation signal sent for task {task_id}"} diff --git a/openrag/routers/openai.py b/openrag/routers/openai.py index e61a7fd2c..6e8b4b8a2 100644 --- a/openrag/routers/openai.py +++ b/openrag/routers/openai.py @@ -1,3 +1,16 @@ +"""OpenAI-compatible RAG endpoints — thin HTTP layer over QueryService. + +Phase 8C.2: the RAG flow (query generation, retrieval, web search, +map-reduce, context/prompt assembly, streaming, and the +``[Sources: N]`` citation filtering) moved to +``services.orchestrators.query_service.QueryService``. This module keeps +HTTP transport only: model→partition resolution, token-limit validation, +the OpenAI ``/models`` listing, request-bound source-link building +(``__prepare_sources`` uses ``request.url_for`` so it stays here and is +handed to the service as a callable), and ``StreamingResponse`` / +``JSONResponse`` wrapping with the SSE error envelope. +""" + import asyncio import json from pathlib import Path @@ -5,19 +18,14 @@ import consts from components.indexer.utils.text_sanitizer import sanitize_text -from components.pipeline import RagPipeline -from components.utils import ( - extract_and_strip_sources_block, - filter_sources_by_citations, - get_num_tokens, - stream_with_source_filtering, -) +from components.utils import get_num_tokens from config import load_config +from di.providers import get_partition_service, get_query_service from fastapi import APIRouter, Body, Depends, HTTPException, Request, status from fastapi.responses import JSONResponse, StreamingResponse -from langchain_core.documents.base import Document from models.openai import OpenAIChatCompletionRequest, OpenAICompletionRequest -from utils.dependencies import get_vectordb +from services.orchestrators.partition_service import PartitionService +from services.orchestrators.query_service import QueryService from utils.exceptions.base import OpenRAGError from utils.logger import get_logger @@ -35,8 +43,6 @@ config = load_config() router = APIRouter() -ragpipe = RagPipeline() - # Cached max model token limit, populated at startup _max_model_tokens: int | None = None @@ -49,14 +55,7 @@ async def _cache_max_model_tokens(): def _make_sse_error(message: str, code: str) -> str: """Format an error as an SSE data chunk for streaming responses.""" - chunk = { - "error": { - "message": message, - "type": "error", - "param": None, - "code": code, - } - } + chunk = {"error": {"message": message, "type": "error", "param": None, "code": code}} return f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n" @@ -81,37 +80,27 @@ def _make_sse_error(message: str, code: str) -> str: response_description="A list of available models in OpenAI format", ) async def list_models( - vectordb=Depends(get_vectordb), user_partitions=Depends(current_user_or_admin_partitions), + partitions: PartitionService = Depends(get_partition_service), ): if [p["partition"] for p in user_partitions] == ["all"]: - user_partitions = await vectordb.list_partitions.remote() + user_partitions = await partitions.list_partitions() logger.debug("Listing models", partition_count=len(user_partitions)) - models = [] - for partition in user_partitions: - model_id = f"{consts.PARTITION_PREFIX}{partition['partition']}" - models.append( - { - "id": model_id, - "object": "model", - "created": partition["created_at"], - "owned_by": "OpenRAG", - } - ) - - models.append( + models = [ { - "id": f"{consts.PARTITION_PREFIX}all", + "id": f"{consts.PARTITION_PREFIX}{partition['partition']}", "object": "model", - "created": 0, + "created": partition["created_at"], "owned_by": "OpenRAG", } - ) + for partition in user_partitions + ] + models.append({"id": f"{consts.PARTITION_PREFIX}all", "object": "model", "created": 0, "owned_by": "OpenRAG"}) return JSONResponse(content={"object": "list", "data": models}) -def __prepare_sources(request: Request, docs: list[Document], web_results: list | None = None): +def __prepare_sources(request: Request, docs: list, web_results: list | None = None): links = [] for doc in docs: doc_metadata = dict(doc.metadata) @@ -144,18 +133,14 @@ def __prepare_sources(request: Request, docs: list[Document], web_results: list def is_direct_llm_model( request: OpenAIChatCompletionRequest | OpenAICompletionRequest, ) -> bool: - """Check if request should use direct LLM (no RAG partition). - - Returns True if model is None, empty, or matches the configured default model. - """ + """True if the request should use the LLM directly (no RAG partition).""" return request.model is None or request.model == "" or request.model == config.llm.model async def _fetch_max_model_tokens() -> int: - """Fetch the maximum model token limit from vLLM's OpenAI server. + """Fetch the max model token limit from vLLM's OpenAI server. - Queries `/v1/models` and looks for `max_model_len` for the configured LLM model. - Falls back to `config.llm_context.max_llm_context_size` (default 8192) if unavailable. + Falls back to ``config.llm_context.max_llm_context_size`` if unavailable. """ default_limit = int(config.llm_context.max_llm_context_size) model_id = config.llm.model @@ -165,17 +150,13 @@ async def _fetch_max_model_tokens() -> int: if model is None: logger.warning(f"No model found for {model_id}. Using default context size.") return default_limit - model_data = model.model_dump() if hasattr(model, "model_dump") else model.dict() max_len = model_data.get("max_model_len") or model_data.get("model_extra", {}).get("max_model_len") - if max_len is None: logger.warning(f"max_model_len not found for {model_id}. Using default context size.") return default_limit - logger.info("Fetched max_model_len from vLLM at startup", model=model_id, max_model_len=int(max_len)) return int(max_len) - except Exception as e: logger.warning("Failed to query /v1/models for max_model_len; using default", error=str(e)) return default_limit @@ -192,15 +173,7 @@ def validate_tokens_limit( request: OpenAIChatCompletionRequest | OpenAICompletionRequest, max_tokens_allowed: int, ) -> tuple[bool, str]: - """Validate if the request respects the maximum token limit. - - Args: - request: The OpenAI request object - max_tokens_allowed: Maximum allowed tokens for the request. - - Returns: - Tuple of (is_valid, error_message) - """ + """Validate if the request respects the maximum token limit.""" try: _length_function = get_num_tokens() @@ -209,15 +182,6 @@ def validate_tokens_limit( default_output_tokens = int(config.llm_context.max_output_tokens) requested_tokens = request.max_tokens or default_output_tokens total_tokens_needed = message_tokens + requested_tokens - - logger.debug( - "Token validation for chat completion", - message_tokens=message_tokens, - requested_tokens=requested_tokens, - total_tokens=total_tokens_needed, - max_allowed=max_tokens_allowed, - ) - if total_tokens_needed > max_tokens_allowed: return False, ( f"Request exceeds maximum token limit. " @@ -232,15 +196,6 @@ def validate_tokens_limit( default_output_tokens = int(config.llm_context.max_output_tokens) requested_tokens = request.max_tokens or default_output_tokens total_tokens_needed = prompt_tokens + requested_tokens - - logger.debug( - "Token validation for completion", - prompt_tokens=prompt_tokens, - requested_tokens=requested_tokens, - total_tokens=total_tokens_needed, - max_allowed=max_tokens_allowed, - ) - if total_tokens_needed > max_tokens_allowed: return False, ( f"Request exceeds maximum token limit. " @@ -251,7 +206,6 @@ def validate_tokens_limit( ) return True, "" - except Exception as e: logger.warning("Error during token validation, skipping check", error=str(e)) return True, "" @@ -262,8 +216,7 @@ def check_tokens_limit( log, ): """Validate token limit and raise HTTPException(413) if exceeded.""" - max_tokens_allowed = get_max_model_tokens() - is_valid, error_message = validate_tokens_limit(request, max_tokens_allowed=max_tokens_allowed) + is_valid, error_message = validate_tokens_limit(request, max_tokens_allowed=get_max_model_tokens()) if not is_valid: log.info("Request exceeds token limit", detail=error_message) raise HTTPException( @@ -289,12 +242,6 @@ def check_tokens_limit( - `stream`: Optional streaming response (true/false) - Standard OpenAI parameters (temperature, max_tokens, etc.) -**RAG Process:** -1. Extracts query from conversation -2. Retrieves relevant documents from specified partition(s) -3. Enriches prompt with document context -4. Generates completion using LLM - **Response:** Returns OpenAI-compatible response with additional `extra` field containing: - `sources`: Array of source documents with metadata and URLs @@ -309,6 +256,8 @@ async def openai_chat_completion( user=Depends(current_user), user_partitions=Depends(current_user_or_admin_partitions_list), _: None = Depends(check_llm_model_availability), + service: QueryService = Depends(get_query_service), + partition_service: PartitionService = Depends(get_partition_service), ): model_name = request.model or config.llm.model log = logger.bind(model=model_name, endpoint="/chat/completions") @@ -320,28 +269,33 @@ async def openai_chat_completion( detail="The last message must be a non-empty user message", ) - log.debug( - "Received chat completion request with messages: {}", - truncate(str(request.messages)), - ) + log.debug("Received chat completion request with messages: {}", truncate(str(request.messages))) if is_direct_llm_model(request): check_tokens_limit(request, log) partitions = None else: - partitions = await get_partition_name(model_name, user_partitions, is_admin=user["is_admin"]) + partitions = await get_partition_name( + model_name, + user_partitions, + partition_service=partition_service, + is_admin=user["is_admin"], + ) log.debug(f"Using partitions: {partitions}") - llm_output, docs, web_results = await ragpipe.chat_completion(partition=partitions, payload=request.model_dump()) - log.debug("RAG chat completion pipeline executed.") - - sources = __prepare_sources(request2, docs, web_results=web_results) + def prep(docs, web): + return __prepare_sources(request2, docs, web) if request.stream: async def stream_response(): try: - async for sse_line in stream_with_source_filtering(llm_output, sources, model_name): + async for sse_line in service.chat_stream( + partitions=partitions, + payload=request.model_dump(), + prepare_sources=prep, + model_name=model_name, + ): yield sse_line except asyncio.CancelledError: log.info("Client disconnected during streaming") @@ -354,18 +308,15 @@ async def stream_response(): yield _make_sse_error("An unexpected error occurred during streaming", "UNEXPECTED_ERROR") return StreamingResponse(stream_response(), media_type="text/event-stream") - else: - chunk = await llm_output.__anext__() - chunk["model"] = model_name - content = chunk.get("choices", [{}])[0].get("message", {}).get("content", "") or "" - clean_content, citations = extract_and_strip_sources_block(content) - chunk["choices"][0]["message"]["content"] = clean_content - - filtered = filter_sources_by_citations(sources, citations) - chunk["extra"] = json.dumps({"sources": filtered}) - log.debug("Returning non-streaming completion chunk.") - return JSONResponse(content=chunk) + chunk = await service.chat( + partitions=partitions, + payload=request.model_dump(), + prepare_sources=prep, + model_name=model_name, + ) + log.debug("Returning non-streaming completion chunk.") + return JSONResponse(content=chunk) @router.post( @@ -384,11 +335,6 @@ async def stream_response(): - `model`: Model/partition to use - Standard OpenAI parameters (temperature, max_tokens, etc.) -**RAG Process:** -1. Retrieves relevant documents from specified partition(s) -2. Enriches prompt with document context -3. Generates completion using LLM - **Response:** Returns OpenAI-compatible response with additional `extra` field containing: - `sources`: Array of source documents with metadata and URLs @@ -402,16 +348,15 @@ async def openai_completion( user=Depends(current_user), user_partitions=Depends(current_user_or_admin_partitions_list), _: None = Depends(check_llm_model_availability), + service: QueryService = Depends(get_query_service), + partition_service: PartitionService = Depends(get_partition_service), ): model_name = request.model or config.llm.model log = logger.bind(model=model_name, endpoint="/completions") if not request.prompt: log.warning("Prompt is missing.") - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail="The prompt is required", - ) + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="The prompt is required") if request.stream: log.warning("Streaming not supported for this endpoint.") @@ -424,20 +369,17 @@ async def openai_completion( check_tokens_limit(request, log) partitions = None else: - partitions = await get_partition_name(model_name, user_partitions, is_admin=user["is_admin"]) - - llm_output, docs = await ragpipe.completions(partition=partitions, payload=request.model_dump()) - log.debug("RAG completion pipeline executed.") - - sources = __prepare_sources(request2, docs) - - complete_response = await llm_output.__anext__() - - text = complete_response.get("choices", [{}])[0].get("text", "") or "" - clean_text, citations = extract_and_strip_sources_block(text) - complete_response["choices"][0]["text"] = clean_text + partitions = await get_partition_name( + model_name, + user_partitions, + partition_service=partition_service, + is_admin=user["is_admin"], + ) - filtered = filter_sources_by_citations(sources, citations) - complete_response["extra"] = json.dumps({"sources": filtered}) + resp = await service.complete( + partitions=partitions, + payload=request.model_dump(), + prepare_sources=lambda docs, _web: __prepare_sources(request2, docs), + ) log.debug("Returning completion response.") - return JSONResponse(content=complete_response) + return JSONResponse(content=resp) diff --git a/openrag/routers/partition.py b/openrag/routers/partition.py index cb4fd3f93..e25935c62 100644 --- a/openrag/routers/partition.py +++ b/openrag/routers/partition.py @@ -1,13 +1,25 @@ +"""Partition routes — thin HTTP layer over :class:`PartitionService`. + +Phase 8B.1: partition CRUD, membership, file/chunk reads and the +relationship queries moved to +``services.orchestrators.partition_service.PartitionService``. This +module keeps HTTP transport only: request-scoped authorization (the +shared ``Depends`` wrappers in ``routers/utils.py``), ``request.url_for`` +link building, and the conflict / not-found guards whose exact +non-bracketed ``{"detail": ...}`` body the legacy endpoints returned via +``HTTPException``. +""" + from typing import Literal from urllib.parse import quote +from di.providers import get_partition_service from fastapi import APIRouter, Depends, Form, HTTPException, Request, Response, status from fastapi.responses import JSONResponse -from utils.dependencies import get_vectordb +from services.orchestrators.partition_service import PartitionService from utils.logger import get_logger from .utils import ( - ROLE_HIERARCHY, partitions_with_details, require_partition_owner, require_partition_viewer, @@ -16,7 +28,7 @@ logger = get_logger() router = APIRouter() -RoleType = Literal[*list(ROLE_HIERARCHY.keys())] +RoleType = Literal["viewer", "editor", "owner"] def _quote_param_value(s: str) -> str: @@ -37,11 +49,11 @@ def _quote_param_value(s: str) -> str: """, ) async def list_existant_partitions( - vectordb=Depends(get_vectordb), partitions=Depends(partitions_with_details), + service: PartitionService = Depends(get_partition_service), ): if len(partitions) == 1 and partitions[0]["partition"] == "all": - partitions = await vectordb.list_partitions.remote() + partitions = await service.list_partitions() logger.debug("Returned list of existing partitions.", partition_count=len(partitions)) return JSONResponse(status_code=status.HTTP_200_OK, content={"partitions": partitions}) @@ -65,11 +77,10 @@ async def list_existant_partitions( ) async def delete_partition( partition: str, - vectordb=Depends(get_vectordb), partition_owner=Depends(require_partition_owner), + service: PartitionService = Depends(get_partition_service), ): - await vectordb.delete_partition.remote(partition) - logger.debug("Partition successfully deleted.") + await service.delete_partition(partition) return Response(status_code=status.HTTP_204_NO_CONTENT) @@ -96,13 +107,10 @@ async def list_files( request: Request, partition: str, limit: int | None = None, - vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), + service: PartitionService = Depends(get_partition_service), ): - log = logger.bind(partition=partition) - file_obj_l = await vectordb.list_partition_files.remote(partition=partition, limit=limit) - file_dicts = file_obj_l.get("files", []) - log.debug("Listed files in partition", file_count=len(file_dicts)) + file_dicts = await service.list_files(partition, limit) def process_file(file_dict): return { @@ -145,19 +153,17 @@ async def get_file( partition: str, file_id: str, limit: int = 2000, - vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), + service: PartitionService = Depends(get_partition_service), ): - if not await vectordb.file_exists.remote(file_id, partition): + if not await service.file_exists(file_id, partition): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"'{file_id}' not found in partition '{partition}'", ) - results = await vectordb.get_file_chunks.remote(partition=partition, file_id=file_id, include_id=True, limit=limit) - - documents = [{"link": str(request.url_for("get_extract", extract_id=doc.metadata["_id"]))} for doc in results] - - metadata = {k: v for k, v in results[0].metadata.items() if k != "_id"} if results else {} + rows = await service.get_file_chunks(partition=partition, file_id=file_id, limit=limit) + documents = [{"link": str(request.url_for("get_extract", extract_id=row["_id"]))} for row in rows] + metadata = {k: v for k, v in rows[0].items() if k != "_id"} if rows else {} return JSONResponse( status_code=status.HTTP_200_OK, @@ -190,17 +196,17 @@ async def list_all_chunks( request: Request, partition: str, include_embedding: bool = True, - vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), + service: PartitionService = Depends(get_partition_service), ): - chunks = await vectordb.list_all_chunk.remote(partition=partition, include_embedding=include_embedding) + items = await service.list_all_chunks(partition=partition, include_embedding=include_embedding) chunks = [ { - "link": str(request.url_for("get_extract", extract_id=chunk.metadata["_id"])), - "content": chunk.page_content, - "metadata": chunk.metadata, + "link": str(request.url_for("get_extract", extract_id=it["metadata"]["_id"])), + "content": it["content"], + "metadata": it["metadata"], } - for chunk in chunks + for it in items ] return JSONResponse(status_code=status.HTTP_200_OK, content={"chunks": chunks}) @@ -224,14 +230,18 @@ async def list_all_chunks( Returns 409 Conflict if partition already exists. """, ) -async def create_partition(request: Request, partition: str, vectordb=Depends(get_vectordb)): - if await vectordb.partition_exists.remote(partition): +async def create_partition( + request: Request, + partition: str, + service: PartitionService = Depends(get_partition_service), +): + if await service.partition_exists(partition): raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail=f"Partition '{partition}' already exists.", ) user_id = request.state.user["id"] - await vectordb.create_partition.remote(partition=partition, user_id=user_id) + await service.create_partition(partition=partition, user_id=user_id) return Response(status_code=status.HTTP_201_CREATED) @@ -259,17 +269,11 @@ async def create_partition(request: Request, partition: str, vectordb=Depends(ge ) async def list_partition_users( partition: str, - vectordb=Depends(get_vectordb), partition_owner=Depends(require_partition_owner), + service: PartitionService = Depends(get_partition_service), ): - """ - List all users who are members of the given partition. - """ - log = logger.bind(partition=partition) - - members = await vectordb.list_partition_members.remote(partition=partition) - - log.debug("Returned list of partition members.", member_count=len(members)) + """List all users who are members of the given partition.""" + members = await service.list_members(partition=partition) return JSONResponse(status_code=status.HTTP_200_OK, content={"members": members}) @@ -298,17 +302,11 @@ async def add_partition_user( partition: str, user_id: int = Form(...), role: RoleType = Form("viewer"), - vectordb=Depends(get_vectordb), partition_owner=Depends(require_partition_owner), + service: PartitionService = Depends(get_partition_service), ): - """ - Add a user as a member of the given partition. - """ - log = logger.bind(partition=partition, user_id=user_id) - - await vectordb.add_partition_member.remote(partition=partition, user_id=user_id, role=role) - - log.debug("User added to partition successfully") + """Add a user as a member of the given partition.""" + await service.add_member(partition=partition, user_id=user_id, role=role) return Response(status_code=status.HTTP_201_CREATED) @@ -335,17 +333,11 @@ async def add_partition_user( async def remove_partition_user( partition: str, user_id: int, - vectordb=Depends(get_vectordb), partition_owner=Depends(require_partition_owner), + service: PartitionService = Depends(get_partition_service), ): - """ - Remove a user from the given partition. - """ - log = logger.bind(partition=partition, user_id=user_id) - - await vectordb.remove_partition_member.remote(partition=partition, user_id=user_id) - - log.debug("User removed from partition successfully") + """Remove a user from the given partition.""" + await service.remove_member(partition=partition, user_id=user_id) return Response(status_code=status.HTTP_204_NO_CONTENT) @@ -374,17 +366,11 @@ async def update_partition_user_role( partition: str, user_id: int, role: RoleType = Form(...), - vectordb=Depends(get_vectordb), partition_owner=Depends(require_partition_owner), + service: PartitionService = Depends(get_partition_service), ): - """ - Update a user's role in the given partition. - """ - log = logger.bind(partition=partition, user_id=user_id, role=role) - - await vectordb.update_partition_member_role.remote(partition=partition, user_id=user_id, new_role=role) - - log.debug("User role updated successfully") + """Update a user's role in the given partition.""" + await service.update_role(partition=partition, user_id=user_id, new_role=role) return Response(status_code=status.HTTP_200_OK) @@ -413,19 +399,13 @@ async def update_partition_user_role( """, ) async def get_related_files( - request: Request, partition: str, relationship_id: str, - vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), + service: PartitionService = Depends(get_partition_service), ): - log = logger.bind(partition=partition, relationship_id=relationship_id) - files = await vectordb.get_files_by_relationship.remote(partition=partition, relationship_id=relationship_id) - log.debug("Listed related files", file_count=len(files)) - return JSONResponse( - status_code=status.HTTP_200_OK, - content={"files": files}, - ) + files = await service.get_related_files(partition=partition, relationship_id=relationship_id) + return JSONResponse(status_code=status.HTTP_200_OK, content={"files": files}) @router.get( @@ -455,26 +435,18 @@ async def get_related_files( """, ) async def get_file_ancestors( - request: Request, partition: str, file_id: str, - max_ancestor_depth: int | None = None, # Optional limit on ancestor depth - vectordb=Depends(get_vectordb), + max_ancestor_depth: int | None = None, partition_viewer=Depends(require_partition_viewer), + service: PartitionService = Depends(get_partition_service), ): - log = logger.bind(partition=partition, file_id=file_id) - - if not await vectordb.file_exists.remote(file_id, partition): + if not await service.file_exists(file_id, partition): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"'{file_id}' not found in partition '{partition}'", ) - - ancestors = await vectordb.get_file_ancestors.remote( + ancestors = await service.get_file_ancestors( partition=partition, file_id=file_id, max_ancestor_depth=max_ancestor_depth ) - log.debug("Listed file ancestors", ancestor_count=len(ancestors)) - return JSONResponse( - status_code=status.HTTP_200_OK, - content={"ancestors": ancestors}, - ) + return JSONResponse(status_code=status.HTTP_200_OK, content={"ancestors": ancestors}) diff --git a/openrag/routers/queue.py b/openrag/routers/queue.py index 576c5bbaa..3f2431d4d 100644 --- a/openrag/routers/queue.py +++ b/openrag/routers/queue.py @@ -1,30 +1,21 @@ -from collections import Counter +"""Queue routes — thin HTTP layer over :class:`JobService`. -from config import load_config +Phase 8D.2: queue aggregation and the ``?task_status=`` filtering moved +to ``services.orchestrators.job_service.JobService``. This module keeps +HTTP transport only: the shared admin/current-user ``Depends`` wrappers +and the ``request.url_for`` link building. +""" + +from di.providers import get_job_service from fastapi import APIRouter, Depends, Request, status from fastapi.responses import JSONResponse -from utils.dependencies import get_task_state_manager +from services.orchestrators.job_service import JobService from .utils import current_user, require_admin -# load config -config = load_config() - -# Create an APIRouter instance router = APIRouter() -def _format_pool_info(worker_info: dict[str, int]) -> dict[str, int]: - """ - Convert SerializerQueue.pool_info() output into a concise dict for the API. - """ - return { - "total_slots": worker_info["total_capacity"], - "pool_size": worker_info["pool_size"], - "max_per_actor": worker_info["max_tasks_per_worker"], - } - - @router.get( "/info", description="""Get queue and worker pool information. @@ -51,25 +42,11 @@ def _format_pool_info(worker_info: dict[str, int]) -> dict[str, int]: Monitor system load and worker utilization. """, ) -async def get_queue_info(admin=Depends(require_admin), task_state_manager=Depends(get_task_state_manager)): - all_states: dict = await task_state_manager.get_all_states.remote() - status_counts = Counter(all_states.values()) - - active_statuses = ["QUEUED", "SERIALIZING", "CHUNKING", "INSERTING"] - active = {s: status_counts.get(s, 0) for s in active_statuses} - - task_summary = { - "active": sum(active.values()), - "active_statuses": active, - "total_cancelled": status_counts.get("CANCELLED", 0), - "total_completed": status_counts.get("COMPLETED", 0), - "total_failed": status_counts.get("FAILED", 0), - } - - worker_info = await task_state_manager.get_pool_info.remote() - workers_block = _format_pool_info(worker_info) - - return {"workers": workers_block, "tasks": task_summary} +async def get_queue_info( + admin=Depends(require_admin), + service: JobService = Depends(get_job_service), +): + return await service.get_queue_info() @router.get( @@ -111,40 +88,30 @@ async def get_queue_info(admin=Depends(require_admin), task_state_manager=Depend async def list_tasks( request: Request, task_status: str | None = None, - task_state_manager=Depends(get_task_state_manager), user=Depends(current_user), + service: JobService = Depends(get_job_service), ): """ - ?task_status=active → QUEUED | SERIALIZING | CHUNKING | INSERTING - ?task_status= → exact match (case-insensitive) - (none) → all tasks """ - # fetch task info - if user.get("is_admin"): - all_info: dict[str, dict] = await task_state_manager.get_all_info.remote() - else: - all_info: dict[str, dict] = await task_state_manager.get_all_user_info.remote(user.get("id")) - - if task_status is None: - filtered = all_info.items() - else: - if task_status.lower() == "active": - active_states = {"QUEUED", "SERIALIZING", "CHUNKING", "INSERTING"} - filtered = [(tid, info) for tid, info in all_info.items() if info["state"] in active_states] - else: - filtered = [(tid, info) for tid, info in all_info.items() if info["state"].lower() == task_status.lower()] - - # format the response + rows = await service.list_tasks( + is_admin=bool(user.get("is_admin")), + user_id=user.get("id"), + task_status=task_status, + ) + tasks = [] - for task_id, info in filtered: + for row in rows: + task_id = row["task_id"] item = { "task_id": task_id, - "state": info["state"], - "details": info["details"], - # include an error URL if applicable + "state": row["state"], + "details": row["details"], **( {"error_url": str(request.url_for("get_task_error", task_id=task_id))} - if info["state"] == "FAILED" + if row["state"] == "FAILED" else {} ), "url": str(request.url_for("get_task_status", task_id=task_id)), diff --git a/openrag/routers/search.py b/openrag/routers/search.py index 05d3595c1..40692828f 100644 --- a/openrag/routers/search.py +++ b/openrag/routers/search.py @@ -1,11 +1,21 @@ +"""Semantic search routes — thin HTTP layer over RetrievalService. + +Phase 8C.1: the retrieval call (was ``indexer.asearch.remote`` + the +legacy ``_expand_with_related_chunks``) moved to +``services.orchestrators.retrieval_service.RetrievalService.search``. +This module keeps HTTP transport only: request-scoped authorization, +partition resolution from the authenticated user, the byte-identical +workspace-not-found 404 guard, ``request.url_for`` links, and response +shaping (domain ``Chunk`` → ``{link, metadata, content}``). +""" + from typing import Annotated -from components.ray_utils import call_ray_actor_with_timeout -from components.retriever import _expand_with_related_chunks -from config import load_config +from di.providers import get_retrieval_service, get_workspace_service from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from fastapi.responses import JSONResponse -from utils.dependencies import get_indexer, get_vectordb +from services.orchestrators.retrieval_service import RetrievalService +from services.orchestrators.workspace_service import WorkspaceService from utils.logger import get_logger from .utils import ( @@ -14,9 +24,6 @@ require_partitions_viewer, ) -_config = load_config() -VECTORDB_TIMEOUT = _config.ray.indexer.vectordb_timeout - logger = get_logger() router = APIRouter() @@ -59,6 +66,25 @@ def __init__( self.filter = filter +def _documents(request: Request, chunks) -> list[dict]: + docs: list[dict] = [] + for c in chunks: + # Restore the legacy response metadata shape. Chunk.from_langchain + # lifts file_id / partition / page / _id out of the free-form + # metadata into typed Chunk fields; to_langchain merges them back so + # the API contract (metadata.file_id, _id, …) matches the + # pre-Phase-8 router that returned the raw Document metadata. + meta = c.to_langchain().metadata + docs.append( + { + "link": str(request.url_for("get_extract", extract_id=meta.get("_id") or c.id)), + "metadata": meta, + "content": c.text, + } + ) + return docs + + @router.get( "", description="""Perform semantic search across multiple partitions. @@ -114,12 +140,11 @@ async def search_multiple_partitions( related_params: Annotated[RelatedDocSearchParams, Depends()], partitions: list[str] | None = Query(default=["all"], description="List of partitions to search"), workspace: str | None = Query(None, description="Workspace ID to filter results"), - indexer=Depends(get_indexer), - vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partitions_viewer), user_partitions=Depends(current_user_or_admin_partitions_list), + service: RetrievalService = Depends(get_retrieval_service), + workspaces: WorkspaceService = Depends(get_workspace_service), ): - # Fetch user partitions if "all" is specified, or all partitions if super admin if partitions == ["all"]: partitions = user_partitions @@ -134,53 +159,29 @@ async def search_multiple_partitions( filter_params = None if workspace: - ws = await call_ray_actor_with_timeout( - vectordb.get_workspace.remote(workspace), - timeout=VECTORDB_TIMEOUT, - task_description=f"get_workspace({workspace})", - ) + ws = await workspaces.get_workspace(workspace) if not ws or ws["partition_name"] not in partitions: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Workspace not found") - filter_params = {"workspace_id": workspace} - results = await indexer.asearch.remote( - query=search_params.text, + results = await service.search( + text=search_params.text, + partitions=partitions, top_k=search_params.top_k, similarity_threshold=search_params.similarity_threshold, - partition=partitions, filter=search_params.filter, filter_params=filter_params, + include_related=related_params.include_related, + include_ancestors=related_params.include_ancestors, + related_limit=related_params.related_limit, + max_ancestor_depth=related_params.max_ancestor_depth, ) - log.info( - "Semantic search on multiple partitions completed.", - result_count=len(results), - ) - - # Expand with related/ancestor chunks if requested - if related_params.include_related or related_params.include_ancestors: - results = await _expand_with_related_chunks( - results=results, - db=vectordb, - include_related=related_params.include_related, - include_ancestors=related_params.include_ancestors, - related_limit=related_params.related_limit, - max_ancestor_depth=related_params.max_ancestor_depth, - ) - log.info( - "Expanded results with related/ancestor chunks.", - expanded_count=len(results), - ) + log.info("Semantic search on multiple partitions completed.", result_count=len(results)) - documents = [ - { - "link": str(request.url_for("get_extract", extract_id=doc.metadata["_id"])), - "metadata": doc.metadata, - "content": doc.page_content, - } - for doc in results - ] - return JSONResponse(status_code=status.HTTP_200_OK, content={"documents": documents}) + return JSONResponse( + status_code=status.HTTP_200_OK, + content={"documents": _documents(request, results)}, + ) @router.get( @@ -229,9 +230,9 @@ async def search_one_partition( search_params: Annotated[CommonSearchParams, Depends()], related_params: Annotated[RelatedDocSearchParams, Depends()], workspace: str | None = Query(None, description="Workspace ID to filter results"), - indexer=Depends(get_indexer), - vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), + service: RetrievalService = Depends(get_retrieval_service), + workspaces: WorkspaceService = Depends(get_workspace_service), ): log = logger.bind( partition=partition, @@ -243,51 +244,29 @@ async def search_one_partition( ) filter_params = None if workspace: - ws = await call_ray_actor_with_timeout( - vectordb.get_workspace.remote(workspace), - timeout=VECTORDB_TIMEOUT, - task_description=f"get_workspace({workspace})", - ) + ws = await workspaces.get_workspace(workspace) if not ws or ws["partition_name"] != partition: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Workspace not found") - filter_params = {"workspace_id": workspace} - results = await indexer.asearch.remote( - query=search_params.text, + results = await service.search( + text=search_params.text, + partitions=partition, top_k=search_params.top_k, similarity_threshold=search_params.similarity_threshold, - partition=partition, filter=search_params.filter, filter_params=filter_params, + include_related=related_params.include_related, + include_ancestors=related_params.include_ancestors, + related_limit=related_params.related_limit, + max_ancestor_depth=related_params.max_ancestor_depth, ) - log.info("Semantic search on single partition completed.", result_count=len(results)) - # Expand with related/ancestor chunks if requested - if related_params.include_related or related_params.include_ancestors: - results = await _expand_with_related_chunks( - results=results, - db=vectordb, - include_related=related_params.include_related, - include_ancestors=related_params.include_ancestors, - related_limit=related_params.related_limit, - max_ancestor_depth=related_params.max_ancestor_depth, - ) - log.info( - "Expanded results with related/ancestor chunks.", - expanded_count=len(results), - ) - - documents = [ - { - "link": str(request.url_for("get_extract", extract_id=doc.metadata["_id"])), - "metadata": doc.metadata, - "content": doc.page_content, - } - for doc in results - ] - return JSONResponse(status_code=status.HTTP_200_OK, content={"documents": documents}) + return JSONResponse( + status_code=status.HTTP_200_OK, + content={"documents": _documents(request, results)}, + ) @router.get( @@ -331,31 +310,25 @@ async def search_file( partition: str, file_id: str, search_params: Annotated[CommonSearchParams, Depends()], - indexer=Depends(get_indexer), - vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), + service: RetrievalService = Depends(get_retrieval_service), ): log = logger.bind(partition=partition, file_id=file_id, query=search_params.text, top_k=search_params.top_k) filter = "file_id == {_file_id}" + (f" AND {search_params.filter}" if search_params.filter else "") params = {"_file_id": file_id} - results = await indexer.asearch.remote( - query=search_params.text, + results = await service.search( + text=search_params.text, + partitions=partition, top_k=search_params.top_k, similarity_threshold=search_params.similarity_threshold, - partition=partition, filter=filter, filter_params=params, ) log.info("Semantic search on specific file completed.", result_count=len(results)) - documents = [ - { - "link": str(request.url_for("get_extract", extract_id=doc.metadata["_id"])), - "metadata": doc.metadata, - "content": doc.page_content, - } - for doc in results - ] - return JSONResponse(status_code=status.HTTP_200_OK, content={"documents": documents}) + return JSONResponse( + status_code=status.HTTP_200_OK, + content={"documents": _documents(request, results)}, + ) diff --git a/openrag/routers/test_auth_router.py b/openrag/routers/test_auth_router.py index 83a5eb57d..fc052d1c5 100644 --- a/openrag/routers/test_auth_router.py +++ b/openrag/routers/test_auth_router.py @@ -1,777 +1,244 @@ -"""Integration tests for the OIDC auth router. - -The router transitively imports ``utils.dependencies``, which spins up Ray -actors at import time (indexer, marker pool, semaphores, …). To avoid that -in a unit-test context, we stub ``utils.dependencies`` in ``sys.modules`` -*before* importing the router, then drive it via FastAPI's ``TestClient``. - -IdP interactions are mocked end-to-end with ``respx`` using a real RSA key -pair so the router exercises actual JWT verification. +"""Transport tests for the thin OIDC auth router (Phase 8A.1). + +Every OIDC business decision (PKCE/state generation, code exchange, user +lookup/provisioning, session creation, logout-URL construction, JWT +verification) now lives in :class:`services.orchestrators.auth_service. +AuthService` and is covered end-to-end by +``services/orchestrators/test_auth_service.py``. + +This module only asserts what the router still owns: the ``AUTH_MODE`` +gate, delegation to the injected service, cookie set/clear, and the +``OIDCFlowError`` → HTTP-response mapping. The service is stubbed via +``dependency_overrides`` so no container / Ray / IdP is needed. + +The router transitively imports ``services.workers.bootstrap`` (Ray actors +at import time) through its dependency graph, so we stub that module in +``sys.modules`` before importing the router. """ from __future__ import annotations import sys -import time import types -from typing import Any import pytest -# --------------------------------------------------------------------------- -# pytest sometimes initialises warning filters before we run — tolerate it. -# --------------------------------------------------------------------------- - -pytest.importorskip("respx") -pytest.importorskip("httpx") -pytest.importorskip("authlib") pytest.importorskip("fastapi") -pytest.importorskip("itsdangerous") -pytest.importorskip("cryptography") +pytest.importorskip("httpx") -import httpx # noqa: E402 -import respx # noqa: E402 -from authlib.jose import JsonWebKey, JsonWebToken # noqa: E402 -from cryptography.fernet import Fernet # noqa: E402 from fastapi import FastAPI # noqa: E402 from fastapi.testclient import TestClient # noqa: E402 -from sqlalchemy.exc import IntegrityError # noqa: E402 - -# --------------------------------------------------------------------------- -# Constants — align with the existing auth unit tests -# --------------------------------------------------------------------------- - -ISSUER = "https://idp.example.com/realms/openrag" -CLIENT_ID = "openrag-client" -CLIENT_SECRET = "test-secret" -REDIRECT_URI = "https://openrag.example.com/auth/callback" -SCOPES = "openid email profile offline_access" - -DISCOVERY_DOC = { - "issuer": ISSUER, - "authorization_endpoint": f"{ISSUER}/protocol/openid-connect/auth", - "token_endpoint": f"{ISSUER}/protocol/openid-connect/token", - "userinfo_endpoint": f"{ISSUER}/protocol/openid-connect/userinfo", - "jwks_uri": f"{ISSUER}/protocol/openid-connect/certs", - "end_session_endpoint": f"{ISSUER}/protocol/openid-connect/logout", -} - - -def _make_rsa_key_pair(): - private = JsonWebKey.generate_key("RSA", 2048, is_private=True) - return private, private.as_dict(is_private=True), private.as_dict() - - -_RSA_PRIVATE, _RSA_PRIVATE_JWK, _RSA_PUBLIC_JWK = _make_rsa_key_pair() -_RSA_PUBLIC_JWK["use"] = "sig" -_RSA_PUBLIC_JWK["alg"] = "RS256" -_RSA_PUBLIC_JWK["kid"] = "test-key-1" -_RSA_PRIVATE_JWK["kid"] = "test-key-1" - -JWKS_RESPONSE = {"keys": [_RSA_PUBLIC_JWK]} - - -def _sign_jwt(payload: dict) -> str: - header = {"alg": "RS256", "kid": "test-key-1"} - # Authlib >=1.0 requires the allowed-algorithms list on JsonWebToken. - jwt = JsonWebToken(["RS256"]) - token = jwt.encode(header, payload, _RSA_PRIVATE) - return token.decode() if isinstance(token, bytes) else token - - -def _id_token_payload( - nonce: str, *, sub: str = "sub-abc", email: str | None = "user@example.com", extra: dict | None = None -) -> dict: - now = int(time.time()) - payload = { - "iss": ISSUER, - "sub": sub, - "aud": CLIENT_ID, - "exp": now + 300, - "iat": now, - "nonce": nonce, - } - if email is not None: - payload["email"] = email - if extra: - payload.update(extra) - return payload - - -def _logout_token_payload(*, sid: str | None = None, sub: str | None = None) -> dict: - now = int(time.time()) - payload: dict[str, Any] = { - "iss": ISSUER, - "aud": CLIENT_ID, - "iat": now, - "jti": "lt-001", - "events": {"http://schemas.openid.net/event/backchannel-logout": {}}, - } - if sid is not None: - payload["sid"] = sid - if sub is not None: - payload["sub"] = sub - return payload - # --------------------------------------------------------------------------- # Stub heavy dependencies BEFORE importing the router # --------------------------------------------------------------------------- -_FERNET_KEY = Fernet.generate_key().decode() - - -class _RayMethodStub: - """Mimics a Ray actor method: ``method.remote(...)`` returns an awaitable.""" - def __init__(self, name: str, fn, call_log: list): - self._name = name - self._fn = fn - self._call_log = call_log +_STUBBED_MODULES = ("utils", "utils.logger", "services.workers.bootstrap") - async def remote(self, *args, **kwargs): - self._call_log.append((self._name, args, kwargs)) - return self._fn(*args, **kwargs) +def _install_dependencies_stub() -> dict[str, types.ModuleType | None]: + previous_modules = {name: sys.modules.get(name) for name in _STUBBED_MODULES} -class _StubVectorDB: - """Minimal Ray-actor stand-in — exposes ``.method.remote(...)`` awaitables.""" - - def __init__(self): - self.calls: list[tuple[str, tuple, dict]] = [] - self._users_by_sub: dict[str, dict] = {} - self._users_by_email: dict[str, dict] = {} - self._users_by_id: dict[int, dict] = {} - self._sessions: dict[int, dict] = {} - self._sessions_by_token: dict[str, int] = {} - self._next_session_id = 1 - self._next_user_id = 1000 - # Bind each underlying impl as an actor-style accessor. - self.get_user_by_external_id = _RayMethodStub( - "get_user_by_external_id", self._impl_get_user_by_external_id, self.calls - ) - self.get_user_by_email = _RayMethodStub("get_user_by_email", self._impl_get_user_by_email, self.calls) - self.update_user_fields = _RayMethodStub("update_user_fields", self._impl_update_user_fields, self.calls) - self.create_user = _RayMethodStub("create_user", self._impl_create_user, self.calls) - self.create_oidc_session = _RayMethodStub("create_oidc_session", self._impl_create_oidc_session, self.calls) - self.get_oidc_session_by_token = _RayMethodStub( - "get_oidc_session_by_token", self._impl_get_oidc_session_by_token, self.calls - ) - self.revoke_oidc_session_by_id = _RayMethodStub( - "revoke_oidc_session_by_id", self._impl_revoke_oidc_session_by_id, self.calls - ) - self.revoke_oidc_sessions_by_sid = _RayMethodStub( - "revoke_oidc_sessions_by_sid", self._impl_revoke_oidc_sessions_by_sid, self.calls - ) - - # Test-only helpers ----------------------------------------------------- - - def add_user( - self, - *, - user_id: int, - email: str | None = None, - external_user_id: str | None = None, - display_name: str | None = None, - ) -> dict: - user = { - "id": user_id, - "email": email, - "external_user_id": external_user_id, - "is_admin": False, - "display_name": display_name or f"user-{user_id}", - } - self._users_by_id[user_id] = user - if external_user_id: - self._users_by_sub[external_user_id] = user - if email: - self._users_by_email[email.strip().lower()] = user - return user - - # Impls ------------------------------------------------------------------ - - def _impl_get_user_by_external_id(self, external_user_id: str): - return self._users_by_sub.get(external_user_id) - - def _impl_get_user_by_email(self, email: str): - if not isinstance(email, str) or not email.strip(): - return None - return self._users_by_email.get(email.strip().lower()) - - def _impl_create_user(self, body): - # Accept either UserCreate or a plain dict, like the real Ray method - # would (Pydantic instances are serializable). - if hasattr(body, "model_dump"): - data = body.model_dump() - else: - data = dict(body) - normalized_email = data.get("email").strip().lower() if data.get("email") else None - # Faithful to the unique ix_users_email index: a second row with the - # same email is rejected, just like Postgres would. - if normalized_email and normalized_email in self._users_by_email: - raise IntegrityError("duplicate key value violates unique constraint", None, Exception()) - user_id = self._next_user_id - self._next_user_id += 1 - user = { - "id": user_id, - "display_name": data.get("display_name"), - "external_user_id": data.get("external_user_id"), - "email": normalized_email, - "is_admin": bool(data.get("is_admin", False)), - "file_quota": data.get("file_quota"), - "file_count": 0, - "token": "or-stub", - } - self._users_by_id[user_id] = user - if data.get("external_user_id"): - self._users_by_sub[data["external_user_id"]] = user - if normalized_email: - self._users_by_email[normalized_email] = user - return user - - def _impl_update_user_fields(self, user_id: int, fields: dict): - user = self._users_by_id.get(user_id) - if user is None: - raise ValueError(f"User {user_id} not found") - _ALLOWED = {"display_name", "email"} - bad = set(fields) - _ALLOWED - if bad: - raise ValueError(f"Cannot update non-whitelisted user fields: {sorted(bad)}") - for k, v in fields.items(): - if v is None: - continue - if k == "email" and isinstance(v, str): - v = v.strip().lower() - user[k] = v - - def _impl_create_oidc_session(self, **kwargs): - sid = self._next_session_id - self._next_session_id += 1 - row = { - "id": sid, - "session_expires_at": kwargs["session_expires_at"], - "id_token_encrypted": kwargs["id_token_encrypted"], - **{k: v for k, v in kwargs.items() if k != "session_token_plain"}, - } - self._sessions[sid] = row - self._sessions_by_token[kwargs["session_token_plain"]] = sid - return row - - def _impl_get_oidc_session_by_token(self, session_token_plain: str): - # Mirror PartitionFileManager.get_oidc_session_by_token semantics: - # reject revoked rows AND rows whose session_expires_at is in the past - # relative to datetime.now(). The expiry check is what makes this stub - # a faithful regression target for the M2 timezone fix. - from datetime import datetime as _dt - - sid = self._sessions_by_token.get(session_token_plain) - if sid is None: - return None - row = self._sessions[sid] - if row.get("revoked_at"): - return None - exp = row.get("session_expires_at") - if isinstance(exp, _dt) and exp < _dt.now(): - return None - return row - - def _impl_revoke_oidc_session_by_id(self, session_id: int): - row = self._sessions.get(session_id) - if row: - row["revoked_at"] = time.time() - - def _impl_revoke_oidc_sessions_by_sid(self, sid: str) -> int: - count = 0 - for row in self._sessions.values(): - if row.get("sid") == sid and not row.get("revoked_at"): - row["revoked_at"] = time.time() - count += 1 - return count - - -_stub_vectordb_singleton = _StubVectorDB() - - -def _install_dependencies_stub(): - """Replace ``utils.dependencies`` with a stub providing only ``get_vectordb``.""" - stub = types.ModuleType("utils.dependencies") - stub.get_vectordb = lambda: _stub_vectordb_singleton + stub = types.ModuleType("services.workers.bootstrap") + stub.actor_creation_map = {} stub.get_task_state_manager = lambda: None stub.get_serializer = lambda: None - stub.get_indexer = lambda: None stub.get_marker_pool = lambda: None - sys.modules["utils.dependencies"] = stub + sys.modules["services.workers.bootstrap"] = stub + + def _logger(): + logger = types.SimpleNamespace( + debug=lambda *args, **kwargs: None, + info=lambda *args, **kwargs: None, + warning=lambda *args, **kwargs: None, + error=lambda *args, **kwargs: None, + exception=lambda *args, **kwargs: None, + ) + logger.bind = lambda *args, **kwargs: logger + return logger + + logger_stub = types.ModuleType("utils.logger") + logger_stub.escape_markup = lambda s: s.replace("\\", "\\\\").replace("<", "\\<").replace(">", "\\>") + logger_stub.mask_email = ( + lambda email: f"{email.partition('@')[0][0]}***@{email.partition('@')[2]}" + if isinstance(email, str) and "@" in email and email.partition("@")[0] + else "***" + ) + logger_stub.get_logger = _logger + sys.modules["utils.logger"] = logger_stub + return previous_modules -_install_dependencies_stub() +def _restore_dependencies_stub(previous_modules: dict[str, types.ModuleType | None]) -> None: + for name, previous_module in previous_modules.items(): + if previous_module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous_module + +_PREVIOUS_MODULES = _install_dependencies_stub() -# Now we can import the router. -import importlib # noqa: E402 +from di.providers import get_auth_service # noqa: E402 +from routers.auth import router as auth_router # noqa: E402 +from services.orchestrators.auth_service import ( # noqa: E402 + SESSION_COOKIE_NAME, + CallbackResult, + LoginRedirect, + OIDCFlowError, +) -# Reset the OIDC client singleton between tests — important when env changes. -from components.auth import deps as _auth_deps # noqa: E402 +_restore_dependencies_stub(_PREVIOUS_MODULES) -# Import the router module, forcing a fresh import. -sys.modules.pop("routers.auth", None) -_auth_router_module = importlib.import_module("routers.auth") -auth_router = _auth_router_module.router +STATE_COOKIE_NAME = "openrag_oidc_state" # --------------------------------------------------------------------------- -# Fixtures +# Stub AuthService # --------------------------------------------------------------------------- -@pytest.fixture -def env_oidc(monkeypatch): - monkeypatch.setenv("AUTH_MODE", "oidc") - monkeypatch.setenv("OIDC_ENDPOINT", ISSUER) - monkeypatch.setenv("OIDC_CLIENT_ID", CLIENT_ID) - monkeypatch.setenv("OIDC_CLIENT_SECRET", CLIENT_SECRET) - monkeypatch.setenv("OIDC_REDIRECT_URI", REDIRECT_URI) - monkeypatch.setenv("OIDC_SCOPES", SCOPES) - monkeypatch.setenv("OIDC_TOKEN_ENCRYPTION_KEY", _FERNET_KEY) - monkeypatch.setenv("OIDC_CLAIM_SOURCE", "id_token") - monkeypatch.setenv("OIDC_POST_LOGOUT_REDIRECT_URI", "/") - monkeypatch.delenv("OIDC_CLAIM_MAPPING", raising=False) - _auth_deps.reset_oidc_client() +class StubAuthService: + def __init__(self) -> None: + self.calls: list[tuple] = [] + self.login_result = LoginRedirect( + authorization_url="https://idp.example.com/auth?response_type=code", + state_cookie_name=STATE_COOKIE_NAME, + state_cookie_value="state-val", + state_cookie_max_age=600, + ) + self.callback_result = CallbackResult( + session_cookie_name=SESSION_COOKIE_NAME, + session_cookie_value="sess-plain", + session_cookie_max_age=300, + next_url="/next", + ) + self.logout_target: str | None = "https://idp.example.com/logout" + self.raise_on: dict[str, OIDCFlowError] = {} + async def start_oidc_login(self, next_url): + self.calls.append(("start_oidc_login", next_url)) + if "login" in self.raise_on: + raise self.raise_on["login"] + return self.login_result -@pytest.fixture -def env_token(monkeypatch): - monkeypatch.setenv("AUTH_MODE", "token") - _auth_deps.reset_oidc_client() + async def handle_oidc_callback(self, *, code, state, state_cookie_raw): + self.calls.append(("handle_oidc_callback", code, state, state_cookie_raw)) + if "callback" in self.raise_on: + raise self.raise_on["callback"] + return self.callback_result + async def handle_backchannel_logout(self, logout_token): + self.calls.append(("handle_backchannel_logout", logout_token)) + if "bcl" in self.raise_on: + raise self.raise_on["bcl"] + return 1 -@pytest.fixture -def fresh_stub_vectordb(): - global _stub_vectordb_singleton - # Re-create so tests see a clean state. - _stub_vectordb_singleton.__init__() - return _stub_vectordb_singleton + async def logout(self, session_cookie_value): + self.calls.append(("logout", session_cookie_value)) + return self.logout_target + + +def _set_cookies(response) -> list[str]: + return [v for k, v in response.headers.multi_items() if k.lower() == "set-cookie"] @pytest.fixture -def client(env_oidc, fresh_stub_vectordb): - """TestClient for the minimal FastAPI app. +def stub() -> StubAuthService: + return StubAuthService() - The OIDCClient singleton uses a shared respx-mocked transport so every - IdP route can be stubbed per-test via ``mock.router.get(...)``. - """ + +@pytest.fixture +def client(stub: StubAuthService) -> TestClient: app = FastAPI() app.include_router(auth_router) + app.dependency_overrides[get_auth_service] = lambda: stub + return TestClient(app) - # Replace the OIDCClient's internal httpx client with one backed by respx. - # respx >= 0.22 removed the top-level MockTransport; use MockRouter + - # httpx.MockTransport(router.handler) instead. - router = respx.MockRouter(assert_all_called=False) - http = httpx.AsyncClient(transport=httpx.MockTransport(router.handler)) - - # Force singleton creation using our mocked http client. - _auth_deps.reset_oidc_client() - _auth_deps._client = _auth_router_module.OIDCClient( - issuer=ISSUER, - client_id=CLIENT_ID, - client_secret=CLIENT_SECRET, - redirect_uri=REDIRECT_URI, - scopes=SCOPES, - http_client=http, - ) - - c = TestClient(app) - c.oidc_router = router # type: ignore[attr-defined] - yield c - -def _setup_discovery(router): - router.get(f"{ISSUER}/.well-known/openid-configuration").mock(return_value=httpx.Response(200, json=DISCOVERY_DOC)) +@pytest.fixture +def oidc_env(monkeypatch): + monkeypatch.setenv("AUTH_MODE", "oidc") -def _setup_jwks(router): - router.get(f"{ISSUER}/protocol/openid-connect/certs").mock(return_value=httpx.Response(200, json=JWKS_RESPONSE)) +@pytest.fixture +def token_env(monkeypatch): + monkeypatch.setenv("AUTH_MODE", "token") # --------------------------------------------------------------------------- -# GET /auth/login +# AUTH_MODE gate — token mode must 400 *before* the service is resolved. +# Regression: previously these returned 503 because Depends(get_auth_service) +# resolved (and failed on a missing container) before the gate ran. # --------------------------------------------------------------------------- -def test_login_rejected_in_token_mode(env_token, fresh_stub_vectordb): +@pytest.mark.parametrize( + ("method", "path", "kwargs"), + [ + ("get", "/auth/login", {}), + ("get", "/auth/callback?code=x&state=y", {}), + ("post", "/auth/backchannel-logout", {"data": {"logout_token": "x"}}), + ("get", "/auth/logout", {}), + ], +) +def test_routes_rejected_in_token_mode(token_env, method, path, kwargs): + """No service override and no container — must still be a clean 400.""" app = FastAPI() app.include_router(auth_router) c = TestClient(app) - r = c.get("/auth/login", follow_redirects=False) + r = getattr(c, method)(path, follow_redirects=False, **kwargs) assert r.status_code == 400 - - -def test_login_redirects_to_idp_with_pkce(client): - _setup_discovery(client.oidc_router) - r = client.get("/auth/login", follow_redirects=False) - assert r.status_code == 302 - loc = r.headers["location"] - assert loc.startswith(f"{ISSUER}/protocol/openid-connect/auth") - assert "code_challenge_method=S256" in loc - assert "state=" in loc - assert "nonce=" in loc - assert "code_challenge=" in loc - # State cookie set - assert "openrag_oidc_state" in r.cookies + assert "AUTH_MODE" in r.json()["detail"] # --------------------------------------------------------------------------- -# GET /auth/callback — failure paths +# GET /auth/login # --------------------------------------------------------------------------- -def test_callback_rejected_in_token_mode(env_token, fresh_stub_vectordb): - app = FastAPI() - app.include_router(auth_router) - c = TestClient(app) - r = c.get("/auth/callback?code=x&state=y", follow_redirects=False) - assert r.status_code == 400 - - -def test_callback_missing_state_cookie(client): - r = client.get("/auth/callback?code=x&state=y", follow_redirects=False) - assert r.status_code == 400 - assert "state cookie" in r.json()["detail"].lower() - +def test_login_redirects_and_sets_state_cookie(oidc_env, client, stub): + r = client.get("/auth/login?next=/dashboard", follow_redirects=False) + assert r.status_code == 302 + assert r.headers["location"] == stub.login_result.authorization_url + assert stub.calls == [("start_oidc_login", "/dashboard")] + assert any(STATE_COOKIE_NAME in c for c in _set_cookies(r)) -def test_callback_state_mismatch(client): - _setup_discovery(client.oidc_router) - # First, obtain a legitimate state cookie via /auth/login. - login_resp = client.get("/auth/login", follow_redirects=False) - assert login_resp.status_code == 302 - # Now call /auth/callback with a *different* state in the query. - r = client.get( - "/auth/callback?code=x&state=WRONG", - follow_redirects=False, - ) - assert r.status_code == 400 - assert "state" in r.json()["detail"].lower() +def test_login_flow_error_maps_to_http(oidc_env, client, stub): + stub.raise_on["login"] = OIDCFlowError("boom", status_code=502) + r = client.get("/auth/login", follow_redirects=False) + assert r.status_code == 502 + assert r.json()["detail"] == "boom" # --------------------------------------------------------------------------- -# GET /auth/callback — success paths +# GET /auth/callback # --------------------------------------------------------------------------- -def _begin_login_and_extract_state(client) -> tuple[str, str]: - """Call /auth/login and return (state, nonce) values from the redirect query.""" - _setup_discovery(client.oidc_router) - r = client.get("/auth/login", follow_redirects=False) - assert r.status_code == 302 - loc = r.headers["location"] - from urllib.parse import parse_qs, urlparse - - qs = parse_qs(urlparse(loc).query) - return qs["state"][0], qs["nonce"][0] - - -def _mock_token_endpoint(router, id_token: str, *, refresh_token: str | None = "rt-1"): - payload = { - "id_token": id_token, - "access_token": "at-1", - "expires_in": 300, - "token_type": "Bearer", - } - if refresh_token is not None: - payload["refresh_token"] = refresh_token - router.post(f"{ISSUER}/protocol/openid-connect/token").mock(return_value=httpx.Response(200, json=payload)) - - -def test_callback_success_by_external_id(client, fresh_stub_vectordb): - fresh_stub_vectordb.add_user(user_id=42, email="user@example.com", external_user_id="sub-abc") - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - id_token = _sign_jwt(_id_token_payload(nonce, sub="sub-abc")) - _mock_token_endpoint(client.oidc_router, id_token) - - r = client.get( - f"/auth/callback?code=authcode&state={state}", - follow_redirects=False, - ) - assert r.status_code == 302, r.text - assert r.headers["location"] == "/" - assert "openrag_session" in r.cookies - # At least one create_oidc_session call recorded. - assert any(c[0] == "create_oidc_session" for c in fresh_stub_vectordb.calls) - - -def test_callback_user_not_registered(client, fresh_stub_vectordb): - """Unknown sub → 403 by default (no email fallback, no auto-provisioning).""" - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - id_token = _sign_jwt(_id_token_payload(nonce, sub="sub-unknown", email="ghost@example.com")) - _mock_token_endpoint(client.oidc_router, id_token) - - r = client.get( - f"/auth/callback?code=c&state={state}", - follow_redirects=False, - ) - assert r.status_code == 403 - assert "not registered" in r.json()["detail"].lower() - # Without OIDC_AUTO_PROVISION_LOGIN, no user must have been created. - assert not any(c[0] == "create_user" for c in fresh_stub_vectordb.calls) - - -def test_callback_auto_provisions_user_when_enabled(client, fresh_stub_vectordb, monkeypatch): - """OIDC_AUTO_PROVISION_LOGIN=true: unknown sub triggers user creation - from ID-token claims, never as admin, and login proceeds (302 + cookie).""" - monkeypatch.setenv("OIDC_AUTO_PROVISION_LOGIN", "true") - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - id_token = _sign_jwt( - _id_token_payload( - nonce, - sub="sub-new-user", - email="alice@example.com", - extra={"name": "Alice Liddell"}, - ) - ) - _mock_token_endpoint(client.oidc_router, id_token) - +def test_callback_success_sets_session_clears_state(oidc_env, client, stub): r = client.get( - f"/auth/callback?code=c&state={state}", + "/auth/callback?code=ac&state=st", + cookies={STATE_COOKIE_NAME: "raw-state"}, follow_redirects=False, ) - assert r.status_code == 302, r.text - assert "openrag_session" in r.cookies - - # create_user was called exactly once with the IdP claims. - create_calls = [c for c in fresh_stub_vectordb.calls if c[0] == "create_user"] - assert len(create_calls) == 1 - body = create_calls[0][1][0] - data = body.model_dump() if hasattr(body, "model_dump") else dict(body) - assert data["external_user_id"] == "sub-new-user" - assert data["display_name"] == "Alice Liddell" - assert data["email"] == "alice@example.com" - assert data["is_admin"] is False # auto-provisioned users are NEVER admin - - -def test_callback_auto_provision_email_collision_returns_409(client, fresh_stub_vectordb, monkeypatch): - """First-login auto-provisioning where the email already belongs to a row - under a different identity must not 500. Matching is external_id-only, so - the existing row isn't found by sub; create_user then hits the unique email - index. The callback recovers with an actionable 409 instead of 500.""" - monkeypatch.setenv("OIDC_AUTO_PROVISION_LOGIN", "true") - # Pre-existing user with this email but a *different* external_user_id. - fresh_stub_vectordb.add_user(user_id=7, email="alice@example.com", external_user_id="sub-old") - - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - id_token = _sign_jwt(_id_token_payload(nonce, sub="sub-new-user", email="alice@example.com", extra={})) - _mock_token_endpoint(client.oidc_router, id_token) - - r = client.get(f"/auth/callback?code=c&state={state}", follow_redirects=False) - - assert r.status_code == 409, r.text - assert "already exists" in r.text - # No session was minted for the failed login. - assert not any(c[0] == "create_oidc_session" for c in fresh_stub_vectordb.calls) - - -def test_callback_auto_provision_falls_back_to_sub_when_no_name(client, fresh_stub_vectordb, monkeypatch): - """When the IdP exposes no readable display name, a deterministic - ``oidc-`` placeholder is used so the UI always has something.""" - monkeypatch.setenv("OIDC_AUTO_PROVISION_LOGIN", "true") - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - id_token = _sign_jwt(_id_token_payload(nonce, sub="abcdef0123456789", email=None, extra={})) - _mock_token_endpoint(client.oidc_router, id_token) - - r = client.get( - f"/auth/callback?code=c&state={state}", - follow_redirects=False, - ) - assert r.status_code == 302, r.text - - create_calls = [c for c in fresh_stub_vectordb.calls if c[0] == "create_user"] - assert len(create_calls) == 1 - body = create_calls[0][1][0] - data = body.model_dump() if hasattr(body, "model_dump") else dict(body) - assert data["display_name"] == "oidc-abcdef01" - assert data["email"] is None - - -def test_callback_auto_provision_disabled_by_default(client, fresh_stub_vectordb, monkeypatch): - """Explicitly setting OIDC_AUTO_PROVISION_LOGIN=false (or unset) keeps the - historical 403 behaviour — non-breaking guard for existing deployments.""" - monkeypatch.setenv("OIDC_AUTO_PROVISION_LOGIN", "false") - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - id_token = _sign_jwt(_id_token_payload(nonce, sub="sub-x", email="x@example.com", extra={"name": "X"})) - _mock_token_endpoint(client.oidc_router, id_token) - - r = client.get(f"/auth/callback?code=c&state={state}", follow_redirects=False) - assert r.status_code == 403 - assert not any(c[0] == "create_user" for c in fresh_stub_vectordb.calls) - - -def test_callback_auto_provision_syncs_existing_user_claims(client, fresh_stub_vectordb, monkeypatch): - """OIDC_AUTO_PROVISION_LOGIN=true also keeps display_name/email of an - already-known user in sync with the IdP claims on every login.""" - monkeypatch.setenv("OIDC_AUTO_PROVISION_LOGIN", "true") - monkeypatch.delenv("OIDC_CLAIM_MAPPING", raising=False) - fresh_stub_vectordb.add_user( - user_id=99, - email="stale@example.com", - external_user_id="sub-existing", - display_name="Stale Name", - ) - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - id_token = _sign_jwt( - _id_token_payload( - nonce, - sub="sub-existing", - email="fresh@example.com", - extra={"name": "Fresh Name"}, - ) - ) - _mock_token_endpoint(client.oidc_router, id_token) - - r = client.get(f"/auth/callback?code=c&state={state}", follow_redirects=False) - assert r.status_code == 302, r.text - - user = fresh_stub_vectordb._users_by_id[99] - assert user["display_name"] == "Fresh Name" - assert user["email"] == "fresh@example.com" - # User row was not re-created — only updated. - assert not any(c[0] == "create_user" for c in fresh_stub_vectordb.calls) - assert sum(1 for c in fresh_stub_vectordb.calls if c[0] == "update_user_fields") == 1 - - -def test_callback_auto_provision_no_db_write_when_claims_match(client, fresh_stub_vectordb, monkeypatch): - """With AUTO_PROVISION_LOGIN on but claims already matching the stored row, - no update_user_fields call is made (the no-op filter avoids DB churn).""" - monkeypatch.setenv("OIDC_AUTO_PROVISION_LOGIN", "true") - monkeypatch.delenv("OIDC_CLAIM_MAPPING", raising=False) - fresh_stub_vectordb.add_user( - user_id=101, - email="same@example.com", - external_user_id="sub-stable", - display_name="Same Name", - ) - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - id_token = _sign_jwt( - _id_token_payload( - nonce, - sub="sub-stable", - email="same@example.com", - extra={"name": "Same Name"}, - ) - ) - _mock_token_endpoint(client.oidc_router, id_token) - - r = client.get(f"/auth/callback?code=c&state={state}", follow_redirects=False) - assert r.status_code == 302, r.text - assert not any(c[0] == "update_user_fields" for c in fresh_stub_vectordb.calls) - - -def test_callback_applies_claim_mapping_from_id_token(client, fresh_stub_vectordb, monkeypatch): - """With OIDC_CLAIM_MAPPING set, claims from the ID token update the user row.""" - monkeypatch.setenv("OIDC_CLAIM_MAPPING", "display_name:name,email:email") - monkeypatch.setenv("OIDC_CLAIM_SOURCE", "id_token") - fresh_stub_vectordb.add_user( - user_id=42, - email="old@example.com", - external_user_id="sub-abc", - display_name="Old Name", - ) - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - id_token = _sign_jwt( - _id_token_payload( - nonce, - sub="sub-abc", - email="dwho@badwolf.org", - extra={"name": "Doctor Who"}, - ) - ) - _mock_token_endpoint(client.oidc_router, id_token) - - r = client.get( - f"/auth/callback?code=c&state={state}", - follow_redirects=False, - ) - assert r.status_code == 302, r.text - user = fresh_stub_vectordb._users_by_id[42] - assert user["display_name"] == "Doctor Who" - # email lowercased by update_user_fields stub - assert user["email"] == "dwho@badwolf.org" - # update_user_fields was called exactly once - assert sum(1 for c in fresh_stub_vectordb.calls if c[0] == "update_user_fields") == 1 - - -def test_callback_applies_claim_mapping_from_userinfo(client, fresh_stub_vectordb, monkeypatch): - """With OIDC_CLAIM_SOURCE=userinfo the claim fetch goes to /userinfo.""" - monkeypatch.setenv("OIDC_CLAIM_MAPPING", "display_name:name,email:email") - monkeypatch.setenv("OIDC_CLAIM_SOURCE", "userinfo") - fresh_stub_vectordb.add_user( - user_id=55, - email=None, - external_user_id="sub-ui", - display_name="legacy", - ) - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - # ID token carries no name/email — the router must pull them from /userinfo. - id_token = _sign_jwt(_id_token_payload(nonce, sub="sub-ui", email=None)) - _mock_token_endpoint(client.oidc_router, id_token) - userinfo_route = client.oidc_router.get(f"{ISSUER}/protocol/openid-connect/userinfo").mock( - return_value=httpx.Response( - 200, - json={"sub": "sub-ui", "name": "UI User", "email": "ui@example.com"}, - ) - ) + assert r.status_code == 302 + assert r.headers["location"] == "/next" + assert stub.calls == [("handle_oidc_callback", "ac", "st", "raw-state")] + cookies = _set_cookies(r) + assert any(SESSION_COOKIE_NAME in c for c in cookies) + # State cookie is cleared (deletion emits a Set-Cookie with Max-Age=0). + assert any(STATE_COOKIE_NAME in c and ("Max-Age=0" in c or "expires=" in c.lower()) for c in cookies) - r = client.get( - f"/auth/callback?code=c&state={state}", - follow_redirects=False, - ) - assert r.status_code == 302, r.text - assert userinfo_route.called - user = fresh_stub_vectordb._users_by_id[55] - assert user["display_name"] == "UI User" - assert user["email"] == "ui@example.com" - - -def test_callback_skips_mapping_when_unset(client, fresh_stub_vectordb, monkeypatch): - """Without OIDC_CLAIM_MAPPING the user row is not touched.""" - monkeypatch.delenv("OIDC_CLAIM_MAPPING", raising=False) - fresh_stub_vectordb.add_user( - user_id=77, - email="tester@example.com", - external_user_id="sub-plain", - display_name="Initial", - ) - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - id_token = _sign_jwt( - _id_token_payload( - nonce, - sub="sub-plain", - email="different@example.com", - extra={"name": "Should Be Ignored"}, - ) - ) - _mock_token_endpoint(client.oidc_router, id_token) - r = client.get( - f"/auth/callback?code=c&state={state}", - follow_redirects=False, - ) - assert r.status_code == 302, r.text - user = fresh_stub_vectordb._users_by_id[77] - # Untouched by the callback when OIDC_CLAIM_MAPPING is empty - assert user["display_name"] == "Initial" - assert user["email"] == "tester@example.com" - # update_user_fields was never called - assert not any(c[0] == "update_user_fields" for c in fresh_stub_vectordb.calls) +def test_callback_flow_error_returns_json_and_clears_state(oidc_env, client, stub): + stub.raise_on["callback"] = OIDCFlowError("bad state", status_code=400) + r = client.get("/auth/callback?code=c&state=s", follow_redirects=False) + assert r.status_code == 400 + assert r.json()["detail"] == "bad state" + assert any(STATE_COOKIE_NAME in c for c in _set_cookies(r)) # --------------------------------------------------------------------------- @@ -779,34 +246,21 @@ def test_callback_skips_mapping_when_unset(client, fresh_stub_vectordb, monkeypa # --------------------------------------------------------------------------- -def test_backchannel_logout_rejects_invalid_token(client): - _setup_discovery(client.oidc_router) - _setup_jwks(client.oidc_router) - r = client.post( - "/auth/backchannel-logout", - data={"logout_token": "not-a-jwt"}, - ) - assert r.status_code == 400 +def test_backchannel_logout_success(oidc_env, client, stub): + r = client.post("/auth/backchannel-logout", data={"logout_token": "lt"}) + assert r.status_code == 200 + assert r.headers["cache-control"] == "no-store" + assert stub.calls == [("handle_backchannel_logout", "lt")] -def test_backchannel_logout_revokes_by_sid(client, fresh_stub_vectordb): - _setup_jwks(client.oidc_router) - _setup_discovery(client.oidc_router) - - # Seed a session to be revoked - fresh_stub_vectordb._sessions[1] = { - "id": 1, - "sid": "sid-target", - "revoked_at": None, - } - token = _sign_jwt(_logout_token_payload(sid="sid-target")) - r = client.post( - "/auth/backchannel-logout", - data={"logout_token": token}, - ) - assert r.status_code == 200 - # The stub increments revoked_at on matching sid - assert fresh_stub_vectordb._sessions[1]["revoked_at"] is not None +def test_backchannel_logout_error_emits_invalid_request(oidc_env, client, stub): + stub.raise_on["bcl"] = OIDCFlowError("nope", error_description="token expired") + r = client.post("/auth/backchannel-logout", data={"logout_token": "lt"}) + assert r.status_code == 400 + body = r.json() + assert body["error"] == "invalid_request" + assert body["error_description"] == "token expired" + assert r.headers["cache-control"] == "no-store" # --------------------------------------------------------------------------- @@ -814,108 +268,21 @@ def test_backchannel_logout_revokes_by_sid(client, fresh_stub_vectordb): # --------------------------------------------------------------------------- -def test_logout_revokes_session_and_deletes_cookie(client, fresh_stub_vectordb): - _setup_discovery(client.oidc_router) - # Seed a session & cookie - session_token = "sess-logout-tok" - fresh_stub_vectordb._sessions[1] = { - "id": 1, - "sid": "sid-1", - "id_token_encrypted": None, # skip decrypt path - "session_expires_at": time.time() + 3600, - "revoked_at": None, - } - fresh_stub_vectordb._sessions_by_token[session_token] = 1 - +def test_logout_redirects_to_idp_and_clears_session(oidc_env, client, stub): r = client.get( "/auth/logout", - cookies={"openrag_session": session_token}, + cookies={SESSION_COOKIE_NAME: "sess"}, follow_redirects=False, ) assert r.status_code == 302 - # Session marked revoked - assert fresh_stub_vectordb._sessions[1]["revoked_at"] is not None - # Cookie cleared in response (max-age=0 or Expires=past) - set_cookie_headers = r.headers.get_list("set-cookie") - assert any("openrag_session=" in h and ("Max-Age=0" in h or "expires=" in h.lower()) for h in set_cookie_headers) + assert r.headers["location"] == stub.logout_target + assert stub.calls == [("logout", "sess")] + assert any(SESSION_COOKIE_NAME in c and ("Max-Age=0" in c or "expires=" in c.lower()) for c in _set_cookies(r)) -def test_logout_rejected_in_token_mode(env_token, fresh_stub_vectordb): - app = FastAPI() - app.include_router(auth_router) - c = TestClient(app) - r = c.get("/auth/logout", follow_redirects=False) - assert r.status_code == 400 - - -# --------------------------------------------------------------------------- -# Skips — scenarios we can add once the full middleware stack is wired (phase 5) -# --------------------------------------------------------------------------- - - -@pytest.mark.skip(reason="requires phase-5 middleware for cookie-based auth on /auth/me") -def test_me_returns_user_info_with_valid_cookie(): - pass - - -# --------------------------------------------------------------------------- -# M2: timezone-consistency regression test -# -# Before the fix, routers/auth.py wrote ``access_token_expires_at`` / -# ``session_expires_at`` via ``datetime.utcnow()`` while every read site -# compared against ``datetime.now()``. On a host whose TZ is east of UTC -# (e.g. Europe/Paris), a newly-issued session thus appeared "already -# expired" by tz_offset hours and ``get_oidc_session_by_token`` returned -# None immediately after the callback. -# --------------------------------------------------------------------------- - - -def test_callback_session_not_prematurely_expired_under_nonutc_tz(client, fresh_stub_vectordb, monkeypatch): - """Callback in a non-UTC timezone must produce an immediately usable session.""" - import os as _os - - # Force a non-UTC timezone for the duration of this test. If the platform - # doesn't support ``time.tzset`` (e.g. Windows CI runners), skip gracefully. - tzset = getattr(time, "tzset", None) - if tzset is None: - pytest.skip("time.tzset not available on this platform; cannot force TZ") - - original_tz = _os.environ.get("TZ") - monkeypatch.setenv("TZ", "Europe/Paris") - tzset() - try: - # Teach the stub to back get_oidc_session_by_token with the same dict we - # created in create_oidc_session (the default stub already does). - fresh_stub_vectordb.add_user(user_id=77, email="tz@example.com", external_user_id="sub-tz") - _setup_jwks(client.oidc_router) - state, nonce = _begin_login_and_extract_state(client) - id_token = _sign_jwt(_id_token_payload(nonce, sub="sub-tz", email="tz@example.com")) - _mock_token_endpoint(client.oidc_router, id_token) - - r = client.get( - f"/auth/callback?code=c&state={state}", - follow_redirects=False, - ) - assert r.status_code == 302, r.text - - # Pull the session cookie value from the response and look it up via - # the stub — this exercises the same staleness comparison the real - # middleware uses at request time. - session_cookie = r.cookies.get("openrag_session") - assert session_cookie, "callback did not set openrag_session cookie" - - fetched = fresh_stub_vectordb._impl_get_oidc_session_by_token(session_cookie) - assert fetched is not None, "Session appeared expired IMMEDIATELY after creation — tz bug (M2)" - - # Additional sanity: session_expires_at must be strictly in the future - # from the perspective of datetime.now() (the read-site clock). - from datetime import datetime as _dt - - session_exp = fetched["session_expires_at"] - assert session_exp > _dt.now(), f"session_expires_at={session_exp} is not in the future vs datetime.now()" - finally: - if original_tz is None: - _os.environ.pop("TZ", None) - else: - _os.environ["TZ"] = original_tz - tzset() +def test_logout_without_idp_target_confirms_in_place(oidc_env, client, stub): + stub.logout_target = None + r = client.get("/auth/logout", follow_redirects=False) + assert r.status_code == 200 + assert r.json()["detail"] == "Logged out" + assert any(SESSION_COOKIE_NAME in c for c in _set_cookies(r)) diff --git a/openrag/routers/test_utils_partition_access.py b/openrag/routers/test_utils_partition_access.py new file mode 100644 index 000000000..8ea1ea40d --- /dev/null +++ b/openrag/routers/test_utils_partition_access.py @@ -0,0 +1,182 @@ +from __future__ import annotations + +import sys +import types + +import pytest +from fastapi import HTTPException + +_STUBBED_MODULES = ( + "utils", + "utils.logger", + "services.workers.bootstrap", + "openai", +) + + +def _install_runtime_stubs() -> dict[str, types.ModuleType | None]: + previous_modules = {name: sys.modules.get(name) for name in _STUBBED_MODULES} + + utils_stub = types.ModuleType("utils") + utils_stub.__path__ = [] + sys.modules["utils"] = utils_stub + + bootstrap_stub = types.ModuleType("services.workers.bootstrap") + bootstrap_stub.get_task_state_manager = lambda: None + sys.modules["services.workers.bootstrap"] = bootstrap_stub + + def _logger(): + logger = types.SimpleNamespace( + debug=lambda *args, **kwargs: None, + info=lambda *args, **kwargs: None, + warning=lambda *args, **kwargs: None, + error=lambda *args, **kwargs: None, + exception=lambda *args, **kwargs: None, + ) + logger.bind = lambda *args, **kwargs: logger + return logger + + logger_stub = types.ModuleType("utils.logger") + logger_stub.escape_markup = lambda s: s.replace("\\", "\\\\").replace("<", "\\<").replace(">", "\\>") + logger_stub.mask_email = ( + lambda email: f"{email.partition('@')[0][0]}***@{email.partition('@')[2]}" + if isinstance(email, str) and "@" in email and email.partition("@")[0] + else "***" + ) + logger_stub.get_logger = _logger + sys.modules["utils.logger"] = logger_stub + + openai_stub = types.ModuleType("openai") + openai_stub.AsyncOpenAI = object + openai_stub.APITimeoutError = TimeoutError + openai_stub.APIConnectionError = ConnectionError + openai_stub.APIError = RuntimeError + sys.modules["openai"] = openai_stub + + return previous_modules + + +def _restore_runtime_stubs(previous_modules: dict[str, types.ModuleType | None]) -> None: + for name, previous_module in previous_modules.items(): + if previous_module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous_module + + +_PREVIOUS_MODULES = _install_runtime_stubs() + +from routers import utils as router_utils # noqa: E402 +from routers.utils import check_user_file_quota, ensure_partition_role, require_task_owner # noqa: E402 +from services.orchestrators.auth_service import AuthService # noqa: E402 + +_restore_runtime_stubs(_PREVIOUS_MODULES) + + +class FakePartitionService: + def __init__(self, existing: set[str]) -> None: + self.existing = existing + self.checked: list[str] = [] + + async def partition_exists(self, partition: str) -> bool: + self.checked.append(partition) + return partition in self.existing + + +class FakeJobService: + def __init__(self, *, details=None, pending_count=0) -> None: + self.details = details + self.pending_count = pending_count + self.detail_checks: list[str] = [] + self.pending_checks: list[int | None] = [] + + async def get_task_details(self, task_id: str): + self.detail_checks.append(task_id) + return self.details + + async def get_user_pending_task_count(self, user_id: int | None) -> int: + self.pending_checks.append(user_id) + return self.pending_count + + +@pytest.mark.asyncio +async def test_ensure_partition_role_allows_unknown_partition_without_membership(): + partition_service = FakePartitionService(existing=set()) + + result = await ensure_partition_role( + partition="new-partition", + user={"id": 1}, + user_partitions=[], + required_role="editor", + auth_service=AuthService, + partition_service=partition_service, + ) + + assert result is True + assert partition_service.checked == ["new-partition"] + + +@pytest.mark.asyncio +async def test_ensure_partition_role_forbids_existing_partition_without_membership(): + partition_service = FakePartitionService(existing={"existing"}) + + with pytest.raises(HTTPException) as exc: + await ensure_partition_role( + partition="existing", + user={"id": 1}, + user_partitions=[], + required_role="viewer", + auth_service=AuthService, + partition_service=partition_service, + ) + + assert exc.value.status_code == 403 + assert exc.value.detail == "Access to partition 'existing' forbidden" + + +@pytest.mark.asyncio +async def test_ensure_partition_role_delegates_membership_role_check_to_auth_service(): + partition_service = FakePartitionService(existing={"p"}) + + with pytest.raises(HTTPException) as exc: + await ensure_partition_role( + partition="p", + user={"id": 1}, + user_partitions=[{"partition": "p", "role": "viewer"}], + required_role="editor", + auth_service=AuthService, + partition_service=partition_service, + ) + + assert exc.value.status_code == 403 + assert exc.value.detail == "Editor role required for partition 'p'" + assert partition_service.checked == [] + + +@pytest.mark.asyncio +async def test_require_task_owner_reads_task_details_through_job_service(): + job_service = FakeJobService(details={"user_id": 7, "filename": "a.pdf"}) + + details = await require_task_owner( + task_id="task-1", + user={"id": 7}, + job_service=job_service, + ) + + assert details == {"user_id": 7, "filename": "a.pdf"} + assert job_service.detail_checks == ["task-1"] + + +@pytest.mark.asyncio +async def test_check_user_file_quota_reads_pending_count_through_job_service(monkeypatch): + monkeypatch.setattr(router_utils, "DEFAULT_FILE_QUOTA", 10) + job_service = FakeJobService(pending_count=2) + + user = await check_user_file_quota( + user={"id": 7, "file_count": 1, "file_quota": 5}, + auth_service=AuthService, + job_service=job_service, + ) + + assert user["id"] == 7 + assert job_service.pending_checks == [7] diff --git a/openrag/routers/tools.py b/openrag/routers/tools.py index de379c473..c2eac0836 100644 --- a/openrag/routers/tools.py +++ b/openrag/routers/tools.py @@ -1,13 +1,24 @@ +"""Tools routes — thin HTTP layer over :class:`ConversionService`. + +Phase 8E: the ``extractText`` serialization moved to +``services.orchestrators.conversion_service.ConversionService`` (the Ray +``DocSerializer`` actor now sits behind the ``FileSerializer`` port). +This module keeps HTTP transport only: the saved-file IO + cleanup, +tool validation/dispatch, and the 4xx/5xx error mapping whose exact +``{"detail": ...}`` body the legacy endpoint returned via +``HTTPException``. +""" + import json from pathlib import Path -import ray -from components.indexer.utils.files import save_file_to_disk, serialize_file -from components.indexer.utils.text_sanitizer import sanitize_extracted_text +from components.indexer.utils.files import save_file_to_disk from config import load_config +from di.providers import get_conversion_service from fastapi import APIRouter, Depends, Form, HTTPException, UploadFile, status from fastapi.responses import JSONResponse from pydantic import BaseModel +from services.orchestrators.conversion_service import ConversionService from utils.logger import get_logger from .utils import ( @@ -92,22 +103,20 @@ async def execute_tool( file: UploadFile = Depends(validate_file_format), tool: str = Depends(validate_tool), metadata: dict = Depends(validate_metadata), + service: ConversionService = Depends(get_conversion_service), ): - save_dir = Path(data_dir) file_path = None try: if tool["name"] == "extractText": - file_path = await save_file_to_disk(file, save_dir, with_random_prefix=True) - metadata.update({"source": str(file_path), "filename": file.filename}) - - task_id = ray.get_runtime_context().get_task_id() + file_path = await save_file_to_disk(file, Path(data_dir), with_random_prefix=True) - logger.debug(f"Execute tool extractText for task {task_id} with file {file.filename}") - doc = await serialize_file(task_id, path=file_path, metadata=metadata) - logger.debug(f"extractText done for task {task_id}") - - # Sanitize the extracted text to remove useless characters and improve quality - sanitized_content = sanitize_extracted_text(doc.page_content) + logger.debug(f"Execute tool extractText with file {file.filename}") + sanitized_content = await service.serialize_file( + file_path=str(file_path), + filename=file.filename, + metadata=metadata, + ) + logger.debug("extractText done") return JSONResponse( status_code=status.HTTP_200_OK, diff --git a/openrag/routers/users.py b/openrag/routers/users.py index 494086c6f..04b56e3d0 100644 --- a/openrag/routers/users.py +++ b/openrag/routers/users.py @@ -1,10 +1,26 @@ +"""User management routes — thin HTTP layer over :class:`UserService`. + +Phase 8A.2: business logic (validation, default-quota rule, existence / +not-found semantics, repo delegation) moved to +``services.orchestrators.user_service.UserService``. This module keeps +HTTP transport only: request-scoped authorization (the shared FastAPI +``Depends`` wrappers in ``routers/utils.py``, retired in a later phase), +the two ``id == 1`` guard rules whose exact ``{"detail": ...}`` body the +legacy endpoints returned via ``HTTPException``, and response shaping. + +``GET /users/info`` stays here unchanged — it computes effective quota +from the ``TaskStateManager`` Ray actor, which orchestrators must not +touch (Phase 8H); it will move to a service once the queue is de-Ray'd. +""" + +from di.providers import get_user_service from fastapi import APIRouter, Depends, HTTPException, Response, status from fastapi.responses import JSONResponse from models.user import UserCreate, UserPublic, UserUpdate -from utils.dependencies import get_task_state_manager, get_vectordb +from services.orchestrators.user_service import UserService from utils.logger import get_logger -from .utils import DEFAULT_FILE_QUOTA, current_user, require_admin, require_admin_or_self +from .utils import current_user, require_admin, require_admin_or_self logger = get_logger() router = APIRouter() @@ -28,9 +44,11 @@ **Note:** User tokens are not included in the response. """, ) -async def list_users(vectordb=Depends(get_vectordb), admin_user=Depends(require_admin)): - users = await vectordb.list_users.remote() - logger.debug("Returned list of users.", user_count=len(users)) +async def list_users( + admin_user=Depends(require_admin), + service: UserService = Depends(get_user_service), +): + users = await service.list_users() return JSONResponse(status_code=status.HTTP_200_OK, content={"users": users}) @@ -59,40 +77,12 @@ async def list_users(vectordb=Depends(get_vectordb), admin_user=Depends(require_ ) async def get_current_user_info( user=Depends(current_user), - task_state_manager=Depends(get_task_state_manager), + service: UserService = Depends(get_user_service), ): """Get current authenticated user info""" - - user_id = user.get("id") - is_admin = user.get("is_admin", False) - - if is_admin: - user_quota = float("inf") - elif DEFAULT_FILE_QUOTA < 0: - user_quota = float("inf") - else: - user_quota = user.get("file_quota", None) - if user_quota is None: - user_quota = DEFAULT_FILE_QUOTA - elif user_quota < 0: - user_quota = float("inf") - - file_count = user.get("file_count", 0) # Get indexed file count from user info - pending_count = await task_state_manager.get_user_pending_task_count.remote( - user_id - ) # Get pending task count from task manager - - total = file_count + pending_count - return JSONResponse( status_code=status.HTTP_200_OK, - content={ - **user, - "file_count": file_count, - "pending_files": pending_count, - "total_files": total, - "file_quota": -1 if user_quota == float("inf") else user_quota, - }, + content=await service.get_current_user_info(user), ) @@ -125,14 +115,11 @@ async def get_current_user_info( ) async def create_user( body: UserCreate, - vectordb=Depends(get_vectordb), admin_user=Depends(require_admin), + service: UserService = Depends(get_user_service), ): - """ - Create a new user and generate a token. - """ - user = await vectordb.create_user.remote(body) - logger.info("Created new user", user_id=user["id"]) + """Create a new user and generate a token.""" + user = await service.create_user(body) return JSONResponse(status_code=status.HTTP_201_CREATED, content=user) @@ -157,11 +144,13 @@ async def create_user( **Note:** User token is not included in the response. """, ) -async def get_user(user_id: int, vectordb=Depends(get_vectordb), admin_user=Depends(require_admin)): - """ - Get details of a specific user (without exposing token). - """ - user = await vectordb.get_user.remote(user_id) +async def get_user( + user_id: int, + admin_user=Depends(require_admin), + service: UserService = Depends(get_user_service), +): + """Get details of a specific user (without exposing token).""" + user = await service.get_user(user_id) return JSONResponse(status_code=status.HTTP_200_OK, content=user) @@ -186,16 +175,18 @@ async def get_user(user_id: int, vectordb=Depends(get_vectordb), admin_user=Depe **Note:** Cannot delete the default admin user (ID: 1). """, ) -async def delete_user(user_id: int, vectordb=Depends(get_vectordb), admin_user=Depends(require_admin)): - """ - Delete a user. - """ +async def delete_user( + user_id: int, + admin_user=Depends(require_admin), + service: UserService = Depends(get_user_service), +): + """Delete a user.""" if user_id == 1: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Cannot delete default admin user.", ) - await vectordb.delete_user.remote(user_id) + await service.delete_user(user_id) return Response(status_code=status.HTTP_204_NO_CONTENT) @@ -225,19 +216,11 @@ async def delete_user(user_id: int, vectordb=Depends(get_vectordb), admin_user=D ) async def regenerate_user_token( user_id: int, - vectordb=Depends(get_vectordb), _auth=Depends(require_admin_or_self), + service: UserService = Depends(get_user_service), ): - """ - Regenerate a user's token. - """ - user = await vectordb.regenerate_user_token.remote(user_id) - if user is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail=f"User '{user_id}' not found", - ) - logger.info("Regenerated user token", user_id=user_id) + """Regenerate a user's token.""" + user = await service.regenerate_token(user_id) return JSONResponse(status_code=status.HTTP_200_OK, content=user) @@ -274,18 +257,14 @@ async def regenerate_user_token( async def update_user( user_id: int, body: UserUpdate, - vectordb=Depends(get_vectordb), admin_user=Depends(require_admin), + service: UserService = Depends(get_user_service), ) -> UserPublic: - """ - Update a user's profile fields. - """ - # Only block if is_admin was explicitly set to False in the request + """Update a user's profile fields.""" + # Only block if is_admin was explicitly set to False in the request. if user_id == 1 and "is_admin" in body.model_fields_set and body.is_admin is False: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Cannot revoke admin privileges from the default admin user.", ) - user = await vectordb.update_user.remote(user_id, body) - logger.info("Updated user info", user_id=user_id) - return user + return await service.update_user(user_id, body) diff --git a/openrag/routers/utils.py b/openrag/routers/utils.py index bfd8a908c..624266e68 100644 --- a/openrag/routers/utils.py +++ b/openrag/routers/utils.py @@ -1,4 +1,3 @@ -import json import os from pathlib import Path from typing import Any @@ -6,15 +5,19 @@ import consts import openai from config import load_config +from core.indexing import validators as core_validators +from core.utils.exceptions import OpenRAGError +from di.providers import get_auth_service, get_job_service, get_partition_service from fastapi import Depends, Form, HTTPException, Request, UploadFile, status from openai import AsyncOpenAI -from utils.dependencies import get_task_state_manager, get_vectordb +from services.orchestrators.auth_service import AuthService +from services.orchestrators.job_service import JobService +from services.orchestrators.partition_service import PartitionService from utils.logger import get_logger # load config config = load_config() logger = get_logger() -task_state_manager = get_task_state_manager() SUPER_ADMIN_MODE = os.getenv("SUPER_ADMIN_MODE", "false").lower() == "true" DATA_DIR = config.paths.data_dir @@ -26,12 +29,6 @@ ACCEPTED_FILE_FORMATS = config.loader.file_loaders.model_dump().keys() DICT_MIMETYPES = config.loader.mimetypes.to_dict() -ROLE_HIERARCHY = { - "viewer": 1, - "editor": 2, - "owner": 3, -} - # File quota per user DEFAULT_FILE_QUOTA = config.rdb.default_file_quota @@ -84,33 +81,36 @@ async def ensure_partition_role( user, user_partitions, required_role: str, + *, + auth_service: AuthService, + partition_service: PartitionService, ): """Ensure the user has at least `required_role` for the partition.""" - # Super-admin bypass - vectordb = get_vectordb() if SUPER_ADMIN_MODE and user.get("is_admin"): return True - # Find membership membership = next((p for p in user_partitions if p["partition"] == partition), None) - if not membership: - # Partition exists but no membership - partition_exists = await vectordb.partition_exists.remote(partition) - if partition_exists: + if await partition_service.partition_exists(partition): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=f"Access to partition '{partition}' forbidden", ) - else: - return True + return True - user_role = membership.get("role") - if ROLE_HIERARCHY[user_role] < ROLE_HIERARCHY[required_role]: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"{required_role.capitalize()} role required for partition '{partition}'", + try: + auth_service.check_partition_access( + user=user, + partition=partition, + user_partitions=user_partitions, + required_role=required_role, + super_admin_mode=SUPER_ADMIN_MODE, ) + except OpenRAGError as exc: + raise HTTPException( + status_code=exc.status_code, + detail=exc.message, + ) from exc return True @@ -119,8 +119,17 @@ async def require_partition_viewer( partition=Depends(request_partition), user=Depends(current_user), user_partitions=Depends(current_user_partitions), + auth_service: AuthService = Depends(get_auth_service), + partition_service: PartitionService = Depends(get_partition_service), ): - await ensure_partition_role(partition, user, user_partitions, "viewer") + await ensure_partition_role( + partition, + user, + user_partitions, + "viewer", + auth_service=auth_service, + partition_service=partition_service, + ) return user @@ -128,8 +137,17 @@ async def require_partition_editor( partition=Depends(request_partition), user=Depends(current_user), user_partitions=Depends(current_user_partitions), + auth_service: AuthService = Depends(get_auth_service), + partition_service: PartitionService = Depends(get_partition_service), ): - await ensure_partition_role(partition, user, user_partitions, "editor") + await ensure_partition_role( + partition, + user, + user_partitions, + "editor", + auth_service=auth_service, + partition_service=partition_service, + ) return user @@ -137,8 +155,17 @@ async def require_partition_owner( partition=Depends(request_partition), user=Depends(current_user), user_partitions=Depends(current_user_partitions), + auth_service: AuthService = Depends(get_auth_service), + partition_service: PartitionService = Depends(get_partition_service), ): - await ensure_partition_role(partition, user, user_partitions, "owner") + await ensure_partition_role( + partition, + user, + user_partitions, + "owner", + auth_service=auth_service, + partition_service=partition_service, + ) return user @@ -146,19 +173,32 @@ async def require_partitions_viewer( partitions=Depends(request_partitions), user=Depends(current_user), user_partitions=Depends(current_user_partitions), + auth_service: AuthService = Depends(get_auth_service), + partition_service: PartitionService = Depends(get_partition_service), ): if SUPER_ADMIN_MODE and user.get("is_admin"): return user if isinstance(partitions, list) and len(partitions) == 1 and partitions[0] == "all": return user for partition in partitions: - await ensure_partition_role(partition, user, user_partitions, "viewer") + await ensure_partition_role( + partition, + user, + user_partitions, + "viewer", + auth_service=auth_service, + partition_service=partition_service, + ) logger.info(f"User has viewer access to partition '{partition}'") return user -async def require_task_owner(task_id=Depends(request_task_id), user=Depends(current_user)): - task_details = await task_state_manager.get_details.remote(task_id) +async def require_task_owner( + task_id=Depends(request_task_id), + user=Depends(current_user), + job_service: JobService = Depends(get_job_service), +): + task_details = await job_service.get_task_details(task_id) if not task_details: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -215,6 +255,8 @@ def require_admin_or_self( async def check_user_file_quota( user=Depends(current_user), + auth_service: AuthService = Depends(get_auth_service), + job_service: JobService = Depends(get_job_service), ): """ Check if user has reached their file quota. @@ -228,92 +270,56 @@ async def check_user_file_quota( - user.file_quota >= 0 → specific limit """ - # Admins have unlimited quota - if user.get("is_admin"): + if user.get("is_admin", False): return user - - if DEFAULT_FILE_QUOTA < 0: # disabled quota checking + if DEFAULT_FILE_QUOTA < 0: return user - - # Determine quota user_quota = user.get("file_quota") - - if user_quota is None: - # Use global quota - user_quota = DEFAULT_FILE_QUOTA - - if user_quota < 0: # unlimited quota + if user_quota is not None and user_quota < 0: return user - # Now user_quota >= 0 - user_id = user.get("id") - indexed_count = user.get("file_count", 0) # Get indexed file count from user info - pending_count = await task_state_manager.get_user_pending_task_count.remote( - user_id - ) # Get pending task count from task manager - - total = indexed_count + pending_count + pending_count = await job_service.get_user_pending_task_count(user_id) logger.debug( "User file quota check", user_id=user_id, - indexed_count=indexed_count, pending_count=pending_count, - user_quota=user_quota, ) - if total >= user_quota: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"File quota exceeded. You have {indexed_count} indexed files and {pending_count} pending tasks. Limit: {user_quota}", + try: + auth_service.validate_file_quota( + user, + pending_task_count=pending_count, + default_quota=DEFAULT_FILE_QUOTA, ) + except OpenRAGError as exc: + raise HTTPException( + status_code=exc.status_code, + detail=exc.message, + ) from exc return user -def is_file_id_valid(file_id: str) -> bool: - return not any(c in file_id for c in FORBIDDEN_CHARS_IN_FILE_ID) - - async def validate_file_id(file_id: str): - if not is_file_id_valid(file_id): - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=f"File ID contains forbidden characters: {', '.join(FORBIDDEN_CHARS_IN_FILE_ID)}", - ) - if not file_id.strip(): - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="File ID cannot be empty.") - return file_id + return core_validators.validate_file_id(file_id, FORBIDDEN_CHARS_IN_FILE_ID) async def validate_metadata(metadata: Any | None = Form(None)): - try: - processed_metadata = metadata or "{}" - processed_metadata = json.loads(processed_metadata) - return processed_metadata - except json.JSONDecodeError: - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid JSON in metadata") + return core_validators.parse_metadata(metadata) async def validate_file_format( file: UploadFile, metadata: dict = Depends(validate_metadata), ): - file_extension = file.filename.split(".")[-1].lower() if "." in file.filename else "" - mimetype = metadata.get("mimetype", None) - - if file_extension not in ACCEPTED_FILE_FORMATS and mimetype not in DICT_MIMETYPES.keys(): - details = ( - f"Unsupported file format: {file_extension} or file mimetype.\n" - f"Supported formats: {', '.join(ACCEPTED_FILE_FORMATS)}\n" - f"Supported mimetypes: {', '.join(DICT_MIMETYPES.keys())}" - ) - raise HTTPException( - status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, - detail=details, - ) - + core_validators.validate_file_format( + filename=file.filename, + accepted_formats=ACCEPTED_FILE_FORMATS, + accepted_mimetypes=DICT_MIMETYPES.keys(), + mimetype=metadata.get("mimetype"), + ) return file @@ -394,9 +400,13 @@ async def check_llm_model_availability(request: Request): ) -async def get_partition_name(model_name, user_partitions, is_admin=False): - vectordb = get_vectordb() - +async def get_partition_name( + model_name, + user_partitions, + *, + partition_service: PartitionService, + is_admin=False, +): partition_prefix = consts.PARTITION_PREFIX if model_name.startswith(consts.LEGACY_PARTITION_PREFIX): # XXX - This is for backward compatibility, but should eventually be removed @@ -408,7 +418,7 @@ async def get_partition_name(model_name, user_partitions, is_admin=False): detail=f"Model not found. Model should respect this format: {consts.PARTITION_PREFIX}partition_name", ) partition = model_name.split(partition_prefix)[1] - if partition != "all" and not await vectordb.partition_exists.remote(partition): + if partition != "all" and not await partition_service.partition_exists(partition): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Partition `{partition}` not found for given model `{model_name}`", diff --git a/openrag/routers/workspaces.py b/openrag/routers/workspaces.py index 829e7d76f..39ade74ff 100644 --- a/openrag/routers/workspaces.py +++ b/openrag/routers/workspaces.py @@ -1,13 +1,21 @@ -"""Workspace management endpoints.""" +"""Workspace management endpoints — thin HTTP layer over WorkspaceService. + +Phase 8B.2: workspace CRUD, file association and the cross-cutting +delete-with-orphan-cleanup moved to +``services.orchestrators.workspace_service.WorkspaceService``. This +module keeps HTTP transport only: request-scoped authorization, request +schema validation, and the guards whose exact non-bracketed +``{"detail": ...}`` body the legacy endpoints returned via +``HTTPException`` (409 duplicate, the workspace-in-partition 404, the +unknown/missing-file 404s, the not-removed 404). +""" -import asyncio import re -from components.ray_utils import call_ray_actor_with_timeout -from config import load_config +from di.providers import get_workspace_service from fastapi import APIRouter, Depends, HTTPException, status from pydantic import BaseModel, ConfigDict, field_validator -from utils.dependencies import get_vectordb +from services.orchestrators.workspace_service import WorkspaceService from utils.logger import get_logger from .utils import require_partition_editor, require_partition_owner, require_partition_viewer @@ -15,9 +23,6 @@ router = APIRouter() logger = get_logger() -_config = load_config() -VECTORDB_TIMEOUT = _config.ray.indexer.vectordb_timeout - _WORKSPACE_ID_RE = re.compile(r"[a-zA-Z0-9_-]+") @@ -43,13 +48,13 @@ class AddFilesRequest(BaseModel): file_ids: list[str] -async def require_workspace_in_partition(partition: str, workspace_id: str, vectordb=Depends(get_vectordb)) -> dict: +async def require_workspace_in_partition( + partition: str, + workspace_id: str, + service: WorkspaceService = Depends(get_workspace_service), +) -> dict: """Validate that a workspace exists and belongs to the given partition.""" - ws = await call_ray_actor_with_timeout( - vectordb.get_workspace.remote(workspace_id), - timeout=VECTORDB_TIMEOUT, - task_description=f"get_workspace({workspace_id})", - ) + ws = await service.get_workspace(workspace_id) if not ws or ws["partition_name"] != partition: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Workspace not found") return ws @@ -63,27 +68,18 @@ async def create_workspace( partition: str, body: CreateWorkspaceRequest, user=Depends(require_partition_editor), - vectordb=Depends(get_vectordb), + service: WorkspaceService = Depends(get_workspace_service), ): - existing = await call_ray_actor_with_timeout( - vectordb.get_workspace.remote(body.workspace_id), - timeout=VECTORDB_TIMEOUT, - task_description=f"get_workspace({body.workspace_id})", - ) - if existing: + if await service.get_workspace(body.workspace_id): raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail=f"Workspace '{body.workspace_id}' already exists.", ) - await call_ray_actor_with_timeout( - vectordb.create_workspace.remote( - workspace_id=body.workspace_id, - partition=partition, - user_id=user["id"], - display_name=body.display_name, - ), - timeout=VECTORDB_TIMEOUT, - task_description=f"create_workspace({body.workspace_id})", + await service.create_workspace( + workspace_id=body.workspace_id, + partition=partition, + user_id=user["id"], + display_name=body.display_name, ) return {"status": "created", "workspace_id": body.workspace_id} @@ -92,13 +88,11 @@ async def create_workspace( "/partition/{partition}/workspaces", dependencies=[Depends(require_partition_viewer)], ) -async def list_workspaces(partition: str, vectordb=Depends(get_vectordb)): - workspaces = await call_ray_actor_with_timeout( - vectordb.list_workspaces.remote(partition), - timeout=VECTORDB_TIMEOUT, - task_description=f"list_workspaces({partition})", - ) - return {"workspaces": workspaces} +async def list_workspaces( + partition: str, + service: WorkspaceService = Depends(get_workspace_service), +): + return {"workspaces": await service.list_workspaces(partition)} @router.get( @@ -114,38 +108,13 @@ async def get_workspace(ws=Depends(require_workspace_in_partition)): dependencies=[Depends(require_partition_owner)], ) async def delete_workspace( - partition: str, workspace_id: str, vectordb=Depends(get_vectordb), _ws=Depends(require_workspace_in_partition) + partition: str, + workspace_id: str, + _ws=Depends(require_workspace_in_partition), + service: WorkspaceService = Depends(get_workspace_service), ): - orphaned = await call_ray_actor_with_timeout( - vectordb.delete_workspace.remote(workspace_id), - timeout=VECTORDB_TIMEOUT, - task_description=f"delete_workspace({workspace_id})", - ) - deleted_count = 0 - failed_file_ids: list[str] = [] - if orphaned: - results = await asyncio.gather( - *[ - call_ray_actor_with_timeout( - vectordb.delete_file.remote(file_id, partition), - timeout=VECTORDB_TIMEOUT, - task_description=f"delete_file({file_id})", - ) - for file_id in orphaned - ], - return_exceptions=True, - ) - for file_id, result in zip(orphaned, results): - if isinstance(result, Exception): - logger.warning("Failed to delete orphaned file from Milvus", file_id=file_id, error=str(result)) - failed_file_ids.append(file_id) - else: - deleted_count += 1 - return { - "status": "deleted", - "orphaned_files_deleted": deleted_count, - "orphaned_files_failed": failed_file_ids, - } + result = await service.delete_workspace(partition, workspace_id) + return {"status": "deleted", **result} @router.post( @@ -156,25 +125,17 @@ async def add_files_to_workspace( partition: str, workspace_id: str, body: AddFilesRequest, - vectordb=Depends(get_vectordb), _ws=Depends(require_workspace_in_partition), + service: WorkspaceService = Depends(get_workspace_service), ): - existing_ids = await call_ray_actor_with_timeout( - vectordb.get_existing_file_ids.remote(partition, body.file_ids), - timeout=VECTORDB_TIMEOUT, - task_description=f"get_existing_file_ids({partition})", - ) + existing_ids = await service.get_existing_file_ids(partition, body.file_ids) unknown_ids = sorted(set(body.file_ids) - set(existing_ids)) if unknown_ids: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"File IDs not found in partition '{partition}': {unknown_ids}", ) - missing = await call_ray_actor_with_timeout( - vectordb.add_files_to_workspace.remote(workspace_id, body.file_ids), - timeout=VECTORDB_TIMEOUT, - task_description=f"add_files_to_workspace({workspace_id})", - ) + missing = await service.add_files(workspace_id, body.file_ids) if missing: # TOCTOU: files were deleted between the pre-check and the insert. raise HTTPException( @@ -189,26 +150,23 @@ async def add_files_to_workspace( dependencies=[Depends(require_partition_viewer)], ) async def list_workspace_files( - workspace_id: str, vectordb=Depends(get_vectordb), _ws=Depends(require_workspace_in_partition) + workspace_id: str, + _ws=Depends(require_workspace_in_partition), + service: WorkspaceService = Depends(get_workspace_service), ): - file_ids = await call_ray_actor_with_timeout( - vectordb.list_workspace_files.remote(workspace_id), - timeout=VECTORDB_TIMEOUT, - task_description=f"list_workspace_files({workspace_id})", - ) - return {"file_ids": file_ids} + return {"file_ids": await service.list_files(workspace_id)} @router.get( "/partition/{partition}/files/{file_id}/workspaces", dependencies=[Depends(require_partition_viewer)], ) -async def list_file_workspaces(partition: str, file_id: str, vectordb=Depends(get_vectordb)): - workspace_ids = await call_ray_actor_with_timeout( - vectordb.get_file_workspaces.remote(file_id, partition), - timeout=VECTORDB_TIMEOUT, - task_description=f"get_file_workspaces({file_id})", - ) +async def list_file_workspaces( + partition: str, + file_id: str, + service: WorkspaceService = Depends(get_workspace_service), +): + workspace_ids = await service.get_file_workspaces(file_id, partition) return {"file_id": file_id, "workspace_ids": workspace_ids} @@ -217,13 +175,12 @@ async def list_file_workspaces(partition: str, file_id: str, vectordb=Depends(ge dependencies=[Depends(require_partition_editor)], ) async def remove_file_from_workspace( - workspace_id: str, file_id: str, vectordb=Depends(get_vectordb), _ws=Depends(require_workspace_in_partition) + workspace_id: str, + file_id: str, + _ws=Depends(require_workspace_in_partition), + service: WorkspaceService = Depends(get_workspace_service), ): - removed = await call_ray_actor_with_timeout( - vectordb.remove_file_from_workspace.remote(workspace_id, file_id), - timeout=VECTORDB_TIMEOUT, - task_description=f"remove_file_from_workspace({workspace_id}, {file_id})", - ) + removed = await service.remove_file(workspace_id, file_id) if not removed: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="File not found in workspace") return {"status": "removed"} diff --git a/openrag/scripts/backup.py b/openrag/scripts/backup.py index 657583de6..cea613eb3 100644 --- a/openrag/scripts/backup.py +++ b/openrag/scripts/backup.py @@ -5,14 +5,38 @@ import sys from typing import IO, Any -from components.indexer.vectordb.utils import PartitionFileManager from pymilvus import Collection, connections +from services.persistence.schema import files as files_table +from services.persistence.schema import partitions as partitions_table +from sqlalchemy import create_engine, select from utils.logger import get_logger +def _list_partitions(conn) -> list[dict]: + rows = conn.execute(select(partitions_table)).all() + return [{"partition": r.partition, "created_at": r.created_at.isoformat()} for r in rows] + + +def _list_partition_files(conn, partition_name: str, limit: int | None = None) -> dict: + q = select(files_table).where(files_table.c.partition_name == partition_name) + if limit is not None: + q = q.limit(limit) + rows = conn.execute(q).all() + file_list = [] + for r in rows: + entry = {"file_id": r.file_id, "partition": r.partition_name} + if r.relationship_id is not None: + entry["relationship_id"] = r.relationship_id + if r.parent_id is not None: + entry["parent_id"] = r.parent_id + entry.update(r.file_metadata or {}) + file_list.append(entry) + return {"files": file_list} + + def dump_rdb_part( out_fh: IO[str], - pfm: PartitionFileManager, + conn, partitions: dict[str, dict[str, Any]], logger: Any, verbose: bool = False, @@ -49,7 +73,7 @@ def dump_rdb_part( ) try: - files = pfm.list_partition_files(part_name) + files = _list_partition_files(conn, part_name) except Exception as e: logger.error(f"Failed while requesting the list of files in partition '{part_name}'\n{e}") raise @@ -223,15 +247,15 @@ def load_openrag_config(logger): logger.info(f"rdb @ {rdb.host}:{rdb.port} | vdb @ {vdb.host}:{vdb.port} | collection: {vdb.collection_name}") # List existing partitions + database_url = ( + f"postgresql://{rdb.user}:{rdb.password}@{rdb.host}:{rdb.port}/partitions_for_collection_{vdb.collection_name}" + ) try: - pfm = PartitionFileManager( - database_url=f"postgresql://{rdb.user}:{rdb.password}@{rdb.host}:{rdb.port}/partitions_for_collection_{vdb.collection_name}", - logger=logger, - ) - - existing_partitions = {item["partition"]: item for item in pfm.list_partitions()} + engine = create_engine(database_url) + with engine.connect() as conn: + existing_partitions = {item["partition"]: item for item in _list_partitions(conn)} except Exception as e: - logger.error(f"Failed while accessing PartitionFileManager at {rdb.host}:{rdb.port}\n{e}") + logger.error(f"Failed while accessing catalog database at {rdb.host}:{rdb.port}\n{e}") raise if args.include_only: @@ -268,8 +292,9 @@ def load_openrag_config(logger): try: with open_output_file(args.output, logger) as out_fh: - # Dump data from RDB (one line per document) - dump_rdb_part(out_fh, pfm, partitions, logger, args.verbose) + with create_engine(database_url).connect() as conn: + # Dump data from RDB (one line per document) + dump_rdb_part(out_fh, conn, partitions, logger, args.verbose) # Dump data from VDB (one line per chunk) dump_vdb_part( diff --git a/openrag/scripts/check_file_counts.py b/openrag/scripts/check_file_counts.py index 59b374d4a..55b047c23 100644 --- a/openrag/scripts/check_file_counts.py +++ b/openrag/scripts/check_file_counts.py @@ -5,12 +5,10 @@ import os import sys -from sqlalchemy import create_engine, func -from sqlalchemy.orm import sessionmaker +from sqlalchemy import create_engine, func, select, update -# Add parent dirs so we can import the models sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) -from components.indexer.vectordb.utils import File, User +from services.persistence.schema import files, users def build_database_url(args): @@ -25,29 +23,27 @@ def build_database_url(args): def check_file_counts(database_url, fix=False): engine = create_engine(database_url) - Session = sessionmaker(bind=engine) - with Session() as session: - # Actual file counts per user from the files table + with engine.connect() as conn: actual_counts = dict( - session.query(File.created_by, func.count(File.id)) - .filter(File.created_by.isnot(None)) - .group_by(File.created_by) - .all() + conn.execute( + select(files.c.created_by, func.count(files.c.id)) + .where(files.c.created_by.isnot(None)) + .group_by(files.c.created_by) + ).all() ) - users = session.query(User).order_by(User.id).all() + users_rows = conn.execute(select(users).order_by(users.c.id)).all() rows = [] has_mismatch = False - for u in users: + for u in users_rows: actual = actual_counts.get(u.id, 0) ok = u.file_count == actual if not ok: has_mismatch = True rows.append((u.id, u.display_name or "", u.file_count, actual, ok)) - # Print table green = "\033[92m" red = "\033[91m" bold = "\033[1m" @@ -58,10 +54,7 @@ def check_file_counts(database_url, fix=False): print("-" * len(header.expandtabs())) for uid, name, stored, actual, ok in rows: - if ok: - status = f"{green}OK{reset}" - else: - status = f"{red}MISMATCH{reset}" + status = f"{green}OK{reset}" if ok else f"{red}MISMATCH{reset}" print(f"{uid:>4} {name:<30} {stored:>8} {actual:>8} {status}") print() @@ -73,9 +66,9 @@ def check_file_counts(database_url, fix=False): if fix: for uid, _, stored, actual, ok in rows: if not ok: - session.query(User).filter(User.id == uid).update({User.file_count: actual}) + conn.execute(update(users).where(users.c.id == uid).values(file_count=actual)) print(f" Fixed user {uid}: {stored} -> {actual}") - session.commit() + conn.commit() print(f"{green}All counts have been fixed.{reset}") else: print(f"Run with {bold}--fix{reset} to correct the values.") @@ -93,8 +86,7 @@ def main(): parser.add_argument("--fix", action="store_true", help="Fix mismatched counts in the database") args = parser.parse_args() - database_url = build_database_url(args) - sys.exit(check_file_counts(database_url, fix=args.fix)) + sys.exit(check_file_counts(build_database_url(args), fix=args.fix)) if __name__ == "__main__": diff --git a/openrag/scripts/restore.py b/openrag/scripts/restore.py index 735aa6f8d..1e13ff889 100644 --- a/openrag/scripts/restore.py +++ b/openrag/scripts/restore.py @@ -5,16 +5,62 @@ import time from typing import IO, Any -import ray -from components.indexer.vectordb import MilvusDB -from components.indexer.vectordb.utils import PartitionFileManager from pymilvus import MilvusClient +from services.persistence.schema import files as files_table +from services.persistence.schema import partition_memberships +from services.persistence.schema import partitions as partitions_table +from services.persistence.schema import users as users_table +from sqlalchemy import create_engine, select +from sqlalchemy.dialects.postgresql import insert as pg_insert from utils.logger import get_logger +def _list_partitions(conn) -> list[dict]: + rows = conn.execute(select(partitions_table)).all() + return [{"partition": r.partition, "created_at": r.created_at.isoformat()} for r in rows] + + +def _add_file_to_partition(conn, file_id: str, partition: str, file_metadata: dict, user_id: int) -> bool: + existing = conn.execute( + select(files_table.c.id).where( + files_table.c.file_id == file_id, + files_table.c.partition_name == partition, + ) + ).first() + if existing: + return False + + partition_exists = conn.execute( + select(partitions_table.c.id).where(partitions_table.c.partition == partition) + ).first() + if not partition_exists: + conn.execute(partitions_table.insert().values(partition=partition)) + conn.execute( + pg_insert(partition_memberships) + .values(partition_name=partition, user_id=user_id, role="owner") + .on_conflict_do_nothing() + ) + + conn.execute( + files_table.insert().values( + file_id=file_id, + partition_name=partition, + file_metadata=file_metadata, + created_by=user_id, + relationship_id=file_metadata.get("relationship_id"), + parent_id=file_metadata.get("parent_id"), + ) + ) + conn.execute( + users_table.update().where(users_table.c.id == user_id).values(file_count=users_table.c.file_count + 1) + ) + conn.commit() + return True + + def read_rdb_section( fh: IO[str], - pfm: PartitionFileManager, + conn, include_only: list[str] | None, added_documents: dict[str, set[str]], existing_partitions: dict[str, Any], @@ -67,10 +113,10 @@ def read_rdb_section( if not dry_run: try: - res = pfm.add_file_to_partition(doc["file_id"], part["name"], doc, user_id) + res = _add_file_to_partition(conn, doc["file_id"], part["name"], doc, user_id) except Exception as e: logger.exception( - f"{type(e)} in add_file_to_partition({doc['file_id']}, {part['name']}, ...)\n" + str(e) + f"{type(e)} in _add_file_to_partition({doc['file_id']}, {part['name']}, ...)\n" + str(e) ) raise else: @@ -252,33 +298,22 @@ def load_openrag_config(logger: Any): logger = get_logger() - try: - # It will create a the Milvus collection if it doesn't exist - vdb_tmp = MilvusDB.options(name="Vectordb", namespace="openrag", lifetime="detached").remote() - - ray.get( - vdb_tmp.__ray_ready__.remote() - ) # ensure the actor is fully initialized and ready: collection and all created if nont existing - print("VectorDB (Milvus) actor fully initialized") - except Exception as e: - logger.exception(f"Failed while trying to create Milvus collection: {e}") - # TODO: stop execution here - rdb, vdb = load_openrag_config(logger) if args.verbose: logger.info(f"rdb @ {rdb.host}:{rdb.port} | vdb @ {vdb.host}:{vdb.port} | collection: {vdb.collection_name}") + database_url = ( + f"postgresql://{rdb.user}:{rdb.password}@{rdb.host}:{rdb.port}/partitions_for_collection_{vdb.collection_name}" + ) + # List existing partitions try: - pfm = PartitionFileManager( - database_url=f"postgresql://{rdb.user}:{rdb.password}@{rdb.host}:{rdb.port}/partitions_for_collection_{vdb.collection_name}", - logger=logger, - ) - - existing_partitions = {item["partition"]: item for item in pfm.list_partitions()} + engine = create_engine(database_url) + with engine.connect() as conn: + existing_partitions = {item["partition"]: item for item in _list_partitions(conn)} except Exception as e: - logger.error(f"Failed while accessing PartitionFileManager at {rdb.host}:{rdb.port}\n{e}") + logger.error(f"Failed while accessing catalog database at {rdb.host}:{rdb.port}\n{e}") raise if args.include_only: @@ -290,7 +325,7 @@ def load_openrag_config(logger: Any): client = MilvusClient(uri=f"http://{vdb.host}:{vdb.port}") try: - with open_backup_file(args.input, logger) as fh: + with open_backup_file(args.input, logger) as fh, create_engine(database_url).connect() as conn: added_documents = {} for line in fh: @@ -299,7 +334,7 @@ def load_openrag_config(logger: Any): if line in ["rdb"]: read_rdb_section( fh, - pfm, + conn, args.include_only, added_documents, existing_partitions, diff --git a/openrag/services/__init__.py b/openrag/services/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/services/auth/__init__.py b/openrag/services/auth/__init__.py new file mode 100644 index 000000000..519dafcb8 --- /dev/null +++ b/openrag/services/auth/__init__.py @@ -0,0 +1,20 @@ +"""Auth service layer - OIDC client, session tokens, state cookie, and deps.""" + +from .deps import get_oidc_client, reset_oidc_client +from .oidc_client import LogoutTokenClaims, OIDCClient, TokenBundle +from .session_tokens import decrypt_token, encrypt_token, hash_session_token, issue_session_token +from .state_cookie import StateCookiePayload, StateCookieSerializer + +__all__ = [ + "OIDCClient", + "TokenBundle", + "LogoutTokenClaims", + "issue_session_token", + "encrypt_token", + "decrypt_token", + "hash_session_token", + "StateCookieSerializer", + "StateCookiePayload", + "get_oidc_client", + "reset_oidc_client", +] diff --git a/openrag/services/auth/deps.py b/openrag/services/auth/deps.py new file mode 100644 index 000000000..b8333c775 --- /dev/null +++ b/openrag/services/auth/deps.py @@ -0,0 +1,85 @@ +"""Lazy, process-local singleton for the OIDCClient. + +Kept in a dedicated module to avoid circular imports between the router +(``openrag/routers/auth.py``) and the application entry point (``openrag/api.py``). + +The OIDC config env vars are resolved here via ``os.getenv`` — the same values +that ``openrag/api.py`` validates at startup. In ``AUTH_MODE=oidc`` mode, these +are guaranteed to be non-empty (api.py refuses to start otherwise), so this +module simply trusts them. +""" + +from __future__ import annotations + +import os +from threading import Lock + +from services.auth.oidc_client import OIDCClient + +_client: OIDCClient | None = None +_lock = Lock() + + +def get_oidc_client() -> OIDCClient: + """Return the shared OIDCClient instance, creating it on first call. + + The instance caches the discovery doc and JWKS, so a single shared client + per worker process is both correct and more efficient than one-per-request. + + Env vars read (all required in AUTH_MODE=oidc): + - OIDC_ENDPOINT + - OIDC_CLIENT_ID + - OIDC_CLIENT_SECRET + - OIDC_REDIRECT_URI + - OIDC_SCOPES (default ``openid email profile offline_access``) + """ + global _client + if _client is not None: + return _client + with _lock: + if _client is not None: + return _client + issuer = os.environ["OIDC_ENDPOINT"] + client_id = os.environ["OIDC_CLIENT_ID"] + client_secret = os.environ["OIDC_CLIENT_SECRET"] + redirect_uri = os.environ["OIDC_REDIRECT_URI"] + scopes = os.getenv("OIDC_SCOPES", "openid email profile offline_access") + _client = OIDCClient( + issuer=issuer, + client_id=client_id, + client_secret=client_secret, + redirect_uri=redirect_uri, + scopes=scopes, + ) + return _client + + +def reset_oidc_client() -> None: + """Test hook — drops the cached client so the next call rebuilds from env. + + Best-effort closes the underlying httpx.AsyncClient to avoid "Unclosed + client session" warnings and leaking connections when tests repeatedly + reset the singleton. If no event loop is running we skip the close call + — the GC will eventually reclaim the socket. + """ + global _client + with _lock: + old = _client + _client = None + if old is None: + return + try: + import asyncio + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.get_event_loop_policy().get_event_loop() + loop.run_until_complete(old.aclose()) + else: + # Schedule close on the running loop without awaiting — caller + # doesn't need to be async. + loop.create_task(old.aclose()) + except Exception: + # Closing is best-effort; never let a reset blow up the caller. + pass diff --git a/openrag/services/auth/oidc_client.py b/openrag/services/auth/oidc_client.py new file mode 100644 index 000000000..30db6850b --- /dev/null +++ b/openrag/services/auth/oidc_client.py @@ -0,0 +1,376 @@ +"""Lightweight OIDC Relying Party client for OpenRAG. + +Wraps Authlib's JWT/JWK primitives with: +- Discovery endpoint caching (1 h TTL) +- JWKS caching with automatic refresh on kid-miss +- PKCE pair generation (S256) +- Authorization URL builder +- Code exchange with ID token verification +- Token refresh (lazy, called by middleware when access_token near expiry) +- Userinfo fetch +- Back-channel logout token verification + +One instance per (issuer, client_id, client_secret) tuple. +The instance is not thread-safe for writes but safe for concurrent reads once +the metadata and JWKS caches are populated. +""" + +import base64 +import hashlib +import secrets +import time +from dataclasses import dataclass +from typing import Any +from urllib.parse import urlencode + +import httpx +from authlib.jose import JsonWebKey, JsonWebToken +from authlib.jose.errors import JoseError + + +@dataclass +class TokenBundle: + """Holds the token set returned by the IdP together with verified ID token claims.""" + + id_token: str + access_token: str + refresh_token: str | None + expires_in: int # seconds + token_type: str # usually "Bearer" + claims: dict[str, Any] # verified claims from id_token + + +@dataclass +class LogoutTokenClaims: + """Verified claims from a back-channel logout token.""" + + iss: str + aud: str | list[str] + sub: str | None + sid: str | None + iat: int + jti: str | None + + +class OIDCClient: + """Lightweight OIDC Relying Party client. + + One instance per (issuer, client_id, client_secret) tuple. + """ + + _DISCOVERY_TTL = 3600 # 1 hour + _JWKS_TTL = 3600 # 1 hour + + def __init__( + self, + *, + issuer: str, + client_id: str, + client_secret: str, + redirect_uri: str, + scopes: str, + http_client: httpx.AsyncClient | None = None, + ): + # Keep the issuer string verbatim (including any trailing "/") — the OIDC + # spec mandates strict byte-for-byte equality between ``self.issuer``, the + # issuer advertised by the discovery document, and the ``iss`` claim in + # tokens. Operators must configure ``OIDC_ENDPOINT`` to match EXACTLY + # what the IdP returns. + self.issuer = issuer + self.client_id = client_id + self.client_secret = client_secret + self.redirect_uri = redirect_uri + self.scopes = scopes + self._http = http_client or httpx.AsyncClient(timeout=10.0) + self._metadata: dict | None = None + self._metadata_fetched_at: float = 0.0 + self._jwks: JsonWebKey | None = None + self._jwks_fetched_at: float = 0.0 + + # ------------------------------------------------------------------ + # Discovery + # ------------------------------------------------------------------ + + async def discover(self) -> dict: + """Fetch and cache the OIDC discovery document. + + Returns the cached document if it is less than _DISCOVERY_TTL seconds old. + Raises ValueError if the returned issuer does not match the configured one. + """ + if self._metadata and (time.time() - self._metadata_fetched_at) < self._DISCOVERY_TTL: + return self._metadata + url = f"{self.issuer.rstrip('/')}/.well-known/openid-configuration" + resp = await self._http.get(url) + resp.raise_for_status() + self._metadata = resp.json() + self._metadata_fetched_at = time.time() + if self._metadata.get("issuer") != self.issuer: + raise ValueError(f"Issuer mismatch: configured {self.issuer!r}, got {self._metadata.get('issuer')!r}") + return self._metadata + + # ------------------------------------------------------------------ + # JWKS + # ------------------------------------------------------------------ + + async def _load_jwks(self, force: bool = False) -> JsonWebKey: + meta = await self.discover() + if not force and self._jwks and (time.time() - self._jwks_fetched_at) < self._JWKS_TTL: + return self._jwks + resp = await self._http.get(meta["jwks_uri"]) + resp.raise_for_status() + self._jwks = JsonWebKey.import_key_set(resp.json()) + self._jwks_fetched_at = time.time() + return self._jwks + + # ------------------------------------------------------------------ + # PKCE helpers + # ------------------------------------------------------------------ + + @staticmethod + def generate_pkce_pair() -> tuple[str, str]: + """Generate a PKCE (code_verifier, code_challenge) pair using S256. + + Returns: + (verifier, challenge) — verifier is 128 url-safe chars, + challenge is the base64url-encoded SHA-256 of the verifier. + """ + verifier = secrets.token_urlsafe(96)[:128] + digest = hashlib.sha256(verifier.encode()).digest() + challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode() + return verifier, challenge + + @staticmethod + def generate_state_and_nonce() -> tuple[str, str]: + """Generate cryptographically random state and nonce values.""" + return secrets.token_urlsafe(32), secrets.token_urlsafe(32) + + # ------------------------------------------------------------------ + # Authorization URL + # ------------------------------------------------------------------ + + async def build_authorization_url(self, *, state: str, nonce: str, code_challenge: str) -> str: + """Build the full authorization URL to redirect the browser to.""" + meta = await self.discover() + params = { + "response_type": "code", + "client_id": self.client_id, + "redirect_uri": self.redirect_uri, + "scope": self.scopes, + "state": state, + "nonce": nonce, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + } + return f"{meta['authorization_endpoint']}?{urlencode(params)}" + + # ------------------------------------------------------------------ + # Code exchange + # ------------------------------------------------------------------ + + async def exchange_code(self, *, code: str, code_verifier: str, expected_nonce: str) -> TokenBundle: + """Exchange an authorization code for tokens. + + Verifies the returned id_token (signature, iss, aud, exp, nonce). + + Args: + code: The authorization code from the IdP callback. + code_verifier: The PKCE verifier corresponding to the challenge sent earlier. + expected_nonce: The nonce value that was sent in the authorization request. + + Returns: + A TokenBundle with verified claims. + """ + meta = await self.discover() + data = { + "grant_type": "authorization_code", + "code": code, + "redirect_uri": self.redirect_uri, + "client_id": self.client_id, + "client_secret": self.client_secret, + "code_verifier": code_verifier, + } + resp = await self._http.post(meta["token_endpoint"], data=data, headers={"Accept": "application/json"}) + resp.raise_for_status() + payload = resp.json() + id_token = payload["id_token"] + claims = await self._verify_id_token(id_token, expected_nonce=expected_nonce) + return TokenBundle( + id_token=id_token, + access_token=payload["access_token"], + refresh_token=payload.get("refresh_token"), + expires_in=int(payload.get("expires_in", 0)), + token_type=payload.get("token_type", "Bearer"), + claims=claims, + ) + + # ------------------------------------------------------------------ + # Token refresh + # ------------------------------------------------------------------ + + async def refresh_access_token(self, refresh_token: str) -> TokenBundle: + """Use the refresh_token to obtain a new access_token. + + If the IdP returns a new id_token, it is re-verified (nonce check skipped + per RFC 8252 §8.2 — nonce is only required during the initial code exchange). + If the IdP omits the refresh_token in the response, the caller's existing + refresh_token is preserved. + + Returns: + A new TokenBundle. + """ + meta = await self.discover() + data = { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": self.client_id, + "client_secret": self.client_secret, + } + resp = await self._http.post(meta["token_endpoint"], data=data, headers={"Accept": "application/json"}) + resp.raise_for_status() + payload = resp.json() + new_id_token = payload.get("id_token") + claims: dict[str, Any] = {} + if new_id_token: + claims = await self._verify_id_token(new_id_token, expected_nonce=None) + return TokenBundle( + id_token=new_id_token or "", + access_token=payload["access_token"], + # Some IdPs omit the refresh_token on rotation — keep the old one. + refresh_token=payload.get("refresh_token", refresh_token), + expires_in=int(payload.get("expires_in", 0)), + token_type=payload.get("token_type", "Bearer"), + claims=claims, + ) + + # ------------------------------------------------------------------ + # Userinfo + # ------------------------------------------------------------------ + + async def fetch_userinfo(self, access_token: str) -> dict: + """Fetch the userinfo endpoint with the given access token.""" + meta = await self.discover() + resp = await self._http.get( + meta["userinfo_endpoint"], + headers={"Authorization": f"Bearer {access_token}"}, + ) + resp.raise_for_status() + return resp.json() + + # ------------------------------------------------------------------ + # ID token verification + # ------------------------------------------------------------------ + + async def _verify_id_token(self, token: str, *, expected_nonce: str | None) -> dict[str, Any]: + """Verify an ID token's signature and standard claims. + + Retries with a fresh JWKS fetch on kid-miss (covers IdP key rotation). + Raises JoseError / ValueError on any validation failure. + """ + jwks = await self._load_jwks() + jwt = JsonWebToken(["RS256", "ES256", "EdDSA", "RS384", "RS512"]) + try: + claims = jwt.decode(token, jwks) + except JoseError: + # Force JWKS refresh in case of kid rotation; retry once. + jwks = await self._load_jwks(force=True) + claims = jwt.decode(token, jwks) + + # Manual validation — avoids authlib version differences around claims.params + decoded: dict[str, Any] = dict(claims) + now = int(time.time()) + + if decoded.get("iss") != self.issuer: + raise ValueError(f"ID token iss mismatch: expected {self.issuer!r}, got {decoded.get('iss')!r}") + + aud = decoded.get("aud") + if isinstance(aud, list): + if self.client_id not in aud: + raise ValueError(f"ID token aud {aud!r} does not contain client_id {self.client_id!r}") + elif aud != self.client_id: + raise ValueError(f"ID token aud {aud!r} != client_id {self.client_id!r}") + + if "exp" not in decoded: + raise ValueError("ID token missing exp claim") + if int(decoded["exp"]) < now: + raise ValueError("ID token has expired") + + if "iat" not in decoded: + raise ValueError("ID token missing iat claim") + + if expected_nonce is not None: + if decoded.get("nonce") != expected_nonce: + raise ValueError("OIDC nonce mismatch") + + return decoded + + # ------------------------------------------------------------------ + # Back-channel logout token verification + # ------------------------------------------------------------------ + + async def verify_logout_token(self, token: str) -> LogoutTokenClaims: + """Verify an OIDC back-channel logout token. + + Validates: + - Signature (with JWKS kid-miss retry) + - Standard claims (iss, aud, iat) + - events claim contains the back-channel-logout URI key + - nonce must NOT be present (spec requirement) + - At least one of sub or sid must be present + + Returns: + LogoutTokenClaims with the verified values. + Raises: + ValueError: on any spec violation. + """ + jwks = await self._load_jwks() + jwt = JsonWebToken(["RS256", "ES256", "EdDSA", "RS384", "RS512"]) + try: + claims = jwt.decode(token, jwks) + except JoseError: + jwks = await self._load_jwks(force=True) + claims = jwt.decode(token, jwks) + + decoded: dict[str, Any] = dict(claims) + now = int(time.time()) + + if decoded.get("iss") != self.issuer: + raise ValueError(f"logout_token iss mismatch: expected {self.issuer!r}, got {decoded.get('iss')!r}") + + aud = decoded.get("aud") + if isinstance(aud, list): + if self.client_id not in aud: + raise ValueError(f"logout_token aud {aud!r} does not contain client_id {self.client_id!r}") + elif aud != self.client_id: + raise ValueError(f"logout_token aud {aud!r} != client_id {self.client_id!r}") + + if "iat" not in decoded: + raise ValueError("logout_token missing iat claim") + if int(decoded.get("exp", now + 1)) < now: + raise ValueError("logout_token has expired") + + events = decoded.get("events") or {} + if "http://schemas.openid.net/event/backchannel-logout" not in events: + raise ValueError("logout_token missing required back-channel-logout event claim") + + if "nonce" in decoded: + raise ValueError("logout_token must not contain nonce") + + if not decoded.get("sub") and not decoded.get("sid"): + raise ValueError("logout_token must contain sub or sid") + + return LogoutTokenClaims( + iss=decoded["iss"], + aud=decoded["aud"], + sub=decoded.get("sub"), + sid=decoded.get("sid"), + iat=int(decoded["iat"]), + jti=decoded.get("jti"), + ) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def aclose(self) -> None: + """Close the underlying HTTP client.""" + await self._http.aclose() diff --git a/openrag/services/auth/refresh.py b/openrag/services/auth/refresh.py new file mode 100644 index 000000000..2f72dfec5 --- /dev/null +++ b/openrag/services/auth/refresh.py @@ -0,0 +1,176 @@ +"""Lazy refresh helper for OIDC access tokens. + +Extracted from ``AuthMiddleware`` (Phase 5) to keep ``api.py`` small and +independently testable. Called per-request when a valid cookie session is +found; a no-op when the access token is still fresh. + +Timezone policy +--------------- +Phase 2 stores all OIDC session timestamps as **naive local time** via +``datetime.now()`` (see ``test_oidc_sessions.py`` and +``PartitionFileManager.get_oidc_session_by_token``). We match that style +everywhere in this module to avoid tz-mismatch bugs when comparing +``access_token_expires_at`` against "now". + +Refresh-token stampede guard (M1) +--------------------------------- +IdPs with refresh_token rotation enabled invalidate the old refresh_token the +first time it is redeemed. Under concurrency, multiple requests can each notice +"my access_token is about to expire" at the same time and race each other to +the token endpoint. The second attempt fails with ``invalid_grant`` and +(without a guard) its session would be revoked mid-flight. + +We mitigate that with two cooperating mechanisms: + +1. A **short-circuit** here: if ``last_refresh_at`` was bumped less than 5 + seconds ago, we assume a sibling request already rotated the tokens, + re-read the row, and reuse those freshly rotated tokens instead of calling + the IdP. +2. A **row-level write lock** in :meth:`PartitionFileManager.update_oidc_session_tokens` + (``SELECT ... FOR UPDATE``) so that only one writer commits at a time on + Postgres. +3. An **error-recovery branch** here: if the IdP does reject our refresh_token + (typically because a sibling raced us and won), we re-read the row once + more and, if the tokens were advanced meanwhile, return the fresh session + rather than giving up. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta +from typing import Any + +from services.auth.deps import get_oidc_client +from services.auth.session_tokens import decrypt_token, encrypt_token +from utils.logger import get_logger + +_REFRESH_BUFFER = timedelta(seconds=60) +_STAMPEDE_WINDOW = timedelta(seconds=5) + +logger = get_logger() + + +def _to_dt(val: Any) -> datetime: + """Coerce a datetime-or-ISO-string into a ``datetime``. + + Ray occasionally ships values across actors in serialised form; accept + either shape so callers never have to care about the transport. + """ + if isinstance(val, datetime): + return val + if isinstance(val, str): + return datetime.fromisoformat(val) + raise TypeError(f"Expected datetime or ISO string, got {type(val).__name__}") + + +async def refresh_session_if_needed( + *, + session: dict[str, Any], + enc_key: str, + auth_service: Any, +) -> dict[str, Any] | None: + """Refresh the IdP access_token if it is within ``_REFRESH_BUFFER`` of expiry. + + Behaviour: + - If the access_token is still valid with the 60s buffer → return ``session`` unchanged. + - Stampede guard: if another request has just refreshed this session + (``last_refresh_at`` within 5s), re-read the row and reuse the fresh + tokens without calling the IdP. + - If near/past expiry AND a ``refresh_token_encrypted`` blob is stored → + call the IdP, persist rotated tokens, return an updated session dict. + - If near/past expiry AND no refresh_token is stored → return ``session`` as-is + when still formally valid, or ``None`` when already expired (caller should + treat as a revoked session). + - If the refresh call raises (typically because a sibling already rotated + the tokens and the IdP now rejects ours) → re-read the row; if a sibling + succeeded, return their fresh session; otherwise ``None``. + + The session dict returned mirrors the DB row shape produced by + ``PartitionFileManager._oidc_session_to_dict``. + """ + now = datetime.now() + access_exp = _to_dt(session["access_token_expires_at"]) + + if access_exp > now + _REFRESH_BUFFER: + return session + + # --- Stampede short-circuit ------------------------------------------- + # If a sibling request just refreshed this same session, re-read the row + # and reuse the freshly rotated tokens. This avoids racing the IdP with a + # refresh_token that the sibling's success has already invalidated. + last_refresh_at = session.get("last_refresh_at") + if last_refresh_at is not None: + try: + last_refresh_at_dt = _to_dt(last_refresh_at) + except TypeError: + last_refresh_at_dt = None + if last_refresh_at_dt is not None and (now - last_refresh_at_dt) < _STAMPEDE_WINDOW: + try: + fresh = await auth_service.get_oidc_session_by_id_for_request(session["id"]) + except Exception as e: + logger.bind(session_id=session.get("id"), error=str(e)).warning( + "Stampede-guard re-read failed; falling through to refresh" + ) + fresh = None + if fresh is not None: + fresh_exp = _to_dt(fresh["access_token_expires_at"]) + if fresh_exp > now + _REFRESH_BUFFER: + return fresh + + refresh_enc = session.get("refresh_token_encrypted") + if not refresh_enc: + # No refresh_token available. + # - If still formally valid (within the 60s buffer window but not yet past exp), + # keep using it. + # - If already expired, caller should treat the session as dead. + return session if access_exp > now else None + + try: + refresh_token = decrypt_token(refresh_enc, enc_key) + client = get_oidc_client() + bundle = await client.refresh_access_token(refresh_token) + except Exception as e: + # Maybe a sibling refreshed between our staleness check and the IdP call + # and the IdP has already invalidated our refresh_token. Re-read the + # row once before giving up: if the tokens were rotated meanwhile, + # treat this as a successful refresh (the sibling's). + logger.bind(session_id=session.get("id"), error=str(e)).warning( + "OIDC refresh_token exchange failed — re-reading session for stampede recovery" + ) + try: + fresh = await auth_service.get_oidc_session_by_id_for_request(session["id"]) + except Exception as re: + logger.bind(session_id=session.get("id"), error=str(re)).error( + "Post-failure re-read of OIDC session failed — invalidating" + ) + return None + if fresh is not None: + fresh_exp = _to_dt(fresh["access_token_expires_at"]) + if fresh_exp > now + _REFRESH_BUFFER: + return fresh + return None + + new_access_exp = now + timedelta(seconds=max(int(bundle.expires_in or 0), 60)) + new_access_enc = encrypt_token(bundle.access_token, enc_key) + new_refresh_enc = encrypt_token(bundle.refresh_token, enc_key) if bundle.refresh_token else refresh_enc + + try: + await auth_service.update_oidc_session_tokens_for_request( + session_id=session["id"], + access_token_encrypted=new_access_enc, + refresh_token_encrypted=new_refresh_enc, + access_token_expires_at=new_access_exp, + ) + except Exception as e: + logger.bind(session_id=session.get("id"), error=str(e)).error( + "Failed to persist refreshed OIDC tokens — invalidating session" + ) + return None + + return { + **session, + "access_token_encrypted": new_access_enc, + "access_token_expires_at": new_access_exp, + "refresh_token_encrypted": new_refresh_enc, + "last_refresh_at": now, + } diff --git a/openrag/services/auth/session_tokens.py b/openrag/services/auth/session_tokens.py new file mode 100644 index 000000000..79cc085e1 --- /dev/null +++ b/openrag/services/auth/session_tokens.py @@ -0,0 +1,64 @@ +"""Session token utilities for OpenRAG OIDC sessions. + +Opaque session tokens are issued at callback and stored hashed (SHA-256) in the DB. +IdP tokens (access_token, refresh_token) are encrypted with Fernet before storage. + +The Fernet key is provided via the OIDC_TOKEN_ENCRYPTION_KEY environment variable. +Generate one with: + python -c 'from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())' +""" + +import hashlib +import secrets + +from cryptography.fernet import Fernet, InvalidToken + + +def issue_session_token() -> tuple[str, str]: + """Generate a new session token. + + Returns: + (plaintext, sha256_hex) — the plaintext is set in the cookie, + the hash is stored in the database. + """ + plain = secrets.token_urlsafe(32) # 43 chars, >= 256 bits entropy + return plain, hash_session_token(plain) + + +def hash_session_token(token: str) -> str: + """Return the SHA-256 hex digest of the session token.""" + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +def _fernet(key: str | bytes) -> Fernet: + try: + return Fernet(key.encode("utf-8") if isinstance(key, str) else key) + except Exception as e: + raise ValueError( + "OIDC_TOKEN_ENCRYPTION_KEY is not a valid Fernet key. " + "Generate one with: python -c 'from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())'" + ) from e + + +def encrypt_token(plaintext: str | None, key: str) -> bytes | None: + """Encrypt a plaintext token string. + + Returns None if plaintext is None (refresh_token may be absent). + """ + if plaintext is None: + return None + return _fernet(key).encrypt(plaintext.encode("utf-8")) + + +def decrypt_token(ciphertext: bytes | None, key: str) -> str | None: + """Decrypt a Fernet-encrypted token. + + Returns None if ciphertext is None. + Raises ValueError on key mismatch or data corruption. + """ + if ciphertext is None: + return None + try: + return _fernet(key).decrypt(ciphertext).decode("utf-8") + except InvalidToken as e: + raise ValueError("Failed to decrypt stored OIDC token — key mismatch or corruption") from e diff --git a/openrag/services/auth/state_cookie.py b/openrag/services/auth/state_cookie.py new file mode 100644 index 000000000..c6e5c566f --- /dev/null +++ b/openrag/services/auth/state_cookie.py @@ -0,0 +1,55 @@ +"""Signed state cookie for OIDC Authorization Code + PKCE flow. + +The cookie transports state/nonce/code_verifier between /auth/login and /auth/callback. +It is signed (not encrypted) using itsdangerous.URLSafeTimedSerializer with HMAC-SHA1. + +The signing key is the OIDC_TOKEN_ENCRYPTION_KEY (a Fernet base64url key, which is +valid arbitrary bytes for HMAC). The consuming code (phase 4 router) will pass the +key to StateCookieSerializer(key). Using the same key for both Fernet encryption and +HMAC signing is safe since itsdangerous derives separate subkeys via HMAC. + +TTL defaults to 600 s (10 minutes) — long enough for a slow user at the IdP login page. +""" + +from dataclasses import asdict, dataclass + +from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer + + +@dataclass +class StateCookiePayload: + state: str + nonce: str + code_verifier: str + next_url: str = "/" + + +class StateCookieSerializer: + """Signs/verifies the short-lived cookie holding OIDC state/nonce/code_verifier. + + TTL defaults to 600 s (10 minutes) — long enough for a slow user at the IdP. + """ + + COOKIE_NAME = "openrag_oidc_state" + DEFAULT_TTL_SECONDS = 600 + + def __init__(self, secret_key: str, salt: str = "openrag-oidc-state-v1"): + self._serializer = URLSafeTimedSerializer(secret_key, salt=salt) + + def dumps(self, payload: StateCookiePayload) -> str: + """Serialize and sign the payload, returning an opaque cookie value.""" + return self._serializer.dumps(asdict(payload)) + + def loads(self, token: str, max_age: int = DEFAULT_TTL_SECONDS) -> StateCookiePayload: + """Verify and deserialize the cookie value. + + Raises: + ValueError: if the cookie is expired or the signature is invalid. + """ + try: + data = self._serializer.loads(token, max_age=max_age) + except SignatureExpired as e: + raise ValueError("OIDC state cookie expired") from e + except BadSignature as e: + raise ValueError("OIDC state cookie signature invalid") from e + return StateCookiePayload(**data) diff --git a/openrag/services/events/__init__.py b/openrag/services/events/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/services/inference/__init__.py b/openrag/services/inference/__init__.py new file mode 100644 index 000000000..dfb018f30 --- /dev/null +++ b/openrag/services/inference/__init__.py @@ -0,0 +1,28 @@ +"""Inference service layer — clients, resilience, and concurrency primitives. + +Importing this package registers implementations in the core registries +(``llm_registry``, ``embedder_registry``, ``vlm_registry``, ``reranker_registry``) +so they can be created via ``registry.create("name", **kwargs)``. +""" + +from ._circuit_breaker import get_breaker, with_circuit_breaker +from ._retry import with_retry +from .distributed_semaphore import DistributedSemaphore, DistributedSemaphoreActor +from .ollama_client import OllamaClient, OllamaEmbedder +from .reranker_clients import InfinityReranker, OpenAIReranker +from .vllm_client import VLLMClient, VLLMEmbedder, VLLMVision + +__all__ = [ + "DistributedSemaphore", + "DistributedSemaphoreActor", + "InfinityReranker", + "OllamaClient", + "OllamaEmbedder", + "OpenAIReranker", + "VLLMClient", + "VLLMEmbedder", + "VLLMVision", + "get_breaker", + "with_circuit_breaker", + "with_retry", +] diff --git a/openrag/services/inference/_circuit_breaker.py b/openrag/services/inference/_circuit_breaker.py new file mode 100644 index 000000000..2d332699c --- /dev/null +++ b/openrag/services/inference/_circuit_breaker.py @@ -0,0 +1,85 @@ +from datetime import timedelta +from functools import wraps + +import httpx +from aiobreaker import CircuitBreaker, CircuitBreakerError, CircuitBreakerListener +from core.utils.exceptions import InferenceConnectionError, LLMParsingError, OpenRAGError +from prometheus_client import Gauge +from utils.logger import get_logger + +logger = get_logger() + +_breakers: dict[str, CircuitBreaker] = {} +_breaker_config: dict[str, tuple[int, float]] = {} + +try: + CIRCUIT_BREAKER_STATE = Gauge( + "openrag_circuit_breaker_state", + "Circuit breaker state (0=closed, 1=open, 2=half-open)", + ["name"], + ) +except ValueError: + from prometheus_client import REGISTRY + + CIRCUIT_BREAKER_STATE = REGISTRY._names_to_collectors["openrag_circuit_breaker_state"] + +_STATE_VALUES = {"ClosedState": 0, "OpenState": 1, "HalfOpenState": 2} + + +def _is_client_error(exc: Exception) -> bool: + if isinstance(exc, httpx.HTTPStatusError): + return 400 <= exc.response.status_code < 500 + if isinstance(exc, OpenRAGError): + return 400 <= exc.status_code < 500 + return False + + +def _is_excluded(exc: Exception) -> bool: + if _is_client_error(exc): + return True + if isinstance(exc, LLMParsingError): + return True + return False + + +class _LoggingListener(CircuitBreakerListener): + def state_change(self, breaker, old, new): + state_name = type(new).__name__ + logger.warning( + "Circuit breaker '{name}' state: {old} -> {new}", + name=breaker.name, + old=type(old).__name__, + new=state_name, + ) + CIRCUIT_BREAKER_STATE.labels(name=breaker.name).set(_STATE_VALUES.get(state_name, -1)) + + +def get_breaker(name: str, fail_max: int = 50, timeout_duration: float = 60.0) -> CircuitBreaker: + requested = (fail_max, timeout_duration) + if name not in _breakers: + _breakers[name] = CircuitBreaker( + fail_max=fail_max, + timeout_duration=timedelta(seconds=timeout_duration), + name=name, + exclude=[_is_excluded], + listeners=[_LoggingListener()], + ) + _breaker_config[name] = requested + elif _breaker_config.get(name) != requested: + raise ValueError(f"Breaker '{name}' already exists with config={_breaker_config[name]}, requested={requested}") + return _breakers[name] + + +def with_circuit_breaker(name: str, fail_max: int = 50, timeout_duration: float = 60.0): + def decorator(fn): + @wraps(fn) + async def wrapper(*args, **kwargs): + breaker = get_breaker(name, fail_max, timeout_duration) + try: + return await breaker.call_async(fn, *args, **kwargs) + except CircuitBreakerError: + raise InferenceConnectionError(f"Circuit open for '{name}'") + + return wrapper + + return decorator diff --git a/openrag/services/inference/_retry.py b/openrag/services/inference/_retry.py new file mode 100644 index 000000000..f74ff6bc1 --- /dev/null +++ b/openrag/services/inference/_retry.py @@ -0,0 +1,44 @@ +import httpx +from core.utils.exceptions import OpenRAGError +from tenacity import ( + RetryCallState, + retry, + retry_if_exception, + stop_after_attempt, + wait_exponential_jitter, +) +from utils.logger import get_logger + +logger = get_logger() + +_RETRYABLE_STATUS_CODES = {429, 502, 503, 504} + + +def _is_retryable(exc: BaseException) -> bool: + if isinstance(exc, (httpx.TimeoutException, httpx.ConnectError)): + return True + if isinstance(exc, httpx.HTTPStatusError): + return exc.response.status_code in _RETRYABLE_STATUS_CODES + if isinstance(exc, OpenRAGError): + return exc.status_code in _RETRYABLE_STATUS_CODES + return False + + +def _log_retry(state: RetryCallState) -> None: + exc = state.outcome.exception() if state.outcome else None + logger.warning( + "Retrying after transient failure (attempt {attempt}/{max}): {exc}", + attempt=state.attempt_number, + max=state.retry_object.stop.max_attempt_number, # type: ignore[union-attr] + exc=repr(exc), + ) + + +def with_retry(max_attempts: int = 3, base_wait: float = 1.0, max_wait: float = 30.0): + return retry( + stop=stop_after_attempt(max_attempts), + wait=wait_exponential_jitter(initial=base_wait, max=max_wait, exp_base=2), + retry=retry_if_exception(_is_retryable), + before_sleep=_log_retry, + reraise=True, + ) diff --git a/openrag/services/inference/distributed_semaphore.py b/openrag/services/inference/distributed_semaphore.py new file mode 100644 index 000000000..c535e2bda --- /dev/null +++ b/openrag/services/inference/distributed_semaphore.py @@ -0,0 +1,61 @@ +"""Ray-based distributed semaphore for cluster-wide concurrency limiting. + +Extracted from ``components/utils.py``. The actor handles acquire/release; +``DistributedSemaphore`` locates (or creates) the actor and wraps it as an +async context manager. +""" + +from __future__ import annotations + +import asyncio + +import ray + + +@ray.remote(max_restarts=5) +class DistributedSemaphoreActor: + def __init__(self, max_concurrent_ops: int): + self.semaphore = asyncio.Semaphore(max_concurrent_ops) + + async def acquire(self): + await self.semaphore.acquire() + + def release(self): + self.semaphore.release() + + +class DistributedSemaphore: + """Async context manager backed by a detached Ray actor. + + The actor is created on first use (get-or-create) and survives across + callers within the same Ray cluster. + """ + + def __init__( + self, + name: str = "llmSemaphore", + namespace: str = "openrag", + max_concurrent_ops: int = 10, + ): + self._name = name + self._namespace = namespace + self._max_concurrent_ops = max_concurrent_ops + + def _get_or_create_actor(self): + try: + return ray.get_actor(self._name, namespace=self._namespace) + except ValueError: + return DistributedSemaphoreActor.options( + name=self._name, + namespace=self._namespace, + lifetime="detached", + ).remote(self._max_concurrent_ops) + + async def __aenter__(self): + semaphore_actor = self._get_or_create_actor() + await semaphore_actor.acquire.remote() + return self + + async def __aexit__(self, exc_type, exc, tb): + semaphore_actor = self._get_or_create_actor() + await semaphore_actor.release.remote() diff --git a/openrag/services/inference/healthcheck.py b/openrag/services/inference/healthcheck.py new file mode 100644 index 000000000..7cf259006 --- /dev/null +++ b/openrag/services/inference/healthcheck.py @@ -0,0 +1,90 @@ +"""Probe inference endpoints for readiness. + +Used at container startup (fail-fast) and by the ``/health_check`` route. +Uses raw ``httpx`` — no OpenAI SDK dependency. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from enum import Enum + +import httpx +from utils.logger import get_logger + +logger = get_logger() + + +class EndpointStatus(str, Enum): + HEALTHY = "healthy" + UNHEALTHY = "unhealthy" + UNREACHABLE = "unreachable" + + +@dataclass +class HealthResult: + url: str + status: EndpointStatus + latency_ms: float = 0.0 + models: list[str] = field(default_factory=list) + http_status: int | None = None + error: str | None = None + + +async def check_endpoint_health(endpoint: str, *, timeout: float = 5.0) -> HealthResult: + """Probe a vLLM / OpenAI-compatible server via ``GET /v1/models``.""" + url = endpoint.rstrip("/") + start = time.monotonic() + try: + async with httpx.AsyncClient(timeout=timeout) as client: + resp = await client.get(f"{url}/v1/models") + latency = (time.monotonic() - start) * 1000 + if resp.status_code == 200: + models = [m["id"] for m in resp.json().get("data", [])] + return HealthResult(url=url, status=EndpointStatus.HEALTHY, latency_ms=latency, models=models) + return HealthResult(url=url, status=EndpointStatus.UNHEALTHY, latency_ms=latency, http_status=resp.status_code) + except (httpx.ConnectError, httpx.TimeoutException) as exc: + latency = (time.monotonic() - start) * 1000 + return HealthResult(url=url, status=EndpointStatus.UNREACHABLE, latency_ms=latency, error=str(exc)) + except Exception as exc: + latency = (time.monotonic() - start) * 1000 + logger.warning("Unexpected error probing endpoint", url=url, error=str(exc)) + return HealthResult(url=url, status=EndpointStatus.UNREACHABLE, latency_ms=latency, error=str(exc)) + + +async def check_infinity(endpoint: str, *, timeout: float = 5.0) -> HealthResult: + """Probe an Infinity reranker server via ``GET /health``.""" + url = endpoint.rstrip("/") + start = time.monotonic() + try: + async with httpx.AsyncClient(timeout=timeout) as client: + resp = await client.get(f"{url}/health") + latency = (time.monotonic() - start) * 1000 + if resp.status_code == 200: + return HealthResult(url=url, status=EndpointStatus.HEALTHY, latency_ms=latency) + return HealthResult(url=url, status=EndpointStatus.UNHEALTHY, latency_ms=latency, http_status=resp.status_code) + except (httpx.ConnectError, httpx.TimeoutException) as exc: + latency = (time.monotonic() - start) * 1000 + return HealthResult(url=url, status=EndpointStatus.UNREACHABLE, latency_ms=latency, error=str(exc)) + except Exception as exc: + latency = (time.monotonic() - start) * 1000 + logger.warning("Unexpected error probing infinity endpoint", url=url, error=str(exc)) + return HealthResult(url=url, status=EndpointStatus.UNREACHABLE, latency_ms=latency, error=str(exc)) + + +async def check_model_available(endpoint: str, model: str, *, timeout: float = 5.0) -> HealthResult: + """Probe an OpenAI-compatible endpoint and verify a specific model is served.""" + result = await check_endpoint_health(endpoint, timeout=timeout) + if result.status != EndpointStatus.HEALTHY: + return result + if model not in result.models: + available = ", ".join(result.models) if result.models else "(none)" + return HealthResult( + url=result.url, + status=EndpointStatus.UNHEALTHY, + latency_ms=result.latency_ms, + models=result.models, + error=f"Model '{model}' not found. Available: {available}", + ) + return result diff --git a/openrag/services/inference/ollama_client.py b/openrag/services/inference/ollama_client.py new file mode 100644 index 000000000..b200546b5 --- /dev/null +++ b/openrag/services/inference/ollama_client.py @@ -0,0 +1,220 @@ +"""Ollama inference clients. + +Ollama exposes an OpenAI-compatible ``/v1`` API since v0.1.24, so these +clients are thin wrappers over the vLLM clients with Ollama-specific defaults +and without vLLM-only fields (``truncate_prompt_tokens``). + +* ``OllamaClient`` → ``LLM`` (chat completions via /v1/chat/completions) +* ``OllamaEmbedder`` → ``Embedder`` (embeddings via /v1/embeddings) +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator + +import httpx +from core.embeddings import Embedder, embedder_registry +from core.llm import LLM, llm_registry +from core.utils.exceptions import ( + EmbeddingAPIError, + EmbeddingResponseError, + InferenceConnectionError, + InferenceError, + InferenceTimeoutError, +) +from utils.logger import get_logger + +from ._circuit_breaker import with_circuit_breaker +from ._retry import with_retry +from .vllm_client import _parse_response + +logger = get_logger() +_ERROR_SNIPPET_LIMIT = 500 + + +def _error_snippet(text: str) -> str: + snippet = " ".join(text.split())[:_ERROR_SNIPPET_LIMIT] + if len(text) > _ERROR_SNIPPET_LIMIT: + return f"{snippet}...(truncated)" + return snippet + + +# --------------------------------------------------------------------------- +# LLM +# --------------------------------------------------------------------------- + + +@llm_registry.register("ollama") +class OllamaClient(LLM): + """Ollama LLM client using the OpenAI-compatible /v1 API. + + *endpoint* should point to the Ollama server root or include the ``/v1`` + prefix, e.g. ``http://localhost:11434/v1``. + """ + + def __init__( + self, + endpoint: str, + model_name: str, + *, + timeout: float = 240.0, + **kwargs, + ) -> None: + self._endpoint = endpoint.rstrip("/") + if not self._endpoint.endswith("/v1"): + self._endpoint = f"{self._endpoint}/v1" + self._model = model_name + self._defaults: dict = kwargs + self._client = httpx.AsyncClient( + timeout=timeout, + headers={"Content-Type": "application/json"}, + ) + + @with_circuit_breaker("llm") + @with_retry(max_attempts=3) + async def generate(self, prompt: str, **kwargs) -> dict: + payload = {**self._defaults, **kwargs, "model": self._model, "prompt": prompt} + payload.pop("metadata", None) + try: + resp = await self._client.post(f"{self._endpoint}/completions", json=payload) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach Ollama at {self._endpoint}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"Ollama request timed out at {self._endpoint}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceError( + f"Ollama error ({exc.response.status_code}): {exc.response.text[:500]}", + status_code=exc.response.status_code, + ) from exc + return _parse_response(resp) + + @with_circuit_breaker("llm") + @with_retry(max_attempts=3) + async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: + payload = {**self._defaults, **kwargs, "model": self._model, "messages": messages, "stream": False} + payload.pop("metadata", None) + try: + resp = await self._client.post(f"{self._endpoint}/chat/completions", json=payload) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach Ollama at {self._endpoint}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"Ollama request timed out at {self._endpoint}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceError( + f"Ollama error ({exc.response.status_code}): {exc.response.text[:500]}", + status_code=exc.response.status_code, + ) from exc + return _parse_response(resp) + + async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIterator[str]: + payload = {**self._defaults, **kwargs, "model": self._model, "messages": messages, "stream": True} + payload.pop("metadata", None) + try: + async with self._client.stream("POST", f"{self._endpoint}/chat/completions", json=payload) as resp: + if resp.status_code >= 400: + await resp.aread() + raise InferenceError( + f"Ollama streaming error ({resp.status_code}): {resp.text[:500]}", + status_code=resp.status_code, + ) + async for line in resp.aiter_lines(): + yield line + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach Ollama at {self._endpoint}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"Ollama streaming request timed out at {self._endpoint}") from exc + + async def aclose(self) -> None: + await self._client.aclose() + + +# --------------------------------------------------------------------------- +# Embedder +# --------------------------------------------------------------------------- + + +@embedder_registry.register("ollama") +class OllamaEmbedder(Embedder): + """Ollama embedding client using the OpenAI-compatible /v1/embeddings API.""" + + def __init__( + self, + endpoint: str, + model_name: str, + *, + dimension: int | None = None, + timeout: float = 60.0, + **_kwargs, + ) -> None: + self._endpoint = endpoint.rstrip("/") + if not self._endpoint.endswith("/v1"): + self._endpoint = f"{self._endpoint}/v1" + self._model = model_name + self._dimension: int | None = dimension + self._client = httpx.AsyncClient(timeout=timeout) + + @with_circuit_breaker("embedder") + @with_retry(max_attempts=3) + async def embed(self, texts: list[str]) -> list[list[float]]: + body = {"model": self._model, "input": texts} + try: + resp = await self._client.post(f"{self._endpoint}/embeddings", json=body) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise EmbeddingAPIError( + f"Cannot reach Ollama embedder at {self._endpoint}", + model_name=self._model, + base_url=self._endpoint, + error=str(exc), + ) from exc + except httpx.TimeoutException as exc: + raise EmbeddingAPIError( + f"Ollama embedder request timed out at {self._endpoint}", + model_name=self._model, + base_url=self._endpoint, + error=str(exc), + ) from exc + except httpx.HTTPStatusError as exc: + raise EmbeddingAPIError( + f"Ollama embedder API error ({exc.response.status_code})", + model_name=self._model, + base_url=self._endpoint, + error=_error_snippet(exc.response.text), + ) from exc + + try: + data = resp.json()["data"] + embeddings = [item["embedding"] for item in sorted(data, key=lambda x: x["index"])] + except (ValueError, KeyError, IndexError, TypeError) as exc: + raise EmbeddingResponseError( + "Unexpected Ollama embedding response format", + model_name=self._model, + base_url=self._endpoint, + error=str(exc), + ) from exc + + if self._dimension is None and embeddings: + self._dimension = len(embeddings[0]) + return embeddings + + async def embed_single(self, text: str) -> list[float]: + result = await self.embed([text]) + if not result: + raise EmbeddingResponseError( + "Empty Ollama embedding response", + model_name=self._model, + base_url=self._endpoint, + error="No vectors returned", + ) + return result[0] + + @property + def dimension(self) -> int: + if self._dimension is None: + raise RuntimeError("Embedding dimension unknown — call embed() first") + return self._dimension + + async def aclose(self) -> None: + await self._client.aclose() diff --git a/openrag/services/inference/parsers/__init__.py b/openrag/services/inference/parsers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/services/inference/parsers/_base_openai_parser.py b/openrag/services/inference/parsers/_base_openai_parser.py new file mode 100644 index 000000000..3b46320e0 --- /dev/null +++ b/openrag/services/inference/parsers/_base_openai_parser.py @@ -0,0 +1,99 @@ +"""Common scaffolding for OpenAI-VLM-backed PDF parsers. + +Provides reusable helpers — PDF rendering, single-page VLM calls under a +semaphore, JSON-fence stripping, picture-bbox cropping — but takes no +opinion on response shape or block layout. Concrete subclasses +implement ``parse()`` and stitch blocks together however suits the +model they target. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from abc import ABC +from io import BytesIO +from typing import Any + +from core.indexing.image_preprocessor import pil_to_png_bytes +from core.indexing.parsers.document_parser import BaseClientParser +from core.models.document import DocumentType +from core.vlm import VLM + +logger = logging.getLogger(__name__) + + +class BaseOpenAIPdfClient(BaseClientParser, ABC): + """OpenAI-compatible VLM-backed PDF parser scaffolding.""" + + def __init__( + self, + vlm: VLM, + *, + scale: float = 1.0, + concurrency_limit: int = 4, + ) -> None: + self._vlm = vlm + self._scale = scale + self._semaphore = asyncio.Semaphore(max(1, concurrency_limit)) + + def supported_types(self) -> list[str]: + return [DocumentType.PDF.value] + + # ----- helpers ----- + + @staticmethod + def _render_pdf_pages(raw_bytes: bytes, scale: float) -> list[Any]: + """Render every PDF page into a PIL Image. Pure-CPU; runs in a thread.""" + import pypdfium2 as pdfium + + pdf = pdfium.PdfDocument(raw_bytes) + try: + return [page.render(scale=scale).to_pil() for page in pdf] + finally: + pdf.close() + + async def _ocr_one(self, page_img: Any, prompt: str) -> str | None: + """Send one page image through the VLM with ``prompt``; return raw text.""" + async with self._semaphore: + try: + png_bytes = pil_to_png_bytes(page_img) + return await self._vlm.caption_image(png_bytes, prompt=prompt) + except Exception as exc: + logger.warning("OpenAI VLM OCR call failed: %s", exc) + return None + + @staticmethod + def _strip_json_fences(raw: str) -> str: + """Strip ```json ... ``` fences and surrounding whitespace from a VLM response.""" + text = raw.strip() + if text.startswith("```"): + text = text.strip("`") + if text.lower().startswith("json"): + text = text[4:].lstrip() + return text + + @staticmethod + def _load_json(raw: str | None) -> Any | None: + """Decode a JSON payload from a VLM response, tolerating fences and whitespace.""" + if not raw: + return None + text = BaseOpenAIPdfClient._strip_json_fences(raw) + try: + return json.loads(text) + except json.JSONDecodeError as exc: + logger.warning("OCR response was not valid JSON: %s", exc) + return None + + @staticmethod + def _crop_to_png_bytes(page_img: Any, bbox: Any) -> bytes | None: + """Crop a region from a PIL page image and return PNG bytes.""" + try: + cropped = page_img.crop(tuple(bbox)) + buf = BytesIO() + cropped.save(buf, format="PNG") + return buf.getvalue() + except Exception as exc: + logger.warning("Failed to crop bbox %s: %s", bbox, exc) + return None diff --git a/openrag/services/inference/parsers/dotsocr.py b/openrag/services/inference/parsers/dotsocr.py new file mode 100644 index 000000000..d65904540 --- /dev/null +++ b/openrag/services/inference/parsers/dotsocr.py @@ -0,0 +1,143 @@ +"""DotsOCR PDF parser — concrete :class:`BaseOpenAIPdfClient` subclass. + +DotsOCR returns a JSON list of layout elements (``Picture``, ``Table``, +``Text``, ``Title`` …) with bounding boxes and text content, sorted by +reading order. + +Block emission: + +- One :class:`TextBlock` per page (1-indexed ``page_number``), holding + every non-``Picture`` element's text joined in reading order. +- One :class:`ImageBlock` per ``Picture`` element, carrying the cropped + PNG bytes. Captioning is left to a downstream stage — the parser + does **not** call the VLM for captions. +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from enum import Enum + +from core.models.document import Document, ImageBlock, ProcessedDocument, TextBlock +from pydantic import BaseModel, RootModel, ValidationError + +from ._base_openai_parser import BaseOpenAIPdfClient + +logger = logging.getLogger(__name__) + + +_DOTSOCR_PROMPT = """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox. + +1. Bbox format: [x1, y1, x2, y2] + +2. Layout Categories: The possible categories are ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title']. + +3. Text Extraction & Formatting Rules: + - Picture: For the 'Picture' category, the text field should be omitted. + - Formula: Format its text as LaTeX. + - Table: Format its text as HTML. + - All Others (Text, Title, etc.): Format their text as Markdown. + +4. Constraints: + - The output text must be the original text from the image, with no translation. + - All layout elements must be sorted according to human reading order. + +5. Final Output: The entire output must be a single JSON object. +""" + + +class DotsOCRCategory(str, Enum): + CAPTION = "Caption" + FOOTNOTE = "Footnote" + FORMULA = "Formula" + LIST_ITEM = "List-item" + PAGE_FOOTER = "Page-footer" + PAGE_HEADER = "Page-header" + PICTURE = "Picture" + SECTION_HEADER = "Section-header" + TABLE = "Table" + TEXT = "Text" + TITLE = "Title" + + +class DotsOCRElement(BaseModel): + """One layout element on a page.""" + + bbox: tuple[float, float, float, float] + category: DotsOCRCategory + text: str = "" + + +class DotsOCRPage(RootModel[list[DotsOCRElement]]): + """One page's DotsOCR output: layout elements in reading order.""" + + def pictures(self) -> list[DotsOCRElement]: + return [e for e in self.root if e.category is DotsOCRCategory.PICTURE] + + def text(self) -> str: + """Join every non-``Picture`` element's text in reading order.""" + return "\n".join( + e.text.strip() for e in self.root if e.category is not DotsOCRCategory.PICTURE and e.text and e.text.strip() + ) + + +class DotsOCRPdfClient(BaseOpenAIPdfClient): + """OpenAI-compatible PDF parser using the DotsOCR layout-aware prompt.""" + + PROMPT: str = _DOTSOCR_PROMPT + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + start = time.time() + try: + page_imgs = await asyncio.to_thread(self._render_pdf_pages, document.raw_bytes, self._scale) + raw_responses = await asyncio.gather(*(self._ocr_one(img, self.PROMPT) for img in page_imgs)) + except Exception: + logger.exception("DotsOCR PDF parse failed (id=%s)", document.id) + raise + + text_blocks: list[TextBlock] = [] + images: list[ImageBlock] = [] + for page_number, (page_img, raw) in enumerate(zip(page_imgs, raw_responses, strict=True), start=1): + page = self._parse_page(raw) + if page is None: + continue + page_text = page.text() + if page_text: + text_blocks.append(TextBlock(text=page_text, page_number=page_number)) + for element in page.pictures(): + png = self._crop_to_png_bytes(page_img, element.bbox) + if png is not None: + images.append(ImageBlock(image_bytes=png, page_number=page_number)) + + logger.info("DotsOCR PDF parsed (id=%s) in %.2fs", document.id, time.time() - start) + + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=dict(document.metadata), + page_count=len(page_imgs), + ) + + @classmethod + def _parse_page(cls, raw: str | None) -> DotsOCRPage | None: + """Validate one page's raw VLM response into a :class:`DotsOCRPage`.""" + payload = cls._load_json(raw) + if payload is None: + return None + # Tolerate ``{"items": [...]}`` envelope as well as a bare list. + if isinstance(payload, dict) and "items" in payload: + payload = payload["items"] + try: + return DotsOCRPage.model_validate(payload) + except ValidationError as exc: + logger.warning("DotsOCR response did not match expected schema: %s", exc) + return None diff --git a/openrag/services/inference/parsers/openai_audio.py b/openrag/services/inference/parsers/openai_audio.py new file mode 100644 index 000000000..54058d0ad --- /dev/null +++ b/openrag/services/inference/parsers/openai_audio.py @@ -0,0 +1,130 @@ +"""OpenAI-compatible audio transcription client. + +Pipeline: + +1. Materialize ``Document.raw_bytes`` to a temporary file via + :meth:`Document.as_temporary_file`. +2. If the file's suffix is in ``direct_upload_suffixes``, send it to the + transcription endpoint as-is. Otherwise, decode through + ``pydub.AudioSegment`` and re-encode as WAV (libsndfile-friendly). +3. Optionally run a caller-provided language detector against the + prepared file (its result is forwarded to the OpenAI ``language`` + param). The detector is a plain async callable so this client stays + free of Ray / model-loader coupling — the wiring layer can plug in a + Whisper actor or any other implementation. +4. Send the file to ``audio.transcriptions.create`` and emit a single + :class:`TextBlock` with the resulting transcript. + +Adapted from the legacy +``components/indexer/loaders/audio/openai.py`` ``AudioTranscriber``; +the new version drops the in-memory ``components.utils`` semaphore (now +per-instance via ``concurrency_limit``) and the embedded WhisperActor +ref-getter (now an injected callable). +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from collections.abc import Awaitable, Callable, Iterable +from pathlib import Path + +from core.indexing.parsers.document_parser import BaseClientParser +from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock +from openai import AsyncOpenAI +from pydub import AudioSegment + +logger = logging.getLogger(__name__) + + +# Suffixes the transcription backend can ingest as-is, avoiding the ~10x +# size inflation from WAV conversion (Scaleway cap: 100 MB; OpenAI: 25 MB). +_DEFAULT_DIRECT_UPLOAD_SUFFIXES: tuple[str, ...] = (".mp3", ".m4a", ".ogg", ".webm", ".wav") + +LanguageDetector = Callable[[Path], Awaitable[str | None]] + + +class OpenAIAudioClient(BaseClientParser): + """OpenAI-compatible audio transcription client.""" + + def __init__( + self, + *, + base_url: str, + api_key: str, + model: str, + timeout: float = 120.0, + direct_upload_suffixes: Iterable[str] = _DEFAULT_DIRECT_UPLOAD_SUFFIXES, + language_detector: LanguageDetector | None = None, + concurrency_limit: int = 1, + ) -> None: + self._client = AsyncOpenAI(base_url=base_url, api_key=api_key, timeout=timeout) + self._model = model + self._direct_upload_suffixes = {s.lower() for s in direct_upload_suffixes} + self._language_detector = language_detector + self._semaphore = asyncio.Semaphore(max(1, concurrency_limit)) + + def supported_types(self) -> list[str]: + return [DocumentType.AUDIO.value, DocumentType.VIDEO.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + start = time.time() + try: + async with document.as_temporary_file() as input_path: + async with self._semaphore: + upload_path, cleanup = await self._prepare_upload(input_path) + try: + language: str | None = None + if self._language_detector is not None: + try: + language = await self._language_detector(upload_path) + except Exception as exc: + logger.warning("Language detection failed: %s", exc) + text = await self._transcribe(upload_path, language=language) + finally: + if cleanup: + await asyncio.to_thread(upload_path.unlink, True) + except Exception: + logger.exception("OpenAI audio transcription failed (id=%s)", document.id) + raise + + logger.info("OpenAI audio transcribed (id=%s) in %.2fs", document.id, time.time() - start) + + text = text.strip() + text_blocks = [TextBlock(text=text, page_number=1)] if text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + metadata=dict(document.metadata), + page_count=1 if text else 0, + ) + + async def _prepare_upload(self, input_path: Path) -> tuple[Path, bool]: + """Return ``(path_to_upload, needs_cleanup)``. + + Files in :attr:`_direct_upload_suffixes` are sent as-is; others + are decoded by ``pydub`` (ffmpeg) and re-exported as WAV next to + the input — the caller unlinks that temporary on the way out. + """ + if input_path.suffix.lower() in self._direct_upload_suffixes: + return input_path, False + + sound = await asyncio.to_thread(AudioSegment.from_file, input_path) + logger.info("Converting audio to WAV (duration=%.1fs)", len(sound) / 1000) + wav_path = input_path.with_suffix(".wav") + await asyncio.to_thread(sound.export, wav_path, format="wav") + return wav_path, True + + async def _transcribe(self, path: Path, *, language: str | None) -> str: + kwargs: dict[str, object] = {"model": self._model, "file": path} + if language: + kwargs["language"] = language + response = await self._client.audio.transcriptions.create(**kwargs) + return response.text or "" diff --git a/openrag/services/inference/parsers/test_openai_audio.py b/openrag/services/inference/parsers/test_openai_audio.py new file mode 100644 index 000000000..e6c665e2c --- /dev/null +++ b/openrag/services/inference/parsers/test_openai_audio.py @@ -0,0 +1,149 @@ +"""Unit tests for :class:`OpenAIAudioClient`. + +``pydub`` is shimmed at import time via ``sys.modules`` so the test +runs on Python 3.13 (where ``audioop`` was dropped from stdlib and +plain ``import pydub`` fails). The mock is good enough for the control +flow we exercise — neither real audio decoding nor a real OpenAI +client is needed. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import pytest + +# ---- shim pydub before importing openai_audio ------------------------------ + +if "pydub" not in sys.modules: + pydub = types.ModuleType("pydub") + pydub.AudioSegment = MagicMock() # type: ignore[attr-defined] + sys.modules["pydub"] = pydub + +from core.models.document import Document, DocumentType # noqa: E402 + +from .openai_audio import OpenAIAudioClient # noqa: E402 + +# ---- shared fixtures ------------------------------------------------------- + + +@pytest.fixture +def mock_openai_client(): + """Build an ``AsyncOpenAI``-shaped mock with an awaitable ``audio.transcriptions.create``.""" + fake = MagicMock() + fake.audio = MagicMock() + fake.audio.transcriptions = MagicMock() + fake.audio.transcriptions.create = AsyncMock() + return fake + + +def _client(mock_openai_client, **overrides) -> OpenAIAudioClient: + defaults = {"base_url": "http://x", "api_key": "k", "model": "whisper-mock"} + client = OpenAIAudioClient(**{**defaults, **overrides}) + # Constructor stores config only; swap in our mock before any call. + client._client = mock_openai_client + return client + + +def _audio_doc(raw: bytes = b"audio-bytes", filename: str = "x.mp3") -> Document: + return Document(filename=filename, content_type=DocumentType.AUDIO, raw_bytes=raw) + + +# ---- _prepare_upload ------------------------------------------------------- + + +class TestPrepareUpload: + @pytest.mark.asyncio + async def test_direct_upload_skips_conversion(self, mock_openai_client): + client = _client(mock_openai_client) + path = Path("/tmp/audio.mp3") + upload, cleanup = await client._prepare_upload(path) + assert upload == path + assert cleanup is False + + @pytest.mark.asyncio + async def test_unsupported_suffix_falls_back_to_wav(self, mock_openai_client, monkeypatch): + from services.inference.parsers import openai_audio as mod + + sound = MagicMock() + sound.__len__ = MagicMock(return_value=1500) + sound.export = MagicMock() + from_file = MagicMock(return_value=sound) + monkeypatch.setattr(mod.AudioSegment, "from_file", from_file) + + client = _client(mock_openai_client) + path = Path("/tmp/audio.flac") + upload, cleanup = await client._prepare_upload(path) + + assert upload == path.with_suffix(".wav") + assert cleanup is True + from_file.assert_called_once_with(path) + sound.export.assert_called_once() + assert sound.export.call_args.kwargs == {"format": "wav"} + + +# ---- parse() --------------------------------------------------------------- + + +class TestParse: + @pytest.mark.asyncio + async def test_empty_raw_bytes_returns_empty(self, mock_openai_client): + result = await _client(mock_openai_client).parse(_audio_doc(raw=b"")) + assert result.text_blocks == [] and result.page_count == 0 + + @pytest.mark.asyncio + async def test_returns_text_block_on_success(self, mock_openai_client): + mock_openai_client.audio.transcriptions.create.return_value = MagicMock(text=" hello world ") + result = await _client(mock_openai_client).parse(_audio_doc()) + + assert len(result.text_blocks) == 1 + assert result.text_blocks[0].text == "hello world" + assert result.text_blocks[0].page_number == 1 + assert result.page_count == 1 + mock_openai_client.audio.transcriptions.create.assert_awaited_once() + kwargs = mock_openai_client.audio.transcriptions.create.await_args.kwargs + assert kwargs["model"] == "whisper-mock" + assert "language" not in kwargs + + @pytest.mark.asyncio + async def test_empty_transcript_yields_no_text_block(self, mock_openai_client): + mock_openai_client.audio.transcriptions.create.return_value = MagicMock(text=" ") + result = await _client(mock_openai_client).parse(_audio_doc()) + assert result.text_blocks == [] and result.page_count == 0 + + @pytest.mark.asyncio + async def test_language_detector_result_forwarded(self, mock_openai_client): + mock_openai_client.audio.transcriptions.create.return_value = MagicMock(text="bonjour") + detector = AsyncMock(return_value="fr") + result = await _client(mock_openai_client, language_detector=detector).parse(_audio_doc()) + + detector.assert_awaited_once() + kwargs = mock_openai_client.audio.transcriptions.create.await_args.kwargs + assert kwargs["language"] == "fr" + assert result.text_blocks[0].text == "bonjour" + + @pytest.mark.asyncio + async def test_language_detector_failure_is_swallowed(self, mock_openai_client): + mock_openai_client.audio.transcriptions.create.return_value = MagicMock(text="ok") + detector = AsyncMock(side_effect=RuntimeError("detector down")) + result = await _client(mock_openai_client, language_detector=detector).parse(_audio_doc()) + + # Transcription proceeds without ``language`` and the call still succeeds. + kwargs = mock_openai_client.audio.transcriptions.create.await_args.kwargs + assert "language" not in kwargs + assert result.text_blocks[0].text == "ok" + + @pytest.mark.asyncio + async def test_transcribe_exception_propagates(self, mock_openai_client): + mock_openai_client.audio.transcriptions.create.side_effect = RuntimeError("api down") + with pytest.raises(RuntimeError, match="api down"): + await _client(mock_openai_client).parse(_audio_doc()) + + +def test_supported_types(mock_openai_client): + types_ = _client(mock_openai_client).supported_types() + assert DocumentType.AUDIO.value in types_ + assert DocumentType.VIDEO.value in types_ diff --git a/openrag/services/inference/reranker_clients.py b/openrag/services/inference/reranker_clients.py new file mode 100644 index 000000000..2e6366f3e --- /dev/null +++ b/openrag/services/inference/reranker_clients.py @@ -0,0 +1,127 @@ +"""Reranker inference clients. + +Two classes — Infinity and OpenAI-compatible — both implementing the +``Reranker`` ABC. Both talk to a ``/rerank`` endpoint with the same +payload shape, differing only in the base URL and transport library +the old code used. Now both use ``httpx`` directly. +""" + +from __future__ import annotations + +import httpx +from core.rerankers import Reranker, reranker_registry +from core.utils.exceptions import InferenceConnectionError, InferenceTimeoutError +from utils.logger import get_logger + +from ._circuit_breaker import with_circuit_breaker +from ._retry import with_retry + +logger = get_logger() + + +@reranker_registry.register("infinity") +class InfinityReranker(Reranker): + """Reranker backed by an Infinity server.""" + + def __init__( + self, + endpoint: str, + model_name: str, + *, + api_key: str = "", + timeout: float = 30.0, + **_kwargs, + ): + self._endpoint = endpoint.rstrip("/") + self._model = model_name + headers = {} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + self._client = httpx.AsyncClient(timeout=timeout, headers=headers) + + @with_circuit_breaker("reranker") + @with_retry(max_attempts=2) + async def rerank(self, query: str, documents: list[str], top_k: int | None = None) -> list[tuple[int, float]]: + top_k = min(top_k, len(documents)) if top_k is not None else len(documents) + try: + resp = await self._client.post( + f"{self._endpoint}/rerank", + json={ + "model": self._model, + "query": query, + "documents": documents, + "top_n": top_k, + "return_documents": False, + "raw_scores": True, + }, + ) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach reranker at {self._endpoint}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"Reranker request timed out at {self._endpoint}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceConnectionError( + f"Reranker at {self._endpoint} returned HTTP {exc.response.status_code}" + ) from exc + try: + results = resp.json()["results"] + return [(r["index"], r["relevance_score"]) for r in results] + except (KeyError, TypeError, ValueError) as exc: + raise InferenceConnectionError(f"Unexpected reranker response format from {self._endpoint}") from exc + + async def aclose(self) -> None: + await self._client.aclose() + + +@reranker_registry.register("openai") +class OpenAIReranker(Reranker): + """Reranker backed by an OpenAI-compatible reranking endpoint.""" + + def __init__( + self, + endpoint: str, + model_name: str, + *, + api_key: str = "", + timeout: float = 30.0, + **_kwargs, + ): + self._endpoint = endpoint.rstrip("/") + self._model = model_name + headers = {} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + self._client = httpx.AsyncClient(timeout=timeout, headers=headers) + + @with_circuit_breaker("reranker") + @with_retry(max_attempts=2) + async def rerank(self, query: str, documents: list[str], top_k: int | None = None) -> list[tuple[int, float]]: + top_k = min(top_k, len(documents)) if top_k is not None else len(documents) + try: + resp = await self._client.post( + f"{self._endpoint}/rerank", + json={ + "model": self._model, + "query": query, + "documents": documents, + "top_n": top_k, + }, + ) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach reranker at {self._endpoint}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"Reranker request timed out at {self._endpoint}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceConnectionError( + f"Reranker at {self._endpoint} returned HTTP {exc.response.status_code}" + ) from exc + try: + results = resp.json()["results"] + return [(r["index"], r["relevance_score"]) for r in results] + except (KeyError, TypeError, ValueError) as exc: + raise InferenceConnectionError(f"Unexpected reranker response format from {self._endpoint}") from exc + + async def aclose(self) -> None: + await self._client.aclose() diff --git a/openrag/services/inference/test_circuit_breaker.py b/openrag/services/inference/test_circuit_breaker.py new file mode 100644 index 000000000..dedb3a58d --- /dev/null +++ b/openrag/services/inference/test_circuit_breaker.py @@ -0,0 +1,116 @@ +import httpx +import pytest +from core.utils.exceptions import InferenceConnectionError, LLMParsingError +from services.inference._circuit_breaker import ( + _breaker_config, + _breakers, + get_breaker, + with_circuit_breaker, +) + + +@pytest.fixture(autouse=True) +def _clean_breakers(): + for breaker in _breakers.values(): + breaker.close() + _breakers.clear() + _breaker_config.clear() + yield + for breaker in _breakers.values(): + breaker.close() + _breakers.clear() + _breaker_config.clear() + + +class TestGetBreaker: + def test_returns_same_instance(self): + b1 = get_breaker("llm") + b2 = get_breaker("llm") + assert b1 is b2 + + def test_different_names_different_instances(self): + b1 = get_breaker("llm") + b2 = get_breaker("embedder") + assert b1 is not b2 + + def test_default_fail_max_is_50(self): + b = get_breaker("test-default") + assert b.fail_max == 50 + + +class TestExclusions: + @pytest.mark.asyncio + async def test_client_4xx_excluded(self): + breaker = get_breaker("test-4xx", fail_max=2, timeout_duration=1.0) + + async def fail_4xx(): + req = httpx.Request("GET", "http://test") + raise httpx.HTTPStatusError("bad request", request=req, response=httpx.Response(400, request=req)) + + for _ in range(5): + with pytest.raises(httpx.HTTPStatusError): + await breaker.call_async(fail_4xx) + + assert "Closed" in type(breaker.state).__name__ + + @pytest.mark.asyncio + async def test_llm_parsing_error_excluded(self): + breaker = get_breaker("test-parse", fail_max=2, timeout_duration=1.0) + + async def fail_parse(): + raise LLMParsingError(raw_response="not json") + + for _ in range(5): + with pytest.raises(LLMParsingError): + await breaker.call_async(fail_parse) + + assert "Closed" in type(breaker.state).__name__ + + @pytest.mark.asyncio + async def test_server_5xx_trips_breaker(self): + breaker = get_breaker("test-5xx", fail_max=2, timeout_duration=1.0) + + async def fail_5xx(): + req = httpx.Request("GET", "http://test") + raise httpx.HTTPStatusError("bad gateway", request=req, response=httpx.Response(502, request=req)) + + with pytest.raises(httpx.HTTPStatusError): + await breaker.call_async(fail_5xx) + + from aiobreaker import CircuitBreakerError + + with pytest.raises(CircuitBreakerError): + await breaker.call_async(fail_5xx) + + assert "Open" in type(breaker.state).__name__ + + +class TestWithCircuitBreaker: + @pytest.mark.asyncio + async def test_passes_through_on_success(self): + @with_circuit_breaker("test-ok", fail_max=3, timeout_duration=1.0) + async def ok(): + return "result" + + assert await ok() == "result" + + @pytest.mark.asyncio + async def test_raises_inference_connection_error_when_open(self): + call_count = 0 + + @with_circuit_breaker("test-open", fail_max=2, timeout_duration=60.0) + async def always_fail(): + nonlocal call_count + call_count += 1 + raise ConnectionError("down") + + with pytest.raises(ConnectionError): + await always_fail() + + with pytest.raises(InferenceConnectionError, match="Circuit open"): + await always_fail() + + with pytest.raises(InferenceConnectionError, match="Circuit open"): + await always_fail() + + assert call_count == 2 diff --git a/openrag/services/inference/test_distributed_semaphore.py b/openrag/services/inference/test_distributed_semaphore.py new file mode 100644 index 000000000..cc77c50cc --- /dev/null +++ b/openrag/services/inference/test_distributed_semaphore.py @@ -0,0 +1,18 @@ +from services.inference.distributed_semaphore import DistributedSemaphore, DistributedSemaphoreActor + + +class TestDistributedSemaphore: + def test_default_params(self): + sem = DistributedSemaphore() + assert sem._name == "llmSemaphore" + assert sem._namespace == "openrag" + assert sem._max_concurrent_ops == 10 + + def test_custom_params(self): + sem = DistributedSemaphore(name="vlm", namespace="test", max_concurrent_ops=5) + assert sem._name == "vlm" + assert sem._namespace == "test" + assert sem._max_concurrent_ops == 5 + + def test_actor_class_exists(self): + assert hasattr(DistributedSemaphoreActor, "remote") diff --git a/openrag/services/inference/test_healthcheck.py b/openrag/services/inference/test_healthcheck.py new file mode 100644 index 000000000..24ff2054e --- /dev/null +++ b/openrag/services/inference/test_healthcheck.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from .healthcheck import ( + EndpointStatus, + check_endpoint_health, + check_infinity, + check_model_available, +) + + +@pytest.fixture +def models_response(): + return httpx.Response(200, json={"data": [{"id": "mistral-small"}, {"id": "bge-m3"}]}) + + +@pytest.fixture +def health_ok(): + return httpx.Response(200, json={"status": "ok"}) + + +class TestCheckOpenAICompatible: + @pytest.mark.asyncio + async def test_healthy(self, models_response): + transport = httpx.MockTransport(lambda req: models_response) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_endpoint_health("http://vllm:8000") + assert result.status == EndpointStatus.HEALTHY + assert "mistral-small" in result.models + assert "bge-m3" in result.models + assert result.latency_ms > 0 + + @pytest.mark.asyncio + async def test_strips_trailing_slash(self, models_response): + transport = httpx.MockTransport(lambda req: models_response) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_endpoint_health("http://vllm:8000/") + assert result.url == "http://vllm:8000" + + @pytest.mark.asyncio + async def test_unhealthy_status(self): + transport = httpx.MockTransport(lambda req: httpx.Response(503)) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_endpoint_health("http://vllm:8000") + assert result.status == EndpointStatus.UNHEALTHY + assert result.http_status == 503 + + @pytest.mark.asyncio + async def test_connection_error(self): + async def raise_connect_error(*a, **kw): + raise httpx.ConnectError("Connection refused") + + client = AsyncMock() + client.get = raise_connect_error + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=False) + with patch("services.inference.healthcheck.httpx.AsyncClient", return_value=client): + result = await check_endpoint_health("http://vllm:8000") + assert result.status == EndpointStatus.UNREACHABLE + assert "Connection refused" in result.error + + @pytest.mark.asyncio + async def test_timeout(self): + async def raise_timeout(*a, **kw): + raise httpx.TimeoutException("timed out") + + client = AsyncMock() + client.get = raise_timeout + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=False) + with patch("services.inference.healthcheck.httpx.AsyncClient", return_value=client): + result = await check_endpoint_health("http://vllm:8000") + assert result.status == EndpointStatus.UNREACHABLE + assert "timed out" in result.error + + +class TestCheckInfinity: + @pytest.mark.asyncio + async def test_healthy(self, health_ok): + transport = httpx.MockTransport(lambda req: health_ok) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_infinity("http://reranker:7997") + assert result.status == EndpointStatus.HEALTHY + assert result.latency_ms > 0 + + @pytest.mark.asyncio + async def test_unhealthy(self): + transport = httpx.MockTransport(lambda req: httpx.Response(500)) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_infinity("http://reranker:7997") + assert result.status == EndpointStatus.UNHEALTHY + assert result.http_status == 500 + + +class TestCheckModelAvailable: + @pytest.mark.asyncio + async def test_model_found(self, models_response): + transport = httpx.MockTransport(lambda req: models_response) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_model_available("http://vllm:8000", "mistral-small") + assert result.status == EndpointStatus.HEALTHY + + @pytest.mark.asyncio + async def test_model_not_found(self, models_response): + transport = httpx.MockTransport(lambda req: models_response) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_model_available("http://vllm:8000", "nonexistent-model") + assert result.status == EndpointStatus.UNHEALTHY + assert "nonexistent-model" in result.error + assert "mistral-small" in result.error + + @pytest.mark.asyncio + async def test_endpoint_unreachable_skips_model_check(self): + async def raise_connect_error(*a, **kw): + raise httpx.ConnectError("refused") + + client = AsyncMock() + client.get = raise_connect_error + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=False) + with patch("services.inference.healthcheck.httpx.AsyncClient", return_value=client): + result = await check_model_available("http://vllm:8000", "mistral-small") + assert result.status == EndpointStatus.UNREACHABLE diff --git a/openrag/services/inference/test_ollama_client.py b/openrag/services/inference/test_ollama_client.py new file mode 100644 index 000000000..f54f42f60 --- /dev/null +++ b/openrag/services/inference/test_ollama_client.py @@ -0,0 +1,367 @@ +from __future__ import annotations + +import json +from unittest.mock import AsyncMock + +import httpx +import pytest +from core.utils.exceptions import ( + EmbeddingAPIError, + EmbeddingResponseError, + InferenceConnectionError, + InferenceError, + InferenceTimeoutError, +) +from services.inference._circuit_breaker import _breakers + +from .ollama_client import OllamaClient, OllamaEmbedder + + +@pytest.fixture(autouse=True) +def _clean_breakers(): + yield + for breaker in _breakers.values(): + breaker.close() + _breakers.clear() + + +def _make_transport(handler): + return httpx.MockTransport(handler) + + +def _chat_response(content: str = "hello") -> httpx.Response: + return httpx.Response(200, json={"choices": [{"message": {"content": content}}]}) + + +def _completions_response(text: str = "result") -> httpx.Response: + return httpx.Response(200, json={"choices": [{"text": text}]}) + + +def _embed_response(vectors: list[list[float]] | None = None) -> httpx.Response: + vectors = vectors or [[0.1, 0.2, 0.3]] + data = [{"index": i, "embedding": v} for i, v in enumerate(vectors)] + return httpx.Response(200, json={"data": data}) + + +# --------------------------------------------------------------------------- +# OllamaClient (LLM) +# --------------------------------------------------------------------------- + + +class TestOllamaClient: + def _make_client(self, handler, endpoint="http://ollama:11434", **kwargs): + client = OllamaClient(endpoint=endpoint, model_name="llama3", **kwargs) + client._client = httpx.AsyncClient(transport=_make_transport(handler)) + return client + + @pytest.mark.asyncio + async def test_chat_returns_full_response(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert "/v1/chat/completions" in str(request.url) + assert body["model"] == "llama3" + assert body["stream"] is False + return _chat_response("world") + + result = await self._make_client(handler).chat([{"role": "user", "content": "hi"}]) + assert result["choices"][0]["message"]["content"] == "world" + + @pytest.mark.asyncio + async def test_generate_returns_full_response(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert "/v1/completions" in str(request.url) + assert "/chat/" not in str(request.url) + assert body["prompt"] == "say something" + return _completions_response("done") + + result = await self._make_client(handler).generate("say something") + assert result["choices"][0]["text"] == "done" + + @pytest.mark.asyncio + async def test_stream_chat_yields_raw_sse_lines(self): + sse_body = ( + 'data: {"choices":[{"delta":{"content":"Hello"}}]}\n' + 'data: {"choices":[{"delta":{"content":" world"}}]}\n' + "data: [DONE]\n" + ) + + def handler(request: httpx.Request) -> httpx.Response: + assert json.loads(request.content)["stream"] is True + return httpx.Response(200, text=sse_body) + + client = self._make_client(handler) + lines = [line async for line in client.stream_chat([{"role": "user", "content": "hi"}])] + assert 'data: {"choices":[{"delta":{"content":"Hello"}}]}' in lines + assert 'data: {"choices":[{"delta":{"content":" world"}}]}' in lines + + @pytest.mark.asyncio + async def test_stream_chat_error_raises(self): + client = self._make_client(lambda req: httpx.Response(503, text="unavailable")) + with pytest.raises(InferenceError): + async for _ in client.stream_chat([{"role": "user", "content": "hi"}]): + pass + + @pytest.mark.asyncio + async def test_chat_connection_error(self): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + client = OllamaClient(endpoint="http://ollama:11434", model_name="llama3") + client._client = AsyncMock() + client._client.post = fail + with pytest.raises(InferenceConnectionError): + await client.chat([{"role": "user", "content": "hi"}]) + + @pytest.mark.asyncio + async def test_chat_timeout(self): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + client = OllamaClient(endpoint="http://ollama:11434", model_name="llama3") + client._client = AsyncMock() + client._client.post = fail + with pytest.raises(InferenceTimeoutError): + await client.chat([{"role": "user", "content": "hi"}]) + + @pytest.mark.asyncio + async def test_chat_http_error_raises_inference_error(self): + client = self._make_client(lambda req: httpx.Response(500, text="server error")) + with pytest.raises(InferenceError): + await client.chat([{"role": "user", "content": "hi"}]) + + @pytest.mark.asyncio + async def test_defaults_forwarded(self): + captured: dict = {} + + def capture(req: httpx.Request) -> httpx.Response: + captured.update(json.loads(req.content)) + return _chat_response() + + await self._make_client(capture, temperature=0.7).chat([{"role": "user", "content": "hi"}]) + assert captured["temperature"] == 0.7 + + @pytest.mark.asyncio + async def test_per_request_kwargs_override_defaults(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["temperature"] == 0.1 + assert body["max_tokens"] == 256 + return _chat_response() + + await self._make_client(handler, temperature=0.9).chat( + [{"role": "user", "content": "hi"}], temperature=0.1, max_tokens=256 + ) + + @pytest.mark.asyncio + async def test_metadata_stripped_from_payload(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert "metadata" not in body + return _chat_response() + + await self._make_client(handler).chat([{"role": "user", "content": "hi"}], metadata={"llm_override": {}}) + + @pytest.mark.asyncio + async def test_metadata_default_stripped_from_generate_payload(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert "metadata" not in body + return _completions_response() + + await self._make_client(handler, metadata={"llm_override": {}}).generate("hi") + + @pytest.mark.asyncio + async def test_metadata_default_stripped_from_chat_payload(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert "metadata" not in body + return _chat_response() + + await self._make_client(handler, metadata={"llm_override": {}}).chat([{"role": "user", "content": "hi"}]) + + @pytest.mark.asyncio + async def test_metadata_default_stripped_from_stream_payload(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert "metadata" not in body + return httpx.Response(200, text="data: [DONE]\n") + + client = self._make_client(handler, metadata={"llm_override": {}}) + lines = [line async for line in client.stream_chat([{"role": "user", "content": "hi"}])] + assert lines == ["data: [DONE]"] + + def test_endpoint_v1_appended_when_missing(self): + client = OllamaClient(endpoint="http://ollama:11434", model_name="llama3") + assert client._endpoint == "http://ollama:11434/v1" + + def test_endpoint_v1_not_doubled(self): + client = OllamaClient(endpoint="http://ollama:11434/v1", model_name="llama3") + assert client._endpoint == "http://ollama:11434/v1" + + def test_trailing_slash_stripped(self): + client = OllamaClient(endpoint="http://ollama:11434/v1/", model_name="llama3") + assert client._endpoint == "http://ollama:11434/v1" + + def test_no_auth_header_sent_by_default(self): + client = OllamaClient(endpoint="http://ollama:11434", model_name="llama3") + assert "Authorization" not in client._client.headers + + @pytest.mark.asyncio + async def test_aclose(self): + client = OllamaClient(endpoint="http://ollama:11434", model_name="llama3") + client._client = AsyncMock() + await client.aclose() + client._client.aclose.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# OllamaEmbedder +# --------------------------------------------------------------------------- + + +class TestOllamaEmbedder: + def _make_embedder(self, handler, endpoint="http://ollama:11434", **kwargs): + embedder = OllamaEmbedder(endpoint=endpoint, model_name="nomic-embed-text", **kwargs) + embedder._client = httpx.AsyncClient(transport=_make_transport(handler)) + return embedder + + @pytest.mark.asyncio + async def test_embed_returns_sorted_vectors(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["model"] == "nomic-embed-text" + assert body["input"] == ["hello", "world"] + return httpx.Response( + 200, + json={"data": [{"index": 1, "embedding": [0.3, 0.4]}, {"index": 0, "embedding": [0.1, 0.2]}]}, + ) + + result = await self._make_embedder(handler).embed(["hello", "world"]) + assert result == [[0.1, 0.2], [0.3, 0.4]] + + @pytest.mark.asyncio + async def test_embed_single(self): + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.5, 0.6, 0.7]}]}) + + result = await self._make_embedder(handler).embed_single("test") + assert result == [0.5, 0.6, 0.7] + + @pytest.mark.asyncio + async def test_dimension_auto_detected(self): + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.1, 0.2, 0.3]}]}) + + embedder = self._make_embedder(handler) + + with pytest.raises(RuntimeError, match="unknown"): + _ = embedder.dimension + + await embedder.embed(["test"]) + assert embedder.dimension == 3 + + def test_dimension_from_init(self): + assert OllamaEmbedder(endpoint="http://ollama:11434", model_name="m", dimension=768).dimension == 768 + + @pytest.mark.asyncio + async def test_no_truncate_prompt_tokens_in_payload(self): + """Ollama doesn't support truncate_prompt_tokens — must never be sent.""" + + def handler(request: httpx.Request) -> httpx.Response: + assert "truncate_prompt_tokens" not in json.loads(request.content) + return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.1]}]}) + + await self._make_embedder(handler).embed(["test"]) + + @pytest.mark.asyncio + async def test_embed_connection_error(self): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + embedder = OllamaEmbedder(endpoint="http://ollama:11434", model_name="nomic-embed-text") + embedder._client = AsyncMock() + embedder._client.post = fail + with pytest.raises(EmbeddingAPIError): + await embedder.embed(["text"]) + + @pytest.mark.asyncio + async def test_embed_timeout(self): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + embedder = OllamaEmbedder(endpoint="http://ollama:11434", model_name="nomic-embed-text") + embedder._client = AsyncMock() + embedder._client.post = fail + with pytest.raises(EmbeddingAPIError): + await embedder.embed(["text"]) + + @pytest.mark.asyncio + async def test_embed_http_error_raises_embedding_api_error(self): + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response(503, text="service unavailable") + + with pytest.raises(EmbeddingAPIError): + await self._make_embedder(handler).embed(["text"]) + + @pytest.mark.asyncio + async def test_embed_http_error_truncates_response_body(self): + body = "secret-token " + ("x" * 1000) + + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response(503, text=body) + + with pytest.raises(EmbeddingAPIError) as exc_info: + await self._make_embedder(handler).embed(["text"]) + + error = exc_info.value.extra["error"] + assert len(error) < len(body) + assert error.endswith("...(truncated)") + + @pytest.mark.asyncio + async def test_embed_bad_response_format(self): + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"wrong": "shape"}) + + with pytest.raises(EmbeddingResponseError): + await self._make_embedder(handler).embed(["text"]) + + @pytest.mark.asyncio + async def test_embed_single_empty_response_raises_embedding_response_error(self): + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": []}) + + with pytest.raises(EmbeddingResponseError, match="Empty Ollama embedding response"): + await self._make_embedder(handler).embed_single("text") + + def test_endpoint_v1_appended_when_missing(self): + embedder = OllamaEmbedder(endpoint="http://ollama:11434", model_name="m") + assert embedder._endpoint == "http://ollama:11434/v1" + + def test_endpoint_v1_not_doubled(self): + embedder = OllamaEmbedder(endpoint="http://ollama:11434/v1", model_name="m") + assert embedder._endpoint == "http://ollama:11434/v1" + + @pytest.mark.asyncio + async def test_aclose(self): + embedder = OllamaEmbedder(endpoint="http://ollama:11434", model_name="m") + embedder._client = AsyncMock() + await embedder.aclose() + embedder._client.aclose.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# Registry integration +# --------------------------------------------------------------------------- + + +class TestRegistryIntegration: + def test_llm_registered(self): + from core.llm import llm_registry + + assert "ollama" in llm_registry + + def test_embedder_registered(self): + from core.embeddings import embedder_registry + + assert "ollama" in embedder_registry diff --git a/openrag/services/inference/test_reranker_clients.py b/openrag/services/inference/test_reranker_clients.py new file mode 100644 index 000000000..860dd94c8 --- /dev/null +++ b/openrag/services/inference/test_reranker_clients.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock + +import httpx +import pytest +from core.utils.exceptions import InferenceConnectionError, InferenceTimeoutError + +from .reranker_clients import InfinityReranker, OpenAIReranker + + +def _rerank_response(results: list[dict] | None = None) -> httpx.Response: + results = results or [ + {"index": 0, "relevance_score": 0.9}, + {"index": 2, "relevance_score": 0.7}, + {"index": 1, "relevance_score": 0.3}, + ] + return httpx.Response(200, json={"results": results}) + + +DOCS = ["doc zero", "doc one", "doc two"] + + +class TestInfinityReranker: + @pytest.fixture + def reranker(self): + return InfinityReranker(endpoint="http://reranker:7997", model_name="gte-reranker") + + @pytest.mark.asyncio + async def test_rerank(self, reranker): + transport = httpx.MockTransport(lambda req: _rerank_response()) + reranker._client = httpx.AsyncClient(transport=transport) + result = await reranker.rerank("query", DOCS) + assert result == [(0, 0.9), (2, 0.7), (1, 0.3)] + + @pytest.mark.asyncio + async def test_rerank_with_top_k(self, reranker): + captured = {} + + def capture(req): + import json + + captured.update(json.loads(req.content)) + return _rerank_response([{"index": 0, "relevance_score": 0.9}]) + + transport = httpx.MockTransport(capture) + reranker._client = httpx.AsyncClient(transport=transport) + result = await reranker.rerank("query", DOCS, top_k=1) + assert captured["top_n"] == 1 + assert len(result) == 1 + + @pytest.mark.asyncio + async def test_top_k_clamped_to_doc_count(self, reranker): + captured = {} + + def capture(req): + import json + + captured.update(json.loads(req.content)) + return _rerank_response() + + transport = httpx.MockTransport(capture) + reranker._client = httpx.AsyncClient(transport=transport) + await reranker.rerank("query", DOCS, top_k=100) + assert captured["top_n"] == 3 + + @pytest.mark.asyncio + async def test_sends_raw_scores(self, reranker): + captured = {} + + def capture(req): + import json + + captured.update(json.loads(req.content)) + return _rerank_response() + + transport = httpx.MockTransport(capture) + reranker._client = httpx.AsyncClient(transport=transport) + await reranker.rerank("query", DOCS) + assert captured["raw_scores"] is True + assert captured["return_documents"] is False + + @pytest.mark.asyncio + async def test_connection_error(self, reranker): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + reranker._client = AsyncMock() + reranker._client.post = fail + with pytest.raises(InferenceConnectionError): + await reranker.rerank("query", DOCS) + + @pytest.mark.asyncio + async def test_timeout(self, reranker): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + reranker._client = AsyncMock() + reranker._client.post = fail + with pytest.raises(InferenceTimeoutError): + await reranker.rerank("query", DOCS) + + @pytest.mark.asyncio + async def test_trailing_slash_stripped(self): + r = InfinityReranker(endpoint="http://reranker:7997/", model_name="m") + assert r._endpoint == "http://reranker:7997" + await r.aclose() + + +class TestOpenAIReranker: + @pytest.fixture + def reranker(self): + return OpenAIReranker(endpoint="http://reranker:8000/v1", model_name="gte-reranker", api_key="k") + + @pytest.mark.asyncio + async def test_rerank(self, reranker): + transport = httpx.MockTransport(lambda req: _rerank_response()) + reranker._client = httpx.AsyncClient(transport=transport) + result = await reranker.rerank("query", DOCS) + assert result == [(0, 0.9), (2, 0.7), (1, 0.3)] + + @pytest.mark.asyncio + async def test_connection_error(self, reranker): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + reranker._client = AsyncMock() + reranker._client.post = fail + with pytest.raises(InferenceConnectionError): + await reranker.rerank("query", DOCS) + + @pytest.mark.asyncio + async def test_timeout(self, reranker): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + reranker._client = AsyncMock() + reranker._client.post = fail + with pytest.raises(InferenceTimeoutError): + await reranker.rerank("query", DOCS) + + +class TestRegistryIntegration: + def test_infinity_registered(self): + from core.rerankers import reranker_registry + + assert "infinity" in reranker_registry + + def test_openai_registered(self): + from core.rerankers import reranker_registry + + assert "openai" in reranker_registry diff --git a/openrag/services/inference/test_retry.py b/openrag/services/inference/test_retry.py new file mode 100644 index 000000000..222e6c5d4 --- /dev/null +++ b/openrag/services/inference/test_retry.py @@ -0,0 +1,127 @@ +import httpx +import pytest +from core.utils.exceptions import OpenRAGError, ServiceUnavailableError +from services.inference._retry import _is_retryable, with_retry + + +class TestIsRetryable: + def test_timeout_exception(self): + assert _is_retryable(httpx.ReadTimeout("timeout")) + + def test_connect_error(self): + assert _is_retryable(httpx.ConnectError("refused")) + + @pytest.mark.parametrize("status", [429, 502, 503, 504]) + def test_retryable_http_status(self, status): + req = httpx.Request("GET", "http://test") + exc = httpx.HTTPStatusError("err", request=req, response=httpx.Response(status, request=req)) + assert _is_retryable(exc) + + @pytest.mark.parametrize("status", [400, 401, 403, 404, 422, 500]) + def test_non_retryable_http_status(self, status): + req = httpx.Request("GET", "http://test") + exc = httpx.HTTPStatusError("err", request=req, response=httpx.Response(status, request=req)) + assert not _is_retryable(exc) + + def test_openrag_error_retryable(self): + assert _is_retryable(ServiceUnavailableError("down")) # 503 + + def test_openrag_error_non_retryable(self): + assert not _is_retryable(OpenRAGError("bad", status_code=404)) + + def test_unrelated_exception(self): + assert not _is_retryable(ValueError("nope")) + + +class TestWithRetry: + @pytest.mark.asyncio + async def test_success_no_retry(self): + call_count = 0 + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def succeed(): + nonlocal call_count + call_count += 1 + return "ok" + + result = await succeed() + assert result == "ok" + assert call_count == 1 + + @pytest.mark.asyncio + async def test_retries_on_connect_error(self): + call_count = 0 + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def fail_connect(): + nonlocal call_count + call_count += 1 + raise httpx.ConnectError("refused") + + with pytest.raises(httpx.ConnectError): + await fail_connect() + + assert call_count == 3 + + @pytest.mark.asyncio + async def test_retries_on_429(self): + call_count = 0 + req = httpx.Request("GET", "http://test") + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def fail_429(): + nonlocal call_count + call_count += 1 + raise httpx.HTTPStatusError("rate limited", request=req, response=httpx.Response(429, request=req)) + + with pytest.raises(httpx.HTTPStatusError): + await fail_429() + + assert call_count == 3 + + @pytest.mark.asyncio + async def test_no_retry_on_400(self): + call_count = 0 + req = httpx.Request("GET", "http://test") + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def fail_400(): + nonlocal call_count + call_count += 1 + raise httpx.HTTPStatusError("bad request", request=req, response=httpx.Response(400, request=req)) + + with pytest.raises(httpx.HTTPStatusError): + await fail_400() + + assert call_count == 1 + + @pytest.mark.asyncio + async def test_succeeds_after_transient_failure(self): + call_count = 0 + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def flaky(): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise httpx.ConnectError("transient") + return "recovered" + + result = await flaky() + assert result == "recovered" + assert call_count == 3 + + @pytest.mark.asyncio + async def test_retries_openrag_503(self): + call_count = 0 + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def fail_503(): + nonlocal call_count + call_count += 1 + raise ServiceUnavailableError("down") + + with pytest.raises(ServiceUnavailableError): + await fail_503() + + assert call_count == 3 diff --git a/openrag/services/inference/test_vllm_client.py b/openrag/services/inference/test_vllm_client.py new file mode 100644 index 000000000..797bb2760 --- /dev/null +++ b/openrag/services/inference/test_vllm_client.py @@ -0,0 +1,439 @@ +from __future__ import annotations + +import json +from unittest.mock import AsyncMock + +import httpx +import pytest +from core.utils.exceptions import ( + EmbeddingAPIError, + EmbeddingResponseError, + InferenceConnectionError, + InferenceError, + InferenceTimeoutError, +) +from services.inference._circuit_breaker import _breakers + +from .vllm_client import VLLMClient, VLLMEmbedder, VLLMVision + + +@pytest.fixture(autouse=True) +def _clean_breakers(): + yield + for breaker in _breakers.values(): + breaker.close() + _breakers.clear() + + +def _make_transport(handler): + return httpx.MockTransport(handler) + + +def _chat_response(content: str = "hello") -> httpx.Response: + return httpx.Response(200, json={"choices": [{"message": {"content": content}}]}) + + +def _completions_response(text: str = "result") -> httpx.Response: + return httpx.Response(200, json={"choices": [{"text": text}]}) + + +def _embed_response(vectors: list[list[float]] | None = None) -> httpx.Response: + vectors = vectors or [[0.1, 0.2, 0.3]] + data = [{"index": i, "embedding": v} for i, v in enumerate(vectors)] + return httpx.Response(200, json={"data": data}) + + +# --------------------------------------------------------------------------- +# VLLMClient (LLM) +# --------------------------------------------------------------------------- + + +class TestVLLMClient: + def _make_client(self, handler, **kwargs): + client = VLLMClient( + endpoint="http://vllm:8000/v1", + model_name="test-model", + api_key="test-key", + temperature=0.3, + **kwargs, + ) + client._client = httpx.AsyncClient(transport=_make_transport(handler)) + return client + + @pytest.mark.asyncio + async def test_chat_returns_full_response(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert "/chat/completions" in str(request.url) + assert body["model"] == "test-model" + assert body["stream"] is False + assert body["temperature"] == 0.3 + return _chat_response("world") + + result = await self._make_client(handler).chat([{"role": "user", "content": "hi"}]) + assert result["choices"][0]["message"]["content"] == "world" + + @pytest.mark.asyncio + async def test_generate_returns_full_response(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert "/completions" in str(request.url) + assert "/chat/" not in str(request.url) + assert body["prompt"] == "say something" + return _completions_response("done") + + result = await self._make_client(handler).generate("say something") + assert result["choices"][0]["text"] == "done" + + @pytest.mark.asyncio + async def test_stream_chat_yields_raw_sse_lines(self): + sse_body = ( + 'data: {"choices":[{"delta":{"content":"Hello"}}]}\n' + 'data: {"choices":[{"delta":{"content":" world"}}]}\n' + "data: [DONE]\n" + ) + + def handler(request: httpx.Request) -> httpx.Response: + assert json.loads(request.content)["stream"] is True + return httpx.Response(200, text=sse_body) + + client = self._make_client(handler) + lines = [line async for line in client.stream_chat([{"role": "user", "content": "hi"}])] + assert 'data: {"choices":[{"delta":{"content":"Hello"}}]}' in lines + assert 'data: {"choices":[{"delta":{"content":" world"}}]}' in lines + + @pytest.mark.asyncio + async def test_stream_chat_error_raises(self): + client = self._make_client(lambda req: httpx.Response(503, text="unavailable")) + with pytest.raises(InferenceError): + async for _ in client.stream_chat([{"role": "user", "content": "hi"}]): + pass + + @pytest.mark.asyncio + async def test_chat_connection_error(self): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + client = VLLMClient(endpoint="http://vllm:8000/v1", model_name="m") + client._client = AsyncMock() + client._client.post = fail + with pytest.raises(InferenceConnectionError): + await client.chat([{"role": "user", "content": "hi"}]) + + @pytest.mark.asyncio + async def test_chat_timeout(self): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + client = VLLMClient(endpoint="http://vllm:8000/v1", model_name="m") + client._client = AsyncMock() + client._client.post = fail + with pytest.raises(InferenceTimeoutError): + await client.chat([{"role": "user", "content": "hi"}]) + + @pytest.mark.asyncio + async def test_defaults_forwarded(self): + captured: dict = {} + + def capture(req: httpx.Request) -> httpx.Response: + captured.update(json.loads(req.content)) + return _chat_response() + + await self._make_client(capture).chat([{"role": "user", "content": "hi"}]) + assert captured["temperature"] == 0.3 + + @pytest.mark.asyncio + async def test_per_request_kwargs_override_defaults(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["temperature"] == 0.9 + assert body["max_tokens"] == 100 + return _chat_response() + + await self._make_client(handler).chat([{"role": "user", "content": "hi"}], temperature=0.9, max_tokens=100) + + @pytest.mark.asyncio + async def test_trailing_slash_stripped(self): + c = VLLMClient(endpoint="http://vllm:8000/v1/", model_name="m") + assert c._endpoint == "http://vllm:8000/v1" + await c.aclose() + + @pytest.mark.asyncio + async def test_aclose(self): + client = VLLMClient(endpoint="http://vllm:8000/v1", model_name="m") + client._client = AsyncMock() + await client.aclose() + client._client.aclose.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# VLLMClientOverrides +# --------------------------------------------------------------------------- + + +class TestVLLMClientOverrides: + """Tests for _resolve_overrides (partition-level model selection).""" + + def _make_client(self): + return VLLMClient( + endpoint="http://default:8000/v1", + model_name="default-model", + api_key="default-key", + ) + + def test_no_override_uses_defaults(self): + client = self._make_client() + kwargs: dict = {} + base_url, model, headers = client._resolve_overrides(kwargs) + assert base_url == "http://default:8000/v1" + assert model == "default-model" + assert headers is None + + def test_llm_override_in_metadata(self): + client = self._make_client() + original_metadata = { + "llm_override": { + "base_url": "http://custom:9000/v1/", + "api_key": "custom-key", + "model": "custom-model", + }, + } + kwargs: dict = {"metadata": original_metadata} + base_url, model, headers = client._resolve_overrides(kwargs) + assert base_url == "http://custom:9000/v1" + assert model == "custom-model" + assert headers is not None + assert headers["Authorization"] == "Bearer custom-key" + # kwargs must not be mutated — retries depend on llm_override surviving. + assert kwargs["metadata"] is original_metadata + assert "llm_override" in kwargs["metadata"] + + def test_llm_override_partial(self): + client = self._make_client() + original_metadata = { + "llm_override": {"model": "override-model"}, + "use_map_reduce": True, + } + kwargs: dict = {"metadata": original_metadata} + base_url, model, headers = client._resolve_overrides(kwargs) + assert base_url == "http://default:8000/v1" + assert model == "override-model" + assert headers is None + assert kwargs["metadata"] is original_metadata + assert kwargs["metadata"] == { + "llm_override": {"model": "override-model"}, + "use_map_reduce": True, + } + + def test_trailing_slash_stripped(self): + client = self._make_client() + kwargs: dict = {"metadata": {"llm_override": {"base_url": "http://custom:9000/v1///"}}} + base_url, _, _ = client._resolve_overrides(kwargs) + assert base_url == "http://custom:9000/v1" + + +# --------------------------------------------------------------------------- +# VLLMEmbedder +# --------------------------------------------------------------------------- + + +class TestVLLMEmbedder: + def _make_embedder(self, handler, **kwargs): + embedder = VLLMEmbedder( + endpoint="http://vllm:8000/v1", + model_name="bge-m3", + api_key="test-key", + **kwargs, + ) + embedder._client = httpx.AsyncClient(transport=_make_transport(handler)) + return embedder + + @pytest.mark.asyncio + async def test_embed_returns_sorted_vectors(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["model"] == "bge-m3" + assert body["input"] == ["hello", "world"] + return httpx.Response( + 200, + json={"data": [{"index": 1, "embedding": [0.3, 0.4]}, {"index": 0, "embedding": [0.1, 0.2]}]}, + ) + + result = await self._make_embedder(handler).embed(["hello", "world"]) + assert result == [[0.1, 0.2], [0.3, 0.4]] + + @pytest.mark.asyncio + async def test_embed_single(self): + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.5, 0.6, 0.7]}]}) + + result = await self._make_embedder(handler).embed_single("test") + assert result == [0.5, 0.6, 0.7] + + @pytest.mark.asyncio + async def test_dimension_auto_detected(self): + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.1, 0.2, 0.3]}]}) + + embedder = self._make_embedder(handler) + + with pytest.raises(RuntimeError, match="unknown"): + _ = embedder.dimension + + await embedder.embed(["test"]) + assert embedder.dimension == 3 + + def test_dimension_from_init(self): + assert VLLMEmbedder(endpoint="http://x", model_name="m", dimension=768).dimension == 768 + + @pytest.mark.asyncio + async def test_truncate_prompt_tokens_included(self): + def handler(request: httpx.Request) -> httpx.Response: + assert json.loads(request.content)["truncate_prompt_tokens"] == 8192 + return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.1]}]}) + + await self._make_embedder(handler, max_model_len=8192).embed(["test"]) + + @pytest.mark.asyncio + async def test_truncate_prompt_tokens_absent_when_none(self): + def handler(request: httpx.Request) -> httpx.Response: + assert "truncate_prompt_tokens" not in json.loads(request.content) + return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.1]}]}) + + await self._make_embedder(handler).embed(["test"]) + + @pytest.mark.asyncio + async def test_embed_connection_error(self): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + embedder = VLLMEmbedder(endpoint="http://vllm:8000/v1", model_name="bge-m3") + embedder._client = AsyncMock() + embedder._client.post = fail + with pytest.raises(EmbeddingAPIError): + await embedder.embed(["text"]) + + @pytest.mark.asyncio + async def test_embed_timeout(self): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + embedder = VLLMEmbedder(endpoint="http://vllm:8000/v1", model_name="bge-m3") + embedder._client = AsyncMock() + embedder._client.post = fail + with pytest.raises(EmbeddingAPIError): + await embedder.embed(["text"]) + + @pytest.mark.asyncio + async def test_embed_bad_response_format(self): + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"wrong": "shape"}) + + with pytest.raises(EmbeddingResponseError): + await self._make_embedder(handler).embed(["text"]) + + +# --------------------------------------------------------------------------- +# VLLMVision +# --------------------------------------------------------------------------- + + +class TestVLLMVision: + def _make_vision(self, handler, **kwargs): + vision = VLLMVision( + endpoint="http://vllm:8000/v1", + model_name="qwen-vl", + api_key="test-key", + **kwargs, + ) + vision._client = httpx.AsyncClient(transport=_make_transport(handler)) + return vision + + @pytest.mark.asyncio + async def test_caption_image(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["model"] == "qwen-vl" + assert body["max_tokens"] == 1024 + msg = body["messages"][0] + assert msg["content"][0]["type"] == "image_url" + assert msg["content"][0]["image_url"]["url"].startswith("data:image/png;base64,") + assert msg["content"][1]["type"] == "text" + return _chat_response("A red car") + + result = await self._make_vision(handler).caption_image(b"\x89PNG\r\n\x1a\n", prompt="What is this?") + assert result == "A red car" + + @pytest.mark.asyncio + async def test_caption_image_default_prompt(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["messages"][0]["content"][1]["text"] == "Describe this image in detail." + return _chat_response("An image") + + await self._make_vision(handler).caption_image(b"\x89PNG\r\n\x1a\n") + + @pytest.mark.asyncio + async def test_caption_images_batch(self): + call_count = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal call_count + call_count += 1 + return _chat_response(f"Caption {call_count}") + + results = await self._make_vision(handler).caption_images_batch([b"img1", b"img2", b"img3"]) + assert len(results) == 3 + assert call_count == 3 + + @pytest.mark.asyncio + async def test_custom_max_tokens(self): + def handler(request: httpx.Request) -> httpx.Response: + assert json.loads(request.content)["max_tokens"] == 512 + return _chat_response("ok") + + await self._make_vision(handler, max_tokens=512).caption_image(b"img") + + @pytest.mark.asyncio + async def test_caption_connection_error(self): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + vision = VLLMVision(endpoint="http://vllm:8000/v1", model_name="qwen-vl") + vision._client = AsyncMock() + vision._client.post = fail + with pytest.raises(InferenceConnectionError): + await vision.caption_image(b"img") + + @pytest.mark.asyncio + async def test_caption_timeout(self): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + vision = VLLMVision(endpoint="http://vllm:8000/v1", model_name="qwen-vl") + vision._client = AsyncMock() + vision._client.post = fail + with pytest.raises(InferenceTimeoutError): + await vision.caption_image(b"img") + + +# --------------------------------------------------------------------------- +# Registry integration +# --------------------------------------------------------------------------- + + +class TestRegistryIntegration: + def test_llm_registered(self): + from core.llm import llm_registry + + assert "vllm" in llm_registry + + def test_embedder_registered(self): + from core.embeddings import embedder_registry + + assert "vllm" in embedder_registry + + def test_vlm_registered(self): + from core.vlm import vlm_registry + + assert "vllm" in vlm_registry diff --git a/openrag/services/inference/vllm_client.py b/openrag/services/inference/vllm_client.py new file mode 100644 index 000000000..1b6fd9e4b --- /dev/null +++ b/openrag/services/inference/vllm_client.py @@ -0,0 +1,323 @@ +"""vLLM / OpenAI-compatible inference clients. + +Three classes grouped by server — all talk to the same OpenAI-compatible API: + +* ``VLLMClient`` → ``LLM`` (chat completions) +* ``VLLMEmbedder`` → ``Embedder`` (embeddings) +* ``VLLMVision`` → ``VLM`` (image captioning via chat completions) + +Each class has its own circuit breaker instance so an embedder outage +doesn't trip the LLM breaker. +""" + +from __future__ import annotations + +import asyncio +import base64 +from collections.abc import AsyncIterator + +import httpx +from core.embeddings import Embedder, embedder_registry +from core.llm import LLM, llm_registry +from core.utils.exceptions import ( + EmbeddingAPIError, + EmbeddingResponseError, + InferenceConnectionError, + InferenceError, + InferenceTimeoutError, +) +from core.vlm import VLM, vlm_registry +from utils.logger import get_logger + +from ._circuit_breaker import with_circuit_breaker +from ._retry import with_retry + +logger = get_logger() + + +def _parse_response(resp: httpx.Response) -> dict: + try: + return resp.json() + except ValueError as e: + raise InferenceError(f"Invalid JSON from inference server ({resp.url}): {e}", status_code=502) from e + + +# --------------------------------------------------------------------------- +# LLM +# --------------------------------------------------------------------------- + + +@llm_registry.register("vllm") +class VLLMClient(LLM): + """OpenAI-compatible LLM client backed by vLLM. + + *endpoint* should include the version prefix, e.g. ``http://vllm:8000/v1``. + A single long-lived ``httpx.AsyncClient`` is reused across requests for + connection pooling. + """ + + def __init__( + self, + endpoint: str, + model_name: str, + *, + api_key: str = "", + timeout: float = 240.0, + **kwargs, + ) -> None: + self._endpoint = endpoint.rstrip("/") + self._model = model_name + self._api_key = api_key + self._defaults: dict = kwargs + headers: dict[str, str] = {"Content-Type": "application/json"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + self._client = httpx.AsyncClient(timeout=timeout, headers=headers) + + def _resolve_overrides(self, kwargs: dict) -> tuple[str, str, dict[str, str] | None]: + """Read ``metadata.llm_override`` from *kwargs* without mutating caller data. + + Pure read: ``kwargs`` is untouched so retries see the original override on + every attempt. The caller strips ``metadata`` from the outbound payload via + ``_payload_kwargs`` — every ``metadata`` key is OpenRAG-internal and never + belongs on the wire. + """ + base_url = self._endpoint + model = self._model + override_headers: dict[str, str] | None = None + + llm_override = (kwargs.get("metadata") or {}).get("llm_override") or {} + if llm_override: + if llm_override.get("base_url"): + base_url = llm_override["base_url"].rstrip("/") + if llm_override.get("model"): + model = llm_override["model"] + if llm_override.get("api_key"): + override_headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {llm_override['api_key']}", + } + + return base_url, model, override_headers + + @with_circuit_breaker("llm") + @with_retry(max_attempts=3) + async def generate(self, prompt: str, **kwargs) -> dict: + base_url, model, headers = self._resolve_overrides(kwargs) + kwargs.pop("metadata", None) + payload = {**self._defaults, **kwargs, "model": model, "prompt": prompt} + try: + resp = await self._client.post(f"{base_url}/completions", json=payload, headers=headers) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach LLM at {base_url}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"LLM request timed out at {base_url}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceError( + f"LLM error ({exc.response.status_code}): {exc.response.text[:500]}", + status_code=exc.response.status_code, + ) from exc + return _parse_response(resp) + + @with_circuit_breaker("llm") + @with_retry(max_attempts=3) + async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: + base_url, model, headers = self._resolve_overrides(kwargs) + kwargs.pop("metadata", None) + payload = {**self._defaults, **kwargs, "model": model, "messages": messages, "stream": False} + try: + resp = await self._client.post(f"{base_url}/chat/completions", json=payload, headers=headers) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach LLM at {base_url}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"LLM request timed out at {base_url}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceError( + f"LLM error ({exc.response.status_code}): {exc.response.text[:500]}", + status_code=exc.response.status_code, + ) from exc + return _parse_response(resp) + + async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIterator[str]: + base_url, model, headers = self._resolve_overrides(kwargs) + kwargs.pop("metadata", None) + payload = {**self._defaults, **kwargs, "model": model, "messages": messages, "stream": True} + try: + async with self._client.stream( + "POST", f"{base_url}/chat/completions", json=payload, headers=headers + ) as resp: + if resp.status_code >= 400: + await resp.aread() + raise InferenceError( + f"LLM streaming error ({resp.status_code}): {resp.text[:500]}", + status_code=resp.status_code, + ) + async for line in resp.aiter_lines(): + yield line + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach LLM at {base_url}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"LLM streaming request timed out at {base_url}") from exc + + async def aclose(self) -> None: + await self._client.aclose() + + +# --------------------------------------------------------------------------- +# Embedder +# --------------------------------------------------------------------------- + + +@embedder_registry.register("vllm") +class VLLMEmbedder(Embedder): + """OpenAI-compatible embedding client backed by vLLM. + + Replaces the sync ``openai.OpenAI`` SDK with an async ``httpx`` client. + """ + + def __init__( + self, + endpoint: str, + model_name: str, + *, + max_model_len: int | None = None, + dimension: int | None = None, + timeout: float = 60.0, + api_key: str = "", + **_kwargs, + ) -> None: + self._endpoint = endpoint.rstrip("/") + self._model = model_name + self._max_model_len = max_model_len + self._dimension: int | None = dimension + headers: dict[str, str] = {} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + self._client = httpx.AsyncClient(timeout=timeout, headers=headers) + + @with_circuit_breaker("embedder") + @with_retry(max_attempts=3) + async def embed(self, texts: list[str]) -> list[list[float]]: + body: dict = {"model": self._model, "input": texts} + if self._max_model_len is not None: + body["truncate_prompt_tokens"] = self._max_model_len + try: + resp = await self._client.post(f"{self._endpoint}/embeddings", json=body) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise EmbeddingAPIError( + f"Cannot reach embedder at {self._endpoint}", + model_name=self._model, + base_url=self._endpoint, + error=str(exc), + ) from exc + except httpx.TimeoutException as exc: + raise EmbeddingAPIError( + f"Embedder request timed out at {self._endpoint}", + model_name=self._model, + base_url=self._endpoint, + error=str(exc), + ) from exc + except httpx.HTTPStatusError as exc: + raise EmbeddingAPIError( + f"Embedder API error ({exc.response.status_code})", + model_name=self._model, + base_url=self._endpoint, + error=exc.response.text, + ) from exc + + try: + data = resp.json()["data"] + embeddings = [item["embedding"] for item in sorted(data, key=lambda x: x["index"])] + except (ValueError, KeyError, IndexError, TypeError) as exc: + raise EmbeddingResponseError( + "Unexpected embedding response format", + model_name=self._model, + base_url=self._endpoint, + error=str(exc), + ) from exc + + if self._dimension is None and embeddings: + self._dimension = len(embeddings[0]) + return embeddings + + async def embed_single(self, text: str) -> list[float]: + result = await self.embed([text]) + return result[0] + + @property + def dimension(self) -> int: + if self._dimension is None: + raise RuntimeError("Embedding dimension unknown — call embed() first") + return self._dimension + + async def aclose(self) -> None: + await self._client.aclose() + + +# --------------------------------------------------------------------------- +# VLM (Vision-Language Model) +# --------------------------------------------------------------------------- + + +@vlm_registry.register("vllm") +class VLLMVision(VLLMClient, VLM): + """OpenAI-compatible vision client backed by vLLM. + + Inherits connection pooling, retry, and circuit breaker from VLLMClient. + Adds image captioning via the same OpenAI-compatible chat/completions endpoint. + """ + + def __init__( + self, + endpoint: str, + model_name: str, + *, + timeout: float = 60.0, + api_key: str = "", + max_tokens: int = 1024, + **kwargs, + ) -> None: + super().__init__(endpoint=endpoint, model_name=model_name, api_key=api_key, timeout=timeout, **kwargs) + self._max_tokens = max_tokens + + @with_circuit_breaker("vlm") + @with_retry(max_attempts=2) + async def caption_image(self, image_bytes: bytes, prompt: str | None = None) -> str: + image_b64 = base64.b64encode(image_bytes).decode() + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{image_b64}"}, + }, + { + "type": "text", + "text": prompt or "Describe this image in detail.", + }, + ], + } + ] + try: + resp = await self._client.post( + f"{self._endpoint}/chat/completions", + json={"model": self._model, "messages": messages, "max_tokens": self._max_tokens}, + ) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach VLM at {self._endpoint}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"VLM request timed out at {self._endpoint}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceError( + f"VLM error ({exc.response.status_code}): {exc.response.text[:500]}", + status_code=exc.response.status_code, + ) from exc + return _parse_response(resp)["choices"][0]["message"]["content"] + + async def caption_images_batch(self, images: list[bytes], prompt: str | None = None) -> list[str]: + return list(await asyncio.gather(*(self.caption_image(img, prompt) for img in images))) diff --git a/openrag/services/orchestrators/__init__.py b/openrag/services/orchestrators/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/services/orchestrators/auth_service.py b/openrag/services/orchestrators/auth_service.py new file mode 100644 index 000000000..e0d73e5c2 --- /dev/null +++ b/openrag/services/orchestrators/auth_service.py @@ -0,0 +1,708 @@ +"""AuthService — OIDC flow + auth-policy orchestration (Phase 8A.1). + +Business logic extracted from ``routers/auth.py`` and the auth helpers in +``routers/utils.py``. The router keeps HTTP transport only (cookies, +redirects, JSON error shaping, the ``AUTH_MODE`` gate); every decision +lives here and is unit-testable with fake repos / OIDC client. + +Compared to the legacy router this service talks to the Phase 7 domain +repositories (``UserRepository``, ``OIDCSessionRepository``) instead of +the Ray ``vectordb`` actor, so it deals in :class:`User` / +:class:`OIDCSession` models rather than ad-hoc dicts. + +The cryptographic / cookie primitives (``OIDCClient``, the state-cookie +serializer, Fernet token (de)encryption, opaque session-token issuance) +still come from ``components.auth`` during the Phase-8 shim period — +those are infrastructure adapters scheduled to move under +``services/auth`` in Phase 9. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import TYPE_CHECKING, Any +from urllib.parse import urlencode, urlparse + +from components.auth import ( + OIDCClient, + StateCookiePayload, + StateCookieSerializer, + decrypt_token, + encrypt_token, + hash_session_token, + issue_session_token, +) +from core.models.user import OIDCSession, User +from core.utils.exceptions import AuthError, OpenRAGError +from utils.logger import get_logger, mask_email + +if TYPE_CHECKING: + from core.config.auth import OIDCConfig + from core.ports.oidc_session_repo import OIDCSessionRepository + from core.ports.partition_membership_repo import PartitionMembershipRepository + from core.ports.user_repo import UserRepository + +logger = get_logger() + +SESSION_COOKIE_NAME = "openrag_session" + + +class OIDCFlowError(OpenRAGError): + """Raised for any recoverable failure inside the OIDC flow. + + Carries the exact ``status_code`` the legacy router used so the thin + router can reproduce the previous HTTP responses verbatim. + ``error_description`` is only set for back-channel logout, where the + OIDC spec wants an ``error_description`` field in the JSON body. + """ + + def __init__( + self, + message: str, + *, + status_code: int = 400, + error_description: str | None = None, + ) -> None: + super().__init__(message, code="OIDC_FLOW_ERROR", status_code=status_code) + self.error_description = error_description + + +@dataclass +class LoginRedirect: + """Everything the router needs to start the Authorization Code flow.""" + + authorization_url: str + state_cookie_name: str + state_cookie_value: str + state_cookie_max_age: int + + +@dataclass +class CallbackResult: + """Everything the router needs to finish login (set session cookie).""" + + session_cookie_name: str + session_cookie_value: str # plaintext — only the SHA-256 hash is stored + session_cookie_max_age: int + next_url: str + + +def _utcnow() -> datetime: + """Naive local ``now`` — matches the DB columns. + + The ``oidc_sessions`` timestamp columns are ``TIMESTAMP WITHOUT TIME + ZONE`` and every read site compares against ``datetime.now()``; using a + tz-aware value here would make freshly-issued sessions look pre-expired + on non-UTC hosts (and asyncpg refuses tz-aware against tz-naive). + """ + return datetime.now() + + +class AuthService: + """Owns the OIDC Authorization-Code + PKCE flow and auth policy.""" + + ROLE_HIERARCHY: dict[str, int] = {"viewer": 1, "editor": 2, "owner": 3} + + # Defence-in-depth: an IdP claim mapping may only ever write these two + # columns. Mirrors the startup validator and the repo whitelist. + _OIDC_CLAIM_MAPPING_ALLOWED_FIELDS = frozenset({"display_name", "email"}) + + def __init__( + self, + *, + user_repo: UserRepository, + oidc_session_repo: OIDCSessionRepository, + membership_repo: PartitionMembershipRepository, + oidc_client: OIDCClient | None, + config: OIDCConfig, + ) -> None: + self._user_repo = user_repo + self._oidc_session_repo = oidc_session_repo + # Retained for the role/membership helpers that later phases will + # route through this service (8B onward); the OIDC flow itself + # does not touch it. + self._membership_repo = membership_repo + self._oidc_client = oidc_client + self._config = config + + # ------------------------------------------------------------------ + # OIDC flow + # ------------------------------------------------------------------ + + async def start_oidc_login(self, next_url: str | None) -> LoginRedirect: + """Generate PKCE + state/nonce and build the IdP authorization URL.""" + client = self._require_client() + state, nonce = OIDCClient.generate_state_and_nonce() + code_verifier, code_challenge = OIDCClient.generate_pkce_pair() + + try: + auth_url = await client.build_authorization_url( + state=state, + nonce=nonce, + code_challenge=code_challenge, + ) + except Exception as e: + logger.error(f"Failed to build OIDC authorization URL: {e}") + raise OIDCFlowError( + "OIDC discovery failed — see server logs.", + status_code=502, + ) from e + + payload = StateCookiePayload( + state=state, + nonce=nonce, + code_verifier=code_verifier, + next_url=self.sanitize_next_url(next_url), + ) + cookie_value = self._state_serializer().dumps(payload) + return LoginRedirect( + authorization_url=auth_url, + state_cookie_name=StateCookieSerializer.COOKIE_NAME, + state_cookie_value=cookie_value, + state_cookie_max_age=StateCookieSerializer.DEFAULT_TTL_SECONDS, + ) + + async def handle_oidc_callback( + self, + *, + code: str | None, + state: str | None, + state_cookie_raw: str | None, + ) -> CallbackResult: + """Validate the IdP redirect, resolve the user, create a session.""" + client = self._require_client() + + if not code or not state: + raise OIDCFlowError("Missing 'code' or 'state' query parameter.") + if not state_cookie_raw: + raise OIDCFlowError("OIDC state cookie missing.") + + try: + payload = self._state_serializer().loads(state_cookie_raw) + except ValueError as e: + logger.warning(f"Invalid OIDC state cookie: {e}") + raise OIDCFlowError("Invalid or expired OIDC state cookie.") from e + + # CSRF: the query ``state`` must match the signed cookie. + if state != payload.state: + logger.warning("OIDC state mismatch between query and cookie") + raise OIDCFlowError("OIDC state mismatch.") + + try: + bundle = await client.exchange_code( + code=code, + code_verifier=payload.code_verifier, + expected_nonce=payload.nonce, + ) + except Exception as e: + # Generic message — IdP URLs / internals must not leak via HTTP. + logger.exception("OIDC code exchange failed") + raise OIDCFlowError("OIDC code exchange failed") from e + + sub = bundle.claims.get("sub") + if not sub: + raise OIDCFlowError("ID token missing 'sub' claim.") + + user = await self._resolve_user(sub, bundle.claims) + user = await self._sync_auto_provisioned(user, sub, bundle.claims) + user = await self._apply_claim_mapping(user, bundle) + + now = _utcnow() + expires_in = max(int(bundle.expires_in or 0), 60) + access_token_expires_at = now + timedelta(seconds=expires_in) + session_expires_at = now + timedelta(days=7) if bundle.refresh_token else access_token_expires_at + + plain, token_hash = issue_session_token() + key = self._config.token_encryption_key + await self._oidc_session_repo.create_session( + OIDCSession( + session_token_hash=token_hash, + user_id=user.id, + sub=sub, + sid=bundle.claims.get("sid"), + id_token_encrypted=encrypt_token(bundle.id_token, key=key), + access_token_encrypted=encrypt_token(bundle.access_token, key=key), + refresh_token_encrypted=encrypt_token(bundle.refresh_token, key=key), + access_token_expires_at=access_token_expires_at, + session_expires_at=session_expires_at, + created_at=now, + ) + ) + + next_url = self.sanitize_next_url(payload.next_url) + max_age = max(int((session_expires_at - now).total_seconds()), 1) + logger.info(f"OIDC login success — user_id={user.id}, sid={bundle.claims.get('sid')!r}, next={next_url!r}") + return CallbackResult( + session_cookie_name=SESSION_COOKIE_NAME, + session_cookie_value=plain, + session_cookie_max_age=max_age, + next_url=next_url, + ) + + async def handle_backchannel_logout(self, logout_token: str) -> int: + """Verify an IdP logout token and revoke the named session(s).""" + client = self._require_client() + try: + claims = await client.verify_logout_token(logout_token) + except ValueError as e: + logger.warning(f"Invalid back-channel logout token: {e}") + raise OIDCFlowError(str(e), error_description=str(e)) from e + except Exception as e: + logger.warning(f"Back-channel logout token verification failed: {e}") + raise OIDCFlowError("invalid_request") from e + + if claims.sid: + count = await self._oidc_session_repo.revoke_by_sid(claims.sid) + logger.info(f"Back-channel logout revoked sessions — sid={claims.sid!r}, count={count}") + return count + + # Policy: sid-less logout tokens are out of scope (still 200 to the + # IdP so it doesn't retry). + logger.warning( + f"Received sid-less back-channel logout token — not supported; " + f"ignoring per implementation policy (sub={claims.sub!r})" + ) + return 0 + + async def logout(self, session_cookie_value: str | None) -> str | None: + """Revoke the local session and build the IdP end-session redirect. + + Returns the redirect target, or ``None`` when neither an IdP + ``end_session_endpoint`` nor a configured post-logout URL exists + (the router then just confirms the logout in place). + """ + client = self._require_client() + + id_token_hint: str | None = None + if session_cookie_value: + session = await self._oidc_session_repo.get_by_token_hash( + hash_session_token(session_cookie_value), + ) + if session: + if session.id_token_encrypted: + try: + id_token_hint = decrypt_token( + session.id_token_encrypted, + key=self._config.token_encryption_key, + ) + except ValueError as e: + logger.warning(f"Failed to decrypt id_token for logout: {e}") + try: + await self._oidc_session_repo.revoke_session(session.id) + except Exception as e: + logger.warning(f"Failed to revoke oidc_session during logout: {e}") + + local_target = self._config.post_logout_redirect_uri or None + redirect_target: str | None = local_target + try: + meta = await client.discover() + end_session = meta.get("end_session_endpoint") + if end_session: + params: dict[str, str] = {"client_id": self._config.client_id} + if local_target: + params["post_logout_redirect_uri"] = local_target + if id_token_hint: + params["id_token_hint"] = id_token_hint + redirect_target = f"{end_session}?{urlencode(params)}" + except Exception as e: + logger.warning(f"OIDC discovery failed during logout, skipping IdP redirect: {e}") + + return redirect_target + + # ------------------------------------------------------------------ + # Request authentication helpers + # ------------------------------------------------------------------ + + async def get_user_for_request(self, user_id: int) -> dict[str, Any] | None: + user = await self._user_repo.get_user(user_id) + return self._user_to_request_dict(user) if user else None + + async def get_user_by_token_for_request(self, token: str) -> dict[str, Any] | None: + user = await self._user_repo.get_user_by_token(hash_session_token(token)) + return self._user_to_request_dict(user) if user else None + + async def list_user_partitions_for_request(self, user_id: int) -> list[dict[str, Any]]: + memberships = await self._membership_repo.list_user_partitions(user_id) + return [ + { + "partition": membership.partition, + "role": membership.role.value, + "created_at": membership.added_at.isoformat() if membership.added_at else None, + } + for membership in memberships + ] + + async def get_oidc_session_by_token_for_request(self, token: str) -> dict[str, Any] | None: + session = await self._oidc_session_repo.get_by_token_hash(hash_session_token(token)) + return self._oidc_session_to_request_dict(session) if session else None + + async def get_oidc_session_by_id_for_request(self, session_id: int) -> dict[str, Any] | None: + session = await self._oidc_session_repo.get_by_id(session_id) + return self._oidc_session_to_request_dict(session) if session else None + + async def update_oidc_session_tokens_for_request( + self, + *, + session_id: int, + access_token_encrypted: bytes, + refresh_token_encrypted: bytes | None, + access_token_expires_at: datetime, + ) -> None: + updates: dict[str, Any] = { + "access_token_encrypted": access_token_encrypted, + "access_token_expires_at": access_token_expires_at, + "last_refresh_at": _utcnow(), + } + if refresh_token_encrypted is not None: + updates["refresh_token_encrypted"] = refresh_token_encrypted + session = await self._oidc_session_repo.update_session(session_id, **updates) + if session is None: + raise ValueError(f"oidc_session id={session_id} does not exist") + + async def revoke_oidc_session_by_id_for_request(self, session_id: int) -> None: + await self._oidc_session_repo.revoke_session(session_id) + + # ------------------------------------------------------------------ + # Auth policy — pure helpers (no I/O) + # ------------------------------------------------------------------ + + @staticmethod + def _uget(user: Any, key: str, default: Any = None) -> Any: + """Read a user attribute whether ``user`` is a dict or :class:`User`. + + The legacy middleware binds ``request.state.user`` as a dict; the + new repos return :class:`User`. Both shapes flow through these + helpers during the shim period. + """ + if isinstance(user, dict): + return user.get(key, default) + return getattr(user, key, default) + + @staticmethod + def _user_to_request_dict(user: User) -> dict[str, Any]: + return { + "id": user.id, + "display_name": user.display_name, + "external_user_id": user.external_user_id, + "email": user.email, + "is_admin": user.is_admin, + "file_quota": user.file_quota, + "file_count": user.file_count, + "memberships": [ + { + "partition": membership.partition, + "role": membership.role.value, + "added_at": membership.added_at.isoformat() if membership.added_at else None, + } + for membership in user.partitions + ], + } + + @staticmethod + def _oidc_session_to_request_dict(session: OIDCSession) -> dict[str, Any]: + return { + "id": session.id, + "user_id": session.user_id, + "sub": session.sub, + "sid": session.sid, + "id_token_encrypted": session.id_token_encrypted, + "access_token_encrypted": session.access_token_encrypted, + "refresh_token_encrypted": session.refresh_token_encrypted, + "access_token_expires_at": session.access_token_expires_at, + "session_expires_at": session.session_expires_at, + "created_at": session.created_at, + "last_refresh_at": session.last_refresh_at, + "revoked_at": session.revoked_at, + } + + @classmethod + def require_admin(cls, user: Any) -> Any: + """Raise :class:`AuthError` (403) unless the user is an admin.""" + if not user or not cls._uget(user, "is_admin", False): + raise AuthError("Admin privileges required", status_code=403) + return user + + @classmethod + def check_partition_access( + cls, + *, + user: Any, + partition: str, + user_partitions: list[dict[str, Any]], + required_role: str, + super_admin_mode: bool = False, + ) -> bool: + """Pure port of ``ensure_partition_role``. + + Unlike the legacy helper this does **not** probe the vector DB for + partition existence — the "unknown partition is allowed" branch was + an I/O side effect. Callers that still need it must validate + existence separately; here, absence of a membership for an + otherwise-known partition is a 403. + """ + if super_admin_mode and cls._uget(user, "is_admin", False): + return True + + membership = next( + (p for p in user_partitions if p.get("partition") == partition), + None, + ) + if not membership: + raise AuthError( + f"Access to partition '{partition}' forbidden", + status_code=403, + ) + + user_role = membership.get("role") + if user_role not in cls.ROLE_HIERARCHY: + raise AuthError( + f"Access to partition '{partition}' forbidden", + status_code=403, + ) + if cls.ROLE_HIERARCHY[user_role] < cls.ROLE_HIERARCHY[required_role]: + raise AuthError( + f"{required_role.capitalize()} role required for partition '{partition}'", + status_code=403, + ) + return True + + @classmethod + def validate_file_quota( + cls, + user: Any, + *, + pending_task_count: int, + default_quota: int, + ) -> None: + """Pure quota check (the pending-task count is supplied by the caller). + + Quota semantics are unchanged from ``check_user_file_quota``: + admins bypass; ``default_quota < 0`` disables; ``file_quota`` of + ``None`` falls back to the default; ``< 0`` means unlimited. + """ + if cls._uget(user, "is_admin", False): + return + if default_quota < 0: + return + + user_quota = cls._uget(user, "file_quota") + if user_quota is None: + user_quota = default_quota + if user_quota < 0: + return + + indexed_count = cls._uget(user, "file_count", 0) or 0 + total = indexed_count + pending_task_count + if total >= user_quota: + raise OpenRAGError( + f"File quota exceeded. You have {indexed_count} indexed files " + f"and {pending_task_count} pending tasks. Limit: {user_quota}", + code="FILE_QUOTA_EXCEEDED", + status_code=403, + ) + + def sanitize_next_url(self, next_url: str | None) -> str: + """Block open redirects. + + Accept a same-origin relative path (``/...`` but not ``//...``) or + an absolute URL whose origin is explicitly whitelisted; fall back + to ``/`` otherwise. + """ + if not next_url: + return "/" + if next_url.startswith("/") and not next_url.startswith("//"): + return next_url + parsed = urlparse(next_url) + if parsed.scheme in ("http", "https") and parsed.netloc: + origin = f"{parsed.scheme}://{parsed.netloc}" + if origin in self._allowed_next_origins(): + return next_url + return "/" + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + def _require_client(self) -> OIDCClient: + if self._oidc_client is None: + raise OIDCFlowError("OIDC is not configured.", status_code=400) + return self._oidc_client + + def _state_serializer(self) -> StateCookieSerializer: + return StateCookieSerializer(secret_key=self._config.token_encryption_key) + + @staticmethod + def _allowed_next_origins() -> set[str]: + """Origins accepted as post-login redirect targets. + + Mirrors the CORS allowlist: localhost dev ports plus + ``INDEXERUI_URL`` so the separately-served indexer-ui can receive + the user back after the flow. + """ + origins = {"http://localhost:3042", "http://localhost:5173"} + indexer_ui = os.getenv("INDEXERUI_URL") + if indexer_ui: + origins.add(indexer_ui.rstrip("/")) + return origins + + @staticmethod + def _display_name_from_claims(claims: dict[str, Any], sub: str) -> str: + """Pick a printable display name from the standard OIDC claims.""" + for key in ("name", "preferred_username"): + value = claims.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + given = claims.get("given_name") or "" + family = claims.get("family_name") or "" + composed = f"{given} {family}".strip() + if composed: + return composed + return f"oidc-{sub[:8]}" + + @classmethod + def _parse_claim_mapping(cls, raw: str) -> dict[str, str]: + """Parse ``OIDC_CLAIM_MAPPING`` (CSV of ``db_field:claim`` pairs). + + Non-whitelisted / malformed entries are dropped silently — the + hard-failure path is the startup validator; at login time we log + and continue rather than break the flow. + """ + raw = (raw or "").strip() + if not raw: + return {} + mapping: dict[str, str] = {} + for pair in raw.split(","): + pair = pair.strip() + if not pair or ":" not in pair: + continue + db_field, claim = pair.split(":", 1) + db_field = db_field.strip() + claim = claim.strip() + if db_field not in cls._OIDC_CLAIM_MAPPING_ALLOWED_FIELDS or not claim: + continue + mapping[db_field] = claim + return mapping + + async def _resolve_user(self, sub: str, claims: dict[str, Any]) -> User: + """Look the user up by ``sub``; auto-provision if configured.""" + user = await self._user_repo.get_user_by_external_id(sub) + if user is not None: + return user + + if not self._config.auto_provision_login: + logger.warning(f"OIDC login rejected — user not registered (sub={sub!r})") + raise OIDCFlowError("User not registered", status_code=403) + + display_name = self._display_name_from_claims(claims, sub) + email = claims.get("email") + try: + user = await self._user_repo.create_user( + User( + display_name=display_name, + external_user_id=sub, + email=email if isinstance(email, str) and email.strip() else None, + is_admin=False, + ) + ) + except Exception as e: + # Concurrent first-login race or DB failure. Re-read by sub first; + # if an email collision caused the insert failure, return an + # actionable conflict instead of an opaque 500. + logger.exception(f"OIDC auto-provisioning failed for sub={sub!r}: {e}") + user = await self._user_repo.get_user_by_external_id(sub) + if user is None: + if isinstance(email, str) and email.strip(): + existing = await self._user_repo.get_user_by_email(email) + if existing is not None: + logger.error( + f"OIDC auto-provisioning blocked for sub={sub!r}: an account with email " + f"{mask_email(email)} already exists under a different identity. Set that " + f"user's external_user_id to this sub to allow login." + ) + raise OIDCFlowError( + "An account with this email already exists. Ask your administrator to " + "link it to your identity provider login.", + status_code=409, + ) from e + raise OIDCFlowError("Failed to provision user", status_code=500) from e + else: + logger.info(f"OIDC user auto-provisioned (id={user.id}, sub={sub!r})") + return user + + async def _sync_auto_provisioned( + self, + user: User, + sub: str, + claims: dict[str, Any], + ) -> User: + """Keep display_name/email in sync with the IdP on every login. + + Only active when ``auto_provision_login`` is on — then the IdP is + the source of truth for these two fields so a rename upstream does + not drift. No-op when the row already matches. + """ + if not self._config.auto_provision_login: + return user + + derived_display = self._display_name_from_claims(claims, sub) + raw_email = claims.get("email") + derived_email = raw_email.strip() if isinstance(raw_email, str) and raw_email.strip() else None + + updates: dict[str, Any] = {} + if derived_display and user.display_name != derived_display: + updates["display_name"] = derived_display + if derived_email is not None and user.email != derived_email: + updates["email"] = derived_email + if not updates: + return user + + try: + refreshed = await self._user_repo.update_user(user.id, **updates) + except Exception as e: + logger.warning(f"OIDC auto-provision sync failed for user_id={user.id}: {e}") + return user + return refreshed or user + + async def _apply_claim_mapping(self, user: User, bundle: Any) -> User: + """Apply the optional ``OIDC_CLAIM_MAPPING`` (display_name/email only).""" + mapping = self._parse_claim_mapping(self._config.claim_mapping) + if not mapping: + return user + + if self._config.claim_source == "userinfo": + try: + claims_for_mapping = await self._require_client().fetch_userinfo(bundle.access_token) + except Exception as e: + logger.warning(f"OIDC userinfo fetch failed: {e}") + raise OIDCFlowError("Failed to fetch userinfo from IdP.") from e + else: + claims_for_mapping = bundle.claims + + updates: dict[str, Any] = {} + for db_field, claim in mapping.items(): + value = claims_for_mapping.get(claim) + if value is None: + continue + if getattr(user, db_field, None) == value: + continue + updates[db_field] = value + if not updates: + return user + + try: + refreshed = await self._user_repo.update_user(user.id, **updates) + except Exception as e: + logger.warning(f"update_user failed for user_id={user.id}: {e}") + return user + return refreshed or user + + +__all__ = [ + "AuthService", + "OIDCFlowError", + "LoginRedirect", + "CallbackResult", + "SESSION_COOKIE_NAME", +] diff --git a/openrag/services/orchestrators/conversion_service.py b/openrag/services/orchestrators/conversion_service.py new file mode 100644 index 000000000..470b8b1d1 --- /dev/null +++ b/openrag/services/orchestrators/conversion_service.py @@ -0,0 +1,94 @@ +"""ConversionService — document extraction + chunk lookup (Phase 8E). + +Business logic extracted from ``routers/tools.py`` (the ``extractText`` +tool) and ``routers/extract.py`` (chunk-by-id lookup). Both were thin +wrappers; this service keeps them Ray-free: + +- serialization runs in the ``DocSerializer`` Ray actor, reached through + the :class:`~core.indexing.serializer.FileSerializer` port (the + container injects the ``SerializerRayShim`` during the shim period); +- chunk lookup goes through the clean :class:`VectorStore` port + (``query_chunks_by_filter`` on the Milvus ``_id``), mirroring how + PartitionService reads chunks — no LangChain ``Document`` leaks out. + +The thin routers keep HTTP transport only: file save + cleanup IO, tool +dispatch, the request-scoped partition authorization, and the guards +whose exact ``{"detail": ...}`` body the legacy endpoints returned via +``HTTPException`` (404 not-found, 403 forbidden, the 4xx/5xx tool-error +mapping). + +Constructor note: the plan's ``ConversionService(config=config)`` is +underspecified — it takes the two ports it actually needs plus the +established ``collection`` extra (the vector-store collection name the +legacy shim read from ``config.vectordb.collection_name``), supplied by +the container from settings so the service stays Ray/config-free (8H). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from core.utils.text import sanitize_extracted_text +from utils.logger import get_logger + +if TYPE_CHECKING: + from core.indexing.serializer import FileSerializer + from core.vector_stores import VectorStore + +logger = get_logger() + + +class ConversionService: + """File-to-text extraction and single-chunk retrieval.""" + + def __init__( + self, + *, + serializer: FileSerializer, + vector_store: VectorStore, + collection: str, + ) -> None: + self._serializer = serializer + self._vector_store = vector_store + self._collection = collection + + async def serialize_file( + self, + *, + file_path: str, + filename: str | None, + metadata: dict, + ) -> str: + """Serialize ``file_path`` to sanitized raw text (``extractText``).""" + metadata = dict(metadata or {}) + metadata.update({"source": str(file_path), "filename": filename}) + content = await self._serializer.serialize(file_path, metadata) + return sanitize_extracted_text(content) + + async def get_chunk(self, chunk_id: str) -> dict | None: + """Return ``{"page_content", "metadata"}`` for a chunk, or ``None``. + + Milvus ``_id`` is Int64; a non-integer id is treated as not + found (the router maps ``None`` to a 404), matching the legacy + ``get_chunk_by_id``. + """ + try: + chunk_id_int = int(chunk_id) + except (ValueError, TypeError): + logger.warning("Invalid chunk_id format - must be an integer", chunk_id=chunk_id) + return None + + rows = await self._vector_store.query_chunks_by_filter( + self._collection, + {"_id": chunk_id_int}, + ) + if not rows: + return None + row = rows[0] + return { + "page_content": row["text"], + "metadata": {k: v for k, v in row.items() if k not in ("text", "vector")}, + } + + +__all__ = ["ConversionService"] diff --git a/openrag/services/orchestrators/indexing_service.py b/openrag/services/orchestrators/indexing_service.py new file mode 100644 index 000000000..19d0e4473 --- /dev/null +++ b/openrag/services/orchestrators/indexing_service.py @@ -0,0 +1,183 @@ +"""IndexingService — file ingest orchestration. + +Business logic extracted from ``routers/indexer.py``: metadata assembly, +existence/workspace checks, and task dispatch. Indexing jobs are routed +through :class:`~core.indexing.dispatcher.IndexingDispatcher` so this +service stays Ray-free. + +The thin router keeps HTTP transport only: file save to disk (IO), +``request.url_for`` link building, the shared ``Depends`` auth wrappers, +and the guards whose exact ``{"detail": ...}`` body the legacy endpoints +returned via ``HTTPException``. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +from components.indexer.utils.files import extract_temporal_fields +from utils.logger import get_logger + +if TYPE_CHECKING: + from core.indexing.dispatcher import IndexingDispatcher + from core.ports.document_repo import DocumentRepository + from core.ports.workspace_repo import WorkspaceRepository + +logger = get_logger() + +# Client-supplied datetime fields lifted into queryable metadata. +TEMPORAL_FIELDS = ["created_at"] + + +def _human_readable_size(size_bytes: int) -> str: + """Bytes → human-readable string (e.g. ``'2.40 MB'``). + + Kept private here rather than imported from ``routers/utils.py`` — + services must not depend on the HTTP layer. + """ + size = float(size_bytes) + for unit in ["B", "KB", "MB", "GB", "TB"]: + if size < 1024: + return f"{size:.2f} {unit}" + size /= 1024 + return f"{size:.2f} PB" + + +class IndexingService: + """File upload/delete/copy/metadata orchestration over the worker layer.""" + + def __init__( + self, + *, + document_repo: DocumentRepository, + workspace_repo: WorkspaceRepository, + dispatcher: IndexingDispatcher, + ) -> None: + self._document_repo = document_repo + self._workspace_repo = workspace_repo + self._dispatcher = dispatcher + + # ------------------------------------------------------------------ + # Lookups (used by the thin router for its byte-identical guards) + # ------------------------------------------------------------------ + + async def file_exists(self, file_id: str, partition: str) -> bool: + try: + return await self._document_repo.file_exists_in_partition( + file_id=file_id, + partition=partition, + ) + except Exception as e: # pragma: no cover - defensive, matches legacy + logger.exception("File existence check failed.", file_id=file_id, partition=partition, error=str(e)) + return False + + async def get_workspace(self, workspace_id: str) -> dict | None: + return await self._workspace_repo.get_workspace_dict(workspace_id) + + # ------------------------------------------------------------------ + # Ingest + # ------------------------------------------------------------------ + + def _build_metadata( + self, + *, + metadata: dict, + file_path: str, + file_id: str, + sanitized_filename: str, + original_filename: str | None, + ) -> dict: + """Assemble the indexing metadata exactly as the legacy router did.""" + metadata = dict(metadata or {}) + metadata.update( + { + "source": str(file_path), + "filename": sanitized_filename, + "original_filename": original_filename, + } + ) + file_stat = Path(file_path).stat() + metadata["file_size"] = _human_readable_size(file_stat.st_size) + metadata["file_id"] = file_id + metadata.update(extract_temporal_fields(metadata, temporal_fields=TEMPORAL_FIELDS)) + return metadata + + async def add_file( + self, + *, + file_path: str, + file_id: str, + partition: str, + metadata: dict, + sanitized_filename: str, + original_filename: str | None, + user: dict | None, + workspace_ids: list[str] | None = None, + replace: bool = False, + ) -> str: + """Assemble metadata and queue an (re)indexing job; return its task id. + + Workspace association happens inside the worker's ``add_file`` + after a successful index — the router only pre-validates the ids. + """ + full_metadata = self._build_metadata( + metadata=metadata, + file_path=file_path, + file_id=file_id, + sanitized_filename=sanitized_filename, + original_filename=original_filename, + ) + return await self._dispatcher.dispatch_indexing( + path=file_path, + metadata=full_metadata, + partition=partition, + user=user, + workspace_ids=workspace_ids, + replace=replace, + ) + + async def delete_file(self, file_id: str, partition: str) -> None: + await self._dispatcher.delete_file(file_id, partition) + + async def update_metadata( + self, + file_id: str, + metadata: dict, + partition: str, + user: dict | None, + ) -> None: + metadata = dict(metadata or {}) + metadata["file_id"] = file_id + await self._dispatcher.update_file_metadata(file_id, metadata, partition, user) + + async def copy_file( + self, + *, + source_file_id: str, + source_partition: str, + target_file_id: str, + target_partition: str, + metadata: dict, + user: dict | None, + ) -> None: + metadata = dict(metadata or {}) + metadata["file_id"] = target_file_id + metadata["partition"] = target_partition + await self._dispatcher.copy_file(source_file_id, metadata, source_partition, user) + + # ------------------------------------------------------------------ + # Task state + # ------------------------------------------------------------------ + + async def get_task_state(self, task_id: str) -> str | None: + return await self._dispatcher.get_task_state(task_id) + + async def get_task_error(self, task_id: str) -> str | None: + return await self._dispatcher.get_task_error(task_id) + + async def cancel_task(self, task_id: str) -> bool: + return await self._dispatcher.cancel_task(task_id) + + +__all__ = ["IndexingService"] diff --git a/openrag/services/orchestrators/job_service.py b/openrag/services/orchestrators/job_service.py new file mode 100644 index 000000000..4ca4978ef --- /dev/null +++ b/openrag/services/orchestrators/job_service.py @@ -0,0 +1,121 @@ +"""JobService — task-queue queries (Phase 8D.2). + +Thin wrapper around the ``TaskStateManager`` Ray actor, extracted from +``routers/queue.py``. Aggregation/filtering (the active-status rollup, +the per-status counts, the ``?task_status=`` filter) is business logic +and lives here; ``request.url_for`` link building stays in the thin +router (HTTP transport). + +This is the one orchestrator that legitimately keeps Ray remote calls +during the shim — 8H verification explicitly excepts JobService +wrapping ``TaskStateManager``. Phase 9 swaps the actor for a DB-backed +job repository (this service is the hook point for that P0 feature). +""" + +from __future__ import annotations + +from collections import Counter +from typing import Any + +_ACTIVE_STATES = ("QUEUED", "SERIALIZING", "CHUNKING", "INSERTING") + + +class JobService: + """Queue/worker introspection over the TaskStateManager actor.""" + + def __init__(self, task_state_manager: Any, timeout: float = 60.0) -> None: + self._tsm = task_state_manager + self._timeout = timeout + + async def _call(self, future: Any, task_description: str) -> Any: + """Route TaskStateManager calls through the centralized Ray helper. + + Direct ``.remote()`` awaits would bypass timeout/cancellation + handling and can stall the queue APIs under Ray degradation. The + canonical helper lives in ``services.workers.ray_utils`` + (``components.ray_utils`` is a backward-compat re-export). + """ + from services.workers.ray_utils import call_ray_actor_with_timeout + + return await call_ray_actor_with_timeout( + future=future, + timeout=self._timeout, + task_description=task_description, + ) + + @staticmethod + def _format_pool_info(worker_info: dict[str, int]) -> dict[str, int]: + """Condense ``SerializerQueue.pool_info()`` into the API shape.""" + return { + "total_slots": worker_info["total_capacity"], + "pool_size": worker_info["pool_size"], + "max_per_actor": worker_info["max_tasks_per_worker"], + } + + async def get_queue_info(self) -> dict: + all_states: dict = await self._call(self._tsm.get_all_states.remote(), "get_all_states") + status_counts = Counter(all_states.values()) + + active = {s: status_counts.get(s, 0) for s in _ACTIVE_STATES} + task_summary = { + "active": sum(active.values()), + "active_statuses": active, + "total_cancelled": status_counts.get("CANCELLED", 0), + "total_completed": status_counts.get("COMPLETED", 0), + "total_failed": status_counts.get("FAILED", 0), + } + + worker_info = await self._call(self._tsm.get_pool_info.remote(), "get_pool_info") + return {"workers": self._format_pool_info(worker_info), "tasks": task_summary} + + async def list_tasks( + self, + *, + is_admin: bool, + user_id: int | None, + task_status: str | None = None, + ) -> list[dict]: + """Return ``{task_id, state, details}`` rows, filtered. + + - admins see every task; regular users only their own + - ``task_status='active'`` → QUEUED|SERIALIZING|CHUNKING|INSERTING + - any other value → exact match (case-insensitive) + - ``None`` → all tasks + + The router decorates each row with the status / error URLs. + """ + if is_admin: + all_info: dict[str, dict] = await self._call(self._tsm.get_all_info.remote(), "get_all_info") + else: + all_info = await self._call(self._tsm.get_all_user_info.remote(user_id), f"get_all_user_info({user_id})") + + if task_status is None: + filtered = list(all_info.items()) + elif task_status.lower() == "active": + active_states = set(_ACTIVE_STATES) + filtered = [(tid, i) for tid, i in all_info.items() if i["state"] in active_states] + else: + filtered = [(tid, i) for tid, i in all_info.items() if i["state"].lower() == task_status.lower()] + + return [{"task_id": tid, "state": i["state"], "details": i["details"]} for tid, i in filtered] + + async def get_user_pending_task_count(self, user_id: int | None) -> int: + """Pending (not-yet-completed) indexing tasks for one user. + + Used by UserService for the quota-usage block of ``/users/info`` + (the legacy router called the actor directly from the handler). + """ + return await self._call( + self._tsm.get_user_pending_task_count.remote(user_id), + f"get_user_pending_task_count({user_id})", + ) + + async def get_task_details(self, task_id: str) -> dict | None: + """Return task details for ownership checks and status routes.""" + return await self._call( + self._tsm.get_details.remote(task_id), + f"get_details({task_id})", + ) + + +__all__ = ["JobService"] diff --git a/openrag/services/orchestrators/partition_service.py b/openrag/services/orchestrators/partition_service.py new file mode 100644 index 000000000..6e5c4441b --- /dev/null +++ b/openrag/services/orchestrators/partition_service.py @@ -0,0 +1,250 @@ +"""PartitionService — partition CRUD, membership, file/chunk reads (Phase 8B.1). + +Business logic extracted from ``routers/partition.py`` and the partition +slice of the legacy Ray ``vectordb`` shim. The service talks to the +Phase 7 repositories and the :class:`VectorStore` port directly; it does +not depend on Ray or pymilvus. + +``delete_partition`` is the one cross-cutting method — it must drop the +partition's vectors from the store *and* the relational rows. It is +performed through the clean :class:`VectorStore` port +(``query_ids_by_filter`` + ``delete``) rather than a Milvus-specific +filter delete, so PartitionService stays backend-agnostic. + +Chunk reads return plain dicts (never LangChain ``Document`` objects — +8H forbids LangChain in orchestrators); the thin router builds the +``request.url_for`` links and final response shape. + +Constructor notes (two args beyond the plan's four, both to preserve +legacy behaviour without widening into Ray/config): ``collection`` (the +vector-store collection name the legacy shim read from +``config.vectordb.collection_name``) and ``user_repo`` (needed to +reproduce the ``VDBUserNotFound`` 404 the legacy ``add_partition_member`` +raised). The container supplies both from settings/the catalog store. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np +from core.utils.exceptions import ( + NotFoundError, + PartitionNotFoundError, + UserNotFoundError, + ValidationError, +) +from utils.logger import get_logger + +if TYPE_CHECKING: + from core.ports.document_repo import DocumentRepository + from core.ports.partition_membership_repo import PartitionMembershipRepository + from core.ports.partition_repo import PartitionRepository + from core.ports.user_repo import UserRepository + from core.vector_stores import VectorStore + +logger = get_logger() + + +class PartitionService: + """Partition lifecycle, membership and read-through orchestration.""" + + def __init__( + self, + *, + partition_repo: PartitionRepository, + membership_repo: PartitionMembershipRepository, + document_repo: DocumentRepository, + vector_store: VectorStore, + user_repo: UserRepository, + collection: str, + ) -> None: + self._partition_repo = partition_repo + self._membership_repo = membership_repo + self._document_repo = document_repo + self._vector_store = vector_store + self._user_repo = user_repo + self._collection = collection + + # ------------------------------------------------------------------ + # Existence guards (mirror the legacy _check_* helpers, core exceptions) + # ------------------------------------------------------------------ + + async def _ensure_partition(self, partition: str) -> None: + if not await self._partition_repo.partition_exists(name=partition): + logger.warning(f"Partition '{partition}' does not exist.") + raise PartitionNotFoundError(f"Partition '{partition}' does not exist.") + + async def _ensure_user_exists(self, user_id: int) -> None: + if not await self._user_repo.user_exists(user_id): + logger.warning(f"User with ID {user_id} does not exist.") + raise UserNotFoundError(f"User with ID {user_id} does not exist.") + + async def _ensure_membership(self, partition: str, user_id: int) -> None: + await self._ensure_partition(partition) + await self._ensure_user_exists(user_id) + if not await self._membership_repo.user_is_partition_member(user_id, partition): + raise NotFoundError( + f"User with ID {user_id} is not a member of partition '{partition}'.", + code="MEMBERSHIP_NOT_FOUND", + ) + + async def file_exists(self, file_id: str, partition: str) -> bool: + try: + return await self._document_repo.file_exists_in_partition( + file_id=file_id, + partition=partition, + ) + except Exception as e: # pragma: no cover - defensive, matches legacy + logger.exception("File existence check failed.", file_id=file_id, partition=partition, error=str(e)) + return False + + # ------------------------------------------------------------------ + # Partition CRUD + # ------------------------------------------------------------------ + + async def partition_exists(self, partition: str) -> bool: + try: + return await self._partition_repo.partition_exists(name=partition) + except Exception as e: # pragma: no cover - defensive, matches legacy + logger.exception("Partition existence check failed.", partition=partition, error=str(e)) + return False + + async def list_partitions(self) -> list[dict]: + return await self._partition_repo.list_partitions() + + async def create_partition(self, partition: str, user_id: int) -> None: + """Create a partition owned by ``user_id``. + + The 409-on-exists check lives in the thin router (it returns a + non-bracketed ``{"detail": ...}`` body that must stay identical); + this raises only if the race is lost between that check and here. + """ + if await self._partition_repo.partition_exists(name=partition): + raise ValidationError( + f"Partition '{partition}' already exists.", + status_code=409, + code="PARTITION_EXISTS", + ) + await self._partition_repo.create_partition(name=partition, user_id=user_id) + logger.info(f"Partition '{partition}' created by user_id {user_id}.") + + async def delete_partition(self, partition: str) -> None: + """Drop a partition's vectors *and* relational rows (cross-cutting).""" + await self._ensure_partition(partition) + ids = await self._vector_store.query_ids_by_filter( + self._collection, + {"partition": partition}, + ) + if ids: + deleted = await self._vector_store.delete(ids, self._collection) + logger.info("Deleted points from partition", partition=partition, count=deleted) + await self._partition_repo.delete_partition(name=partition) + logger.info("Partition successfully deleted.", partition=partition) + + # ------------------------------------------------------------------ + # File / chunk reads + # ------------------------------------------------------------------ + + async def list_files(self, partition: str, limit: int | None = None) -> list[dict]: + await self._ensure_partition(partition) + result = await self._document_repo.list_partition_files(partition=partition, limit=limit) + return result.get("files", []) + + async def get_file_chunks(self, partition: str, file_id: str, limit: int = 2000) -> list[dict]: + """Return chunk rows (``_id`` kept, ``text`` dropped) for one file. + + The router builds the extract links and strips ``_id`` from the + surfaced metadata, exactly as before. + """ + if not await self.file_exists(file_id, partition): + raise NotFoundError( + f"'{file_id}' not found in partition '{partition}'", + code="FILE_NOT_FOUND", + ) + rows = await self._vector_store.query_chunks_by_filter( + self._collection, + {"partition": partition, "file_id": file_id}, + output_fields=["*"], + ) + if len(rows) > limit: + rows = rows[:limit] + return [{k: v for k, v in row.items() if k != "text"} for row in rows] + + async def list_all_chunks(self, partition: str, include_embedding: bool = True) -> list[dict]: + """Return ``{"content", "metadata"}`` dicts for every chunk.""" + await self._ensure_partition(partition) + excluded = {"text"} if include_embedding else {"text", "vector"} + output_fields = ["*", "vector"] if include_embedding else ["*"] + rows = await self._vector_store.query_chunks_by_filter( + self._collection, + {"partition": partition}, + output_fields=output_fields, + ) + + def _meta(row: dict[str, Any]) -> dict[str, Any]: + meta: dict[str, Any] = {} + for k, v in row.items(): + if k in excluded: + continue + if k == "vector": + # Legacy surfaced the embedding as a flat string. + v = str(np.array(v).flatten().tolist()) + meta[k] = v + return meta + + return [{"content": row.get("text"), "metadata": _meta(row)} for row in rows] + + # ------------------------------------------------------------------ + # Membership + # ------------------------------------------------------------------ + + async def list_members(self, partition: str) -> list[dict]: + await self._ensure_partition(partition) + return await self._membership_repo.list_partition_members(partition) + + async def add_member(self, partition: str, user_id: int, role: str) -> None: + await self._ensure_partition(partition) + await self._ensure_user_exists(user_id) + await self._membership_repo.add_partition_member(partition, user_id, role) + logger.info(f"User_id {user_id} added to partition '{partition}'.") + + async def remove_member(self, partition: str, user_id: int) -> None: + await self._ensure_membership(partition, user_id) + await self._membership_repo.remove_partition_member(partition, user_id) + logger.info(f"User_id {user_id} removed from partition '{partition}'.") + + async def update_role(self, partition: str, user_id: int, new_role: str) -> None: + await self._ensure_membership(partition, user_id) + await self._membership_repo.update_partition_member_role(partition, user_id, new_role) + logger.info(f"User_id {user_id} role updated to '{new_role}' in partition '{partition}'.") + + # ------------------------------------------------------------------ + # Document relationships + # ------------------------------------------------------------------ + + async def get_related_files(self, partition: str, relationship_id: str) -> list[dict]: + return await self._document_repo.get_files_by_relationship( + partition=partition, + relationship_id=relationship_id, + ) + + async def get_file_ancestors( + self, + partition: str, + file_id: str, + max_ancestor_depth: int | None = None, + ) -> list[dict]: + if not await self.file_exists(file_id, partition): + raise NotFoundError( + f"'{file_id}' not found in partition '{partition}'", + code="FILE_NOT_FOUND", + ) + return await self._document_repo.get_file_ancestors( + partition=partition, + file_id=file_id, + max_ancestor_depth=max_ancestor_depth, + ) + + +__all__ = ["PartitionService"] diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py new file mode 100644 index 000000000..fda3495a5 --- /dev/null +++ b/openrag/services/orchestrators/query_service.py @@ -0,0 +1,419 @@ +"""QueryService — RAG orchestration (Phase 8C.2). + +Rebuilt from ``components/pipeline.py:RagPipeline`` + ``map_reduce.py``. +The hardest single extraction in Phase 8: query generation, retrieval, +web search, map-reduce, context formatting, system-prompt assembly, and +streaming all lived tangled in ``RagPipeline``. + +Two logged decisions (REFACTORING_DECISION_LOG Phase 8): + +* **Structured output** — the legacy used a LangChain structured-output + chain for ``SearchQueries`` (query generation) and ``SummarizedChunk`` + (map-reduce). 8H bans LangChain in orchestrators, so QueryService uses + the injected core ``LLM`` with a + JSON-instructed prompt + ``response_format=json_object`` and + ``json.loads`` into the Pydantic model, keeping the legacy fallbacks + (retry → raw user query; relevancy=False on parse failure). +* **Streaming + citations live here; the router is pure transport.** + ``chat_stream`` drives the proven + ``components.utils.stream_with_source_filtering`` (100-char buffer that + strips the ``[Sources: N]`` tag before it reaches the client); + ``chat`` / ``complete`` return the finalized OpenAI dict with the + citation-filtered ``extra`` sources. The router only maps the + partition, builds request-bound source links (``prepare_sources`` + callable — keeps ``request.url_for`` in transport), and wraps + ``StreamingResponse`` / ``JSONResponse``. + +Imports from ``components.*`` (pure helpers / prompts / websearch) are +allowed during the Phase-8 shim (legacy layer, unchecked by the guard; +no LangChain symbol is imported into this file → 8H clean). ``Chunk`` is +converted to LangChain ``Document`` via ``Chunk.to_langchain()`` at the +boundary so the existing ``format_context`` / source helpers are reused +verbatim (no langchain import in this module). +""" + +from __future__ import annotations + +import asyncio +import copy +import json +from collections.abc import AsyncIterator, Callable +from datetime import datetime +from enum import Enum +from typing import TYPE_CHECKING, Any + +from components.prompts import ( + QUERY_CONTEXTUALIZER_PROMPT, + SPOKEN_STYLE_ANSWER_PROMPT, + SYS_PROMPT_TMPLT, +) +from components.utils import ( + SOURCE_SEPARATOR, + detect_language, + extract_and_strip_sources_block, + filter_sources_by_citations, + format_context, + format_web_context, + get_llm_semaphore, + stream_with_source_filtering, +) +from core.models.query import Query, SearchQueries +from utils.logger import get_logger + +if TYPE_CHECKING: + from core.config.root import Settings + from core.llm.llm import LLM + from services.orchestrators.retrieval_service import RetrievalService + from services.orchestrators.workspace_service import WorkspaceService + +logger = get_logger() + +PrepareSources = Callable[[list, list], list] + +_MAP_SYSTEM_PROMPT = """You are an AI assistant specialized in extracting and synthesizing relevant information from text. + +Your task: +1. Analyze the provided text in relation to the user's question +2. Extract only the essential information that directly addresses the query +3. Preserve necessary context (key words, project names, dates) so the summary is self-understandable + +Respond with a JSON object exactly matching this schema: +{"relevancy": , "summary": ""} +Set relevancy=false (and summary="") if the text has no relevant content for the query.""" + +_MAP_USER_PROMPT = """Here is a text: +{content} + +From this document, identify and comprehensively summarize the information useful for answering the following question: +{query}""" + +_QUERY_JSON_HINT = ( + "\n\nRespond ONLY with a JSON object of the form " + '{"query_list": [{"query": "", "temporal_filters": null}]}.' +) + + +class RAGMODE(Enum): + SIMPLERAG = "SimpleRag" + CHATBOTRAG = "ChatBotRag" + + +class QueryService: + """End-to-end RAG: query-gen → retrieve (+web) → map-reduce → answer.""" + + def __init__( + self, + *, + retrieval_service: RetrievalService, + llm: LLM, + config: Settings, + web_search_service: Any | None, + workspace_service: WorkspaceService, + ) -> None: + self._retrieval = retrieval_service + self._llm = llm + self._web = web_search_service + self._workspace = workspace_service + + self._rag_mode = config.rag.mode + self._chat_history_depth = config.rag.chat_history_depth + self._max_contextualized_query_len = config.rag.max_contextualized_query_len + self._max_context_tokens = config.reranker.top_k * config.chunker.chunk_size + + mr = config.map_reduce + self._mr_initial = mr.initial_batch_size + self._mr_expansion = mr.expansion_batch_size + self._mr_max = mr.max_total_documents + + # ------------------------------------------------------------------ + # Query generation (was RagPipeline.generate_query — no LangChain) + # ------------------------------------------------------------------ + + async def generate_query(self, messages: list[dict]) -> SearchQueries: + last_user = messages[-1]["content"] + if RAGMODE(self._rag_mode) is RAGMODE.SIMPLERAG: + return SearchQueries(query_list=[Query(query=last_user)]) + + chat_history = "".join(f"{m['role']}: {m['content']}\n" for m in messages) + prompt = QUERY_CONTEXTUALIZER_PROMPT.format( + query_language=detect_language(last_user), + current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S"), + ) + llm_messages = [ + {"role": "system", "content": prompt + _QUERY_JSON_HINT}, + {"role": "user", "content": f"Here is the chat history: \n{chat_history}\n"}, + ] + params = { + "max_completion_tokens": self._max_contextualized_query_len, + "response_format": {"type": "json_object"}, + } + for attempt in (1, 2): + try: + resp = await self._llm.chat(llm_messages, **params) + content = resp["choices"][0]["message"]["content"] + return SearchQueries.model_validate_json(_json_slice(content)) + except Exception as exc: + if attempt == 1: + logger.warning("Query generation parse error — retrying", error=str(exc)) + else: + logger.warning( + "Query generation failed twice — falling back to raw user query", + error=str(exc), + ) + return SearchQueries(query_list=[Query(query=last_user)]) + + # ------------------------------------------------------------------ + # Map-reduce (was map_reduce.RAGMapReduce — no LangChain) + # ------------------------------------------------------------------ + + async def _infer_relevancy(self, query: str, doc) -> tuple[bool, str]: + async with get_llm_semaphore(): + try: + resp = await self._llm.chat( + [ + {"role": "system", "content": _MAP_SYSTEM_PROMPT}, + {"role": "user", "content": _MAP_USER_PROMPT.format(query=query, content=doc.page_content)}, + ], + max_tokens=512, + temperature=0.3, + response_format={"type": "json_object"}, + ) + data = json.loads(_json_slice(resp["choices"][0]["message"]["content"])) + return bool(data.get("relevancy", False)), str(data.get("summary", "") or "") + except Exception as e: + logger.error("Error during chunk relevancy inference", error=str(e)) + return False, "" + + async def _map_reduce(self, query: str, docs: list) -> list: + """LLM relevancy filter + summarisation, batched with early stop.""" + + async def _batch(chunks: list, summaries: list) -> bool: + outputs = await asyncio.gather(*[self._infer_relevancy(query, c) for c in chunks]) + terminate = all(not rel for rel, _ in outputs[-self._mr_expansion :]) + for (rel, summary), chunk in zip(outputs, chunks, strict=True): + if rel: + summaries.append(_summary_doc(chunk, summary)) + return terminate + + summaries: list = [] + initial, remaining = docs[: self._mr_initial], docs[self._mr_initial :] + terminate = await _batch(initial, summaries) + if terminate or not remaining or len(summaries) >= self._mr_max: + return summaries + + for i in range(0, len(remaining), self._mr_expansion): + n = min(self._mr_expansion, self._mr_max - len(summaries)) + if n <= 0: + break + terminate = await _batch(remaining[i : i + n], summaries) + if terminate or len(summaries) >= self._mr_max: + break + logger.debug("Map reduce completed", relevant_chunks_count=len(summaries)) + return summaries + + # ------------------------------------------------------------------ + # Preparation (was RagPipeline._prepare_for_chat_completion) + # ------------------------------------------------------------------ + + async def _prepare_chat(self, partition: list[str] | None, payload: dict): + messages = payload["messages"][-self._chat_history_depth :] + queries = await self.generate_query(messages) + + metadata = payload.get("metadata") or {} + use_map_reduce = metadata.get("use_map_reduce", False) + spoken_style = metadata.get("spoken_style_answer", False) + use_websearch = metadata.get("websearch", False) + workspace = metadata.get("workspace") + + top_k = self._mr_max if use_map_reduce else None + + if workspace: + ws = await self._workspace.get_workspace(workspace) + if not ws or ("all" not in partition and ws["partition_name"] not in partition): + logger.warning("Workspace not found in partition(s) — ignoring", workspace=workspace) + workspace = None + filter_params = {"workspace_id": workspace} if workspace else None + + web_results: list = [] + if partition is not None and use_websearch: + doc_lists, web_lists = await self._gather_rag_and_web(queries.query_list, partition, top_k, filter_params) + chunks = self._retrieval.fuse(doc_lists, top_k=top_k) + web_results = _dedupe_web(web_lists) + elif partition is not None: + chunks = await self._retrieval.retrieve_multi( + partitions=partition, search_queries=queries, top_k=top_k, filter_params=filter_params + ) + else: + web_results = _dedupe_web(await asyncio.gather(*[self._web.search(q.query) for q in queries.query_list])) + chunks = [] + + if not chunks and not web_results and partition is None: + return payload, [], [] + + docs = [c.to_langchain() for c in chunks] + if use_map_reduce and docs: + docs = await self._map_reduce(" ".join(q.query for q in queries.query_list), docs) + + web_formatted, web_tokens = "", 0 + if web_results: + web_formatted, _, web_tokens = format_web_context( + web_results, start_index=1, max_tokens=self._web.max_tokens + ) + context, included = format_context(docs, max_context_tokens=self._max_context_tokens - web_tokens) + docs = [docs[i] for i in included] + + if web_results: + if docs: + web_formatted, _, _ = format_web_context( + web_results, start_index=len(docs) + 1, max_tokens=self._web.max_tokens + ) + else: + context = "" + context = f"{context}{SOURCE_SEPARATOR}{web_formatted}" if context else web_formatted + + new_messages = copy.deepcopy(messages) + tmpl = SPOKEN_STYLE_ANSWER_PROMPT if spoken_style else SYS_PROMPT_TMPLT + new_messages.insert( + 0, + { + "role": "system", + "content": tmpl.format( + context=context, current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S") + ), + }, + ) + payload["messages"] = new_messages + return payload, docs, web_results + + async def _gather_rag_and_web(self, query_list, partition, top_k, filter_params): + rag = self._retrieval.retrieve_per_query( + partitions=partition, queries=query_list, top_k=top_k, filter_params=filter_params + ) + web = asyncio.gather(*[self._web.search(q.query) for q in query_list]) + doc_lists, web_lists = await asyncio.gather(rag, web) + return doc_lists, web_lists + + async def _prepare_completions(self, partition: list[str], payload: dict): + prompt = payload["prompt"] + queries = await self.generate_query([{"role": "user", "content": prompt}]) + chunks = await self._retrieval.retrieve_multi(partitions=partition, search_queries=queries) + docs = [c.to_langchain() for c in chunks] + context, included = format_context(docs, max_context_tokens=self._max_context_tokens) + docs = [docs[i] for i in included] + if docs: + payload["prompt"] = ( + f"Given the content\n{context}\nComplete the following prompt: {prompt}\n" + "At the very end of your response, on a new line, list which source numbers " + "you used: [Sources: 1, 3]" + ) + return payload, docs + + # ------------------------------------------------------------------ + # Public API (router = transport) + # ------------------------------------------------------------------ + + async def chat( + self, + *, + partitions: list[str] | None, + payload: dict, + prepare_sources: PrepareSources, + model_name: str, + ) -> dict: + """Non-streaming chat completion → finalized OpenAI dict.""" + metadata = payload.get("metadata") or {} + if partitions is None and not metadata.get("websearch", False): + docs, web_results = [], [] + else: + payload, docs, web_results = await self._prepare_chat(partitions, payload) + sources = prepare_sources(docs, web_results) + + chunk = await self._llm.chat(payload["messages"], **_sampling(payload)) + chunk["model"] = model_name + content = chunk.get("choices", [{}])[0].get("message", {}).get("content", "") or "" + clean, citations = extract_and_strip_sources_block(content) + chunk["choices"][0]["message"]["content"] = clean + chunk["extra"] = json.dumps({"sources": filter_sources_by_citations(sources, citations)}) + return chunk + + async def chat_stream( + self, + *, + partitions: list[str] | None, + payload: dict, + prepare_sources: PrepareSources, + model_name: str, + ) -> AsyncIterator[str]: + """Streaming chat completion → SSE strings with filtered sources.""" + metadata = payload.get("metadata") or {} + if partitions is None and not metadata.get("websearch", False): + docs, web_results = [], [] + else: + payload, docs, web_results = await self._prepare_chat(partitions, payload) + sources = prepare_sources(docs, web_results) + + llm_stream = self._llm.stream_chat(payload["messages"], **_sampling(payload)) + async for sse_line in stream_with_source_filtering(llm_stream, sources, model_name): + yield sse_line + + async def complete( + self, + *, + partitions: list[str] | None, + payload: dict, + prepare_sources: PrepareSources, + ) -> dict: + """Non-streaming text completion → finalized OpenAI dict.""" + if partitions is None: + docs = [] + else: + payload, docs = await self._prepare_completions(partitions, payload) + sources = prepare_sources(docs, []) + + resp = await self._llm.generate(payload["prompt"], **_sampling(payload, key="prompt")) + text = resp.get("choices", [{}])[0].get("text", "") or "" + clean, citations = extract_and_strip_sources_block(text) + resp["choices"][0]["text"] = clean + resp["extra"] = json.dumps({"sources": filter_sources_by_citations(sources, citations)}) + return resp + + +# --------------------------------------------------------------------------- # +# Helpers +# --------------------------------------------------------------------------- # + + +def _json_slice(text: str) -> str: + """Best-effort extract the first JSON object from an LLM response.""" + start = text.find("{") + end = text.rfind("}") + return text[start : end + 1] if start != -1 and end > start else text + + +def _summary_doc(chunk, summary: str): + """A summarised copy of a LangChain Document (page_content replaced).""" + return chunk.__class__(page_content=summary, metadata=chunk.metadata) + + +def _dedupe_web(web_lists: list[list]) -> list: + seen: set[str] = set() + out: list = [] + for r in (r for lst in web_lists for r in lst): + if r.url not in seen: + seen.add(r.url) + out.append(r) + return out + + +def _sampling(payload: dict, key: str = "messages") -> dict: + """Sampling kwargs handed to the core LLM (everything but the body). + + Mirrors the legacy ``_LLMShim``: strip the transport keys; the core + ``VLLMClient`` consumes ``metadata`` (llm_override) and the rest as + OpenAI sampling params. + """ + drop = {key, "stream", "model"} + return {k: v for k, v in payload.items() if k not in drop} + + +__all__ = ["QueryService", "RAGMODE"] diff --git a/openrag/services/orchestrators/retrieval_service.py b/openrag/services/orchestrators/retrieval_service.py new file mode 100644 index 000000000..43d509c5f --- /dev/null +++ b/openrag/services/orchestrators/retrieval_service.py @@ -0,0 +1,226 @@ +"""RetrievalService — retrieval orchestration (Phase 8C.1). + +Wraps the clean ``core.retrieval`` pipeline (strategy + optional reranker ++ related/ancestor expansion + RRF fusion). The legacy +``components/retriever.py`` and ``RetrieverPipeline`` were Phase-5 shims +over this same core; this service is the real composition seam. + +Searcher backing (logged decision, Phase 8C): the core retriever talks +to a ``RetrievalSearcher`` port. The only implementation today is +``MilvusRayShim`` (Ray ``Vectordb`` actor — embeds + hybrid-searches +internally). Per the dev-workflow doc, Ray cleanup is Phase 9, and +orchestrators may call Ray actors *behind a port* during the Phase-8 +shim. So the searcher is injected (Ray stays behind the port); this +file has no Ray remote-call and no Ray import (8H stays satisfied). A +clean ``VectorStore``-backed searcher replaces it in Phase 9. + +Constructor deviates from the plan's prescribed +``(vector_store, embedder_factory, reranker_factory, llm_factory, +document_repo, config)`` for the same reason: with the Ray-shim searcher, +the vector store / embedder / document repo are unused (the shim does +embedding + related/ancestor itself). The container injects the already +built ``searcher`` / ``reranker`` / ``llm`` plus ``config``. +""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +from core.retrieval.pipeline import RetrieverPipeline +from core.retrieval.retriever import ( + HyDeRetriever, + MultiQueryRetriever, + SingleRetriever, + _expand_with_related_chunks, +) +from core.retrieval.rrf import rrf_reranking +from utils.logger import get_logger + +if TYPE_CHECKING: + from core.config.root import Settings + from core.llm.llm import LLM + from core.models.chunk import Chunk + from core.models.query import Query, SearchQueries + from core.rerankers.reranker import Reranker + from core.retrieval.searcher import RetrievalSearcher + +logger = get_logger() + + +def _chunk_key(c: Chunk): + return c.id or id(c) + + +class RetrievalService: + """Retrieval pipeline orchestration (search, single/multi retrieve).""" + + def __init__( + self, + *, + searcher: RetrievalSearcher, + reranker: Reranker | None, + llm: LLM | None, + config: Settings, + ) -> None: + self._searcher = searcher + rcfg = config.retriever + common = { + "searcher": searcher, + "top_k": rcfg.top_k, + "similarity_threshold": rcfg.similarity_threshold, + "with_surrounding_chunks": rcfg.with_surrounding_chunks, + "include_related": rcfg.include_related, + "include_ancestors": rcfg.include_ancestors, + "related_limit": rcfg.related_limit, + "max_ancestor_depth": rcfg.max_ancestor_depth, + } + rtype = rcfg.type + if rtype == "multiQuery": + from components.prompts import MULTI_QUERY_PROMPT + + retriever = MultiQueryRetriever( + llm=llm, + multi_query_template=MULTI_QUERY_PROMPT, + k_queries=rcfg.k_queries, + **common, + ) + elif rtype == "hyde": + from components.prompts import HYDE_PROMPT + + retriever = HyDeRetriever( + llm=llm, + hyde_template=HYDE_PROMPT, + combine=rcfg.combine, + **common, + ) + else: + retriever = SingleRetriever(**common) + + self._pipeline = RetrieverPipeline( + retriever=retriever, + reranker=reranker if config.reranker.enabled else None, + reranker_top_k=config.reranker.top_k, + allow_filterless_fallback=rcfg.allow_filterless_fallback, + ) + logger.debug( + "RetrievalService ready", + retriever=rtype, + reranker_enabled=config.reranker.enabled and reranker is not None, + ) + + # ------------------------------------------------------------------ + # Raw semantic search (powers routers/search.py — was indexer.asearch) + # ------------------------------------------------------------------ + + async def search( + self, + *, + text: str, + partitions: str | list[str], + top_k: int, + similarity_threshold: float, + filter: str | None = None, + filter_params: dict | None = None, + include_related: bool = False, + include_ancestors: bool = False, + related_limit: int = 20, + max_ancestor_depth: int | None = None, + ) -> list[Chunk]: + """One similarity search, then optional related/ancestor expansion. + + Faithful port of ``indexer.asearch`` + the legacy + ``_expand_with_related_chunks``: a single ``searcher.search`` (no + query generation / reranking / RRF — those belong to QueryService). + """ + parts = [partitions] if isinstance(partitions, str) else list(partitions) + chunks = await self._searcher.search( + query=text, + partition=parts, + top_k=top_k, + filter=filter, + filter_params=filter_params, + similarity_threshold=similarity_threshold, + with_surrounding_chunks=True, + ) + if include_related or include_ancestors: + chunks = await _expand_with_related_chunks( + searcher=self._searcher, + results=chunks, + include_related=include_related, + include_ancestors=include_ancestors, + related_limit=related_limit, + max_ancestor_depth=max_ancestor_depth, + ) + return chunks + + # ------------------------------------------------------------------ + # Pipeline retrieval (powers QueryService — 8C.2) + # ------------------------------------------------------------------ + + async def retrieve( + self, + *, + partitions: list[str], + query: Query, + top_k: int | None = None, + filter_params: dict | None = None, + ) -> list[Chunk]: + """Single ``Query`` through retrieve → expand → rerank.""" + return await self._pipeline.retrieve_docs( + partition=partitions, + query=query, + top_k=top_k, + filter_params=filter_params, + ) + + async def retrieve_multi( + self, + *, + partitions: list[str], + search_queries: SearchQueries, + top_k: int | None = None, + filter_params: dict | None = None, + ) -> list[Chunk]: + """Every sub-query in parallel, fused with RRF.""" + return await self._pipeline.get_relevant_docs( + partition=partitions, + search_queries=search_queries, + top_k=top_k, + filter_params=filter_params, + ) + + async def retrieve_per_query( + self, + *, + partitions: list[str], + queries: list[Query], + top_k: int | None = None, + filter_params: dict | None = None, + ) -> list[list[Chunk]]: + """Per-sub-query ranked lists (NOT fused). + + QueryService's combined web-search path interleaves these with web + searches concurrently, then fuses; exposing the un-fused lists + lets it run one ``asyncio.gather`` over both. + """ + return await asyncio.gather( + *[ + self._pipeline.retrieve_docs( + partition=partitions, + query=q, + top_k=top_k, + filter_params=filter_params, + ) + for q in queries + ] + ) + + @staticmethod + def fuse(doc_lists: list[list[Chunk]], top_k: int | None = None) -> list[Chunk]: + """RRF-fuse per-query ranked lists (same fusion the pipeline uses).""" + fused = rrf_reranking(doc_lists, key_fn=_chunk_key) + return fused[:top_k] if top_k is not None else fused + + +__all__ = ["RetrievalService"] diff --git a/openrag/services/orchestrators/test_auth_service.py b/openrag/services/orchestrators/test_auth_service.py new file mode 100644 index 000000000..c69d8e479 --- /dev/null +++ b/openrag/services/orchestrators/test_auth_service.py @@ -0,0 +1,422 @@ +"""Unit tests for :class:`AuthService` (Phase 8A.1). + +The OIDC primitives (PKCE/state generation, state-cookie signing, Fernet +token encryption, opaque session-token issuance) are exercised for real; +only the IdP-facing :class:`OIDCClient` and the persistence repos are +faked. Each test asserts behaviour the legacy router used to own. +""" + +from __future__ import annotations + +import pytest +from components.auth import StateCookieSerializer, hash_session_token +from components.auth.oidc_client import LogoutTokenClaims, TokenBundle +from core.config.auth import OIDCConfig +from core.models.user import User +from cryptography.fernet import Fernet +from services.orchestrators.auth_service import AuthService, OIDCFlowError + +KEY = Fernet.generate_key().decode() + + +# --------------------------------------------------------------------------- # +# Fakes +# --------------------------------------------------------------------------- # + + +class FakeUserRepo: + def __init__(self, users: dict[str, User] | None = None): + self._by_ext = users or {} + self._by_email = { + user.email.strip().lower(): user + for user in self._by_ext.values() + if isinstance(user.email, str) and user.email.strip() + } + self.created: list[User] = [] + self.updated: list[tuple[int, dict]] = [] + self._next_id = 100 + + async def get_user_by_external_id(self, external_id: str) -> User | None: + return self._by_ext.get(external_id) + + async def get_user_by_email(self, email: str) -> User | None: + return self._by_email.get(email.strip().lower()) + + async def create_user(self, user: User) -> User: + normalized_email = user.email.strip().lower() if user.email else None + if normalized_email and normalized_email in self._by_email: + raise ValueError("duplicate key value violates unique constraint") + self._next_id += 1 + user.id = self._next_id + user.email = normalized_email + self.created.append(user) + if user.external_user_id: + self._by_ext[user.external_user_id] = user + if user.email: + self._by_email[user.email] = user + return user + + async def update_user(self, user_id: int, **fields): + self.updated.append((user_id, fields)) + for u in self._by_ext.values(): + if u.id == user_id: + for k, v in fields.items(): + if k == "email" and isinstance(v, str): + v = v.strip().lower() + setattr(u, k, v) + if u.email: + self._by_email[u.email] = u + return u + return None + + +class FakeSessionRepo: + def __init__(self): + self.created = [] + self.revoked_ids: list[int] = [] + self.revoked_sids: list[str] = [] + self._by_hash = {} + + async def create_session(self, session): + self.created.append(session) + self._by_hash[session.session_token_hash] = session + return session + + async def get_by_token_hash(self, token_hash: str): + return self._by_hash.get(token_hash) + + async def revoke_session(self, session_id: int) -> bool: + self.revoked_ids.append(session_id) + return True + + async def revoke_by_sid(self, sid: str) -> int: + self.revoked_sids.append(sid) + return 3 + + +class FakeOIDCClient: + def __init__(self, *, bundle=None, logout_claims=None, meta=None, userinfo=None): + self._bundle = bundle + self._logout_claims = logout_claims + self._meta = meta or {} + self._userinfo = userinfo or {} + self.exchange_calls: list[dict] = [] + + async def build_authorization_url(self, *, state, nonce, code_challenge): + return f"https://idp.example/auth?state={state}&cc={code_challenge}" + + async def exchange_code(self, *, code, code_verifier, expected_nonce): + self.exchange_calls.append({"code": code, "cv": code_verifier, "nonce": expected_nonce}) + if isinstance(self._bundle, Exception): + raise self._bundle + return self._bundle + + async def verify_logout_token(self, token): + if isinstance(self._logout_claims, Exception): + raise self._logout_claims + return self._logout_claims + + async def discover(self): + return self._meta + + async def fetch_userinfo(self, access_token): + return self._userinfo + + +def _cfg(**over) -> OIDCConfig: + base = { + "enabled": True, + "client_id": "openrag", + "token_encryption_key": KEY, + "claim_source": "id_token", + "claim_mapping": "", + "post_logout_redirect_uri": "", + "auto_provision_login": False, + } + base.update(over) + return OIDCConfig(**base) + + +def _service(*, user_repo=None, session_repo=None, client=None, cfg=None) -> AuthService: + return AuthService( + user_repo=user_repo or FakeUserRepo(), + oidc_session_repo=session_repo or FakeSessionRepo(), + membership_repo=object(), + oidc_client=client, + config=cfg or _cfg(), + ) + + +def _bundle(**over) -> TokenBundle: + base = { + "id_token": "idtok", + "access_token": "acctok", + "refresh_token": "reftok", + "expires_in": 3600, + "token_type": "Bearer", + "claims": {"sub": "kc-alice", "sid": "sess-1"}, + } + base.update(over) + return TokenBundle(**base) + + +# --------------------------------------------------------------------------- # +# start_oidc_login +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_login_builds_url_and_roundtrips_state_cookie(): + svc = _service(client=FakeOIDCClient()) + result = await svc.start_oidc_login("/dashboard") + + assert result.authorization_url.startswith("https://idp.example/auth?state=") + payload = StateCookieSerializer(secret_key=KEY).loads(result.state_cookie_value) + assert payload.next_url == "/dashboard" + # The state in the signed cookie must match the one in the auth URL. + assert f"state={payload.state}" in result.authorization_url + + +@pytest.mark.asyncio +async def test_login_sanitizes_open_redirect(): + svc = _service(client=FakeOIDCClient()) + result = await svc.start_oidc_login("//evil.com/phish") + payload = StateCookieSerializer(secret_key=KEY).loads(result.state_cookie_value) + assert payload.next_url == "/" + + +@pytest.mark.asyncio +async def test_login_without_client_raises(): + svc = _service(client=None) + with pytest.raises(OIDCFlowError) as ei: + await svc.start_oidc_login("/") + assert ei.value.status_code == 400 + + +# --------------------------------------------------------------------------- # +# handle_oidc_callback +# --------------------------------------------------------------------------- # + + +async def _login_and_get_state(svc) -> tuple[str, str]: + """Run login, return (state, signed_cookie_value).""" + result = await svc.start_oidc_login("/home") + payload = StateCookieSerializer(secret_key=KEY).loads(result.state_cookie_value) + return payload.state, result.state_cookie_value + + +@pytest.mark.asyncio +async def test_callback_happy_path_creates_session(): + user = User(id=7, display_name="Alice", external_user_id="kc-alice") + urepo = FakeUserRepo({"kc-alice": user}) + srepo = FakeSessionRepo() + client = FakeOIDCClient(bundle=_bundle()) + svc = _service(user_repo=urepo, session_repo=srepo, client=client) + + state, cookie = await _login_and_get_state(svc) + result = await svc.handle_oidc_callback(code="abc", state=state, state_cookie_raw=cookie) + + assert result.next_url == "/home" + assert len(srepo.created) == 1 + sess = srepo.created[0] + assert sess.user_id == 7 + assert sess.sub == "kc-alice" + assert sess.sid == "sess-1" + # Only the hash is persisted; the plaintext cookie must hash to it. + assert sess.session_token_hash == hash_session_token(result.session_cookie_value) + # IdP tokens are stored encrypted, not in the clear. + assert sess.access_token_encrypted not in (None, b"acctok") + + +@pytest.mark.asyncio +async def test_callback_missing_code_or_state(): + svc = _service(client=FakeOIDCClient(bundle=_bundle())) + with pytest.raises(OIDCFlowError, match="Missing 'code' or 'state'"): + await svc.handle_oidc_callback(code=None, state="x", state_cookie_raw="y") + + +@pytest.mark.asyncio +async def test_callback_state_mismatch(): + svc = _service(client=FakeOIDCClient(bundle=_bundle())) + _, cookie = await _login_and_get_state(svc) + with pytest.raises(OIDCFlowError, match="state mismatch"): + await svc.handle_oidc_callback(code="abc", state="not-the-state", state_cookie_raw=cookie) + + +@pytest.mark.asyncio +async def test_callback_unregistered_user_rejected(): + svc = _service(user_repo=FakeUserRepo({}), client=FakeOIDCClient(bundle=_bundle())) + state, cookie = await _login_and_get_state(svc) + with pytest.raises(OIDCFlowError) as ei: + await svc.handle_oidc_callback(code="abc", state=state, state_cookie_raw=cookie) + assert ei.value.status_code == 403 + assert ei.value.message == "User not registered" + + +@pytest.mark.asyncio +async def test_callback_auto_provisions_when_enabled(): + urepo = FakeUserRepo({}) + svc = _service( + user_repo=urepo, + client=FakeOIDCClient(bundle=_bundle(claims={"sub": "kc-bob", "name": "Bob", "email": "bob@x.io"})), + cfg=_cfg(auto_provision_login=True), + ) + state, cookie = await _login_and_get_state(svc) + result = await svc.handle_oidc_callback(code="abc", state=state, state_cookie_raw=cookie) + + assert len(urepo.created) == 1 + created = urepo.created[0] + assert created.external_user_id == "kc-bob" + assert created.display_name == "Bob" + assert created.is_admin is False + assert result.next_url == "/home" + + +@pytest.mark.asyncio +async def test_callback_auto_provision_email_collision_returns_conflict(): + existing = User(id=7, display_name="Existing", external_user_id="kc-old", email="alice@example.com") + urepo = FakeUserRepo({"kc-old": existing}) + svc = _service( + user_repo=urepo, + client=FakeOIDCClient(bundle=_bundle(claims={"sub": "kc-new", "email": "alice@example.com"})), + cfg=_cfg(auto_provision_login=True), + ) + + state, cookie = await _login_and_get_state(svc) + with pytest.raises(OIDCFlowError) as ei: + await svc.handle_oidc_callback(code="abc", state=state, state_cookie_raw=cookie) + + assert ei.value.status_code == 409 + assert "already exists" in ei.value.message + assert not urepo.created + + +@pytest.mark.asyncio +async def test_callback_code_exchange_failure_is_masked(): + svc = _service(client=FakeOIDCClient(bundle=RuntimeError("idp 500"))) + state, cookie = await _login_and_get_state(svc) + with pytest.raises(OIDCFlowError, match="OIDC code exchange failed"): + await svc.handle_oidc_callback(code="abc", state=state, state_cookie_raw=cookie) + + +@pytest.mark.asyncio +async def test_claim_mapping_from_userinfo_updates_user(): + user = User(id=9, display_name="Old", external_user_id="kc-c", email="old@x.io") + urepo = FakeUserRepo({"kc-c": user}) + client = FakeOIDCClient( + bundle=_bundle(claims={"sub": "kc-c", "sid": "s"}), + userinfo={"mail": "new@x.io"}, + ) + svc = _service( + user_repo=urepo, + client=client, + cfg=_cfg(claim_source="userinfo", claim_mapping="email:mail"), + ) + state, cookie = await _login_and_get_state(svc) + await svc.handle_oidc_callback(code="abc", state=state, state_cookie_raw=cookie) + + assert urepo.updated and urepo.updated[0][1] == {"email": "new@x.io"} + + +# --------------------------------------------------------------------------- # +# backchannel logout / logout +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_backchannel_logout_revokes_by_sid(): + srepo = FakeSessionRepo() + claims = LogoutTokenClaims(iss="i", aud="openrag", sub="s", sid="sess-9", iat=0, jti=None) + svc = _service(session_repo=srepo, client=FakeOIDCClient(logout_claims=claims)) + count = await svc.handle_backchannel_logout("tok") + assert count == 3 + assert srepo.revoked_sids == ["sess-9"] + + +@pytest.mark.asyncio +async def test_backchannel_logout_invalid_token_carries_description(): + svc = _service(client=FakeOIDCClient(logout_claims=ValueError("bad aud"))) + with pytest.raises(OIDCFlowError) as ei: + await svc.handle_backchannel_logout("tok") + assert ei.value.error_description == "bad aud" + + +@pytest.mark.asyncio +async def test_backchannel_logout_sidless_is_noop_200(): + srepo = FakeSessionRepo() + claims = LogoutTokenClaims(iss="i", aud="openrag", sub="s", sid=None, iat=0, jti=None) + svc = _service(session_repo=srepo, client=FakeOIDCClient(logout_claims=claims)) + assert await svc.handle_backchannel_logout("tok") == 0 + assert srepo.revoked_sids == [] + + +@pytest.mark.asyncio +async def test_logout_revokes_and_builds_end_session_url(): + # Seed a real session via the callback path so the stored id_token is + # encrypted with the configured key. + user = User(id=5, external_user_id="kc-d") + urepo = FakeUserRepo({"kc-d": user}) + srepo = FakeSessionRepo() + client = FakeOIDCClient( + bundle=_bundle(claims={"sub": "kc-d", "sid": "z"}), + meta={"end_session_endpoint": "https://idp.example/logout"}, + ) + svc = _service(user_repo=urepo, session_repo=srepo, client=client) + state, cookie = await _login_and_get_state(svc) + cb = await svc.handle_oidc_callback(code="abc", state=state, state_cookie_raw=cookie) + + target = await svc.logout(cb.session_cookie_value) + assert target.startswith("https://idp.example/logout?") + assert "client_id=openrag" in target + assert "id_token_hint=idtok" in target + assert srepo.revoked_ids == [srepo.created[0].id] + + +@pytest.mark.asyncio +async def test_logout_no_end_session_returns_none(): + svc = _service(client=FakeOIDCClient(meta={})) + assert await svc.logout(None) is None + + +# --------------------------------------------------------------------------- # +# pure auth-policy helpers +# --------------------------------------------------------------------------- # + + +def test_require_admin(): + assert AuthService.require_admin({"is_admin": True}) == {"is_admin": True} + with pytest.raises(OIDCFlowError.__bases__[0]): # OpenRAGError subclass (AuthError) + AuthService.require_admin({"is_admin": False}) + + +def test_check_partition_access_role_hierarchy(): + parts = [{"partition": "p1", "role": "viewer"}] + assert AuthService.check_partition_access( + user={"is_admin": False}, partition="p1", user_partitions=parts, required_role="viewer" + ) + with pytest.raises(Exception): + AuthService.check_partition_access( + user={"is_admin": False}, partition="p1", user_partitions=parts, required_role="owner" + ) + + +def test_check_partition_access_super_admin_bypass(): + assert AuthService.check_partition_access( + user={"is_admin": True}, + partition="anything", + user_partitions=[], + required_role="owner", + super_admin_mode=True, + ) + + +def test_validate_file_quota(): + # Admin bypass. + AuthService.validate_file_quota({"is_admin": True}, pending_task_count=99, default_quota=1) + # Disabled globally. + AuthService.validate_file_quota({"file_count": 50}, pending_task_count=50, default_quota=-1) + # Specific limit exceeded (3 indexed + 2 pending >= 5). + with pytest.raises(Exception): + AuthService.validate_file_quota({"file_count": 3, "file_quota": 5}, pending_task_count=2, default_quota=10) + # Under the limit is fine. + AuthService.validate_file_quota({"file_count": 1, "file_quota": 5}, pending_task_count=1, default_quota=10) diff --git a/openrag/services/orchestrators/test_conversion_service.py b/openrag/services/orchestrators/test_conversion_service.py new file mode 100644 index 000000000..f0700da0e --- /dev/null +++ b/openrag/services/orchestrators/test_conversion_service.py @@ -0,0 +1,83 @@ +"""Unit tests for :class:`ConversionService` (Phase 8E).""" + +from __future__ import annotations + +import pytest +from services.orchestrators.conversion_service import ConversionService + + +class FakeSerializer: + def __init__(self, *, content=" raw\x00 text "): + self._content = content + self.calls: list[tuple[str, dict]] = [] + + async def serialize(self, path: str, metadata: dict) -> str: + self.calls.append((path, metadata)) + return self._content + + +class FakeVectorStore: + def __init__(self, *, rows=None): + self._rows = rows if rows is not None else [] + self.queries: list[tuple[str, dict]] = [] + + async def query_chunks_by_filter(self, collection, filters, output_fields=None): + self.queries.append((collection, filters)) + return list(self._rows) + + +def _service(*, serializer=None, store=None): + return ConversionService( + serializer=serializer or FakeSerializer(), + vector_store=store or FakeVectorStore(), + collection="chunks", + ) + + +@pytest.mark.asyncio +async def test_serialize_file_merges_metadata_and_sanitizes(): + ser = FakeSerializer(content="hello\x00world") + svc = _service(serializer=ser) + + out = await svc.serialize_file( + file_path="/tmp/x.pdf", + filename="x.pdf", + metadata={"author": "a"}, + ) + + # null byte stripped by sanitize_extracted_text + assert "\x00" not in out + assert "hello" in out and "world" in out + path, md = ser.calls[0] + assert path == "/tmp/x.pdf" + assert md == {"author": "a", "source": "/tmp/x.pdf", "filename": "x.pdf"} + + +@pytest.mark.asyncio +async def test_get_chunk_returns_page_content_and_metadata(): + store = FakeVectorStore(rows=[{"text": "chunk body", "vector": [0.1], "partition": "p1", "file_id": "f1"}]) + svc = _service(store=store) + + chunk = await svc.get_chunk("42") + + assert chunk == { + "page_content": "chunk body", + "metadata": {"partition": "p1", "file_id": "f1"}, + } + # queried Milvus _id as int + assert store.queries == [("chunks", {"_id": 42})] + + +@pytest.mark.asyncio +async def test_get_chunk_invalid_id_returns_none_without_query(): + store = FakeVectorStore(rows=[{"text": "x"}]) + svc = _service(store=store) + + assert await svc.get_chunk("not-an-int") is None + assert store.queries == [] + + +@pytest.mark.asyncio +async def test_get_chunk_missing_returns_none(): + svc = _service(store=FakeVectorStore(rows=[])) + assert await svc.get_chunk("7") is None diff --git a/openrag/services/orchestrators/test_indexing_service.py b/openrag/services/orchestrators/test_indexing_service.py new file mode 100644 index 000000000..f719c6342 --- /dev/null +++ b/openrag/services/orchestrators/test_indexing_service.py @@ -0,0 +1,202 @@ +"""Unit tests for :class:`IndexingService` (Phase 8D.1).""" + +from __future__ import annotations + +import pytest +from services.orchestrators.indexing_service import IndexingService + + +class FakeDocumentRepo: + def __init__(self, *, exists: bool = False, raise_on_check: bool = False): + self._exists = exists + self._raise = raise_on_check + + async def file_exists_in_partition(self, file_id: str, partition: str) -> bool: + if self._raise: + raise RuntimeError("boom") + return self._exists + + +class FakeWorkspaceRepo: + def __init__(self, *, workspace=None): + self._workspace = workspace + + async def get_workspace_dict(self, workspace_id: str): + return self._workspace + + +class FakeDispatcher: + def __init__(self): + self.dispatched: list[dict] = [] + self.deleted: list[tuple[str, str]] = [] + self.updated: list[tuple] = [] + self.copied: list[tuple] = [] + self.cancelled: list[str] = [] + self.cancel_result = True + + async def dispatch_indexing(self, *, path, metadata, partition, user, workspace_ids, replace): + self.dispatched.append( + { + "path": path, + "metadata": metadata, + "partition": partition, + "user": user, + "workspace_ids": workspace_ids, + "replace": replace, + } + ) + return "task-abc" + + async def delete_file(self, file_id, partition): + self.deleted.append((file_id, partition)) + + async def update_file_metadata(self, file_id, metadata, partition, user): + self.updated.append((file_id, metadata, partition, user)) + + async def copy_file(self, file_id, metadata, partition, user): + self.copied.append((file_id, metadata, partition, user)) + + async def get_task_state(self, task_id): + return "QUEUED" + + async def get_task_error(self, task_id): + return "trace" + + async def cancel_task(self, task_id): + self.cancelled.append(task_id) + return self.cancel_result + + +def _service(*, doc=None, ws=None, disp=None): + return IndexingService( + document_repo=doc or FakeDocumentRepo(), + workspace_repo=ws or FakeWorkspaceRepo(), + dispatcher=disp or FakeDispatcher(), + ) + + +@pytest.mark.asyncio +async def test_file_exists_passthrough(): + svc = _service(doc=FakeDocumentRepo(exists=True)) + assert await svc.file_exists("f1", "p1") is True + + +@pytest.mark.asyncio +async def test_file_exists_swallows_errors(): + svc = _service(doc=FakeDocumentRepo(raise_on_check=True)) + assert await svc.file_exists("f1", "p1") is False + + +@pytest.mark.asyncio +async def test_get_workspace_passthrough(): + ws = {"workspace_id": "w1", "partition_name": "p1"} + svc = _service(ws=FakeWorkspaceRepo(workspace=ws)) + assert await svc.get_workspace("w1") == ws + + +@pytest.mark.asyncio +async def test_add_file_builds_metadata_and_dispatches(tmp_path): + f = tmp_path / "doc.txt" + f.write_text("hello world") + disp = FakeDispatcher() + svc = _service(disp=disp) + + task_id = await svc.add_file( + file_path=str(f), + file_id="f1", + partition="p1", + metadata={"author": "alice"}, + sanitized_filename="doc.txt", + original_filename="Doc Original.txt", + user={"id": 7}, + workspace_ids=["w1"], + ) + + assert task_id == "task-abc" + assert len(disp.dispatched) == 1 + sent = disp.dispatched[0] + assert sent["path"] == str(f) + assert sent["partition"] == "p1" + assert sent["workspace_ids"] == ["w1"] + assert sent["replace"] is False + md = sent["metadata"] + assert md["author"] == "alice" + assert md["source"] == str(f) + assert md["filename"] == "doc.txt" + assert md["original_filename"] == "Doc Original.txt" + assert md["file_id"] == "f1" + assert md["file_size"] == "11.00 B" + + +@pytest.mark.asyncio +async def test_replace_sets_replace_flag(tmp_path): + f = tmp_path / "doc.txt" + f.write_text("x") + disp = FakeDispatcher() + svc = _service(disp=disp) + await svc.add_file( + file_path=str(f), + file_id="f1", + partition="p1", + metadata={}, + sanitized_filename="doc.txt", + original_filename="doc.txt", + user=None, + replace=True, + ) + assert disp.dispatched[0]["replace"] is True + + +@pytest.mark.asyncio +async def test_delete_file_delegates(): + disp = FakeDispatcher() + svc = _service(disp=disp) + await svc.delete_file("f1", "p1") + assert disp.deleted == [("f1", "p1")] + + +@pytest.mark.asyncio +async def test_update_metadata_injects_file_id(): + disp = FakeDispatcher() + svc = _service(disp=disp) + await svc.update_metadata("f1", {"author": "bob"}, "p1", {"id": 1}) + file_id, md, partition, user = disp.updated[0] + assert file_id == "f1" + assert md == {"author": "bob", "file_id": "f1"} + assert partition == "p1" + assert user == {"id": 1} + + +@pytest.mark.asyncio +async def test_copy_file_sets_target_fields(): + disp = FakeDispatcher() + svc = _service(disp=disp) + await svc.copy_file( + source_file_id="src", + source_partition="p-src", + target_file_id="dst", + target_partition="p-dst", + metadata={"k": "v"}, + user={"id": 2}, + ) + file_id, md, partition, user = disp.copied[0] + assert file_id == "src" + assert partition == "p-src" + assert md == {"k": "v", "file_id": "dst", "partition": "p-dst"} + assert user == {"id": 2} + + +@pytest.mark.asyncio +async def test_task_state_and_error_passthrough(): + svc = _service() + assert await svc.get_task_state("t1") == "QUEUED" + assert await svc.get_task_error("t1") == "trace" + + +@pytest.mark.asyncio +async def test_cancel_task_passthrough(): + disp = FakeDispatcher() + disp.cancel_result = False + svc = _service(disp=disp) + assert await svc.cancel_task("t1") is False + assert disp.cancelled == ["t1"] diff --git a/openrag/services/orchestrators/test_job_service.py b/openrag/services/orchestrators/test_job_service.py new file mode 100644 index 000000000..5f8533b6d --- /dev/null +++ b/openrag/services/orchestrators/test_job_service.py @@ -0,0 +1,133 @@ +"""Unit tests for :class:`JobService` (Phase 8D.2).""" + +from __future__ import annotations + +import sys +import types + +import pytest +from services.orchestrators.job_service import JobService + + +@pytest.fixture(autouse=True) +def _stub_ray_utils(monkeypatch): + async def _call_ray_actor_with_timeout(*, future, timeout, task_description): + return await future + + ray_utils = types.ModuleType("services.workers.ray_utils") + ray_utils.call_ray_actor_with_timeout = _call_ray_actor_with_timeout + monkeypatch.setitem(sys.modules, "services.workers.ray_utils", ray_utils) + + +class _Remote: + """Mimics a Ray actor method: ``actor.method.remote(...)`` awaitable.""" + + def __init__(self, fn): + self._fn = fn + + def remote(self, *args, **kwargs): + async def _coro(): + return self._fn(*args, **kwargs) + + return _coro() + + +class FakeTSM: + def __init__(self, *, states=None, info=None, pool=None): + self._states = states or {} + self._info = info or {} + self._pool = pool or {"total_capacity": 8, "pool_size": 2, "max_tasks_per_worker": 4} + self.get_all_states = _Remote(lambda: dict(self._states)) + self.get_pool_info = _Remote(lambda: dict(self._pool)) + self.get_all_info = _Remote(lambda: dict(self._info)) + self.get_all_user_info = _Remote(lambda uid: {k: v for k, v in self._info.items() if v.get("user") == uid}) + self.get_details = _Remote(lambda task_id: self._info.get(task_id, {}).get("details")) + self.get_user_pending_task_count = _Remote( + lambda user_id: sum(1 for info in self._info.values() if info.get("user_id") == user_id) + ) + + +@pytest.mark.asyncio +async def test_get_queue_info_rolls_up_states(): + tsm = FakeTSM( + states={ + "a": "QUEUED", + "b": "CHUNKING", + "c": "COMPLETED", + "d": "FAILED", + "e": "CANCELLED", + } + ) + out = await JobService(tsm).get_queue_info() + + assert out["workers"] == {"total_slots": 8, "pool_size": 2, "max_per_actor": 4} + tasks = out["tasks"] + assert tasks["active"] == 2 + assert tasks["active_statuses"] == {"QUEUED": 1, "SERIALIZING": 0, "CHUNKING": 1, "INSERTING": 0} + assert tasks["total_completed"] == 1 + assert tasks["total_failed"] == 1 + assert tasks["total_cancelled"] == 1 + + +@pytest.mark.asyncio +async def test_list_tasks_admin_sees_all(): + info = { + "t1": {"state": "QUEUED", "details": {"f": 1}, "user": 1}, + "t2": {"state": "COMPLETED", "details": {"f": 2}, "user": 2}, + } + rows = await JobService(FakeTSM(info=info)).list_tasks(is_admin=True, user_id=1) + assert {r["task_id"] for r in rows} == {"t1", "t2"} + assert rows[0]["details"] == {"f": 1} + + +@pytest.mark.asyncio +async def test_list_tasks_user_scoped(): + info = { + "t1": {"state": "QUEUED", "details": {}, "user": 1}, + "t2": {"state": "QUEUED", "details": {}, "user": 2}, + } + rows = await JobService(FakeTSM(info=info)).list_tasks(is_admin=False, user_id=1) + assert [r["task_id"] for r in rows] == ["t1"] + + +@pytest.mark.asyncio +async def test_list_tasks_active_filter(): + info = { + "t1": {"state": "QUEUED", "details": {}, "user": 1}, + "t2": {"state": "COMPLETED", "details": {}, "user": 1}, + "t3": {"state": "INSERTING", "details": {}, "user": 1}, + } + rows = await JobService(FakeTSM(info=info)).list_tasks(is_admin=True, user_id=1, task_status="active") + assert sorted(r["task_id"] for r in rows) == ["t1", "t3"] + + +@pytest.mark.asyncio +async def test_list_tasks_exact_status_case_insensitive(): + info = { + "t1": {"state": "FAILED", "details": {}, "user": 1}, + "t2": {"state": "COMPLETED", "details": {}, "user": 1}, + } + rows = await JobService(FakeTSM(info=info)).list_tasks(is_admin=True, user_id=1, task_status="failed") + assert [r["task_id"] for r in rows] == ["t1"] + + +@pytest.mark.asyncio +async def test_get_task_details_uses_task_state_manager(): + info = {"t1": {"details": {"user_id": 7, "filename": "a.pdf"}}} + + details = await JobService(FakeTSM(info=info)).get_task_details("t1") + + assert details == {"user_id": 7, "filename": "a.pdf"} + + +@pytest.mark.asyncio +async def test_get_user_pending_task_count_uses_task_state_manager(): + info = { + "t1": {"user_id": 7}, + "t2": {"user_id": 8}, + "t3": {"user_id": 7}, + } + + pending = await JobService(FakeTSM(info=info)).get_user_pending_task_count(7) + + assert pending == 2 diff --git a/openrag/services/orchestrators/test_partition_service.py b/openrag/services/orchestrators/test_partition_service.py new file mode 100644 index 000000000..c62667645 --- /dev/null +++ b/openrag/services/orchestrators/test_partition_service.py @@ -0,0 +1,293 @@ +"""Unit tests for :class:`PartitionService` (Phase 8B.1).""" + +from __future__ import annotations + +import pytest +from core.utils.exceptions import NotFoundError, PartitionNotFoundError, UserNotFoundError, ValidationError +from services.orchestrators.partition_service import PartitionService + + +class FakePartitionRepo: + def __init__(self, existing: set[str] | None = None): + self._existing = existing if existing is not None else set() + self.created: list[tuple[str, int]] = [] + self.deleted: list[str] = [] + + async def partition_exists(self, name: str) -> bool: + return name in self._existing + + async def list_partitions(self) -> list[dict]: + return [{"partition": p} for p in sorted(self._existing)] + + async def create_partition(self, name: str, user_id: int | None = None) -> dict: + self._existing.add(name) + self.created.append((name, user_id)) + return {"partition": name} + + async def delete_partition(self, name: str) -> bool: + self.deleted.append(name) + self._existing.discard(name) + return True + + +class FakeMembershipRepo: + def __init__(self, members: set[tuple[int, str]] | None = None): + self._members = members or set() + self.added: list[tuple[str, int, str]] = [] + self.removed: list[tuple[str, int]] = [] + self.role_updates: list[tuple[str, int, str]] = [] + + async def user_is_partition_member(self, user_id: int, partition: str) -> bool: + return (user_id, partition) in self._members + + async def list_partition_members(self, partition: str) -> list[dict]: + return [{"user_id": u, "role": "viewer"} for (u, p) in self._members if p == partition] + + async def add_partition_member(self, partition: str, user_id: int, role: str) -> bool: + self.added.append((partition, user_id, role)) + return True + + async def remove_partition_member(self, partition: str, user_id: int) -> bool: + self.removed.append((partition, user_id)) + return True + + async def update_partition_member_role(self, partition: str, user_id: int, new_role: str) -> bool: + self.role_updates.append((partition, user_id, new_role)) + return True + + +class FakeDocumentRepo: + def __init__(self, files: set[tuple[str, str]] | None = None, listing: dict | None = None): + self._files = files or set() + self._listing = listing if listing is not None else {} + + async def file_exists_in_partition(self, file_id: str, partition: str) -> bool: + return (file_id, partition) in self._files + + async def list_partition_files(self, partition: str, limit=None) -> dict: + return self._listing + + async def get_files_by_relationship(self, partition: str, relationship_id: str) -> list[dict]: + return [{"file_id": "a", "relationship_id": relationship_id}] + + async def get_file_ancestors(self, partition: str, file_id: str, max_ancestor_depth=None) -> list[dict]: + return [{"file_id": "root"}, {"file_id": file_id}] + + +class FakeVectorStore: + def __init__(self, ids=None, rows=None): + self._ids = ids or [] + self._rows = rows or [] + self.deleted_ids: list[str] = [] + + async def query_ids_by_filter(self, collection, filters): + return list(self._ids) + + async def delete(self, ids, collection="default") -> int: + self.deleted_ids.extend(ids) + return len(ids) + + async def query_chunks_by_filter(self, collection, filters, output_fields=None): + return list(self._rows) + + +class FakeUserRepo: + def __init__(self, existing: set[int] | None = None): + self._existing = existing if existing is not None else set() + + async def user_exists(self, user_id: int) -> bool: + return user_id in self._existing + + +def _svc( + *, + prepo=None, + mrepo=None, + drepo=None, + vstore=None, + urepo=None, + collection="vdb", +) -> PartitionService: + return PartitionService( + partition_repo=prepo or FakePartitionRepo(), + membership_repo=mrepo or FakeMembershipRepo(), + document_repo=drepo or FakeDocumentRepo(), + vector_store=vstore or FakeVectorStore(), + user_repo=urepo or FakeUserRepo(), + collection=collection, + ) + + +# --------------------------------------------------------------------------- # +# CRUD +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_create_partition_conflict_raises_409(): + prepo = FakePartitionRepo(existing={"p1"}) + with pytest.raises(ValidationError) as ei: + await _svc(prepo=prepo).create_partition("p1", 1) + assert ei.value.status_code == 409 + + +@pytest.mark.asyncio +async def test_create_partition_success(): + prepo = FakePartitionRepo() + await _svc(prepo=prepo).create_partition("new", 7) + assert prepo.created == [("new", 7)] + + +@pytest.mark.asyncio +async def test_delete_partition_missing_raises_404(): + with pytest.raises(PartitionNotFoundError): + await _svc(prepo=FakePartitionRepo(existing=set())).delete_partition("ghost") + + +@pytest.mark.asyncio +async def test_delete_partition_drops_vectors_then_rows(): + prepo = FakePartitionRepo(existing={"p1"}) + vstore = FakeVectorStore(ids=["c1", "c2"]) + await _svc(prepo=prepo, vstore=vstore).delete_partition("p1") + assert vstore.deleted_ids == ["c1", "c2"] + assert prepo.deleted == ["p1"] + + +@pytest.mark.asyncio +async def test_delete_partition_no_vectors_still_deletes_rows(): + prepo = FakePartitionRepo(existing={"p1"}) + vstore = FakeVectorStore(ids=[]) + await _svc(prepo=prepo, vstore=vstore).delete_partition("p1") + assert vstore.deleted_ids == [] + assert prepo.deleted == ["p1"] + + +# --------------------------------------------------------------------------- # +# file / chunk reads +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_list_files_missing_partition_404(): + with pytest.raises(PartitionNotFoundError): + await _svc(prepo=FakePartitionRepo(set())).list_files("nope") + + +@pytest.mark.asyncio +async def test_list_files_empty_listing_returns_empty_list(): + svc = _svc(prepo=FakePartitionRepo({"p"}), drepo=FakeDocumentRepo(listing={})) + assert await svc.list_files("p") == [] + + +@pytest.mark.asyncio +async def test_get_file_chunks_missing_file_404(): + svc = _svc(drepo=FakeDocumentRepo(files=set())) + with pytest.raises(NotFoundError) as ei: + await svc.get_file_chunks("p", "f") + assert ei.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_get_file_chunks_strips_text_keeps_id_and_caps_limit(): + rows = [{"_id": str(i), "text": "body", "page": i} for i in range(5)] + svc = _svc( + drepo=FakeDocumentRepo(files={("f", "p")}), + vstore=FakeVectorStore(rows=rows), + ) + out = await svc.get_file_chunks("p", "f", limit=3) + assert len(out) == 3 + assert all("text" not in r for r in out) + assert all("_id" in r for r in out) + + +@pytest.mark.asyncio +async def test_list_all_chunks_excludes_vector_when_no_embedding(): + rows = [{"text": "t", "_id": "1", "vector": [0.1, 0.2]}] + svc = _svc(prepo=FakePartitionRepo({"p"}), vstore=FakeVectorStore(rows=rows)) + out = await svc.list_all_chunks("p", include_embedding=False) + assert out[0]["content"] == "t" + assert "vector" not in out[0]["metadata"] + assert "text" not in out[0]["metadata"] + + +@pytest.mark.asyncio +async def test_list_all_chunks_stringifies_vector_when_included(): + rows = [{"text": "t", "_id": "1", "vector": [0.1, 0.2]}] + svc = _svc(prepo=FakePartitionRepo({"p"}), vstore=FakeVectorStore(rows=rows)) + out = await svc.list_all_chunks("p", include_embedding=True) + assert isinstance(out[0]["metadata"]["vector"], str) + + +# --------------------------------------------------------------------------- # +# membership +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_list_members_missing_partition_404(): + with pytest.raises(PartitionNotFoundError): + await _svc(prepo=FakePartitionRepo(set())).list_members("x") + + +@pytest.mark.asyncio +async def test_add_member_checks_partition_and_user(): + mrepo = FakeMembershipRepo() + svc = _svc( + prepo=FakePartitionRepo({"p"}), + mrepo=mrepo, + urepo=FakeUserRepo({9}), + ) + await svc.add_member("p", 9, "editor") + assert mrepo.added == [("p", 9, "editor")] + + +@pytest.mark.asyncio +async def test_add_member_unknown_user_404(): + svc = _svc(prepo=FakePartitionRepo({"p"}), urepo=FakeUserRepo(set())) + with pytest.raises(UserNotFoundError): + await svc.add_member("p", 123, "viewer") + + +@pytest.mark.asyncio +async def test_remove_member_requires_existing_membership(): + svc = _svc( + prepo=FakePartitionRepo({"p"}), + mrepo=FakeMembershipRepo(members=set()), + urepo=FakeUserRepo({9}), + ) + with pytest.raises(NotFoundError) as ei: + await svc.remove_member("p", 9) + assert ei.value.code == "MEMBERSHIP_NOT_FOUND" + + +@pytest.mark.asyncio +async def test_update_role_success(): + mrepo = FakeMembershipRepo(members={(9, "p")}) + svc = _svc(prepo=FakePartitionRepo({"p"}), mrepo=mrepo, urepo=FakeUserRepo({9})) + await svc.update_role("p", 9, "owner") + assert mrepo.role_updates == [("p", 9, "owner")] + + +# --------------------------------------------------------------------------- # +# relationships +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_get_related_files_delegates(): + out = await _svc().get_related_files("p", "rel-1") + assert out == [{"file_id": "a", "relationship_id": "rel-1"}] + + +@pytest.mark.asyncio +async def test_get_file_ancestors_missing_file_404(): + svc = _svc(drepo=FakeDocumentRepo(files=set())) + with pytest.raises(NotFoundError): + await svc.get_file_ancestors("p", "f") + + +@pytest.mark.asyncio +async def test_get_file_ancestors_success(): + svc = _svc(drepo=FakeDocumentRepo(files={("f", "p")})) + out = await svc.get_file_ancestors("p", "f") + assert out[-1]["file_id"] == "f" diff --git a/openrag/services/orchestrators/test_query_service.py b/openrag/services/orchestrators/test_query_service.py new file mode 100644 index 000000000..99f533636 --- /dev/null +++ b/openrag/services/orchestrators/test_query_service.py @@ -0,0 +1,235 @@ +"""Unit tests for :class:`QueryService` (Phase 8C.2). + +The Ray-backed LLM semaphore and the model-file-backed language detector +are monkeypatched (both are infra concerns exercised in integration, not +here). Retrieval is faked; the real ``format_context`` / +``stream_with_source_filtering`` helpers run against real ``Chunk`` → +``Document`` conversions. +""" + +from __future__ import annotations + +import json +from contextlib import asynccontextmanager +from types import SimpleNamespace + +import pytest +import services.orchestrators.query_service as qs +from core.models.chunk import Chunk +from services.orchestrators.query_service import QueryService + + +@pytest.fixture(autouse=True) +def _patch_infra(monkeypatch): + @asynccontextmanager + async def _noop_sem(): + yield + + monkeypatch.setattr(qs, "get_llm_semaphore", _noop_sem) + monkeypatch.setattr(qs, "detect_language", lambda _t: "en") + + +class FakeLLM: + def __init__(self, *, chat_responses=None, gen_text="answer", stream_lines=None): + self._chat_responses = list(chat_responses or []) + self._gen_text = gen_text + self._stream_lines = stream_lines or ['data: {"choices":[{"delta":{"content":"hi"}}]}\n\n', "data: [DONE]\n\n"] + self.chat_calls: list = [] + + async def chat(self, messages, **kwargs): + self.chat_calls.append((messages, kwargs)) + if self._chat_responses: + content = self._chat_responses.pop(0) + else: + content = "final answer" + return {"choices": [{"message": {"content": content}}]} + + async def generate(self, prompt, **kwargs): + return {"choices": [{"text": self._gen_text}]} + + async def stream_chat(self, messages, **kwargs): + for line in self._stream_lines: + yield line + + +class FakeRetrieval: + def __init__(self, chunks=None): + self._chunks = chunks if chunks is not None else [Chunk(id="c1", text="ctx", metadata={"_id": "c1"})] + + async def retrieve_multi(self, **kwargs): + return list(self._chunks) + + async def retrieve_per_query(self, *, queries, **kwargs): + return [list(self._chunks) for _ in queries] + + @staticmethod + def fuse(doc_lists, top_k=None): + return doc_lists[0] if doc_lists else [] + + +class FakeWeb: + max_tokens = 2000 + + def __init__(self, results=None): + self._results = results or [] + + async def search(self, query): + return list(self._results) + + +class FakeWorkspace: + async def get_workspace(self, wid): + return None + + +def _config(mode="SimpleRag"): + return SimpleNamespace( + rag=SimpleNamespace(mode=mode, chat_history_depth=4, max_contextualized_query_len=512), + reranker=SimpleNamespace(top_k=5), + chunker=SimpleNamespace(chunk_size=512), + map_reduce=SimpleNamespace(initial_batch_size=2, expansion_batch_size=2, max_total_documents=4), + ) + + +def _svc(*, llm=None, retrieval=None, web=None, mode="SimpleRag") -> QueryService: + return QueryService( + retrieval_service=retrieval or FakeRetrieval(), + llm=llm or FakeLLM(), + config=_config(mode), + web_search_service=web or FakeWeb(), + workspace_service=FakeWorkspace(), + ) + + +# --------------------------------------------------------------------------- # +# generate_query +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_generate_query_simplerag_uses_last_message(): + sq = await _svc(mode="SimpleRag").generate_query([{"role": "user", "content": "what is X?"}]) + assert [q.query for q in sq.query_list] == ["what is X?"] + + +@pytest.mark.asyncio +async def test_generate_query_chatbotrag_parses_json(): + payload = json.dumps({"query_list": [{"query": "rewritten", "temporal_filters": None}]}) + svc = _svc(llm=FakeLLM(chat_responses=[payload]), mode="ChatBotRag") + sq = await svc.generate_query([{"role": "user", "content": "hi"}]) + assert sq.query_list[0].query == "rewritten" + + +@pytest.mark.asyncio +async def test_generate_query_chatbotrag_falls_back_on_garbage(): + svc = _svc(llm=FakeLLM(chat_responses=["not json", "still not json"]), mode="ChatBotRag") + sq = await svc.generate_query([{"role": "user", "content": "raw question"}]) + assert sq.query_list[0].query == "raw question" # fallback to raw user query + + +# --------------------------------------------------------------------------- # +# chat / complete +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_chat_direct_mode_skips_retrieval(): + retrieval = FakeRetrieval() + called = {"n": 0} + + async def _spy(**kwargs): + called["n"] += 1 + return [] + + retrieval.retrieve_multi = _spy + svc = _svc(retrieval=retrieval, llm=FakeLLM(chat_responses=["hello [Sources: none]"])) + out = await svc.chat( + partitions=None, + payload={"messages": [{"role": "user", "content": "hi"}], "metadata": {}}, + prepare_sources=lambda d, w: [{"source_type": "document"}], + model_name="m1", + ) + assert called["n"] == 0 # no retrieval in direct mode + assert out["model"] == "m1" + assert out["choices"][0]["message"]["content"] == "hello" # sources tag stripped + assert json.loads(out["extra"])["sources"] == [] # [Sources: none] → no sources + + +@pytest.mark.asyncio +async def test_chat_with_partition_retrieves_and_filters_sources(): + svc = _svc(llm=FakeLLM(chat_responses=["answer [Sources: 1]"])) + sources = [{"source_type": "document", "n": 1}, {"source_type": "document", "n": 2}] + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "q"}], "metadata": {}}, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + filtered = json.loads(out["extra"])["sources"] + assert filtered == [{"source_type": "document", "n": 1}] # only cited source 1 + + +@pytest.mark.asyncio +async def test_complete_strips_and_filters(): + svc = _svc(llm=FakeLLM(gen_text="text body [Sources: none]")) + out = await svc.complete( + partitions=None, + payload={"prompt": "do x"}, + prepare_sources=lambda d, w: [{"x": 1}], + ) + assert out["choices"][0]["text"] == "text body" + assert json.loads(out["extra"])["sources"] == [] + + +@pytest.mark.asyncio +async def test_chat_stream_yields_sse_and_done(): + svc = _svc(llm=FakeLLM()) + lines = [] + async for line in svc.chat_stream( + partitions=None, + payload={"messages": [{"role": "user", "content": "hi"}], "metadata": {}}, + prepare_sources=lambda d, w: [], + model_name="m", + ): + lines.append(line) + assert any("[DONE]" in ln for ln in lines) + + +# --------------------------------------------------------------------------- # +# map-reduce +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_map_reduce_keeps_relevant_drops_irrelevant(): + rel = json.dumps({"relevancy": True, "summary": "kept"}) + irr = json.dumps({"relevancy": False, "summary": ""}) + svc = _svc(llm=FakeLLM(chat_responses=[rel, irr])) + docs = [ + Chunk(id="a", text="A", metadata={"_id": "a"}).to_langchain(), + Chunk(id="b", text="B", metadata={"_id": "b"}).to_langchain(), + ] + out = await svc._map_reduce("q", docs) + assert len(out) == 1 + assert out[0].page_content == "kept" + + +# --------------------------------------------------------------------------- # +# helpers +# --------------------------------------------------------------------------- # + + +def test_json_slice_extracts_object(): + assert qs._json_slice('noise {"a": 1} trailing') == '{"a": 1}' + + +def test_dedupe_web_preserves_first_seen(): + a = SimpleNamespace(url="u1") + b = SimpleNamespace(url="u1") + c = SimpleNamespace(url="u2") + assert qs._dedupe_web([[a, b], [c]]) == [a, c] + + +def test_sampling_strips_transport_keys(): + out = qs._sampling({"messages": [], "stream": True, "model": "m", "temperature": 0.5}) + assert out == {"temperature": 0.5} diff --git a/openrag/services/orchestrators/test_retrieval_service.py b/openrag/services/orchestrators/test_retrieval_service.py new file mode 100644 index 000000000..ed4b570f1 --- /dev/null +++ b/openrag/services/orchestrators/test_retrieval_service.py @@ -0,0 +1,171 @@ +"""Unit tests for :class:`RetrievalService` (Phase 8C.1). + +The Ray-backed searcher is faked (the service is constructed with a +``RetrievalSearcher`` stub, exactly as the container will inject +``MilvusRayShim``). Default config uses the ``single`` retriever with +the reranker disabled, so the core pipeline path is exercised end-to-end +without inference services. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from core.models.chunk import Chunk +from core.models.query import Query, SearchQueries +from services.orchestrators.retrieval_service import RetrievalService + + +def _chunk(cid: str, text: str = "t") -> Chunk: + return Chunk(id=cid, text=text, metadata={"_id": cid}) + + +class FakeSearcher: + def __init__(self): + self.search_calls: list[dict] = [] + self.search_result: list[Chunk] = [] + self.related_result: list[Chunk] = [] + self.ancestor_result: list[Chunk] = [] + + async def search(self, **kwargs): + self.search_calls.append(kwargs) + return list(self.search_result) + + async def multi_query_search(self, **kwargs): + return list(self.search_result) + + async def get_related_chunks(self, **kwargs): + return list(self.related_result) + + async def get_ancestor_chunks(self, **kwargs): + return list(self.ancestor_result) + + +def _config(rtype: str = "single", reranker_enabled: bool = False) -> SimpleNamespace: + return SimpleNamespace( + retriever=SimpleNamespace( + type=rtype, + top_k=6, + similarity_threshold=0.5, + with_surrounding_chunks=False, + include_related=False, + include_ancestors=False, + related_limit=10, + max_ancestor_depth=None, + allow_filterless_fallback=True, + k_queries=3, + combine=False, + ), + reranker=SimpleNamespace(enabled=reranker_enabled, top_k=5), + ) + + +def _svc(searcher, *, rtype="single", reranker_enabled=False) -> RetrievalService: + return RetrievalService( + searcher=searcher, + reranker=None, + llm=None, + config=_config(rtype, reranker_enabled), + ) + + +# --------------------------------------------------------------------------- # +# search() — powers routers/search.py +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_search_normalizes_str_partition_and_passes_params(): + s = FakeSearcher() + s.search_result = [_chunk("1"), _chunk("2")] + out = await _svc(s).search( + text="hello", + partitions="p1", + top_k=7, + similarity_threshold=0.8, + filter="file_id == 'x'", + filter_params={"a": 1}, + ) + assert [c.id for c in out] == ["1", "2"] + call = s.search_calls[0] + assert call["partition"] == ["p1"] # str normalized to list + assert call["query"] == "hello" + assert call["top_k"] == 7 + assert call["similarity_threshold"] == 0.8 + assert call["filter"] == "file_id == 'x'" + assert call["filter_params"] == {"a": 1} + assert call["with_surrounding_chunks"] is True + + +@pytest.mark.asyncio +async def test_search_no_expansion_when_flags_off(): + s = FakeSearcher() + s.search_result = [_chunk("1")] + s.related_result = [_chunk("rel")] + out = await _svc(s).search(text="q", partitions=["p"], top_k=5, similarity_threshold=0.5) + assert [c.id for c in out] == ["1"] # related NOT included + + +@pytest.mark.asyncio +async def test_search_expands_related_when_requested(): + s = FakeSearcher() + # The core expand helper only fetches related chunks for source + # chunks that carry both a partition and a relationship_id. + src = Chunk(id="1", text="t", partition="p", metadata={"_id": "1", "relationship_id": "r1"}) + s.search_result = [src] + s.related_result = [_chunk("rel")] + out = await _svc(s).search( + text="q", + partitions=["p"], + top_k=5, + similarity_threshold=0.5, + include_related=True, + related_limit=3, + ) + ids = {c.id for c in out} + assert "1" in ids and "rel" in ids + + +# --------------------------------------------------------------------------- # +# retrieve / retrieve_multi / fuse — powers QueryService (8C.2) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_retrieve_single_query_via_pipeline(): + s = FakeSearcher() + s.search_result = [_chunk("a"), _chunk("b")] + out = await _svc(s).retrieve(partitions=["p"], query=Query(query="hi")) + assert [c.id for c in out] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_retrieve_multi_fuses_subqueries(): + s = FakeSearcher() + s.search_result = [_chunk("a"), _chunk("b")] + sq = SearchQueries(query_list=[Query(query="q1"), Query(query="q2")]) + out = await _svc(s).retrieve_multi(partitions=["p"], search_queries=sq) + assert {c.id for c in out} == {"a", "b"} + + +@pytest.mark.asyncio +async def test_retrieve_per_query_returns_unfused_lists(): + s = FakeSearcher() + s.search_result = [_chunk("a")] + out = await _svc(s).retrieve_per_query(partitions=["p"], queries=[Query(query="q1"), Query(query="q2")]) + assert len(out) == 2 + assert all(lst[0].id == "a" for lst in out) + + +def test_fuse_rrf_merges_and_dedupes(): + a, b, c = _chunk("a"), _chunk("b"), _chunk("c") + fused = RetrievalService.fuse([[a, b], [b, c]]) + ids = [x.id for x in fused] + assert set(ids) == {"a", "b", "c"} + assert ids[0] == "b" # appears in both lists -> highest RRF score + + +def test_fuse_respects_top_k(): + a, b, c = _chunk("a"), _chunk("b"), _chunk("c") + assert len(RetrievalService.fuse([[a, b], [b, c]], top_k=2)) == 2 diff --git a/openrag/services/orchestrators/test_user_service.py b/openrag/services/orchestrators/test_user_service.py new file mode 100644 index 000000000..462c1c9e9 --- /dev/null +++ b/openrag/services/orchestrators/test_user_service.py @@ -0,0 +1,300 @@ +"""Unit tests for :class:`UserService` (Phase 8A.2).""" + +from __future__ import annotations + +import pytest +from core.models.user import User +from core.utils.exceptions import UserNotFoundError, ValidationError +from models.user import UserCreate, UserUpdate +from services.orchestrators.user_service import UserService + + +class FakeUserRepo: + def __init__(self, existing: set[int] | None = None): + self._existing = existing if existing is not None else set() + self.created: list[dict] = [] + self.deleted: list[int] = [] + self.regenerated: list[int] = [] + self.regen_results: dict[int, dict] = {} + self.updated: list[tuple[int, dict]] = [] + self._users: dict[int, User] = {} + + async def user_exists(self, user_id: int) -> bool: + return user_id in self._existing + + async def create_legacy_user(self, *, display_name, external_user_id, email, is_admin, file_quota): + rec = { + "id": 42, + "display_name": display_name, + "external_user_id": external_user_id, + "email": email, + "token": "or-deadbeef", + "is_admin": is_admin, + "file_quota": file_quota, + "file_count": 0, + } + self.created.append(rec) + return rec + + async def list_users_dict(self): + return [{"id": 1, "display_name": "Admin"}] + + async def get_user_dict_by_id(self, user_id: int): + return {"id": user_id, "display_name": "U"} + + async def delete_user(self, user_id: int) -> bool: + self.deleted.append(user_id) + return True + + async def regenerate_user_token(self, user_id: int): + self.regenerated.append(user_id) + return self.regen_results.get(user_id) + + async def update_user(self, user_id: int, **fields): + self.updated.append((user_id, fields)) + return self._users.get(user_id) + + +class FakePartitionService: + def __init__(self): + self.deleted: list[str] = [] + + async def delete_partition(self, partition: str) -> None: + self.deleted.append(partition) + + +class FakeMembershipRepo: + def __init__(self, owned: dict[int, list[dict]] | None = None): + self._owned = owned or {} + + async def list_user_partitions_dict(self, user_id: int) -> list[dict]: + return self._owned.get(user_id, []) + + +class FakeJobService: + def __init__(self, pending: int = 0): + self._pending = pending + self.calls: list[int | None] = [] + + async def get_user_pending_task_count(self, user_id: int | None) -> int: + self.calls.append(user_id) + return self._pending + + +def _svc( + repo: FakeUserRepo, + *, + default_quota: int = 10, + partition_service: FakePartitionService | None = None, + membership_repo: FakeMembershipRepo | None = None, + job_service: FakeJobService | None = None, +) -> UserService: + return UserService( + user_repo=repo, + auth_service=object(), + default_file_quota=default_quota, + partition_service=partition_service or FakePartitionService(), + membership_repo=membership_repo or FakeMembershipRepo(), + job_service=job_service or FakeJobService(), + ) + + +# --------------------------------------------------------------------------- # +# create_user +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_create_user_passes_through_explicit_quota(): + repo = FakeUserRepo() + svc = _svc(repo, default_quota=10) + out = await svc.create_user(UserCreate(display_name="Bob", file_quota=3)) + assert out["token"] == "or-deadbeef" + assert repo.created[0]["file_quota"] == 3 + + +@pytest.mark.asyncio +async def test_create_user_applies_default_quota_when_none_and_default_positive(): + repo = FakeUserRepo() + svc = _svc(repo, default_quota=7) + await svc.create_user(UserCreate(display_name="Bob", file_quota=None)) + assert repo.created[0]["file_quota"] == 7 + + +@pytest.mark.asyncio +async def test_create_user_no_default_when_default_not_positive(): + repo = FakeUserRepo() + svc = _svc(repo, default_quota=-1) + await svc.create_user(UserCreate(display_name="Bob", file_quota=None)) + assert repo.created[0]["file_quota"] is None + + +@pytest.mark.asyncio +async def test_create_user_rejects_bad_email(): + svc = _svc(FakeUserRepo()) + with pytest.raises(ValidationError) as ei: + await svc.create_user(UserCreate(display_name="Bob", email="not-an-email")) + assert ei.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_create_user_rejects_overlong_display_name(): + svc = _svc(FakeUserRepo()) + with pytest.raises(ValidationError): + await svc.create_user(UserCreate(display_name="x" * 256)) + + +# --------------------------------------------------------------------------- # +# read / delete / regenerate / update +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_list_users_delegates(): + assert await _svc(FakeUserRepo()).list_users() == [{"id": 1, "display_name": "Admin"}] + + +@pytest.mark.asyncio +async def test_get_user_missing_raises_404(): + svc = _svc(FakeUserRepo(existing=set())) + with pytest.raises(UserNotFoundError) as ei: + await svc.get_user(9) + assert ei.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_get_user_existing_returns_dict(): + svc = _svc(FakeUserRepo(existing={9})) + assert await svc.get_user(9) == {"id": 9, "display_name": "U"} + + +@pytest.mark.asyncio +async def test_delete_user_missing_raises_and_no_repo_call(): + repo = FakeUserRepo(existing=set()) + with pytest.raises(UserNotFoundError): + await _svc(repo).delete_user(5) + assert repo.deleted == [] + + +@pytest.mark.asyncio +async def test_delete_user_existing_no_owned_partitions(): + repo = FakeUserRepo(existing={5}) + ps = FakePartitionService() + await _svc(repo, partition_service=ps).delete_user(5) + assert repo.deleted == [5] + assert ps.deleted == [] + + +@pytest.mark.asyncio +async def test_delete_user_cascades_owned_partitions_first(): + repo = FakeUserRepo(existing={5}) + ps = FakePartitionService() + mem = FakeMembershipRepo( + { + 5: [ + {"partition": "p_owned", "role": "owner"}, + {"partition": "p_viewer", "role": "viewer"}, # not cascaded + ] + } + ) + await _svc(repo, partition_service=ps, membership_repo=mem).delete_user(5) + assert ps.deleted == ["p_owned"] # only owner-role partitions + assert repo.deleted == [5] + + +@pytest.mark.asyncio +async def test_regenerate_token_missing_user_404(): + repo = FakeUserRepo(existing=set()) + with pytest.raises(UserNotFoundError): + await _svc(repo).regenerate_token(3) + + +@pytest.mark.asyncio +async def test_regenerate_token_repo_returns_none_404(): + repo = FakeUserRepo(existing={3}) # exists but repo regen returns None + with pytest.raises(UserNotFoundError): + await _svc(repo).regenerate_token(3) + + +@pytest.mark.asyncio +async def test_regenerate_token_success(): + repo = FakeUserRepo(existing={3}) + repo.regen_results[3] = {"id": 3, "token": "or-new"} + out = await _svc(repo).regenerate_token(3) + assert out == {"id": 3, "token": "or-new"} + assert repo.regenerated == [3] + + +@pytest.mark.asyncio +async def test_update_user_missing_404(): + repo = FakeUserRepo(existing=set()) + with pytest.raises(UserNotFoundError): + await _svc(repo).update_user(2, UserUpdate(display_name="X")) + + +@pytest.mark.asyncio +async def test_update_user_returns_legacy_dict_shape(): + repo = FakeUserRepo(existing={2}) + repo._users[2] = User(id=2, display_name="New", email="a@b.io", is_admin=True, file_quota=5, file_count=4) + out = await _svc(repo).update_user(2, UserUpdate(display_name="New")) + assert set(out) == { + "id", + "display_name", + "external_user_id", + "email", + "is_admin", + "created_at", + "file_quota", + "file_count", + } + assert out["id"] == 2 and out["display_name"] == "New" and out["file_count"] == 4 + assert isinstance(out["created_at"], str) # iso-formatted + + +@pytest.mark.asyncio +async def test_update_user_validates_email(): + repo = FakeUserRepo(existing={2}) + with pytest.raises(ValidationError): + await _svc(repo).update_user(2, UserUpdate(email="bogus")) + + +# --------------------------------------------------------------------------- # +# get_current_user_info — quota-usage block (8F: moved out of the router) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_current_user_info_specific_quota_and_pending(): + job = FakeJobService(pending=3) + svc = _svc(FakeUserRepo(), default_quota=10, job_service=job) + + out = await svc.get_current_user_info({"id": 7, "is_admin": False, "file_quota": 5, "file_count": 4}) + + assert out["file_count"] == 4 + assert out["pending_files"] == 3 + assert out["total_files"] == 7 + assert out["file_quota"] == 5 + assert out["id"] == 7 # original fields preserved + assert job.calls == [7] + + +@pytest.mark.asyncio +async def test_current_user_info_admin_is_unlimited(): + svc = _svc(FakeUserRepo(), default_quota=10, job_service=FakeJobService(pending=1)) + out = await svc.get_current_user_info({"id": 1, "is_admin": True, "file_count": 2}) + assert out["file_quota"] == -1 + assert out["total_files"] == 3 + + +@pytest.mark.asyncio +async def test_current_user_info_none_quota_falls_back_to_default(): + svc = _svc(FakeUserRepo(), default_quota=8, job_service=FakeJobService()) + out = await svc.get_current_user_info({"id": 2, "is_admin": False, "file_count": 0}) + assert out["file_quota"] == 8 + + +@pytest.mark.asyncio +async def test_current_user_info_negative_default_is_unlimited(): + svc = _svc(FakeUserRepo(), default_quota=-1, job_service=FakeJobService()) + out = await svc.get_current_user_info({"id": 2, "is_admin": False, "file_quota": 3, "file_count": 0}) + assert out["file_quota"] == -1 diff --git a/openrag/services/orchestrators/test_workspace_service.py b/openrag/services/orchestrators/test_workspace_service.py new file mode 100644 index 000000000..368e0609e --- /dev/null +++ b/openrag/services/orchestrators/test_workspace_service.py @@ -0,0 +1,153 @@ +"""Unit tests for :class:`WorkspaceService` (Phase 8B.2).""" + +from __future__ import annotations + +import pytest +from services.orchestrators.workspace_service import WorkspaceService + + +class FakeWorkspaceRepo: + def __init__(self, *, workspace=None, orphaned=None): + self._workspace = workspace + self._orphaned = orphaned if orphaned is not None else [] + self.created: list[tuple] = [] + self.added: list[tuple[str, list[str]]] = [] + self.removed: list[tuple[str, str]] = [] + self.removed_from_all: list[tuple[str, str]] = [] + self.deleted: list[str] = [] + + async def get_workspace_dict(self, workspace_id: str): + return self._workspace + + async def list_workspaces_dict(self, partition: str) -> list[dict]: + return [{"workspace_id": "w1", "partition_name": partition}] + + async def create_workspace_legacy(self, workspace_id, partition, user_id, display_name): + self.created.append((workspace_id, partition, user_id, display_name)) + + async def get_existing_file_ids(self, partition: str, file_ids): + return [f for f in file_ids if f != "ghost"] + + async def add_files_to_workspace(self, workspace_id: str, file_ids): + self.added.append((workspace_id, file_ids)) + return [] + + async def remove_file_from_workspace(self, workspace_id: str, file_id: str) -> bool: + self.removed.append((workspace_id, file_id)) + return True + + async def list_workspace_files(self, workspace_id: str) -> list[str]: + return ["f1", "f2"] + + async def get_file_workspaces(self, file_id: str, partition: str) -> list[str]: + return ["w1", "w2"] + + async def delete_workspace(self, workspace_id: str) -> list[str]: + self.deleted.append(workspace_id) + return list(self._orphaned) + + async def remove_file_from_all_workspaces(self, file_id: str, partition: str) -> None: + self.removed_from_all.append((file_id, partition)) + + +class FakeDocumentRepo: + def __init__(self, *, fail_on: set[str] | None = None): + self._fail_on = fail_on or set() + self.removed: list[tuple[str, str]] = [] + + async def remove_file_from_partition(self, file_id: str, partition: str) -> bool: + if file_id in self._fail_on: + raise RuntimeError(f"boom:{file_id}") + self.removed.append((file_id, partition)) + return True + + +class FakeVectorStore: + def __init__(self, ids_by_file=None): + self._ids_by_file = ids_by_file or {} + self.deleted: list[list[str]] = [] + + async def query_ids_by_filter(self, collection, filters): + return list(self._ids_by_file.get(filters.get("file_id"), [])) + + async def delete(self, ids, collection="default") -> int: + self.deleted.append(list(ids)) + return len(ids) + + +def _svc(*, wrepo=None, drepo=None, vstore=None, collection="vdb") -> WorkspaceService: + return WorkspaceService( + workspace_repo=wrepo or FakeWorkspaceRepo(), + document_repo=drepo or FakeDocumentRepo(), + vector_store=vstore or FakeVectorStore(), + collection=collection, + ) + + +# --------------------------------------------------------------------------- # +# delegations +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_create_workspace_delegates(): + wrepo = FakeWorkspaceRepo() + await _svc(wrepo=wrepo).create_workspace("w1", "p", 5, "Disp") + assert wrepo.created == [("w1", "p", 5, "Disp")] + + +@pytest.mark.asyncio +async def test_get_existing_file_ids_filters(): + out = await _svc().get_existing_file_ids("p", ["a", "ghost", "b"]) + assert set(out) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_remove_file_and_list_and_workspaces(): + wrepo = FakeWorkspaceRepo() + svc = _svc(wrepo=wrepo) + assert await svc.remove_file("w1", "f1") is True + assert await svc.list_files("w1") == ["f1", "f2"] + assert await svc.get_file_workspaces("f1", "p") == ["w1", "w2"] + + +# --------------------------------------------------------------------------- # +# cross-cutting delete_workspace +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_delete_workspace_no_orphans(): + wrepo = FakeWorkspaceRepo(orphaned=[]) + drepo = FakeDocumentRepo() + vstore = FakeVectorStore() + out = await _svc(wrepo=wrepo, drepo=drepo, vstore=vstore).delete_workspace("p", "w1") + assert out == {"orphaned_files_deleted": 0, "orphaned_files_failed": []} + assert wrepo.deleted == ["w1"] + assert vstore.deleted == [] + assert drepo.removed == [] + + +@pytest.mark.asyncio +async def test_delete_workspace_cleans_orphans_vectors_and_rows(): + wrepo = FakeWorkspaceRepo(orphaned=["fA", "fB"]) + drepo = FakeDocumentRepo() + vstore = FakeVectorStore(ids_by_file={"fA": ["c1", "c2"], "fB": []}) + out = await _svc(wrepo=wrepo, drepo=drepo, vstore=vstore).delete_workspace("p", "w1") + + assert out == {"orphaned_files_deleted": 2, "orphaned_files_failed": []} + # fA had chunks -> a delete call; fB had none -> no delete call. + assert vstore.deleted == [["c1", "c2"]] + assert set(drepo.removed) == {("fA", "p"), ("fB", "p")} + assert set(wrepo.removed_from_all) == {("fA", "p"), ("fB", "p")} + + +@pytest.mark.asyncio +async def test_delete_workspace_collects_per_file_failures(): + wrepo = FakeWorkspaceRepo(orphaned=["good", "bad"]) + drepo = FakeDocumentRepo(fail_on={"bad"}) + vstore = FakeVectorStore(ids_by_file={"good": ["c1"], "bad": ["c2"]}) + out = await _svc(wrepo=wrepo, drepo=drepo, vstore=vstore).delete_workspace("p", "w1") + + assert out["orphaned_files_deleted"] == 1 + assert out["orphaned_files_failed"] == ["bad"] diff --git a/openrag/services/orchestrators/user_service.py b/openrag/services/orchestrators/user_service.py new file mode 100644 index 000000000..f2b941330 --- /dev/null +++ b/openrag/services/orchestrators/user_service.py @@ -0,0 +1,214 @@ +"""UserService — user CRUD orchestration (Phase 8A.2). + +Business logic extracted from ``routers/users.py``. The legacy router was +already mostly delegation (each endpoint issued one Ray ``vectordb`` +user call); this service owns the parts that were *not* pure +delegation: input validation, the default-quota rule, and the +existence / not-found semantics. It talks to the Phase 7 +:class:`UserRepository` directly instead of the Ray ``vectordb`` actor. + +Response shape is kept identical to the legacy endpoints — the repo's +``*_dict`` helpers reproduce the exact ``PartitionFileManager`` dict +contract the Ray actor used to expose, so existing clients are +unaffected. + +Owner-partition cascade (restored in Phase 8B): the legacy Ray +``delete_user`` deleted every partition the user owned (Milvus + +Postgres) before removing the row. That cross-cutting delete is owned by +PartitionService; :meth:`delete_user` composes it so the behaviour +matches the legacy endpoint again. +""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING, Any + +from core.utils.exceptions import UserNotFoundError, ValidationError +from utils.logger import get_logger + +if TYPE_CHECKING: + from core.ports.partition_membership_repo import PartitionMembershipRepository + from core.ports.user_repo import UserRepository + from models.user import UserCreate, UserUpdate + from services.orchestrators.auth_service import AuthService + from services.orchestrators.job_service import JobService + from services.orchestrators.partition_service import PartitionService + +logger = get_logger() + +# Pragmatic, permissive address shape — the IdP / caller is the real +# source of truth; this only rejects obviously malformed input. +_EMAIL_RE = re.compile(r"^[^@\s]+@[^@\s]+\.[^@\s]+$") +_MAX_DISPLAY_NAME = 255 + + +class UserService: + """User account CRUD — validation + repo delegation.""" + + def __init__( + self, + *, + user_repo: UserRepository, + auth_service: AuthService, + default_file_quota: int, + partition_service: PartitionService, + membership_repo: PartitionMembershipRepository, + job_service: JobService, + ) -> None: + self._user_repo = user_repo + # Injected per the Phase 8 prescribed signature: the place future + # phases consolidate authz (e.g. require_admin) once the shared + # FastAPI Depends wrappers in routers/utils.py are retired. + self._auth_service = auth_service + # Legacy ``file_quota_per_user`` (config.rdb.default_file_quota). + # Only applied as a creation default when > 0, matching the old + # ``vectordb.create_user`` behaviour. + self._default_file_quota = default_file_quota + # 8B: reinstating the owner-partition cascade dropped in 8A.2. + # delete_user must also delete every partition the user owns + # (Milvus + Postgres) — that cross-cutting delete is owned by + # PartitionService, so UserService composes it. + self._partition_service = partition_service + self._membership_repo = membership_repo + # /users/info reports quota usage = indexed files + pending + # indexing tasks. The pending count comes from the TaskStateManager + # actor, which JobService wraps (8H excepts JobService, not + # UserService) — so UserService composes JobService rather than + # holding a Ray handle itself. + self._job_service = job_service + + # ------------------------------------------------------------------ + # Validation + # ------------------------------------------------------------------ + + @staticmethod + def _validate_profile(display_name: Any, email: Any) -> None: + if isinstance(display_name, str) and len(display_name) > _MAX_DISPLAY_NAME: + raise ValidationError( + f"display_name exceeds {_MAX_DISPLAY_NAME} characters.", + status_code=400, + ) + if isinstance(email, str) and email.strip() and not _EMAIL_RE.match(email.strip()): + raise ValidationError(f"Invalid email address: {email!r}", status_code=400) + + async def _ensure_exists(self, user_id: int) -> None: + if not await self._user_repo.user_exists(user_id): + logger.warning(f"User with ID {user_id} does not exist.") + raise UserNotFoundError(f"User with ID {user_id} does not exist.") + + # ------------------------------------------------------------------ + # CRUD + # ------------------------------------------------------------------ + + async def create_user(self, body: UserCreate) -> dict: + """Create a user, returning the legacy dict (token shown once).""" + fields = body.model_dump() + self._validate_profile(fields.get("display_name"), fields.get("email")) + + file_quota = fields.get("file_quota") + if self._default_file_quota > 0 and file_quota is None: + file_quota = self._default_file_quota + + user = await self._user_repo.create_legacy_user( + display_name=fields.get("display_name"), + external_user_id=fields.get("external_user_id"), + email=fields.get("email"), + is_admin=fields.get("is_admin", False), + file_quota=file_quota, + ) + logger.info("Created new user", user_id=user["id"]) + return user + + async def get_current_user_info(self, user: dict) -> dict: + """Augment the authenticated user with the quota-usage block. + + ``user`` is the request-state dict the auth middleware set (no DB + fetch — identical to the legacy ``/users/info`` handler). Quota + rule, byte-for-byte: admins and a negative global default mean + unlimited; a ``None`` per-user quota falls back to the global + default; a negative per-user quota means unlimited. ``file_quota`` + is surfaced as ``-1`` when unlimited. + """ + is_admin = user.get("is_admin", False) + if is_admin or self._default_file_quota < 0: + user_quota: float | int = float("inf") + else: + user_quota = user.get("file_quota", None) + if user_quota is None: + user_quota = self._default_file_quota + elif user_quota < 0: + user_quota = float("inf") + + file_count = user.get("file_count", 0) + pending_count = await self._job_service.get_user_pending_task_count(user.get("id")) + total = file_count + pending_count + + return { + **user, + "file_count": file_count, + "pending_files": pending_count, + "total_files": total, + "file_quota": -1 if user_quota == float("inf") else user_quota, + } + + async def list_users(self) -> list[dict]: + users = await self._user_repo.list_users_dict() + logger.debug("Returned list of users.", user_count=len(users)) + return users + + async def get_user(self, user_id: int) -> dict: + await self._ensure_exists(user_id) + user = await self._user_repo.get_user_dict_by_id(user_id) + if user is None: + raise UserNotFoundError(f"User '{user_id}' not found") + return user + + async def delete_user(self, user_id: int) -> None: + """Delete a user, cascading partitions the user owns first. + + Mirrors the legacy Ray ``delete_user``: every partition where the + user holds the ``owner`` role is deleted (vectors + relational + rows, via PartitionService) before the user row is removed. + """ + await self._ensure_exists(user_id) + owned = [ + p["partition"] + for p in await self._membership_repo.list_user_partitions_dict(user_id) + if p.get("role") == "owner" + ] + for partition in owned: + await self._partition_service.delete_partition(partition) + await self._user_repo.delete_user(user_id) + logger.info("Deleted user", user_id=user_id, cascaded_partitions=len(owned)) + + async def regenerate_token(self, user_id: int) -> dict: + await self._ensure_exists(user_id) + user = await self._user_repo.regenerate_user_token(user_id) + if user is None: + raise UserNotFoundError(f"User '{user_id}' not found") + logger.info("Regenerated user token", user_id=user_id) + return user + + async def update_user(self, user_id: int, body: UserUpdate) -> dict: + await self._ensure_exists(user_id) + updates = body.model_dump(exclude_unset=True) + self._validate_profile(updates.get("display_name"), updates.get("email")) + + user = await self._user_repo.update_user(user_id, **updates) + if user is None: + raise UserNotFoundError(f"User '{user_id}' not found") + logger.info("Updated user info", user_id=user_id) + return { + "id": user.id, + "display_name": user.display_name, + "external_user_id": user.external_user_id, + "email": user.email, + "is_admin": user.is_admin, + "created_at": user.created_at.isoformat() if user.created_at else None, + "file_quota": user.file_quota, + "file_count": user.file_count, + } + + +__all__ = ["UserService"] diff --git a/openrag/services/orchestrators/workspace_service.py b/openrag/services/orchestrators/workspace_service.py new file mode 100644 index 000000000..c77457f10 --- /dev/null +++ b/openrag/services/orchestrators/workspace_service.py @@ -0,0 +1,154 @@ +"""WorkspaceService — workspace CRUD + file association (Phase 8B.2). + +Business logic extracted from ``routers/workspaces.py`` and the +workspace slice of the legacy Ray ``vectordb`` shim. The simple +endpoints were already 1:1 repo delegations; the substantive extraction +is :meth:`delete_workspace`, the cross-cutting op that drops the +workspace, then deletes every file orphaned by that removal from *both* +the vector store and the relational catalog (the legacy router looped the Ray vectordb +delete-file call itself). + +The thin router keeps the HTTP guards whose exact non-bracketed +``{"detail": ...}`` body must stay identical (409 on duplicate, the +``require_workspace_in_partition`` 404, the unknown/missing-file 404s). + +Constructor note: ``collection`` (vector-store collection name) is one +arg beyond the plan's three — the legacy ``delete_file`` read it from +``config.vectordb.collection_name``; the container supplies it from +settings so the service stays Ray/config-free (8H). +""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +from utils.logger import get_logger + +if TYPE_CHECKING: + from core.ports.document_repo import DocumentRepository + from core.ports.workspace_repo import WorkspaceRepository + from core.vector_stores import VectorStore + +logger = get_logger() + + +class WorkspaceService: + """Workspace lifecycle, file association and orphan cleanup.""" + + def __init__( + self, + *, + workspace_repo: WorkspaceRepository, + document_repo: DocumentRepository, + vector_store: VectorStore, + collection: str, + ) -> None: + self._workspace_repo = workspace_repo + self._document_repo = document_repo + self._vector_store = vector_store + self._collection = collection + + # ------------------------------------------------------------------ + # CRUD / lookups (thin repo delegations) + # ------------------------------------------------------------------ + + async def get_workspace(self, workspace_id: str) -> dict | None: + return await self._workspace_repo.get_workspace_dict(workspace_id) + + async def list_workspaces(self, partition: str) -> list[dict]: + return await self._workspace_repo.list_workspaces_dict(partition) + + async def create_workspace( + self, + workspace_id: str, + partition: str, + user_id: int | None = None, + display_name: str | None = None, + ) -> None: + """Create a workspace. + + The 409-on-exists guard lives in the thin router (byte-identical + non-bracketed body); this is the plain repo create. + """ + await self._workspace_repo.create_workspace_legacy( + workspace_id, + partition, + user_id, + display_name, + ) + + async def get_existing_file_ids(self, partition: str, file_ids: list[str]) -> list[str]: + return list(await self._workspace_repo.get_existing_file_ids(partition, file_ids)) + + async def add_files(self, workspace_id: str, file_ids: list[str]) -> list[str]: + """Associate files; returns any file_ids that were not found.""" + return await self._workspace_repo.add_files_to_workspace(workspace_id, file_ids) + + async def remove_file(self, workspace_id: str, file_id: str) -> bool: + return await self._workspace_repo.remove_file_from_workspace(workspace_id, file_id) + + async def list_files(self, workspace_id: str) -> list[str]: + return await self._workspace_repo.list_workspace_files(workspace_id) + + async def get_file_workspaces(self, file_id: str, partition: str) -> list[str]: + return await self._workspace_repo.get_file_workspaces(file_id, partition) + + # ------------------------------------------------------------------ + # Cross-cutting: delete workspace + clean up orphaned files + # ------------------------------------------------------------------ + + async def delete_workspace(self, partition: str, workspace_id: str) -> dict: + """Delete the workspace, then fully delete any files it orphaned. + + ``workspace_repo.delete_workspace`` removes the workspace and its + associations and returns the file_ids that are no longer + referenced by *any* workspace. Each of those is deleted from the + vector store and the relational catalog — concurrently, with + per-file failures collected rather than raised, matching the + legacy router's ``asyncio.gather(..., return_exceptions=True)``. + """ + orphaned = await self._workspace_repo.delete_workspace(workspace_id) + + deleted_count = 0 + failed_file_ids: list[str] = [] + if orphaned: + results = await asyncio.gather( + *[self._delete_file(file_id, partition) for file_id in orphaned], + return_exceptions=True, + ) + for file_id, result in zip(orphaned, results, strict=True): + if isinstance(result, Exception): + logger.warning( + "Failed to delete orphaned file from vector store", + file_id=file_id, + error=str(result), + ) + failed_file_ids.append(file_id) + else: + deleted_count += 1 + + return { + "orphaned_files_deleted": deleted_count, + "orphaned_files_failed": failed_file_ids, + } + + async def _delete_file(self, file_id: str, partition: str) -> None: + """Port of the legacy ``vectordb.delete_file``. + + Drops the file's chunks from the vector store (via the clean + port: query ids by filter + delete), then detaches it from every + workspace and removes the relational file row. + """ + ids = await self._vector_store.query_ids_by_filter( + self._collection, + {"partition": partition, "file_id": file_id}, + ) + if ids: + await self._vector_store.delete(ids, self._collection) + await self._workspace_repo.remove_file_from_all_workspaces(file_id, partition) + await self._document_repo.remove_file_from_partition(file_id=file_id, partition=partition) + logger.info("Deleted orphaned file", file_id=file_id, partition=partition) + + +__all__ = ["WorkspaceService"] diff --git a/openrag/services/persistence/__init__.py b/openrag/services/persistence/__init__.py new file mode 100644 index 000000000..30477f4b6 --- /dev/null +++ b/openrag/services/persistence/__init__.py @@ -0,0 +1,63 @@ +"""Postgres persistence adapter — connection manager, schema, repositories. + +This package contains the asyncpg-based Postgres adapter that replaces the +synchronous SQLAlchemy ORM in ``components/indexer/vectordb/utils.py``. + +Public entry points: + - :class:`connection.ConnectionManager` — pool lifecycle + migrations (7A.1) + - :mod:`schema` — metadata-only Alembic target (7A.1) + - Repositories (7A.2): + - Real (decomposed from ``PartitionFileManager``): + ``PgDocumentRepository``, ``PgUserRepository``, + ``PgPartitionRepository``, ``PgPartitionMembershipRepository``, + ``PgOIDCSessionRepository``, ``PgWorkspaceRepository``. + - Stubs (post-refactoring features — raise + :class:`StubRepositoryError` on every call): + ``PgJobRepository``, ``PgChunkRepository``, + ``PgPromptRepository``, ``PgConversationRepository``, + ``PgAuditLogRepository``, ``PgIdempotencyRepository``, + ``PgEntityRepository``, ``PgTopicTagRepository``, + ``PgModelEndpointRepository``, ``PgPresetRepository``. +""" + +from services.persistence._stubs import StubRepositoryError +from services.persistence.audit_log_repo import PgAuditLogRepository +from services.persistence.chunk_repo import PgChunkRepository +from services.persistence.connection import ConnectionManager +from services.persistence.conversation_repo import PgConversationRepository +from services.persistence.document_repo import PgDocumentRepository +from services.persistence.entity_repo import PgEntityRepository +from services.persistence.idempotency_repo import PgIdempotencyRepository +from services.persistence.job_repo import PgJobRepository +from services.persistence.model_endpoint_repo import PgModelEndpointRepository +from services.persistence.oidc_session_repo import PgOIDCSessionRepository +from services.persistence.partition_membership_repo import PgPartitionMembershipRepository +from services.persistence.partition_repo import PgPartitionRepository +from services.persistence.preset_repo import PgPresetRepository +from services.persistence.prompt_repo import PgPromptRepository +from services.persistence.schema import metadata +from services.persistence.topic_tag_repo import PgTopicTagRepository +from services.persistence.user_repo import PgUserRepository +from services.persistence.workspace_repo import PgWorkspaceRepository + +__all__ = [ + "ConnectionManager", + "PgAuditLogRepository", + "PgChunkRepository", + "PgConversationRepository", + "PgDocumentRepository", + "PgEntityRepository", + "PgIdempotencyRepository", + "PgJobRepository", + "PgModelEndpointRepository", + "PgOIDCSessionRepository", + "PgPartitionMembershipRepository", + "PgPartitionRepository", + "PgPresetRepository", + "PgPromptRepository", + "PgTopicTagRepository", + "PgUserRepository", + "PgWorkspaceRepository", + "StubRepositoryError", + "metadata", +] diff --git a/openrag/services/persistence/_stubs.py b/openrag/services/persistence/_stubs.py new file mode 100644 index 000000000..2e2eaf9ba --- /dev/null +++ b/openrag/services/persistence/_stubs.py @@ -0,0 +1,59 @@ +"""Shared stub primitives for repositories with no current source code. + +Phase 7A.2 ships placeholders for ports that have no equivalent in the +legacy :class:`components.indexer.vectordb.utils.PartitionFileManager`. +These are deliberate scaffolds — the architecture is hexagonal so that +adding the real implementation later is a one-file change. Until then +every stub method raises :class:`StubRepositoryError`, a distinctive +exception subclass that grep-finds easily when the post-refactoring +features come online. + +Why not silent fallbacks (return ``None`` / empty list)? Because that +hides bugs: an orchestrator that quietly retrieves zero rows from a +"feature does not exist" repo behaves indistinguishably from a real +empty repo. A loud exception forces the caller to opt in to the gap. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import asyncpg + + +class StubRepositoryError(NotImplementedError): + """Raised by every stub repository method. + + Subclass of :class:`NotImplementedError` so existing ``except + NotImplementedError`` blocks still catch it, but distinguishable + from a third-party library's NotImplementedError when tracing + production issues. + """ + + +def stub_not_implemented(feature: str) -> StubRepositoryError: + """Build a uniformly-phrased error so tracebacks are self-explanatory.""" + return StubRepositoryError( + f"{feature} is on the post-refactoring roadmap — see REFACTORING/Phase 7A.2 stubs.", + ) + + +class _StubRepositoryBase: + """Common asyncpg pool plumbing for stubs. + + A stub still receives a ``pool_getter`` callable so that when the + real implementation lands the constructor signature does not need + to change — only the method bodies. + """ + + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter + + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() + + +__all__ = ["StubRepositoryError", "stub_not_implemented", "_StubRepositoryBase"] diff --git a/openrag/services/persistence/audit_log_repo.py b/openrag/services/persistence/audit_log_repo.py new file mode 100644 index 000000000..8c49205d1 --- /dev/null +++ b/openrag/services/persistence/audit_log_repo.py @@ -0,0 +1,41 @@ +"""Stub :class:`AuditLogRepository`. + +Audit logging is a post-refactoring P2 feature (enterprise compliance): +an append-only record of "who did what when" against the catalog. When +that lands the implementation is straightforward — one INSERT per +sensitive API call, paginated SELECTs from an admin route — but no +table exists today so every method raises. +""" + +from __future__ import annotations + +from typing import Any + +from core.ports.audit_log_repo import AuditLogRepository +from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented + + +class PgAuditLogRepository(_StubRepositoryBase, AuditLogRepository): + """TODO: real impl once the ``audit_log`` table is added.""" + + async def insert( + self, + user_id: int | None, + action: str, + resource_type: str, + resource_id: str | None = None, + details_json: dict | None = None, + request_id: str | None = None, + ) -> None: + raise stub_not_implemented("Audit log") + + async def query( + self, + filters: dict[str, Any], + offset: int = 0, + limit: int = 50, + ) -> list[dict]: + raise stub_not_implemented("Audit log") + + +__all__ = ["PgAuditLogRepository"] diff --git a/openrag/services/persistence/chunk_repo.py b/openrag/services/persistence/chunk_repo.py new file mode 100644 index 000000000..1388ccf86 --- /dev/null +++ b/openrag/services/persistence/chunk_repo.py @@ -0,0 +1,39 @@ +"""Stub :class:`ChunkRepository`. + +Chunks currently live exclusively in Milvus — the vector store owns the +text, the embedding, and the BM25 sparse index. A Postgres-side chunk +table (with ``tsvector`` for FTS) is a post-refactoring feature: it +would unlock keyword-search routes that don't round-trip through +Milvus, plus easier full-table backups. Until that lands every method +raises :class:`StubRepositoryError`. +""" + +from __future__ import annotations + +from core.ports.chunk_repo import ChunkRepository +from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented + + +class PgChunkRepository(_StubRepositoryBase, ChunkRepository): + """TODO: real impl once the ``chunks`` table is added.""" + + async def bulk_insert(self, chunks: list[dict]) -> int: + raise stub_not_implemented("Postgres-side chunk storage") + + async def get_by_ids(self, chunk_ids: list[str]) -> list[dict]: + raise stub_not_implemented("Postgres-side chunk storage") + + async def get_by_document_id(self, document_id: str) -> list[dict]: + raise stub_not_implemented("Postgres-side chunk storage") + + async def delete_by_document_id(self, document_id: str) -> int: + raise stub_not_implemented("Postgres-side chunk storage") + + async def delete_by_partition(self, partition: str) -> int: + raise stub_not_implemented("Postgres-side chunk storage") + + async def bm25_search(self, query_text: str, partition: str, top_k: int = 20) -> list[dict]: + raise stub_not_implemented("Postgres-side BM25 / tsvector FTS") + + +__all__ = ["PgChunkRepository"] diff --git a/openrag/services/persistence/connection.py b/openrag/services/persistence/connection.py new file mode 100644 index 000000000..514085624 --- /dev/null +++ b/openrag/services/persistence/connection.py @@ -0,0 +1,210 @@ +"""Asynchronous Postgres connection manager. + +Owns a single :mod:`asyncpg` pool that every repository in +``services/persistence/`` borrows from. Lifecycle: + +1. :meth:`ConnectionManager.initialize` creates the pool with a bounded + exponential-backoff retry — the database container is often slower to + accept connections than the app at startup. +2. :meth:`ConnectionManager.run_migrations` upgrades the schema to ``head`` + by invoking Alembic synchronously. Safe to call after ``initialize()``. +3. :meth:`ConnectionManager.shutdown` closes the pool. + +The pool is exposed via the :attr:`pool` property; repositories receive a +``pool_getter`` callable so they always see the live pool even when the +manager is reinitialised in tests. +""" + +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from typing import TYPE_CHECKING + +import asyncpg +from utils.logger import get_logger + +if TYPE_CHECKING: + from core.config.infrastructure import RDBConfig + + +logger = get_logger() + + +_MIGRATIONS_DIR = Path(__file__).parent / "migrations" / "alembic" +_ALEMBIC_INI = _MIGRATIONS_DIR / "alembic.ini" + +_RETRY_ATTEMPTS = 5 +_RETRY_BASE_DELAY_SECONDS = 1.0 +_RETRY_MAX_DELAY_SECONDS = 16.0 + + +async def _init_connection(conn: asyncpg.Connection) -> None: + """Per-connection init: decode ``json`` / ``jsonb`` columns into Python. + + asyncpg returns JSON values as strings by default. Repositories prefer + receiving dicts so they don't have to ``json.loads`` every row read; we + register codecs for both flavours so the column type doesn't matter. + """ + await conn.set_type_codec( + "json", + encoder=json.dumps, + decoder=json.loads, + schema="pg_catalog", + ) + await conn.set_type_codec( + "jsonb", + encoder=json.dumps, + decoder=json.loads, + schema="pg_catalog", + ) + + +class ConnectionManager: + """Lifecycle wrapper around an :class:`asyncpg.Pool`.""" + + def __init__(self, config: RDBConfig) -> None: + if not config.database: + raise ValueError( + "RDBConfig.database is required for ConnectionManager — " + "set POSTGRES_DATABASE or wire it from the collection name." + ) + # Passed to asyncpg as discrete kwargs rather than a single DSN + # string: a password containing URL-significant characters + # (``@ : / ? # %``) would corrupt a string DSN and silently point + # the pool at the wrong host or fail auth. + self._conn_kwargs = { + "host": config.host, + "port": config.port, + "user": config.user, + "password": config.password, + "database": config.database, + } + # Password-free string used only for logs / error messages. + self._dsn_log = f"postgresql://{config.user}@{config.host}:{config.port}/{config.database}" + self._min_size = config.pool_min_size + self._max_size = config.pool_max_size + self._command_timeout = config.command_timeout + self._pool: asyncpg.Pool | None = None + + @property + def pool(self) -> asyncpg.Pool: + if self._pool is None: + raise RuntimeError( + "ConnectionManager.initialize() has not been called", + ) + return self._pool + + async def initialize(self) -> None: + """Open the pool with bounded exponential-backoff retry. + + Retries up to ``_RETRY_ATTEMPTS`` times with delays of 1, 2, 4, 8, 16 + seconds. Re-raises the final exception if every attempt fails. + """ + if self._pool is not None: + return + + await asyncio.to_thread(self._ensure_database_exists) + + last_exc: Exception | None = None + for attempt in range(1, _RETRY_ATTEMPTS + 1): + try: + self._pool = await asyncpg.create_pool( + **self._conn_kwargs, + min_size=self._min_size, + max_size=self._max_size, + command_timeout=self._command_timeout, + init=_init_connection, + ) + logger.info(f"Connected to Postgres at {self._dsn_log} (pool={self._min_size}..{self._max_size})") + return + except (OSError, asyncpg.PostgresError) as exc: + last_exc = exc + if attempt == _RETRY_ATTEMPTS: + break + delay = min( + _RETRY_BASE_DELAY_SECONDS * (2 ** (attempt - 1)), + _RETRY_MAX_DELAY_SECONDS, + ) + logger.warning( + f"Postgres connection attempt {attempt}/{_RETRY_ATTEMPTS} failed ({exc}); retrying in {delay:.1f}s" + ) + await asyncio.sleep(delay) + + assert last_exc is not None + raise RuntimeError( + f"Failed to connect to Postgres at {self._dsn_log} after {_RETRY_ATTEMPTS} attempts: {last_exc}", + ) from last_exc + + def _ensure_database_exists(self) -> None: + """Create the configured database before opening the asyncpg pool.""" + from sqlalchemy import URL + from sqlalchemy_utils import create_database, database_exists + + url = URL.create( + drivername="postgresql", + username=self._conn_kwargs["user"], + password=self._conn_kwargs["password"], + host=self._conn_kwargs["host"], + port=self._conn_kwargs["port"], + database=self._conn_kwargs["database"], + ) + if not database_exists(url): + create_database(url) + logger.info(f"Created Postgres database `{self._conn_kwargs['database']}`.") + + async def shutdown(self) -> None: + if self._pool is None: + return + await self._pool.close() + self._pool = None + + async def run_migrations(self) -> None: + """Upgrade the schema to ``head`` via Alembic. + + Alembic uses a synchronous SQLAlchemy engine, so the call is offloaded + to the default executor to avoid blocking the event loop. + """ + await asyncio.to_thread(self._run_migrations_sync) + + def _run_migrations_sync(self) -> None: + # Imported lazily so the module can be imported without alembic + # installed (e.g. for type-checking in tooling). + from alembic import command + from alembic.config import Config + from sqlalchemy import URL + + if not _ALEMBIC_INI.exists(): + raise FileNotFoundError( + f"Alembic config not found at {_ALEMBIC_INI}; did the migrations move?", + ) + + # Build the URL through SQLAlchemy so credentials with special + # characters are percent-encoded correctly. Alembic stores the + # value in a ConfigParser and reads it back with ``%``-interpolation, + # so escape literal ``%`` as ``%%`` to survive that round-trip. + db_url = URL.create( + "postgresql", + username=self._conn_kwargs["user"], + password=self._conn_kwargs["password"], + host=self._conn_kwargs["host"], + port=self._conn_kwargs["port"], + database=self._conn_kwargs["database"], + ).render_as_string(hide_password=False) + + cfg = Config(str(_ALEMBIC_INI)) + # ``script_location`` in alembic.ini is relative to the .ini file via + # %(here)s, so this resolves to services/persistence/migrations/alembic/. + cfg.set_main_option("script_location", str(_MIGRATIONS_DIR)) + cfg.set_main_option("sqlalchemy.url", db_url.replace("%", "%%")) + logger.info(f"Running Alembic migrations against {self._dsn_log}") + try: + command.upgrade(cfg, "head") + except Exception as exc: + logger.error(f"Alembic migration failed for {self._dsn_log}: {exc}") + raise RuntimeError(f"Failed to run Alembic migrations against {self._dsn_log}: {exc}") from exc + logger.info(f"Alembic migrations completed for {self._dsn_log}") + + +__all__ = ["ConnectionManager"] diff --git a/openrag/services/persistence/conversation_repo.py b/openrag/services/persistence/conversation_repo.py new file mode 100644 index 000000000..91e824d20 --- /dev/null +++ b/openrag/services/persistence/conversation_repo.py @@ -0,0 +1,42 @@ +"""Stub :class:`ConversationRepository`. + +OpenRAG does not persist chat history today — Chainlit handles transient +session state and nothing is written to Postgres. The post-refactoring +P2 feature is DB-persisted conversations + messages, used both for +resume-on-reconnect UX and as a source for fine-tuning datasets. +""" + +from __future__ import annotations + +from core.models.conversation import Conversation, Message +from core.ports.conversation_repo import ConversationRepository +from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented + + +class PgConversationRepository(_StubRepositoryBase, ConversationRepository): + """TODO: real impl once chat persistence ships.""" + + async def create_conversation(self, conversation: Conversation) -> Conversation: + raise stub_not_implemented("Chat persistence") + + async def get_conversation(self, conversation_id: str) -> Conversation | None: + raise stub_not_implemented("Chat persistence") + + async def list_conversations( + self, + user_id: int, + partition: str | None = None, + ) -> list[Conversation]: + raise stub_not_implemented("Chat persistence") + + async def delete_conversation(self, conversation_id: str) -> bool: + raise stub_not_implemented("Chat persistence") + + async def add_message(self, message: Message) -> Message: + raise stub_not_implemented("Chat persistence") + + async def list_messages(self, conversation_id: str) -> list[Message]: + raise stub_not_implemented("Chat persistence") + + +__all__ = ["PgConversationRepository"] diff --git a/openrag/services/persistence/document_repo.py b/openrag/services/persistence/document_repo.py new file mode 100644 index 000000000..3a6c7edbe --- /dev/null +++ b/openrag/services/persistence/document_repo.py @@ -0,0 +1,567 @@ +"""Postgres implementation of :class:`DocumentRepository`. + +Backed by the ``files`` table — the canonical record of every file indexed +into OpenRAG. The legacy :class:`components.indexer.vectordb.utils.PartitionFileManager` +exposed eight methods here that are decomposed onto this class: +``add_file_to_partition``, ``remove_file_from_partition``, +``update_file_metadata_in_db``, ``update_file_in_partition``, +``list_partition_files``, ``file_exists_in_partition``, +``get_files_by_relationship``, ``get_file_ancestors``. + +The new port methods (``create_document`` / ``get_document`` / ...) take the +clean :class:`DocumentRecord` domain model — those are what Phase 8 +orchestrators will call. The legacy method names are kept on the concrete +class (not on the ABC) so the Phase 7C shim can delegate to them unchanged. +Both sets read and write the same rows. + +Status / error_message / filename / created_at fields exist on +``DocumentRecord`` but not on the ``files`` table — they are derived from +``file_metadata`` or set to safe defaults. Tracking these in their own +columns is a post-refactoring feature. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +from core.models.catalog import DocumentRecord, DocumentStatus +from core.ports.document_repo import DocumentRepository + +if TYPE_CHECKING: + import asyncpg + +# Note on JSON: ``ConnectionManager.initialize`` registers a json/jsonb codec +# on every connection, so reading a JSON column yields a Python dict and +# binding a dict to a JSON parameter is encoded transparently. The repo +# therefore never calls ``json.dumps`` itself. + + +class PgDocumentRepository(DocumentRepository): + """asyncpg-backed implementation of :class:`DocumentRepository`.""" + + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter + + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() + + # ── DocumentRepository port methods ────────────────────────────── + + async def create_document(self, doc: DocumentRecord) -> DocumentRecord: + """Insert a document row keyed by (file_id, partition). + + The port-level ``DocumentRecord.id`` is treated as the natural + ``file_id`` — the legacy schema uses an integer surrogate PK but + every caller identifies documents by ``file_id``. If ``doc.id`` is + a default UUID and ``doc.file_id`` is also set, the explicit + ``file_id`` wins. + """ + file_id = doc.file_id or doc.id + metadata = dict(doc.metadata or {}) + if doc.filename and "filename" not in metadata: + metadata["filename"] = doc.filename + if doc.status and doc.status != DocumentStatus.QUEUED: + metadata["status"] = doc.status.value + if doc.error_message: + metadata["error_message"] = doc.error_message + await self.pool.execute( + """ + INSERT INTO files (file_id, partition_name, file_metadata, + created_by, relationship_id, parent_id) + VALUES ($1, $2, $3::json, $4, $5, $6) + """, + file_id, + doc.partition, + metadata, + doc.created_by, + doc.relationship_id, + doc.parent_id, + ) + return doc.model_copy(update={"file_id": file_id, "metadata": metadata}) + + async def get_document(self, document_id: str) -> DocumentRecord | None: + """Fetch a document by ``file_id`` (any partition). + + The current schema does not enforce ``file_id`` uniqueness across + partitions, so this returns the first match. Callers that need + partition-scoped lookup should use + :meth:`file_exists_in_partition` or the legacy + :meth:`list_partition_files`. + """ + row = await self.pool.fetchrow( + "SELECT * FROM files WHERE file_id = $1 LIMIT 1", + document_id, + ) + return self._row_to_document(row) if row else None + + async def list_documents( + self, + partition: str | list[str] | None = None, + status: str | None = None, + offset: int = 0, + limit: int = 50, + ) -> list[DocumentRecord]: + clauses: list[str] = [] + params: list[Any] = [] + if isinstance(partition, str): + params.append(partition) + clauses.append(f"partition_name = ${len(params)}") + elif isinstance(partition, list) and partition: + params.append(partition) + clauses.append(f"partition_name = ANY(${len(params)}::text[])") + if status: + # status is stored inside file_metadata; we filter on the JSON path + params.append(status) + clauses.append(f"file_metadata->>'status' = ${len(params)}") + where = f"WHERE {' AND '.join(clauses)}" if clauses else "" + params.extend([limit, offset]) + rows = await self.pool.fetch( + f"SELECT * FROM files {where} ORDER BY id DESC LIMIT ${len(params) - 1} OFFSET ${len(params)}", + *params, + ) + return [self._row_to_document(r) for r in rows] + + async def update_document(self, document_id: str, **fields: Any) -> DocumentRecord | None: + """Patch a document row by ``file_id``. + + Accepts the port's domain field names — ``metadata``, + ``status``, ``error_message``, ``relationship_id``, ``parent_id``, + ``filename``. ``status`` / ``error_message`` / ``filename`` are + folded into ``file_metadata`` since the schema has no dedicated + columns for them. + """ + if not fields: + return await self.get_document(document_id) + async with self.pool.acquire() as conn: + row = await conn.fetchrow( + "SELECT * FROM files WHERE file_id = $1 LIMIT 1", + document_id, + ) + if row is None: + return None + metadata = dict(row["file_metadata"] or {}) + sets: list[str] = [] + params: list[Any] = [] + + for json_only in ("filename", "status", "error_message"): + if json_only in fields: + value = fields.pop(json_only) + if json_only == "status" and hasattr(value, "value"): + value = value.value + metadata[json_only] = value + if "metadata" in fields: + merged = fields.pop("metadata") or {} + metadata.update(merged) + # We always rewrite file_metadata so JSON-only updates are persisted. + params.append(metadata) + sets.append(f"file_metadata = ${len(params)}::json") + + for column in ("relationship_id", "parent_id", "created_by"): + if column in fields: + params.append(fields.pop(column)) + sets.append(f"{column} = ${len(params)}") + if "partition" in fields: + params.append(fields.pop("partition")) + sets.append(f"partition_name = ${len(params)}") + + # Silently ignore any unknown keys to match Pydantic-style flexibility. + params.append(row["id"]) + await conn.execute( + f"UPDATE files SET {', '.join(sets)} WHERE id = ${len(params)}", + *params, + ) + updated = await conn.fetchrow("SELECT * FROM files WHERE id = $1", row["id"]) + return self._row_to_document(updated) if updated else None + + async def delete_document(self, document_id: str) -> bool: + """Delete a document by ``file_id`` across any partition. + + Decrements the uploader's ``file_count`` (clamped at zero) to keep + quota accounting honest, matching :meth:`remove_file_from_partition`. + """ + async with self.pool.acquire() as conn: + async with conn.transaction(): + row = await conn.fetchrow( + "SELECT id, created_by FROM files WHERE file_id = $1 LIMIT 1", + document_id, + ) + if row is None: + return False + await conn.execute("DELETE FROM files WHERE id = $1", row["id"]) + if row["created_by"] is not None: + await conn.execute( + "UPDATE users SET file_count = GREATEST(file_count - 1, 0) WHERE id = $1", + row["created_by"], + ) + return True + + async def delete_documents_by_partition(self, partition: str) -> int: + """Bulk-delete every file in a partition and decrement uploader counts.""" + async with self.pool.acquire() as conn: + async with conn.transaction(): + uploader_rows = await conn.fetch( + """ + SELECT created_by, COUNT(*)::int AS n + FROM files + WHERE partition_name = $1 AND created_by IS NOT NULL + GROUP BY created_by + """, + partition, + ) + result = await conn.execute( + "DELETE FROM files WHERE partition_name = $1", + partition, + ) + for r in uploader_rows: + await conn.execute( + "UPDATE users SET file_count = GREATEST(file_count - $1, 0) WHERE id = $2", + r["n"], + r["created_by"], + ) + # asyncpg returns 'DELETE ' as the command tag. + try: + return int(result.split()[-1]) + except (ValueError, IndexError): + return 0 + + async def count_documents( + self, + partition: str | list[str] | None = None, + status: str | None = None, + ) -> int: + clauses: list[str] = [] + params: list[Any] = [] + if isinstance(partition, str): + params.append(partition) + clauses.append(f"partition_name = ${len(params)}") + elif isinstance(partition, list) and partition: + params.append(partition) + clauses.append(f"partition_name = ANY(${len(params)}::text[])") + if status: + params.append(status) + clauses.append(f"file_metadata->>'status' = ${len(params)}") + where = f"WHERE {' AND '.join(clauses)}" if clauses else "" + return await self.pool.fetchval(f"SELECT COUNT(*)::int FROM files {where}", *params) + + async def file_exists_in_partition(self, file_id: str, partition: str) -> bool: + return await self.pool.fetchval( + "SELECT EXISTS (SELECT 1 FROM files WHERE file_id = $1 AND partition_name = $2)", + file_id, + partition, + ) + + # ── Legacy method names used by the Phase 7C shim ──────────────── + # These are NOT on the ABC. Phase 8 orchestrators must not depend on + # them — they exist solely so the shim can keep every legacy caller + # working unchanged until Phase 9 deletes the actor. Mark TODO so they + # are easy to grep for and remove later. + + async def add_file_to_partition( # noqa: PLR0913 — legacy signature pinned + self, + file_id: str, + partition: str, + file_metadata: dict | None = None, + user_id: int | None = None, + relationship_id: str | None = None, + parent_id: str | None = None, + ) -> bool: + """TODO(phase-9): remove. Mirror of legacy ``add_file_to_partition``. + + Creates the partition row on first use (legacy behaviour). Returns + ``False`` if a row with the same (file_id, partition) already exists. + """ + async with self.pool.acquire() as conn: + async with conn.transaction(): + existing = await conn.fetchval( + "SELECT 1 FROM files WHERE file_id = $1 AND partition_name = $2", + file_id, + partition, + ) + if existing: + return False + + # Auto-create partition + first-owner membership when missing — + # legacy side-effect documented in the phase-7 spec. + created = await conn.fetchval( + """ + INSERT INTO partitions (partition, created_at) + VALUES ($1, NOW()) + ON CONFLICT (partition) DO NOTHING + RETURNING 1 + """, + partition, + ) + if created and user_id is not None: + await conn.execute( + """ + INSERT INTO partition_memberships (partition_name, user_id, role, added_at) + VALUES ($1, $2, 'owner', NOW()) + ON CONFLICT (partition_name, user_id) DO NOTHING + """, + partition, + user_id, + ) + + await conn.execute( + """ + INSERT INTO files (file_id, partition_name, file_metadata, + created_by, relationship_id, parent_id) + VALUES ($1, $2, $3::json, $4, $5, $6) + """, + file_id, + partition, + file_metadata or {}, + user_id, + relationship_id, + parent_id, + ) + if user_id is not None: + await conn.execute( + "UPDATE users SET file_count = file_count + 1 WHERE id = $1", + user_id, + ) + return True + + async def remove_file_from_partition(self, file_id: str, partition: str) -> bool: + """TODO(phase-9): remove. Mirror of legacy ``remove_file_from_partition``.""" + async with self.pool.acquire() as conn: + async with conn.transaction(): + row = await conn.fetchrow( + "SELECT id, created_by FROM files WHERE file_id = $1 AND partition_name = $2", + file_id, + partition, + ) + if row is None: + return False + await conn.execute("DELETE FROM files WHERE id = $1", row["id"]) + if row["created_by"] is not None: + await conn.execute( + "UPDATE users SET file_count = GREATEST(file_count - 1, 0) WHERE id = $1", + row["created_by"], + ) + return True + + async def update_file_metadata_in_db( + self, + file_id: str, + partition: str, + file_metadata: dict, + ) -> bool: + """TODO(phase-9): remove. Updates ``file_metadata`` + syncs structured columns. + + Mirrors the legacy behaviour: when the new metadata blob contains + ``relationship_id`` or ``parent_id`` keys, the dedicated columns are + rewritten too so the JSON never diverges from the structured fields. + """ + rel_id = file_metadata.get("relationship_id") if "relationship_id" in file_metadata else None + parent_id = file_metadata.get("parent_id") if "parent_id" in file_metadata else None + sets = ["file_metadata = $1::json"] + params: list[Any] = [file_metadata] + if "relationship_id" in file_metadata: + params.append(rel_id) + sets.append(f"relationship_id = ${len(params)}") + if "parent_id" in file_metadata: + params.append(parent_id) + sets.append(f"parent_id = ${len(params)}") + params.extend([file_id, partition]) + result = await self.pool.execute( + f""" + UPDATE files SET {", ".join(sets)} + WHERE file_id = ${len(params) - 1} AND partition_name = ${len(params)} + """, + *params, + ) + return result.endswith(" 1") + + _UNSET = object() + + async def update_file_in_partition( + self, + file_id: str, + partition: str, + file_metadata: dict | None = None, + relationship_id: object = _UNSET, + parent_id: object = _UNSET, + ) -> bool: + """TODO(phase-9): remove. PUT-style in-place update. + + Preserves the underlying ``files.id`` so workspace FK rows stay + valid. Pass ``relationship_id=None`` / ``parent_id=None`` + explicitly to clear; omit the kwarg to leave the column alone. + """ + sets: list[str] = [] + params: list[Any] = [] + if file_metadata is not None: + params.append(file_metadata) + sets.append(f"file_metadata = ${len(params)}::json") + if relationship_id is not self._UNSET: + params.append(relationship_id) + sets.append(f"relationship_id = ${len(params)}") + if parent_id is not self._UNSET: + params.append(parent_id) + sets.append(f"parent_id = ${len(params)}") + if not sets: + # Match legacy: report whether the row exists at all. + return await self.file_exists_in_partition(file_id, partition) + params.extend([file_id, partition]) + result = await self.pool.execute( + f""" + UPDATE files SET {", ".join(sets)} + WHERE file_id = ${len(params) - 1} AND partition_name = ${len(params)} + """, + *params, + ) + return result.endswith(" 1") + + async def list_partition_files( + self, + partition: str, + limit: int | None = None, + ) -> dict: + """TODO(phase-9): remove. Returns ``{"files": [...]}`` shape used by routers.""" + sql = "SELECT * FROM files WHERE partition_name = $1" + params: list[Any] = [partition] + if limit is not None: + params.append(limit) + sql += f" LIMIT ${len(params)}" + rows = await self.pool.fetch(sql, *params) + if not rows: + return {} + return {"files": [self._row_to_dict(r) for r in rows]} + + async def get_files_by_relationship( + self, + partition: str, + relationship_id: str, + ) -> list[dict]: + """TODO(phase-9): remove.""" + rows = await self.pool.fetch( + "SELECT * FROM files WHERE partition_name = $1 AND relationship_id = $2", + partition, + relationship_id, + ) + return [self._row_to_dict(r) for r in rows] + + async def get_file_ids_by_relationship( + self, + partition: str, + relationship_id: str, + ) -> list[str]: + """TODO(phase-9): remove.""" + rows = await self.pool.fetch( + "SELECT file_id FROM files WHERE partition_name = $1 AND relationship_id = $2", + partition, + relationship_id, + ) + return [r["file_id"] for r in rows] + + async def get_file_ancestors( + self, + partition: str, + file_id: str, + max_ancestor_depth: int | None = None, + ) -> list[dict]: + """TODO(phase-9): remove. Recursive CTE walking ``parent_id`` upward. + + Returns a list ordered from root → self (depth DESC). When + ``max_ancestor_depth`` is given, the recursion stops once the + accumulated depth meets the cap. + """ + depth_filter = "" + params: list[Any] = [file_id, partition] + if max_ancestor_depth is not None: + params.append(max_ancestor_depth) + depth_filter = f"AND a.depth < ${len(params)}" + rows = await self.pool.fetch( + f""" + WITH RECURSIVE ancestors AS ( + SELECT id, file_id, partition_name, parent_id, file_metadata, + relationship_id, 0 AS depth + FROM files + WHERE file_id = $1 AND partition_name = $2 + AND relationship_id IS NOT NULL + UNION ALL + SELECT f.id, f.file_id, f.partition_name, f.parent_id, + f.file_metadata, f.relationship_id, a.depth + 1 + FROM files f + INNER JOIN ancestors a + ON f.file_id = a.parent_id + AND f.partition_name = a.partition_name + AND f.relationship_id IS NOT NULL + {depth_filter} + ) + SELECT * FROM ancestors ORDER BY depth DESC + """, + *params, + ) + out: list[dict] = [] + for r in rows: + metadata = r["file_metadata"] or {} + out.append( + { + "file_id": r["file_id"], + "partition": r["partition_name"], + "parent_id": r["parent_id"], + "relationship_id": r["relationship_id"], + "depth": r["depth"], + **metadata, + }, + ) + return out + + async def get_ancestor_file_ids( + self, + partition: str, + file_id: str, + max_ancestor_depth: int | None = None, + ) -> list[str]: + """TODO(phase-9): remove.""" + ancestors = await self.get_file_ancestors(partition, file_id, max_ancestor_depth) + return [a["file_id"] for a in ancestors] + + # ── Row → domain helpers ───────────────────────────────────────── + + @staticmethod + def _row_to_dict(row: asyncpg.Record) -> dict: + """Replica of the legacy ``File.to_dict()`` ORM shape. + + The legacy routers consume this exact shape (``partition``, + ``file_id``, ``relationship_id``, ``parent_id`` plus every metadata + key flattened in). Used by the shim's pass-through calls. + """ + metadata = row["file_metadata"] or {} + return { + "partition": row["partition_name"], + "file_id": row["file_id"], + "relationship_id": row["relationship_id"], + "parent_id": row["parent_id"], + **metadata, + } + + @staticmethod + def _row_to_document(row: asyncpg.Record) -> DocumentRecord: + metadata = dict(row["file_metadata"] or {}) + status_raw = metadata.pop("status", None) + error_message = metadata.pop("error_message", None) + filename = metadata.pop("filename", "") or "" + try: + status = DocumentStatus(status_raw) if status_raw else DocumentStatus.QUEUED + except ValueError: + status = DocumentStatus.QUEUED + return DocumentRecord( + id=row["file_id"], + file_id=row["file_id"], + filename=filename, + partition=row["partition_name"], + metadata=metadata, + status=status, + error_message=error_message, + created_by=row["created_by"], + relationship_id=row["relationship_id"], + parent_id=row["parent_id"], + ) + + +__all__ = ["PgDocumentRepository"] diff --git a/openrag/services/persistence/entity_repo.py b/openrag/services/persistence/entity_repo.py new file mode 100644 index 000000000..3de095a32 --- /dev/null +++ b/openrag/services/persistence/entity_repo.py @@ -0,0 +1,37 @@ +"""Stub :class:`EntityRepository`. + +Entity extraction (canonical name + aliases per partition) is a +post-refactoring NER feature with no current implementation. The port +shape is pinned so that when an extraction pipeline is added, only this +file changes. +""" + +from __future__ import annotations + +from core.ports.entity_repo import EntityRepository +from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented + + +class PgEntityRepository(_StubRepositoryBase, EntityRepository): + """TODO: real impl once the ``entities`` table is added.""" + + async def upsert( + self, + partition: str, + entity_type: str, + canonical_name: str, + aliases: list[str], + ) -> str: + raise stub_not_implemented("NER / entity storage") + + async def search(self, partition: str, query: str, top_k: int = 10) -> list[dict]: + raise stub_not_implemented("NER / entity storage") + + async def get_by_document(self, document_id: str) -> list[dict]: + raise stub_not_implemented("NER / entity storage") + + async def delete_by_document(self, document_id: str) -> int: + raise stub_not_implemented("NER / entity storage") + + +__all__ = ["PgEntityRepository"] diff --git a/openrag/services/persistence/idempotency_repo.py b/openrag/services/persistence/idempotency_repo.py new file mode 100644 index 000000000..5ea535dd2 --- /dev/null +++ b/openrag/services/persistence/idempotency_repo.py @@ -0,0 +1,32 @@ +"""Stub :class:`IdempotencyRepository`. + +Idempotency support is a post-refactoring P3 feature: cache the +(method, path, body-hash) of a request so retries from a flaky client +return the original response instead of double-applying. When that +lands the implementation is a tiny table keyed by a SHA-256 hash with +a TTL on cleanup; no support exists today. +""" + +from __future__ import annotations + +from core.ports.idempotency_repo import IdempotencyRepository +from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented + + +class PgIdempotencyRepository(_StubRepositoryBase, IdempotencyRepository): + """TODO: real impl once the ``idempotency_keys`` table is added.""" + + async def get_by_hash(self, key_hash: str) -> dict | None: + raise stub_not_implemented("Idempotency keys") + + async def store( + self, + key_hash: str, + http_method: str, + status_code: int, + response_body: bytes, + ) -> None: + raise stub_not_implemented("Idempotency keys") + + +__all__ = ["PgIdempotencyRepository"] diff --git a/openrag/services/persistence/job_repo.py b/openrag/services/persistence/job_repo.py new file mode 100644 index 000000000..55cbbf3ce --- /dev/null +++ b/openrag/services/persistence/job_repo.py @@ -0,0 +1,41 @@ +"""Stub :class:`JobRepository` — see ``_stubs.py`` for the rationale. + +Job state is currently tracked in-memory by the +:class:`components.indexer.indexer.TaskStateManager` Ray actor. The +post-refactoring P0 feature is to persist jobs to Postgres so they +survive restarts and become visible to operators. When that lands, +swap the body of each method for an asyncpg implementation against a +new ``jobs`` table — the port shape is already pinned by Phase 4. +""" + +from __future__ import annotations + +from typing import Any + +from core.models.catalog import IndexationJob +from core.ports.job_repo import JobRepository +from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented + + +class PgJobRepository(_StubRepositoryBase, JobRepository): + """TODO: real impl once the ``jobs`` table is added. See REFACTORING P0 plan.""" + + async def create_job(self, job: IndexationJob) -> IndexationJob: + raise stub_not_implemented("DB-backed job tracking") + + async def get_job(self, job_id: str) -> IndexationJob | None: + raise stub_not_implemented("DB-backed job tracking") + + async def list_jobs( + self, + status: str | None = None, + offset: int = 0, + limit: int = 50, + ) -> list[IndexationJob]: + raise stub_not_implemented("DB-backed job tracking") + + async def update_job(self, job_id: str, **fields: Any) -> IndexationJob | None: + raise stub_not_implemented("DB-backed job tracking") + + +__all__ = ["PgJobRepository"] diff --git a/openrag/services/persistence/migrations/__init__.py b/openrag/services/persistence/migrations/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/scripts/migrations/alembic/README b/openrag/services/persistence/migrations/alembic/README similarity index 100% rename from openrag/scripts/migrations/alembic/README rename to openrag/services/persistence/migrations/alembic/README diff --git a/openrag/services/persistence/migrations/alembic/__init__.py b/openrag/services/persistence/migrations/alembic/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/scripts/migrations/alembic/alembic.ini b/openrag/services/persistence/migrations/alembic/alembic.ini similarity index 100% rename from openrag/scripts/migrations/alembic/alembic.ini rename to openrag/services/persistence/migrations/alembic/alembic.ini diff --git a/openrag/scripts/migrations/alembic/env.py b/openrag/services/persistence/migrations/alembic/env.py similarity index 64% rename from openrag/scripts/migrations/alembic/env.py rename to openrag/services/persistence/migrations/alembic/env.py index 88a6c4f78..d7954ab5b 100644 --- a/openrag/scripts/migrations/alembic/env.py +++ b/openrag/services/persistence/migrations/alembic/env.py @@ -7,8 +7,8 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from alembic import context -from components.indexer.vectordb.models import Base from config import load_config +from services.persistence.schema import metadata as target_metadata from sqlalchemy import URL, engine_from_config, pool rag_config = load_config() @@ -23,30 +23,35 @@ if config.config_file_name is not None: fileConfig(config.config_file_name) -# override the SQLALCHEMY URL with an environment variable -rdb_user = rag_config.rdb.user -rdb_password = rag_config.rdb.password -rdb_port = rag_config.rdb.port -rdb_host = rag_config.rdb.host - -collection_name = rag_config.vectordb.collection_name - -database_url = URL.create( - drivername="postgresql", - username=rdb_user, - password=rdb_password, - host=rdb_host, - port=rdb_port, - database=f"partitions_for_collection_{collection_name}", -) -config.set_main_option("sqlalchemy.url", database_url.render_as_string(hide_password=False)) - -# add your model's MetaData object here -# for 'autogenerate' support -# from myapp import mymodel -# target_metadata = mymodel.Base.metadata - -target_metadata = Base.metadata +# When the caller (typically ``ConnectionManager.run_migrations``) already +# wired a DSN into ``sqlalchemy.url``, defer to it. The default in +# ``alembic.ini`` is the placeholder ``driver://user:pass@localhost/dbname`` +# from Alembic's template; we treat that and an empty value as "fall back to +# the OpenRAG config". +preset_url = config.get_main_option("sqlalchemy.url") or "" +if (not preset_url) or preset_url.startswith("driver://"): + rdb_user = rag_config.rdb.user + rdb_password = rag_config.rdb.password + rdb_port = rag_config.rdb.port + rdb_host = rag_config.rdb.host + + collection_name = rag_config.vectordb.collection_name + + database_url = URL.create( + drivername="postgresql", + username=rdb_user, + password=rdb_password, + host=rdb_host, + port=rdb_port, + database=f"partitions_for_collection_{collection_name}", + ) + config.set_main_option( + "sqlalchemy.url", + database_url.render_as_string(hide_password=False), + ) + +# Metadata target for autogenerate is imported from +# `services.persistence.schema` (metadata-only Table definitions). # other values from the config, defined by the needs of env.py, # can be acquired: diff --git a/openrag/scripts/migrations/alembic/schema_helpers.py b/openrag/services/persistence/migrations/alembic/schema_helpers.py similarity index 100% rename from openrag/scripts/migrations/alembic/schema_helpers.py rename to openrag/services/persistence/migrations/alembic/schema_helpers.py diff --git a/openrag/scripts/migrations/alembic/script.py.mako b/openrag/services/persistence/migrations/alembic/script.py.mako similarity index 100% rename from openrag/scripts/migrations/alembic/script.py.mako rename to openrag/services/persistence/migrations/alembic/script.py.mako diff --git a/openrag/scripts/migrations/alembic/versions/4add4d260575_initial_migration.py b/openrag/services/persistence/migrations/alembic/versions/4add4d260575_initial_migration.py similarity index 100% rename from openrag/scripts/migrations/alembic/versions/4add4d260575_initial_migration.py rename to openrag/services/persistence/migrations/alembic/versions/4add4d260575_initial_migration.py diff --git a/openrag/services/persistence/migrations/alembic/versions/__init__.py b/openrag/services/persistence/migrations/alembic/versions/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/scripts/migrations/alembic/versions/a1b2c3d4e5f6_add_document_relationships.py b/openrag/services/persistence/migrations/alembic/versions/a1b2c3d4e5f6_add_document_relationships.py similarity index 100% rename from openrag/scripts/migrations/alembic/versions/a1b2c3d4e5f6_add_document_relationships.py rename to openrag/services/persistence/migrations/alembic/versions/a1b2c3d4e5f6_add_document_relationships.py diff --git a/openrag/scripts/migrations/alembic/versions/c224d4befe71_add_file_count_and_file_quota.py b/openrag/services/persistence/migrations/alembic/versions/c224d4befe71_add_file_count_and_file_quota.py similarity index 100% rename from openrag/scripts/migrations/alembic/versions/c224d4befe71_add_file_count_and_file_quota.py rename to openrag/services/persistence/migrations/alembic/versions/c224d4befe71_add_file_count_and_file_quota.py diff --git a/openrag/scripts/migrations/alembic/versions/cd642e4502d8_create_users_memberships_tables.py b/openrag/services/persistence/migrations/alembic/versions/cd642e4502d8_create_users_memberships_tables.py similarity index 100% rename from openrag/scripts/migrations/alembic/versions/cd642e4502d8_create_users_memberships_tables.py rename to openrag/services/persistence/migrations/alembic/versions/cd642e4502d8_create_users_memberships_tables.py diff --git a/openrag/scripts/migrations/alembic/versions/cd9b84278028_merge_heads.py b/openrag/services/persistence/migrations/alembic/versions/cd9b84278028_merge_heads.py similarity index 100% rename from openrag/scripts/migrations/alembic/versions/cd9b84278028_merge_heads.py rename to openrag/services/persistence/migrations/alembic/versions/cd9b84278028_merge_heads.py diff --git a/openrag/scripts/migrations/alembic/versions/e7f8a9b0c1d2_add_workspaces.py b/openrag/services/persistence/migrations/alembic/versions/e7f8a9b0c1d2_add_workspaces.py similarity index 100% rename from openrag/scripts/migrations/alembic/versions/e7f8a9b0c1d2_add_workspaces.py rename to openrag/services/persistence/migrations/alembic/versions/e7f8a9b0c1d2_add_workspaces.py diff --git a/openrag/scripts/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.py b/openrag/services/persistence/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.py similarity index 100% rename from openrag/scripts/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.py rename to openrag/services/persistence/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.py diff --git a/openrag/scripts/migrations/alembic/versions/f5b6c918f741_add_oidc_auth.py b/openrag/services/persistence/migrations/alembic/versions/f5b6c918f741_add_oidc_auth.py similarity index 100% rename from openrag/scripts/migrations/alembic/versions/f5b6c918f741_add_oidc_auth.py rename to openrag/services/persistence/migrations/alembic/versions/f5b6c918f741_add_oidc_auth.py diff --git a/openrag/scripts/migrations/milvus/1.add_created_at_temporal_fields.py b/openrag/services/persistence/migrations/milvus/1.add_created_at_temporal_fields.py similarity index 94% rename from openrag/scripts/migrations/milvus/1.add_created_at_temporal_fields.py rename to openrag/services/persistence/migrations/milvus/1.add_created_at_temporal_fields.py index 563f0193a..df22de6fe 100644 --- a/openrag/scripts/migrations/milvus/1.add_created_at_temporal_fields.py +++ b/openrag/services/persistence/migrations/milvus/1.add_created_at_temporal_fields.py @@ -13,28 +13,28 @@ Usage — prefer the generic runner (from repo root, inside the container): docker compose run --no-deps --rm --entrypoint "" openrag \\ - uv run python scripts/migrations/milvus/migrate.py [--dry-run] [--downgrade] [--target N] + uv run python services/persistence/migrations/milvus/migrate.py [--dry-run] [--downgrade] [--target N] Or run this script directly: # Dry-run first (inspect only, no changes): docker compose run --no-deps --rm --build --entrypoint "" openrag \\ - uv run python scripts/migrations/milvus/1.add_temporal_fields.py --dry-run + uv run python services/persistence/migrations/milvus/1.add_created_at_temporal_fields.py --dry-run # Apply: docker compose run --no-deps --rm --build --entrypoint "" openrag \\ - uv run python scripts/migrations/milvus/1.add_temporal_fields.py + uv run python services/persistence/migrations/milvus/1.add_created_at_temporal_fields.py # Roll back indexes and reset version (fields cannot be dropped in Milvus): docker compose run --no-deps --rm --entrypoint "" openrag \\ - uv run python scripts/migrations/milvus/1.add_temporal_fields.py --downgrade + uv run python services/persistence/migrations/milvus/1.add_created_at_temporal_fields.py --downgrade """ import argparse import sys -from components.indexer.vectordb.vectordb import SCHEMA_VERSION_PROPERTY_KEY from config import load_config from pymilvus import DataType, MilvusClient +from services.storage.milvus_store import SCHEMA_VERSION_PROPERTY_KEY from utils.logger import get_logger TARGET_VERSION = 1 # The schema version this migration brings the collection to. diff --git a/openrag/services/persistence/migrations/milvus/__init__.py b/openrag/services/persistence/migrations/milvus/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/scripts/migrations/milvus/migrate.py b/openrag/services/persistence/migrations/milvus/migrate.py similarity index 94% rename from openrag/scripts/migrations/milvus/migrate.py rename to openrag/services/persistence/migrations/milvus/migrate.py index 33df55c8c..c732d1833 100644 --- a/openrag/scripts/migrations/milvus/migrate.py +++ b/openrag/services/persistence/migrations/milvus/migrate.py @@ -9,19 +9,19 @@ # Dry-run — inspect what would change, no writes: docker compose run --no-deps --rm --build --entrypoint "" openrag \\ - uv run python scripts/migrations/milvus/migrate.py --dry-run + uv run python services/persistence/migrations/milvus/migrate.py --dry-run # Upgrade to latest: docker compose run --no-deps --rm --entrypoint "" openrag \\ - uv run python scripts/migrations/milvus/migrate.py + uv run python services/persistence/migrations/milvus/migrate.py # Upgrade to a specific version: docker compose run --no-deps --rm --entrypoint "" openrag \\ - uv run python scripts/migrations/milvus/migrate.py --target 2 + uv run python services/persistence/migrations/milvus/migrate.py --target 2 # Downgrade to version 0 (resets version property, drops indexes): docker compose run --no-deps --rm --entrypoint "" openrag \\ - uv run python scripts/migrations/milvus/migrate.py --downgrade --target 0 + uv run python services/persistence/migrations/milvus/migrate.py --downgrade --target 0 Convention — each migration module must expose: TARGET_VERSION: int # the version this script brings the DB to @@ -36,9 +36,9 @@ from pathlib import Path from types import ModuleType -from components.indexer.vectordb.vectordb import SCHEMA_VERSION_PROPERTY_KEY from config import load_config from pymilvus import MilvusClient +from services.storage.milvus_store import SCHEMA_VERSION_PROPERTY_KEY from utils.logger import get_logger logger = get_logger() diff --git a/openrag/services/persistence/model_endpoint_repo.py b/openrag/services/persistence/model_endpoint_repo.py new file mode 100644 index 000000000..7e73e65b4 --- /dev/null +++ b/openrag/services/persistence/model_endpoint_repo.py @@ -0,0 +1,31 @@ +"""Stub :class:`ModelEndpointRepository`. + +Model endpoints (embedder URLs, LLM URLs, reranker URLs etc.) are +configured in Hydra YAML today — runtime can't add/swap them without a +restart. A DB-backed registry is a post-refactoring P1 feature so +operators can repoint endpoints from an admin UI. +""" + +from __future__ import annotations + +from core.ports.model_endpoint_repo import ModelEndpointRepository +from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented + + +class PgModelEndpointRepository(_StubRepositoryBase, ModelEndpointRepository): + """TODO: real impl once the ``model_endpoints`` table is added.""" + + async def get(self, name: str, model_type: str) -> dict | None: + raise stub_not_implemented("DB-backed model endpoints") + + async def list_all(self, model_type: str | None = None) -> list[dict]: + raise stub_not_implemented("DB-backed model endpoints") + + async def upsert(self, name: str, model_type: str, config: dict) -> dict: + raise stub_not_implemented("DB-backed model endpoints") + + async def delete(self, name: str, model_type: str) -> bool: + raise stub_not_implemented("DB-backed model endpoints") + + +__all__ = ["PgModelEndpointRepository"] diff --git a/openrag/services/persistence/oidc_session_repo.py b/openrag/services/persistence/oidc_session_repo.py new file mode 100644 index 000000000..482d97728 --- /dev/null +++ b/openrag/services/persistence/oidc_session_repo.py @@ -0,0 +1,399 @@ +"""Postgres implementation of :class:`OIDCSessionRepository`. + +Backs the ``oidc_sessions`` table — the persistence side of the OIDC +authorization-code + PKCE flow. The legacy +:class:`components.indexer.vectordb.utils.PartitionFileManager` exposed +seven OIDC methods (``create_oidc_session``, ``get_oidc_session_by_token``, +``get_oidc_session_by_id``, ``update_oidc_session_tokens``, +``revoke_oidc_sessions_by_sid``, ``revoke_oidc_session_by_id``, +``cleanup_expired_oidc_sessions``) that all land on this class. + +Tokens (``id_token``, ``access_token``, ``refresh_token``) are stored +**Fernet-encrypted** as ``BYTEA``. Encryption and decryption are the +caller's responsibility (see ``components.auth.crypto``) — the repo +treats the bytes as opaque. The plain session cookie value is hashed +(SHA-256) at the caller before being passed in. + +Hard expiry is enforced at read-time: a row is hidden once +``session_expires_at`` is in the past or ``revoked_at`` is set, matching +the legacy behaviour. A periodic call to :meth:`delete_expired` keeps +the table bounded; an explicit retention window (7 days past expiry) +mirrors the legacy ``cleanup_expired_oidc_sessions``. +""" + +from __future__ import annotations + +import hashlib +from collections.abc import Callable +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING, Any + +from core.models.user import OIDCSession +from core.ports.oidc_session_repo import OIDCSessionRepository + +if TYPE_CHECKING: + import asyncpg + + +def _hash_token(token: str) -> str: + """SHA-256 hex digest of an opaque session token.""" + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +# Same retention window the legacy code used — keep expired rows around +# for a week so we can post-mortem broken sessions, then prune. +_EXPIRED_RETENTION = timedelta(days=7) + + +class PgOIDCSessionRepository(OIDCSessionRepository): + """asyncpg-backed implementation of :class:`OIDCSessionRepository`.""" + + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter + + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() + + # ── OIDCSessionRepository port methods ─────────────────────────── + + async def create_session(self, session: OIDCSession) -> OIDCSession: + """Insert a new OIDC session row. + + ``session.session_token_hash`` is treated as already SHA-256 + hashed — the auth service is responsible for hashing the cookie + value before calling this. If it arrives unhashed the row will + still insert but no future lookup will find it. + """ + row = await self.pool.fetchrow( + """ + INSERT INTO oidc_sessions ( + session_token_hash, user_id, sub, sid, + id_token_encrypted, access_token_encrypted, refresh_token_encrypted, + access_token_expires_at, session_expires_at, + created_at, last_refresh_at, revoked_at + ) VALUES ( + $1, $2, $3, $4, + $5, $6, $7, + $8, $9, + COALESCE($10, NOW()), $11, $12 + ) + RETURNING * + """, + session.session_token_hash, + session.user_id, + session.sub, + session.sid, + session.id_token_encrypted, + session.access_token_encrypted, + session.refresh_token_encrypted, + session.access_token_expires_at, + session.session_expires_at, + session.created_at, + session.last_refresh_at, + session.revoked_at, + ) + return self._row_to_session(row) + + async def get_by_token_hash(self, token_hash: str) -> OIDCSession | None: + """Lookup by token hash, filtering out revoked / expired rows.""" + row = await self.pool.fetchrow( + """ + SELECT * FROM oidc_sessions + WHERE session_token_hash = $1 + AND revoked_at IS NULL + AND session_expires_at >= NOW() + """, + token_hash, + ) + return self._row_to_session(row) if row else None + + async def get_by_id(self, session_id: int) -> OIDCSession | None: + """Lookup by primary key, filtering out revoked / expired rows.""" + row = await self.pool.fetchrow( + """ + SELECT * FROM oidc_sessions + WHERE id = $1 + AND revoked_at IS NULL + AND session_expires_at >= NOW() + """, + session_id, + ) + return self._row_to_session(row) if row else None + + async def get_by_sid(self, sid: str) -> list[OIDCSession]: + rows = await self.pool.fetch( + "SELECT * FROM oidc_sessions WHERE sid = $1 ORDER BY created_at", + sid, + ) + return [self._row_to_session(r) for r in rows] + + async def update_session(self, session_id: int, **fields: Any) -> OIDCSession | None: + """Patch a session row by primary key. + + Whitelist matches the columns the auth flow legitimately + rotates: ``access_token_encrypted``, ``refresh_token_encrypted``, + ``id_token_encrypted``, ``access_token_expires_at``, + ``session_expires_at``, ``last_refresh_at``, ``revoked_at``, + ``sid``. Unknown keys are ignored. + """ + allowed = { + "access_token_encrypted", + "refresh_token_encrypted", + "id_token_encrypted", + "access_token_expires_at", + "session_expires_at", + "last_refresh_at", + "revoked_at", + "sid", + } + sets: list[str] = [] + params: list[Any] = [] + for key, value in fields.items(): + if key not in allowed: + continue + params.append(value) + sets.append(f"{key} = ${len(params)}") + if not sets: + row = await self.pool.fetchrow( + "SELECT * FROM oidc_sessions WHERE id = $1", + session_id, + ) + return self._row_to_session(row) if row else None + params.append(session_id) + row = await self.pool.fetchrow( + f"UPDATE oidc_sessions SET {', '.join(sets)} WHERE id = ${len(params)} RETURNING *", + *params, + ) + return self._row_to_session(row) if row else None + + async def revoke_session(self, session_id: int) -> bool: + """Mark a single session revoked (RP-initiated logout).""" + result = await self.pool.execute( + """ + UPDATE oidc_sessions SET revoked_at = NOW() + WHERE id = $1 AND revoked_at IS NULL + """, + session_id, + ) + return result.endswith(" 1") + + async def revoke_by_sid(self, sid: str) -> int: + """Mark every non-revoked session with this OIDC ``sid`` revoked. + + Used by the OIDC back-channel logout flow: the IdP POSTs a + signed logout token whose ``sid`` claim names the session(s) to + terminate; we mark them revoked and return the affected count. + """ + result = await self.pool.execute( + """ + UPDATE oidc_sessions SET revoked_at = NOW() + WHERE sid = $1 AND revoked_at IS NULL + """, + sid, + ) + try: + return int(result.split()[-1]) + except (ValueError, IndexError): + return 0 + + async def revoke_by_user(self, user_id: int) -> int: + """Revoke every non-revoked session belonging to a user.""" + result = await self.pool.execute( + """ + UPDATE oidc_sessions SET revoked_at = NOW() + WHERE user_id = $1 AND revoked_at IS NULL + """, + user_id, + ) + try: + return int(result.split()[-1]) + except (ValueError, IndexError): + return 0 + + async def delete_expired(self) -> int: + """Hard-delete rows expired more than 7 days ago. Returns count.""" + # The ``session_expires_at`` column is TIMESTAMP WITHOUT TIME ZONE, + # so the cutoff has to be tz-naive UTC — asyncpg refuses to bind a + # tz-aware value against a tz-naive column. + cutoff = datetime.now(UTC).replace(tzinfo=None) - _EXPIRED_RETENTION + result = await self.pool.execute( + "DELETE FROM oidc_sessions WHERE session_expires_at < $1", + cutoff, + ) + try: + return int(result.split()[-1]) + except (ValueError, IndexError): + return 0 + + # ── Legacy method names used by the Phase 7C shim ──────────────── + + async def create_oidc_session( # noqa: PLR0913 — legacy signature pinned + self, + *, + user_id: int, + sub: str, + sid: str | None, + session_token_plain: str, + id_token_encrypted: bytes | None, + access_token_encrypted: bytes | None, + refresh_token_encrypted: bytes | None, + access_token_expires_at: datetime, + session_expires_at: datetime, + ) -> dict: + """TODO(phase-9): remove. Legacy interface — hashes the plaintext cookie.""" + token_hash = _hash_token(session_token_plain) + row = await self.pool.fetchrow( + """ + INSERT INTO oidc_sessions ( + session_token_hash, user_id, sub, sid, + id_token_encrypted, access_token_encrypted, refresh_token_encrypted, + access_token_expires_at, session_expires_at, created_at + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, NOW()) + RETURNING * + """, + token_hash, + user_id, + sub, + sid, + id_token_encrypted, + access_token_encrypted, + refresh_token_encrypted, + access_token_expires_at, + session_expires_at, + ) + return self._row_to_dict(row) + + async def get_oidc_session_by_token(self, session_token_plain: str) -> dict | None: + """TODO(phase-9): remove. Returns ``None`` for revoked / expired rows.""" + row = await self.pool.fetchrow( + """ + SELECT * FROM oidc_sessions + WHERE session_token_hash = $1 + AND revoked_at IS NULL + AND session_expires_at >= NOW() + """, + _hash_token(session_token_plain), + ) + return self._row_to_dict(row) if row else None + + async def get_oidc_session_by_id(self, session_id: int) -> dict | None: + """TODO(phase-9): remove. Same hidden-row rules as the legacy code.""" + row = await self.pool.fetchrow( + """ + SELECT * FROM oidc_sessions + WHERE id = $1 + AND revoked_at IS NULL + AND session_expires_at >= NOW() + """, + session_id, + ) + return self._row_to_dict(row) if row else None + + async def update_oidc_session_tokens( + self, + *, + session_id: int, + access_token_encrypted: bytes, + refresh_token_encrypted: bytes | None, + access_token_expires_at: datetime, + ) -> None: + """TODO(phase-9): remove. Atomic token-rotation with row lock. + + ``SELECT ... FOR UPDATE`` serialises concurrent refresh callers on + the same session row — Postgres only. The wider stampede guard in + :mod:`components.auth.refresh` short-circuits before this is even + called in the common case. + """ + async with self.pool.acquire() as conn: + async with conn.transaction(): + row = await conn.fetchrow( + "SELECT id FROM oidc_sessions WHERE id = $1 FOR UPDATE", + session_id, + ) + if row is None: + raise ValueError(f"oidc_session id={session_id} does not exist") + if refresh_token_encrypted is not None: + await conn.execute( + """ + UPDATE oidc_sessions + SET access_token_encrypted = $2, + refresh_token_encrypted = $3, + access_token_expires_at = $4, + last_refresh_at = NOW() + WHERE id = $1 + """, + session_id, + access_token_encrypted, + refresh_token_encrypted, + access_token_expires_at, + ) + else: + await conn.execute( + """ + UPDATE oidc_sessions + SET access_token_encrypted = $2, + access_token_expires_at = $3, + last_refresh_at = NOW() + WHERE id = $1 + """, + session_id, + access_token_encrypted, + access_token_expires_at, + ) + + async def revoke_oidc_sessions_by_sid(self, sid: str) -> int: + """TODO(phase-9): remove. Alias for :meth:`revoke_by_sid`.""" + return await self.revoke_by_sid(sid) + + async def revoke_oidc_session_by_id(self, session_id: int) -> None: + """TODO(phase-9): remove. Returns nothing; legacy contract.""" + await self.revoke_session(session_id) + + async def cleanup_expired_oidc_sessions(self) -> int: + """TODO(phase-9): remove. Alias for :meth:`delete_expired`.""" + return await self.delete_expired() + + # ── Helpers ────────────────────────────────────────────────────── + + @staticmethod + def _row_to_session(row: asyncpg.Record) -> OIDCSession: + return OIDCSession( + id=row["id"], + session_token_hash=row["session_token_hash"], + user_id=row["user_id"], + sid=row["sid"], + sub=row["sub"], + id_token_encrypted=row["id_token_encrypted"], + access_token_encrypted=row["access_token_encrypted"], + refresh_token_encrypted=row["refresh_token_encrypted"], + access_token_expires_at=row["access_token_expires_at"], + session_expires_at=row["session_expires_at"], + created_at=row["created_at"], + last_refresh_at=row["last_refresh_at"], + revoked_at=row["revoked_at"], + ) + + @staticmethod + def _row_to_dict(row: asyncpg.Record) -> dict: + """Legacy dict shape — matches PartitionFileManager._oidc_session_to_dict. + + Encrypted blobs are passed through untouched; the caller decrypts. + """ + return { + "id": row["id"], + "user_id": row["user_id"], + "sub": row["sub"], + "sid": row["sid"], + "id_token_encrypted": row["id_token_encrypted"], + "access_token_encrypted": row["access_token_encrypted"], + "refresh_token_encrypted": row["refresh_token_encrypted"], + "access_token_expires_at": row["access_token_expires_at"], + "session_expires_at": row["session_expires_at"], + "created_at": row["created_at"], + "last_refresh_at": row["last_refresh_at"], + "revoked_at": row["revoked_at"], + } + + +__all__ = ["PgOIDCSessionRepository", "_hash_token"] diff --git a/openrag/services/persistence/partition_membership_repo.py b/openrag/services/persistence/partition_membership_repo.py new file mode 100644 index 000000000..3df72165f --- /dev/null +++ b/openrag/services/persistence/partition_membership_repo.py @@ -0,0 +1,255 @@ +"""Postgres implementation of :class:`PartitionMembershipRepository`. + +Backs the ``partition_memberships`` table. Split out of +:class:`~services.persistence.user_repo.PgUserRepository` so the catalog +matches the 7A.2 one-repo-per-entity layout. + +Two parallel APIs live here on the same table: + +* The clean port methods (:meth:`assign_partition`, :meth:`list_user_partitions`, + …) typed with the :class:`~core.models.user.UserPartition` / + :class:`~core.models.user.PartitionRole` domain models. +* The legacy ``*_partition_member`` / ``*_dict`` methods consumed by the + Phase 7C shim (``vectordb_shims.py``). These return plain dicts to match + the old ``PartitionFileManager`` signatures and are removed in Phase 9 + once the shim is deleted (each carries a ``TODO(phase-9)``). + +``PgUserRepository`` still reads this table directly (``_fetch_memberships``) +to hydrate the ``User`` aggregate's ``partitions`` field — that is a +read-only denormalisation inside the user aggregate boundary, not membership +management, so it stays there. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from core.models.user import PartitionRole, UserPartition +from core.ports.partition_membership_repo import PartitionMembershipRepository + +if TYPE_CHECKING: + import asyncpg + + +class PgPartitionMembershipRepository(PartitionMembershipRepository): + """asyncpg-backed implementation of :class:`PartitionMembershipRepository`.""" + + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter + + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() + + # ── Partition memberships ──────────────────────────────────────── + + async def assign_partition(self, assignment: UserPartition) -> UserPartition: + """Idempotent upsert of (partition, user_id) → role. + + Returns the row as actually persisted (re-reads the DB so the + timestamp reflects what's on disk). + """ + await self.pool.execute( + """ + INSERT INTO partition_memberships (partition_name, user_id, role, added_at) + VALUES ($1, $2, $3, NOW()) + ON CONFLICT (partition_name, user_id) + DO UPDATE SET role = EXCLUDED.role + """, + assignment.partition, + assignment.user_id, + assignment.role.value, + ) + row = await self.pool.fetchrow( + """ + SELECT * FROM partition_memberships + WHERE partition_name = $1 AND user_id = $2 + """, + assignment.partition, + assignment.user_id, + ) + return self._row_to_user_partition(row) + + async def remove_partition(self, user_id: int, partition: str) -> bool: + result = await self.pool.execute( + """ + DELETE FROM partition_memberships + WHERE user_id = $1 AND partition_name = $2 + """, + user_id, + partition, + ) + return result.endswith(" 1") + + async def list_user_partitions(self, user_id: int) -> list[UserPartition]: + rows = await self.pool.fetch( + "SELECT * FROM partition_memberships WHERE user_id = $1 ORDER BY added_at", + user_id, + ) + return [self._row_to_user_partition(r) for r in rows] + + async def list_partition_users(self, partition: str) -> list[UserPartition]: + rows = await self.pool.fetch( + """ + SELECT * FROM partition_memberships + WHERE partition_name = $1 + ORDER BY added_at + """, + partition, + ) + return [self._row_to_user_partition(r) for r in rows] + + async def update_partition_role( + self, + user_id: int, + partition: str, + role: PartitionRole, + ) -> bool: + result = await self.pool.execute( + """ + UPDATE partition_memberships SET role = $3 + WHERE user_id = $1 AND partition_name = $2 + """, + user_id, + partition, + role.value, + ) + return result.endswith(" 1") + + async def count_partition_users(self, partition: str) -> int: + return await self.pool.fetchval( + "SELECT COUNT(*)::int FROM partition_memberships WHERE partition_name = $1", + partition, + ) + + # ── Legacy method names used by the Phase 7C shim ──────────────── + + async def list_partition_members(self, partition: str) -> list[dict]: + """TODO(phase-9): remove. Returns empty list when the partition does not exist.""" + exists = await self.pool.fetchval( + "SELECT 1 FROM partitions WHERE partition = $1", + partition, + ) + if not exists: + return [] + rows = await self.pool.fetch( + """ + SELECT * FROM partition_memberships + WHERE partition_name = $1 + ORDER BY added_at + """, + partition, + ) + return [ + { + "user_id": r["user_id"], + "role": r["role"], + "added_at": r["added_at"].isoformat() if r["added_at"] else None, + } + for r in rows + ] + + async def add_partition_member(self, partition: str, user_id: int, role: str) -> bool: + """TODO(phase-9): remove. Creates the partition row on first use.""" + async with self.pool.acquire() as conn: + async with conn.transaction(): + await conn.execute( + """ + INSERT INTO partitions (partition, created_at) + VALUES ($1, NOW()) + ON CONFLICT (partition) DO NOTHING + """, + partition, + ) + await conn.execute( + """ + INSERT INTO partition_memberships + (partition_name, user_id, role, added_at) + VALUES ($1, $2, $3, NOW()) + ON CONFLICT (partition_name, user_id) + DO UPDATE SET role = EXCLUDED.role + """, + partition, + user_id, + role, + ) + return True + + async def remove_partition_member(self, partition: str, user_id: int) -> bool: + """TODO(phase-9): remove.""" + result = await self.pool.execute( + """ + DELETE FROM partition_memberships + WHERE partition_name = $1 AND user_id = $2 + """, + partition, + user_id, + ) + return result.endswith(" 1") + + async def update_partition_member_role( + self, + partition: str, + user_id: int, + new_role: str, + ) -> bool: + """TODO(phase-9): remove.""" + result = await self.pool.execute( + """ + UPDATE partition_memberships SET role = $3 + WHERE partition_name = $1 AND user_id = $2 + """, + partition, + user_id, + new_role, + ) + return result.endswith(" 1") + + async def user_is_partition_member(self, user_id: int, partition: str) -> bool: + """TODO(phase-9): remove.""" + return await self.pool.fetchval( + """ + SELECT EXISTS ( + SELECT 1 FROM partition_memberships + WHERE user_id = $1 AND partition_name = $2 + ) + """, + user_id, + partition, + ) + + async def list_user_partitions_dict(self, user_id: int) -> list[dict]: + """TODO(phase-9): remove. Legacy ``Partition.to_dict()``-style rows.""" + rows = await self.pool.fetch( + """ + SELECT p.partition, p.created_at, m.role + FROM partitions p + JOIN partition_memberships m + ON m.partition_name = p.partition + WHERE m.user_id = $1 + """, + user_id, + ) + return [ + { + "partition": r["partition"], + "created_at": r["created_at"].isoformat() if r["created_at"] else None, + "role": r["role"], + } + for r in rows + ] + + # ── Helpers ────────────────────────────────────────────────────── + + @staticmethod + def _row_to_user_partition(row: asyncpg.Record) -> UserPartition: + return UserPartition( + user_id=row["user_id"], + partition=row["partition_name"], + role=PartitionRole(row["role"]), + added_at=row["added_at"], + ) + + +__all__ = ["PgPartitionMembershipRepository"] diff --git a/openrag/services/persistence/partition_repo.py b/openrag/services/persistence/partition_repo.py new file mode 100644 index 000000000..b1b29cbf5 --- /dev/null +++ b/openrag/services/persistence/partition_repo.py @@ -0,0 +1,164 @@ +"""Postgres implementation of :class:`PartitionRepository`. + +Manages the ``partitions`` table — the global registry of document +collections. The legacy +:class:`components.indexer.vectordb.utils.PartitionFileManager` exposed +``create_partition``, ``delete_partition``, ``list_partitions``, +``partition_exists``, ``get_partition_file_count``, ``get_total_file_count`` +here; all six map onto this class. + +Deleting a partition cascades to ``files``, ``partition_memberships``, +and ``workspaces`` via the FK ``ON DELETE CASCADE`` rules in the schema. +Per-uploader ``file_count`` is decremented in application code (no SQL +trigger) so the books stay balanced. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from core.ports.partition_repo import PartitionRepository + +if TYPE_CHECKING: + import asyncpg + + +class PgPartitionRepository(PartitionRepository): + """asyncpg-backed implementation of :class:`PartitionRepository`.""" + + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter + + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() + + # ── PartitionRepository port methods ───────────────────────────── + + async def create_partition(self, name: str, user_id: int | None = None) -> dict: + """Insert a partition row; idempotent on the unique constraint. + + When ``user_id`` is provided the caller is granted ``owner`` on + creation. Existing partitions are returned unchanged with no + membership churn — matches the legacy "already exists, log and + skip" behaviour. + """ + async with self.pool.acquire() as conn: + async with conn.transaction(): + row = await conn.fetchrow( + "SELECT * FROM partitions WHERE partition = $1", + name, + ) + if row is None: + row = await conn.fetchrow( + """ + INSERT INTO partitions (partition, created_at) + VALUES ($1, NOW()) + RETURNING * + """, + name, + ) + if user_id is not None: + await conn.execute( + """ + INSERT INTO partition_memberships + (partition_name, user_id, role, added_at) + VALUES ($1, $2, 'owner', NOW()) + ON CONFLICT (partition_name, user_id) DO NOTHING + """, + name, + user_id, + ) + return self._row_to_dict(row) + + async def get_partition(self, name: str) -> dict | None: + row = await self.pool.fetchrow( + "SELECT * FROM partitions WHERE partition = $1", + name, + ) + return self._row_to_dict(row) if row else None + + async def list_partitions(self) -> list[dict]: + rows = await self.pool.fetch("SELECT * FROM partitions ORDER BY created_at") + return [self._row_to_dict(r) for r in rows] + + async def delete_partition(self, name: str) -> bool: + """Delete a partition + its files, memberships, and workspaces. + + ``files.partition_name`` has no ``ON DELETE CASCADE`` (the legacy + ORM relied on SQLAlchemy's Python-side cascade), so we delete file + rows explicitly before the partition. ``workspace_files.file_id`` + cascades, so workspace links clean up with the files. + ``partition_memberships`` and ``workspaces`` cascade from the + partition row. + + Mirrors the legacy bookkeeping: before deleting we count files per + uploader and decrement each uploader's ``file_count`` by that + amount (clamped at zero) so quotas stay accurate. + """ + async with self.pool.acquire() as conn: + async with conn.transaction(): + exists = await conn.fetchval( + "SELECT 1 FROM partitions WHERE partition = $1", + name, + ) + if not exists: + return False + uploader_counts = await conn.fetch( + """ + SELECT created_by, COUNT(*)::int AS n + FROM files + WHERE partition_name = $1 AND created_by IS NOT NULL + GROUP BY created_by + """, + name, + ) + await conn.execute( + "DELETE FROM files WHERE partition_name = $1", + name, + ) + await conn.execute( + "DELETE FROM partitions WHERE partition = $1", + name, + ) + for r in uploader_counts: + await conn.execute( + "UPDATE users SET file_count = GREATEST(file_count - $1, 0) WHERE id = $2", + r["n"], + r["created_by"], + ) + return True + + async def partition_exists(self, name: str) -> bool: + return await self.pool.fetchval( + "SELECT EXISTS (SELECT 1 FROM partitions WHERE partition = $1)", + name, + ) + + # ── Legacy method names used by the Phase 7C shim ──────────────── + + async def get_partition_file_count(self, partition: str) -> int: + """TODO(phase-9): remove.""" + return await self.pool.fetchval( + "SELECT COUNT(*)::int FROM files WHERE partition_name = $1", + partition, + ) + + async def get_total_file_count(self) -> int: + """TODO(phase-9): remove.""" + return await self.pool.fetchval("SELECT COUNT(*)::int FROM files") + + # ── Row → dict helper ──────────────────────────────────────────── + + @staticmethod + def _row_to_dict(row: asyncpg.Record) -> dict: + """Shape mirrors the legacy ``Partition.to_dict()`` ORM helper.""" + created = row["created_at"] + return { + "partition": row["partition"], + "created_at": created.isoformat() if created else None, + } + + +__all__ = ["PgPartitionRepository"] diff --git a/openrag/services/persistence/preset_repo.py b/openrag/services/persistence/preset_repo.py new file mode 100644 index 000000000..be2caede0 --- /dev/null +++ b/openrag/services/persistence/preset_repo.py @@ -0,0 +1,31 @@ +"""Stub :class:`PresetRepository`. + +Pipeline presets — named bundles of chunker/embedder/retriever config — +are P0 on the post-refactoring roadmap. They are the mechanism that +will let each partition pick its own pipeline configuration without +operators touching YAML. No table exists today. +""" + +from __future__ import annotations + +from core.ports.preset_repo import PresetRepository +from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented + + +class PgPresetRepository(_StubRepositoryBase, PresetRepository): + """TODO: real impl once the ``presets`` table is added — see REFACTORING P0 plan.""" + + async def get(self, name: str, preset_type: str) -> dict | None: + raise stub_not_implemented("Per-partition pipeline presets") + + async def list_all(self, preset_type: str | None = None) -> list[dict]: + raise stub_not_implemented("Per-partition pipeline presets") + + async def upsert(self, name: str, preset_type: str, config: dict) -> dict: + raise stub_not_implemented("Per-partition pipeline presets") + + async def delete(self, name: str, preset_type: str) -> bool: + raise stub_not_implemented("Per-partition pipeline presets") + + +__all__ = ["PgPresetRepository"] diff --git a/openrag/services/persistence/prompt_repo.py b/openrag/services/persistence/prompt_repo.py new file mode 100644 index 000000000..6c143546b --- /dev/null +++ b/openrag/services/persistence/prompt_repo.py @@ -0,0 +1,43 @@ +"""Stub :class:`PromptRepository`. + +Prompts are currently disk-based templates (``components/prompts/``). +The post-refactoring P1 feature is DB-stored, per-partition, +versionable prompts that operators can edit without redeploying. When +that lands the on-disk templates become the seed for the table and +this stub becomes a real asyncpg implementation against a ``prompts`` +table. +""" + +from __future__ import annotations + +from core.models.prompt import Prompt +from core.ports.prompt_repo import PromptRepository +from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented + + +class PgPromptRepository(_StubRepositoryBase, PromptRepository): + """TODO: real impl once the ``prompts`` table is added.""" + + async def create_prompt(self, prompt: Prompt) -> Prompt: + raise stub_not_implemented("DB-stored prompts") + + async def get_prompt(self, prompt_id: str) -> Prompt | None: + raise stub_not_implemented("DB-stored prompts") + + async def get_by_type(self, prompt_type: str) -> list[Prompt]: + raise stub_not_implemented("DB-stored prompts") + + async def get_active(self, prompt_type: str) -> Prompt | None: + raise stub_not_implemented("DB-stored prompts") + + async def list_prompts(self) -> list[Prompt]: + raise stub_not_implemented("DB-stored prompts") + + async def update_prompt(self, prompt_id: str, content: str) -> Prompt | None: + raise stub_not_implemented("DB-stored prompts") + + async def delete_prompt(self, prompt_id: str) -> bool: + raise stub_not_implemented("DB-stored prompts") + + +__all__ = ["PgPromptRepository"] diff --git a/openrag/services/persistence/schema.py b/openrag/services/persistence/schema.py new file mode 100644 index 000000000..22ab3ad7b --- /dev/null +++ b/openrag/services/persistence/schema.py @@ -0,0 +1,205 @@ +"""Metadata-only table definitions for the Postgres catalog. + +Defines the same 7 tables that ``components/indexer/vectordb/models.py`` declares +with the SQLAlchemy ORM, but as :class:`sqlalchemy.Table` objects bound to a +single :class:`sqlalchemy.MetaData`. The new persistence layer talks to +Postgres through ``asyncpg`` with raw SQL; this module exists solely so that +Alembic's autogenerate has a metadata target to diff against, and so the +on-startup ``metadata.create_all()`` path keeps working until phase 9 retires +the legacy actor. + +Column types, defaults, foreign keys, unique constraints, check constraints +and indexes must stay identical to the ORM models — Alembic will treat any +divergence as a pending schema change. +""" + +from datetime import datetime + +from sqlalchemy import ( + JSON, + Boolean, + CheckConstraint, + Column, + DateTime, + ForeignKey, + Index, + Integer, + LargeBinary, + MetaData, + String, + Table, + UniqueConstraint, +) + +metadata = MetaData() + + +partitions = Table( + "partitions", + metadata, + Column("id", Integer, primary_key=True), + Column("partition", String, unique=True, nullable=False, index=True), + Column("created_at", DateTime, default=datetime.now, nullable=False, index=True), +) + + +files = Table( + "files", + metadata, + Column("id", Integer, primary_key=True), + Column("file_id", String, nullable=False, index=True), + Column( + "partition_name", + String, + ForeignKey("partitions.partition"), + nullable=False, + index=True, + ), + Column("file_metadata", JSON, nullable=True, default=dict), + Column( + "created_by", + Integer, + ForeignKey("users.id", ondelete="SET NULL"), + nullable=True, + index=True, + ), + Column("relationship_id", String, nullable=True, index=True), + Column("parent_id", String, nullable=True, index=True), + UniqueConstraint("file_id", "partition_name", name="uix_file_id_partition"), + Index("ix_partition_file", "partition_name", "file_id"), + Index("ix_relationship_partition", "relationship_id", "partition_name"), + Index("ix_parent_partition", "parent_id", "partition_name"), +) + + +users = Table( + "users", + metadata, + Column("id", Integer, primary_key=True), + Column("external_user_id", String, unique=True, nullable=True, index=True), + Column("display_name", String, nullable=True), + Column("email", String, unique=True, nullable=True, index=True), + Column("token", String, unique=True, nullable=True, index=True), + Column("is_admin", Boolean, default=False, nullable=False), + Column("created_at", DateTime, default=datetime.now, nullable=False), + Column("file_quota", Integer, nullable=True, default=None), + Column("file_count", Integer, nullable=False, default=0), +) + + +oidc_sessions = Table( + "oidc_sessions", + metadata, + Column("id", Integer, primary_key=True), + Column( + "session_token_hash", + String(64), + unique=True, + nullable=False, + index=True, + ), + Column( + "user_id", + Integer, + ForeignKey("users.id", ondelete="CASCADE"), + nullable=False, + index=True, + ), + Column("sid", String, nullable=True, index=True), + Column("sub", String, nullable=False), + Column("id_token_encrypted", LargeBinary, nullable=True), + Column("access_token_encrypted", LargeBinary, nullable=True), + Column("refresh_token_encrypted", LargeBinary, nullable=True), + Column("access_token_expires_at", DateTime, nullable=False), + Column("session_expires_at", DateTime, nullable=False), + Column("created_at", DateTime, default=datetime.now, nullable=False), + Column("last_refresh_at", DateTime, nullable=True), + Column("revoked_at", DateTime, nullable=True), + Index("ix_oidc_sessions_user_sub", "user_id", "sub"), +) + + +partition_memberships = Table( + "partition_memberships", + metadata, + Column("id", Integer, primary_key=True), + Column( + "partition_name", + String, + ForeignKey("partitions.partition", ondelete="CASCADE"), + nullable=False, + index=True, + ), + Column( + "user_id", + Integer, + ForeignKey("users.id", ondelete="CASCADE"), + nullable=False, + index=True, + ), + Column("role", String, nullable=False), + Column("added_at", DateTime, default=datetime.now, nullable=False), + UniqueConstraint("partition_name", "user_id", name="uix_partition_user"), + CheckConstraint( + "role IN ('owner','editor','viewer')", + name="ck_membership_role", + ), + Index("ix_user_partition", "user_id", "partition_name"), +) + + +workspaces = Table( + "workspaces", + metadata, + Column("id", Integer, primary_key=True), + Column("workspace_id", String, unique=True, nullable=False, index=True), + Column( + "partition_name", + String, + ForeignKey("partitions.partition", ondelete="CASCADE"), + nullable=False, + index=True, + ), + Column( + "created_by", + Integer, + ForeignKey("users.id", ondelete="SET NULL"), + nullable=True, + ), + Column("display_name", String, nullable=True), + Column("created_at", DateTime, default=datetime.now), +) + + +workspace_files = Table( + "workspace_files", + metadata, + Column("id", Integer, primary_key=True), + Column( + "workspace_id", + String, + ForeignKey("workspaces.workspace_id", ondelete="CASCADE"), + nullable=False, + index=True, + ), + Column( + "file_id", + Integer, + ForeignKey("files.id", ondelete="CASCADE"), + nullable=False, + index=True, + ), + UniqueConstraint("workspace_id", "file_id", name="uix_workspace_file"), +) + + +__all__ = [ + "metadata", + "partitions", + "files", + "users", + "oidc_sessions", + "partition_memberships", + "workspaces", + "workspace_files", +] diff --git a/openrag/services/persistence/test_connection.py b/openrag/services/persistence/test_connection.py new file mode 100644 index 000000000..8863b6259 --- /dev/null +++ b/openrag/services/persistence/test_connection.py @@ -0,0 +1,38 @@ +from unittest.mock import Mock + +from services.persistence.connection import ConnectionManager + + +class RDBConfigStub: + host = "db" + port = 5432 + user = "root" + password = "root_password" + database = "partitions_for_collection_test" + pool_min_size = 1 + pool_max_size = 4 + command_timeout = 10 + + +def test_ensure_database_exists_creates_missing_database(monkeypatch): + manager = ConnectionManager(RDBConfigStub()) + create_database = Mock() + + monkeypatch.setattr("sqlalchemy_utils.database_exists", lambda url: False) + monkeypatch.setattr("sqlalchemy_utils.create_database", create_database) + + manager._ensure_database_exists() + + create_database.assert_called_once() + + +def test_ensure_database_exists_skips_existing_database(monkeypatch): + manager = ConnectionManager(RDBConfigStub()) + create_database = Mock() + + monkeypatch.setattr("sqlalchemy_utils.database_exists", lambda url: True) + monkeypatch.setattr("sqlalchemy_utils.create_database", create_database) + + manager._ensure_database_exists() + + create_database.assert_not_called() diff --git a/openrag/components/indexer/vectordb/test_delete_workspace.py b/openrag/services/persistence/test_delete_workspace.py similarity index 100% rename from openrag/components/indexer/vectordb/test_delete_workspace.py rename to openrag/services/persistence/test_delete_workspace.py diff --git a/openrag/services/persistence/topic_tag_repo.py b/openrag/services/persistence/topic_tag_repo.py new file mode 100644 index 000000000..c6398004a --- /dev/null +++ b/openrag/services/persistence/topic_tag_repo.py @@ -0,0 +1,30 @@ +"""Stub :class:`TopicTagRepository`. + +Topic/tag attachment per document is a future feature — useful for +faceted search and for "show me docs about X" UIs. No table exists +today. +""" + +from __future__ import annotations + +from core.ports.topic_tag_repo import TopicTagRepository +from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented + + +class PgTopicTagRepository(_StubRepositoryBase, TopicTagRepository): + """TODO: real impl once the ``topic_tags`` table is added.""" + + async def bulk_insert(self, tags: list[dict]) -> int: + raise stub_not_implemented("Topic / tag storage") + + async def get_by_document(self, document_id: str) -> list[dict]: + raise stub_not_implemented("Topic / tag storage") + + async def delete_by_document(self, document_id: str) -> int: + raise stub_not_implemented("Topic / tag storage") + + async def search(self, partition: str, tag: str, top_k: int = 10) -> list[dict]: + raise stub_not_implemented("Topic / tag storage") + + +__all__ = ["PgTopicTagRepository"] diff --git a/openrag/services/persistence/user_repo.py b/openrag/services/persistence/user_repo.py new file mode 100644 index 000000000..b2d606d0f --- /dev/null +++ b/openrag/services/persistence/user_repo.py @@ -0,0 +1,451 @@ +"""Postgres implementation of :class:`UserRepository`. + +Backs the ``users`` table (and, when the post-refactoring ``api_keys`` +table lands, ``api_keys``). The legacy +:class:`components.indexer.vectordb.utils.PartitionFileManager` exposed +eleven user-shaped methods (``create_user``, ``get_user_by_id``, +``get_user_by_token``, ``delete_user``, ``update_user``, ``list_users``, +``regenerate_user_token``, ``user_exists``, ``get_user_by_external_id``, +``update_user_fields``, ``_ensure_admin_user``) which map onto this class. +The six partition-membership methods moved to +:class:`~services.persistence.partition_membership_repo.PgPartitionMembershipRepository` +(7A.2 one-repo-per-entity layout). This class still *reads* +``partition_memberships`` via :meth:`_fetch_memberships` to hydrate the +``User`` aggregate's ``partitions`` field — a read-only denormalisation +inside the user aggregate boundary, not membership management. + +Notes on the schema vs. the port: + +* The port :class:`~openrag.core.models.user.User` model carries + ``password_hash``, ``is_active`` and ``updated_at`` fields that have + no column today. They are treated as ``None`` / ``True`` / ``created_at`` + respectively at the boundary so the domain shape stays useful for + callers. +* The ``UserRepository`` port also defines four ``api_key_*`` methods. + OpenRAG currently stores one hashed token in ``users.token`` — a real + ``api_keys`` table is on the post-refactoring roadmap. Until then the + api-key methods raise :class:`NotImplementedError` to signal the gap + loudly rather than silently returning empty lists. +""" + +from __future__ import annotations + +import hashlib +import secrets +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +from core.models.user import ApiKey, PartitionRole, User, UserPartition +from core.ports.user_repo import UserRepository + +if TYPE_CHECKING: + import asyncpg + + +def _hash_token(token: str) -> str: + """SHA-256 hex digest of a token string. + + Matches the legacy :meth:`PartitionFileManager.hash_token` so existing + rows continue to validate against the same hash. Exposed at module + level so callers (e.g. auth middleware) can hash before lookup + without instantiating the repo. + """ + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +class PgUserRepository(UserRepository): + """asyncpg-backed implementation of :class:`UserRepository`.""" + + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter + + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() + + # ── User CRUD ──────────────────────────────────────────────────── + + async def create_user(self, user: User) -> User: + """Insert a user row and return it with its assigned PK. + + ``User.password_hash`` is dropped because the column does not + exist yet — when password auth lands we'll add the column and + wire it here. ``token`` / ``token_hash`` should be set out-of-band + via :meth:`set_user_token`; this method does NOT generate one. + """ + row = await self.pool.fetchrow( + """ + INSERT INTO users (display_name, external_user_id, email, + is_admin, file_quota, file_count, created_at) + VALUES ($1, $2, $3, $4, $5, $6, COALESCE($7, NOW())) + RETURNING * + """, + user.display_name, + user.external_user_id, + (user.email.strip().lower() if user.email else None), + user.is_admin, + user.file_quota, + user.file_count, + user.created_at, + ) + return self._row_to_user(row) + + async def get_user(self, user_id: int) -> User | None: + row = await self.pool.fetchrow("SELECT * FROM users WHERE id = $1", user_id) + if row is None: + return None + memberships = await self._fetch_memberships(user_id) + return self._row_to_user(row, memberships) + + async def get_user_by_email(self, email: str) -> User | None: + row = await self.pool.fetchrow( + "SELECT * FROM users WHERE email = $1", + email.strip().lower(), + ) + if row is None: + return None + memberships = await self._fetch_memberships(row["id"]) + return self._row_to_user(row, memberships) + + async def get_user_by_token(self, token_hash: str) -> User | None: + """Lookup by the SHA-256 hash of the bearer token. + + The auth middleware hashes the raw token before calling this, so + the repo never sees plaintext. + """ + row = await self.pool.fetchrow( + "SELECT * FROM users WHERE token = $1", + token_hash, + ) + if row is None: + return None + memberships = await self._fetch_memberships(row["id"]) + return self._row_to_user(row, memberships) + + async def get_user_by_external_id(self, external_id: str) -> User | None: + row = await self.pool.fetchrow( + "SELECT * FROM users WHERE external_user_id = $1", + external_id, + ) + if row is None: + return None + memberships = await self._fetch_memberships(row["id"]) + return self._row_to_user(row, memberships) + + async def list_users(self, offset: int = 0, limit: int = 50) -> list[User]: + rows = await self.pool.fetch( + "SELECT * FROM users ORDER BY id LIMIT $1 OFFSET $2", + limit, + offset, + ) + return [self._row_to_user(r) for r in rows] + + async def update_user(self, user_id: int, **fields: Any) -> User | None: + """Patch fields on a user row. + + Silently ignores unknown columns — keeps the call site forgiving + when the domain model carries fields the schema does not have + yet (``password_hash``, ``is_active``, ``updated_at``). + """ + allowed = { + "display_name", + "external_user_id", + "email", + "is_admin", + "file_quota", + "file_count", + "token", + } + sets: list[str] = [] + params: list[Any] = [] + for key, value in fields.items(): + if key not in allowed: + continue + if key == "email" and isinstance(value, str): + value = value.strip().lower() + params.append(value) + sets.append(f"{key} = ${len(params)}") + if not sets: + return await self.get_user(user_id) + params.append(user_id) + row = await self.pool.fetchrow( + f"UPDATE users SET {', '.join(sets)} WHERE id = ${len(params)} RETURNING *", + *params, + ) + if row is None: + return None + memberships = await self._fetch_memberships(user_id) + return self._row_to_user(row, memberships) + + async def delete_user(self, user_id: int) -> bool: + result = await self.pool.execute("DELETE FROM users WHERE id = $1", user_id) + return result.endswith(" 1") + + async def count_users(self) -> int: + return await self.pool.fetchval("SELECT COUNT(*)::int FROM users") + + # ── API keys (stub — table not yet shipped) ────────────────────── + + async def create_api_key(self, key: ApiKey) -> ApiKey: + raise NotImplementedError( + "api_keys table is on the post-refactoring roadmap; use users.token until then.", + ) + + async def get_api_keys_by_prefix(self, prefix: str) -> list[ApiKey]: + raise NotImplementedError( + "api_keys table is on the post-refactoring roadmap; use users.token until then.", + ) + + async def list_api_keys_for_user(self, user_id: int) -> list[ApiKey]: + raise NotImplementedError( + "api_keys table is on the post-refactoring roadmap; use users.token until then.", + ) + + async def delete_api_key(self, key_id: str) -> bool: + raise NotImplementedError( + "api_keys table is on the post-refactoring roadmap; use users.token until then.", + ) + + # ── Legacy method names used by the Phase 7C shim ──────────────── + + async def create_legacy_user( + self, + display_name: str | None, + external_user_id: str | None, + email: str | None, + is_admin: bool, + file_quota: int | None, + ) -> dict: + """TODO(phase-9): remove. Mirror of legacy ``create_user``. + + Generates a plain ``or-`` token, stores its hash, returns the + plaintext exactly once (it is never persisted unhashed). + """ + plaintext = f"or-{secrets.token_hex(16)}" + token_hash = _hash_token(plaintext) + row = await self.pool.fetchrow( + """ + INSERT INTO users (display_name, external_user_id, email, + token, is_admin, file_quota, file_count, created_at) + VALUES ($1, $2, $3, $4, $5, $6, 0, NOW()) + RETURNING * + """, + display_name, + external_user_id, + (email.strip().lower() if email else None), + token_hash, + is_admin, + file_quota, + ) + return { + "id": row["id"], + "display_name": row["display_name"], + "external_user_id": row["external_user_id"], + "email": row["email"], + "token": plaintext, + "is_admin": row["is_admin"], + "file_quota": row["file_quota"], + "file_count": row["file_count"], + } + + async def regenerate_user_token(self, user_id: int) -> dict | None: + """TODO(phase-9): remove. Rotate ``users.token`` and surface the plaintext.""" + plaintext = f"or-{secrets.token_hex(16)}" + token_hash = _hash_token(plaintext) + row = await self.pool.fetchrow( + """ + UPDATE users SET token = $2 + WHERE id = $1 + RETURNING * + """, + user_id, + token_hash, + ) + if row is None: + return None + return { + "id": row["id"], + "display_name": row["display_name"], + "external_user_id": row["external_user_id"], + "token": plaintext, + "is_admin": row["is_admin"], + "file_quota": row["file_quota"], + "file_count": row["file_count"], + } + + async def get_user_by_token_plain(self, token: str) -> dict | None: + """TODO(phase-9): remove. Hash + lookup + serialise to legacy dict shape.""" + return await self.get_user_dict_by_id( + await self.pool.fetchval( + "SELECT id FROM users WHERE token = $1", + _hash_token(token), + ), + ) + + async def get_user_dict_by_id(self, user_id: int | None) -> dict | None: + """TODO(phase-9): remove. Legacy dict shape with ``memberships`` list.""" + if user_id is None: + return None + row = await self.pool.fetchrow("SELECT * FROM users WHERE id = $1", user_id) + if row is None: + return None + memberships = await self.pool.fetch( + "SELECT * FROM partition_memberships WHERE user_id = $1 ORDER BY added_at", + user_id, + ) + return { + "id": row["id"], + "display_name": row["display_name"], + "external_user_id": row["external_user_id"], + "email": row["email"], + "is_admin": row["is_admin"], + "file_quota": row["file_quota"], + "file_count": row["file_count"], + "memberships": [ + { + "partition": m["partition_name"], + "role": m["role"], + "added_at": m["added_at"].isoformat() if m["added_at"] else None, + } + for m in memberships + ], + } + + async def get_user_by_external_id_dict(self, external_user_id: str) -> dict | None: + """TODO(phase-9): remove. Legacy dict shape, lookup by OIDC sub claim.""" + row = await self.pool.fetchrow( + "SELECT id FROM users WHERE external_user_id = $1", + external_user_id, + ) + return await self.get_user_dict_by_id(row["id"]) if row else None + + async def list_users_dict(self) -> list[dict]: + """TODO(phase-9): remove. Legacy list shape used by /users/ endpoint.""" + rows = await self.pool.fetch("SELECT * FROM users ORDER BY id") + return [ + { + "id": r["id"], + "display_name": r["display_name"], + "external_user_id": r["external_user_id"], + "is_admin": r["is_admin"], + "file_quota": r["file_quota"], + "file_count": r["file_count"], + "created_at": r["created_at"].isoformat() if r["created_at"] else None, + } + for r in rows + ] + + async def user_exists(self, user_id: int) -> bool: + """TODO(phase-9): remove.""" + return await self.pool.fetchval( + "SELECT EXISTS (SELECT 1 FROM users WHERE id = $1)", + user_id, + ) + + # Whitelist mirrored from the legacy PartitionFileManager. Three + # layers (startup validator, claim parser, repo) all enforce the + # same set as defence-in-depth against an OIDC claim mapping that + # would otherwise let a remote IdP rewrite arbitrary user columns. + _OIDC_WRITABLE_USER_FIELDS = frozenset({"display_name", "email"}) + + async def update_user_fields(self, user_id: int, fields: dict[str, Any]) -> None: + """TODO(phase-9): remove. Strict-whitelist update for the OIDC claim mapper.""" + if not fields: + return + bad = set(fields) - self._OIDC_WRITABLE_USER_FIELDS + if bad: + raise ValueError(f"Cannot update non-whitelisted user fields: {sorted(bad)}") + cleaned = {k: v for k, v in fields.items() if v is not None} + if not cleaned: + return + if "email" in cleaned and isinstance(cleaned["email"], str): + cleaned["email"] = cleaned["email"].strip().lower() + sets: list[str] = [] + params: list[Any] = [] + for key, value in cleaned.items(): + params.append(value) + sets.append(f"{key} = ${len(params)}") + params.append(user_id) + result = await self.pool.execute( + f"UPDATE users SET {', '.join(sets)} WHERE id = ${len(params)}", + *params, + ) + if not result.endswith(" 1"): + raise ValueError(f"User {user_id} not found") + + async def ensure_admin_user(self, admin_token: str | None) -> str: + """TODO(phase-9): remove. Bootstrap mirror of the legacy admin-bootstrap. + + Ensures ``users.id = 1`` exists with ``is_admin = TRUE`` and the + token hash matching ``admin_token``. Generates a token if none + is supplied. Returns whichever plaintext token is now valid. + """ + plaintext = admin_token or f"or-{secrets.token_hex(16)}" + token_hash = _hash_token(plaintext) + async with self.pool.acquire() as conn: + async with conn.transaction(): + existing = await conn.fetchrow("SELECT id FROM users WHERE id = 1") + if existing is None: + await conn.execute( + """ + INSERT INTO users (id, display_name, token, is_admin, file_count, created_at) + VALUES (1, 'Admin', $1, TRUE, 0, NOW()) + """, + token_hash, + ) + # Keep the sequence ahead of the explicit id=1 insert so + # subsequent `INSERT INTO users` calls don't collide. + await conn.execute( + "SELECT setval(pg_get_serial_sequence('users','id'), GREATEST(1, (SELECT MAX(id) FROM users)))" + ) + else: + await conn.execute( + "UPDATE users SET is_admin = TRUE, token = $1 WHERE id = 1", + token_hash, + ) + return plaintext + + # ── Helpers ────────────────────────────────────────────────────── + + async def _fetch_memberships(self, user_id: int) -> list[UserPartition]: + rows = await self.pool.fetch( + "SELECT * FROM partition_memberships WHERE user_id = $1 ORDER BY added_at", + user_id, + ) + return [self._row_to_user_partition(r) for r in rows] + + @staticmethod + def _row_to_user( + row: asyncpg.Record, + memberships: list[UserPartition] | None = None, + ) -> User: + # `password_hash`, `is_active`, `updated_at` do not exist on the + # current schema; fall back to safe defaults so the domain model + # stays consistent even though the underlying row is narrower. + created = row["created_at"] + return User( + id=row["id"], + display_name=row["display_name"], + external_user_id=row["external_user_id"], + email=row["email"], + password_hash=None, + is_admin=row["is_admin"], + is_active=True, + file_quota=row["file_quota"], + file_count=row["file_count"], + created_at=created, + updated_at=created, + partitions=memberships or [], + ) + + @staticmethod + def _row_to_user_partition(row: asyncpg.Record) -> UserPartition: + return UserPartition( + user_id=row["user_id"], + partition=row["partition_name"], + role=PartitionRole(row["role"]), + added_at=row["added_at"], + ) + + +__all__ = ["PgUserRepository", "_hash_token"] diff --git a/openrag/services/persistence/workspace_repo.py b/openrag/services/persistence/workspace_repo.py new file mode 100644 index 000000000..28fefb742 --- /dev/null +++ b/openrag/services/persistence/workspace_repo.py @@ -0,0 +1,354 @@ +"""Postgres implementation of :class:`WorkspaceRepository`. + +Backs the ``workspaces`` table and the ``workspace_files`` many-to-many +join. The legacy +:class:`components.indexer.vectordb.utils.PartitionFileManager` exposed +ten workspace methods that all map onto this class: +``create_workspace``, ``list_workspaces``, ``get_workspace``, +``delete_workspace``, ``add_files_to_workspace``, +``remove_file_from_workspace``, ``list_workspace_files``, +``get_file_workspaces``, ``get_existing_file_ids``, +``remove_file_from_all_workspaces``. + +The join references the canonical ``files.id`` integer PK (not the +opaque ``file_id`` string), so deletion cascades correctly — when a +``files`` row goes away the workspace_files entries it backed go with +it without any application-side bookkeeping. Conversely the workspace +APIs accept and emit the human ``file_id`` form; the repo translates at +the boundary. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from core.models.workspace import Workspace +from core.ports.workspace_repo import WorkspaceRepository + +if TYPE_CHECKING: + import asyncpg + + +class PgWorkspaceRepository(WorkspaceRepository): + """asyncpg-backed implementation of :class:`WorkspaceRepository`.""" + + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter + + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() + + # ── WorkspaceRepository port methods ───────────────────────────── + + async def create_workspace(self, workspace: Workspace) -> Workspace: + row = await self.pool.fetchrow( + """ + INSERT INTO workspaces (workspace_id, partition_name, + created_by, display_name, created_at) + VALUES ($1, $2, $3, $4, COALESCE($5, NOW())) + RETURNING * + """, + workspace.workspace_id, + workspace.partition, + workspace.created_by, + workspace.display_name, + workspace.created_at, + ) + return self._row_to_workspace(row) + + async def get_workspace(self, workspace_id: str) -> Workspace | None: + row = await self.pool.fetchrow( + "SELECT * FROM workspaces WHERE workspace_id = $1", + workspace_id, + ) + return self._row_to_workspace(row) if row else None + + async def list_workspaces(self, partition: str) -> list[Workspace]: + rows = await self.pool.fetch( + """ + SELECT * FROM workspaces + WHERE partition_name = $1 + ORDER BY created_at + """, + partition, + ) + return [self._row_to_workspace(r) for r in rows] + + async def delete_workspace(self, workspace_id: str) -> list[str]: + """Delete the workspace and return the orphaned ``file_id`` list. + + An orphan = a file currently in this workspace and in no other. + Because ``workspace_files.file_id`` is an integer FK to + ``files.id``, every workspace_files row already has a backing + files row; the orphan check therefore reduces to + "file_id NOT IN (other workspaces' file_ids)". + + Returning the orphans (rather than auto-deleting them) keeps the + deletion of the underlying file optional — the legacy router + loops over the list and calls the indexer's file-delete path so + the Milvus side is cleaned up too. + """ + async with self.pool.acquire() as conn: + async with conn.transaction(): + orphan_rows = await conn.fetch( + """ + SELECT f.file_id + FROM workspace_files wf + JOIN files f ON f.id = wf.file_id + WHERE wf.workspace_id = $1 + AND wf.file_id NOT IN ( + SELECT file_id FROM workspace_files + WHERE workspace_id <> $1 + ) + """, + workspace_id, + ) + await conn.execute( + "DELETE FROM workspaces WHERE workspace_id = $1", + workspace_id, + ) + return [r["file_id"] for r in orphan_rows] + + async def add_files_to_workspace( + self, + workspace_id: str, + file_ids: list[str], + ) -> list[str]: + """Attach files identified by their string ``file_id`` to a workspace. + + Returns the list of supplied ``file_ids`` that do not exist in + the workspace's partition — callers surface these to the user as + "not found". + """ + if not file_ids: + return [] + async with self.pool.acquire() as conn: + async with conn.transaction(): + workspace = await conn.fetchrow( + "SELECT partition_name FROM workspaces WHERE workspace_id = $1", + workspace_id, + ) + if workspace is None: + return list(file_ids) + partition = workspace["partition_name"] + resolved = await conn.fetch( + """ + SELECT file_id, id FROM files + WHERE file_id = ANY($1::text[]) AND partition_name = $2 + """, + file_ids, + partition, + ) + id_map = {r["file_id"]: r["id"] for r in resolved} + missing = [fid for fid in file_ids if fid not in id_map] + if id_map: + # Insert each row separately with ON CONFLICT DO NOTHING. + # asyncpg has no native bulk-with-conflict; the row count + # is bounded by file_ids so the loop is fine here. + for file_pk in id_map.values(): + await conn.execute( + """ + INSERT INTO workspace_files (workspace_id, file_id) + VALUES ($1, $2) + ON CONFLICT ON CONSTRAINT uix_workspace_file DO NOTHING + """, + workspace_id, + file_pk, + ) + return missing + + async def remove_file_from_workspace( + self, + workspace_id: str, + file_id: str, + ) -> bool: + async with self.pool.acquire() as conn: + async with conn.transaction(): + workspace = await conn.fetchrow( + "SELECT partition_name FROM workspaces WHERE workspace_id = $1", + workspace_id, + ) + if workspace is None: + return False + file_pk = await conn.fetchval( + """ + SELECT id FROM files + WHERE file_id = $1 AND partition_name = $2 + """, + file_id, + workspace["partition_name"], + ) + if file_pk is None: + return False + result = await conn.execute( + """ + DELETE FROM workspace_files + WHERE workspace_id = $1 AND file_id = $2 + """, + workspace_id, + file_pk, + ) + try: + return int(result.split()[-1]) > 0 + except (ValueError, IndexError): + return False + + async def list_workspace_files(self, workspace_id: str) -> list[str]: + rows = await self.pool.fetch( + """ + SELECT f.file_id + FROM workspace_files wf + JOIN files f ON f.id = wf.file_id + WHERE wf.workspace_id = $1 + """, + workspace_id, + ) + return [r["file_id"] for r in rows] + + async def get_file_workspaces( + self, + file_id: str, + partition: str, + ) -> list[str]: + """Workspaces containing ``file_id``, scoped to ``partition``. + + Scoping is necessary because a given ``file_id`` string is unique + only within a partition — the underlying ``files`` rows are + distinct PKs across partitions. + """ + rows = await self.pool.fetch( + """ + SELECT wf.workspace_id + FROM workspace_files wf + JOIN files f ON f.id = wf.file_id + JOIN workspaces w ON w.workspace_id = wf.workspace_id + WHERE f.file_id = $1 + AND f.partition_name = $2 + AND w.partition_name = $2 + """, + file_id, + partition, + ) + return [r["workspace_id"] for r in rows] + + async def get_existing_file_ids( + self, + partition: str, + file_ids: list[str], + ) -> set[str]: + if not file_ids: + return set() + rows = await self.pool.fetch( + """ + SELECT file_id FROM files + WHERE file_id = ANY($1::text[]) AND partition_name = $2 + """, + file_ids, + partition, + ) + return {r["file_id"] for r in rows} + + async def remove_file_from_all_workspaces( + self, + file_id: str, + partition: str, + ) -> None: + """Detach a file from every workspace in its partition. + + Called from the file-delete path so workspace integrity is + restored before the underlying file row goes away. A no-op when + the file does not exist in the partition. + """ + async with self.pool.acquire() as conn: + async with conn.transaction(): + file_pk = await conn.fetchval( + """ + SELECT id FROM files + WHERE file_id = $1 AND partition_name = $2 + """, + file_id, + partition, + ) + if file_pk is None: + return + await conn.execute( + """ + DELETE FROM workspace_files + WHERE file_id = $1 + AND workspace_id IN ( + SELECT workspace_id FROM workspaces + WHERE partition_name = $2 + ) + """, + file_pk, + partition, + ) + + # ── Legacy method names used by the Phase 7C shim ──────────────── + + async def create_workspace_legacy( + self, + workspace_id: str, + partition: str, + user_id: int | None, + display_name: str | None = None, + ) -> None: + """TODO(phase-9): remove. Positional-arg mirror of legacy ``create_workspace``.""" + await self.pool.execute( + """ + INSERT INTO workspaces (workspace_id, partition_name, + created_by, display_name, created_at) + VALUES ($1, $2, $3, $4, NOW()) + """, + workspace_id, + partition, + user_id, + display_name, + ) + + async def list_workspaces_dict(self, partition: str) -> list[dict]: + """TODO(phase-9): remove. Legacy router-facing dict shape.""" + rows = await self.pool.fetch( + """ + SELECT * FROM workspaces + WHERE partition_name = $1 + ORDER BY created_at + """, + partition, + ) + return [self._row_to_dict(r) for r in rows] + + async def get_workspace_dict(self, workspace_id: str) -> dict | None: + """TODO(phase-9): remove. Legacy router-facing dict shape.""" + row = await self.pool.fetchrow( + "SELECT * FROM workspaces WHERE workspace_id = $1", + workspace_id, + ) + return self._row_to_dict(row) if row else None + + # ── Helpers ────────────────────────────────────────────────────── + + @staticmethod + def _row_to_workspace(row: asyncpg.Record) -> Workspace: + return Workspace( + workspace_id=row["workspace_id"], + partition=row["partition_name"], + display_name=row["display_name"], + created_by=row["created_by"], + created_at=row["created_at"], + ) + + @staticmethod + def _row_to_dict(row: asyncpg.Record) -> dict: + return { + "workspace_id": row["workspace_id"], + "partition_name": row["partition_name"], + "display_name": row["display_name"], + "created_by": row["created_by"], + "created_at": str(row["created_at"]) if row["created_at"] else None, + } + + +__all__ = ["PgWorkspaceRepository"] diff --git a/openrag/services/storage/__init__.py b/openrag/services/storage/__init__.py new file mode 100644 index 000000000..19a3f9090 --- /dev/null +++ b/openrag/services/storage/__init__.py @@ -0,0 +1,21 @@ +"""Storage adapters — concrete :class:`CatalogStore` and :class:`VectorStore`. + +Phase 7 splits the legacy ``MilvusDB`` Ray god object into two clean +adapters that orchestrators consume through the core port ABCs: + +* :class:`postgres_store.PostgresStore` — composes the asyncpg + :class:`ConnectionManager` with every repository implementation under + :mod:`services.persistence`, satisfying + :class:`core.ports.catalog_store.CatalogStore`. +* :class:`milvus_store.MilvusVectorStore` — Milvus 2.6 backed vector ops + satisfying :class:`core.vector_stores.VectorStore`. + +The Ray actor that callers know today lives at +:mod:`services.storage.milvus_ray_shim` and will be folded into the +new stores during Phase 7C. +""" + +from services.storage.milvus_store import MilvusVectorStore +from services.storage.postgres_store import PostgresStore + +__all__ = ["MilvusVectorStore", "PostgresStore"] diff --git a/openrag/services/storage/milvus_store.py b/openrag/services/storage/milvus_store.py new file mode 100644 index 000000000..9b22f2027 --- /dev/null +++ b/openrag/services/storage/milvus_store.py @@ -0,0 +1,1066 @@ +"""Milvus 2.6 vector store adapter implementing :class:`VectorStore`. + +Scope: + Pure vector operations against a single Milvus collection (the one named + in ``config.vectordb.collection_name``). Embedding, metadata persistence, + surrounding-chunk hydration, workspace resolution, and cross-store + orchestration live elsewhere. + +Collection model: + OpenRAG uses **one shared Milvus collection with a partition_key field**. + The ``collection`` argument on the :class:`VectorStore` ABC therefore maps + to the **``partition`` row-value** tagged on each entity, not to a Milvus + collection name. ``ensure_collection`` / ``drop_collection`` operate at + partition-row granularity. + +Client split (Milvus 2.6): + ``AsyncMilvusClient`` covers the data plane (``insert``, ``search``, + ``hybrid_search``, ``query``, ``delete``, ``upsert``). The admin/lifecycle + plane (``has_collection``, ``create_collection``, ``load_collection``, + ``alter_collection_properties``, ``describe_collection``, + ``query_iterator``, ``prepare_index_params``) is sync-only, so the sync + :class:`MilvusClient` is kept alongside. + +Hybrid BM25: + Milvus 2.6 native ``Function(FunctionType.BM25)`` computes the sparse + vector server-side from the ``text`` field at both insert and query time. + Hybrid is config-driven, not a separate entry point: :meth:`search` + dispatches to :meth:`_hybrid_search` when ``config.hybrid_search`` is on + and :meth:`_dense_search` otherwise. The ``query_text`` argument carries + the raw query Milvus's server-side BM25 ``Function`` needs alongside the + dense embedding; the dense-only path ignores it. +""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Iterator +from contextlib import contextmanager +from datetime import UTC, datetime +from typing import Any + +from pymilvus import ( + AnnSearchRequest, + AsyncMilvusClient, + DataType, + Function, + FunctionType, + MilvusClient, + MilvusException, + RRFRanker, +) + +from openrag.core.config.infrastructure import VectorDBConfig +from openrag.core.models.chunk import Chunk +from openrag.core.utils.exceptions import ( + UnexpectedVDBError, + VDBConnectionError, + VDBCreateOrLoadCollectionError, + VDBDeleteError, + VDBInsertError, + VDBSchemaMigrationRequiredError, + VDBSearchError, +) +from openrag.core.vector_stores import VectorStore + +# --------------------------------------------------------------------------- +# Module constants — lifted verbatim from the legacy MilvusDB so the schema +# is bit-for-bit identical and existing collections load without migration. +# --------------------------------------------------------------------------- + +#: Milvus VARCHAR upper bound used for ``text`` / ``partition`` / ``file_id``. +MAX_LENGTH = 65_535 + +#: Custom collection property holding the schema version integer. +SCHEMA_VERSION_PROPERTY_KEY = "openrag.schema_version" + +#: Scalar time fields that get an ``STL_SORT`` index. +INDEXED_TIME_FIELDS = ["created_at"] + +#: Dense ANN search params for the HNSW/COSINE index on ``vector``. ``ef`` +#: governs the search-time candidate pool size and trades recall for latency. +DEFAULT_DENSE_SEARCH_PARAMS: dict[str, Any] = { + "metric_type": "COSINE", + "params": {"ef": 64}, +} + +#: COSINE upper bound for range search. With ``metric_type="COSINE"`` Milvus +#: keeps hits whose similarity is in ``(radius, range_filter]``; cosine +#: similarity maxes at 1.0, so this is the inclusive ceiling and +#: ``similarity_threshold`` supplies the exclusive ``radius`` floor. +COSINE_RANGE_FILTER_MAX = 1.0 + +#: BM25 search params for the SPARSE_INVERTED_INDEX on ``sparse``. +#: ``drop_ratio_build`` matches the legacy MilvusDB tuning. +DEFAULT_BM25_SEARCH_PARAMS: dict[str, Any] = { + "metric_type": "BM25", + "params": {"drop_ratio_build": 0.2}, +} + +#: Native Milvus 2.6 RRF fusion constant — k=100 matches the legacy MilvusDB +#: tuning and the rank-fusion literature default. +RRF_K = 100 + +#: Entity-level keys to strip from search-result records — ``vector`` is +#: noisy and large; ``text`` stays in the payload (callers need it). +_SEARCH_RESULT_DROPPED_KEYS = frozenset({"vector"}) + +#: BM25 analyzer params for the ``text`` field — standard tokenizer plus +#: OpenRAG-specific stop words so chunk-boundary / image-placeholder markers +#: don't pollute lexical scores. +analyzer_params: dict[str, Any] = { + "tokenizer": "standard", + "filter": [ + { + "type": "stop", + "stop_words": [ + "", + "", + "[Image Placeholder]", + "_english_", + "_french_", + "[CHUNK_START]", + "[CHUNK_END]", + "[CONTEXT]", + ], + } + ], +} + + +class MilvusVectorStore(VectorStore): + """Milvus 2.6 implementation of :class:`VectorStore`. + + The store is constructed cheaply (no I/O); the collection is materialised + on the first :meth:`initialize` call. ``initialize`` is idempotent and + takes the embedding dimension as an argument so the schema does not need + to import the embedder. + """ + + def __init__(self, config: VectorDBConfig) -> None: + self._config = config + self._collection_name = config.collection_name + self._hybrid = config.hybrid_search + self._uri = f"http://{config.host}:{config.port}" + + try: + self._client = MilvusClient(uri=self._uri) + self._async_client = AsyncMilvusClient(uri=self._uri) + except MilvusException as e: + raise VDBConnectionError( + f"Failed to connect to Milvus: {e!s}", + db_url=self._uri, + db_type="Milvus", + ) from e + + self._embedding_dimension: int | None = None + self._loaded = False + self._load_lock = asyncio.Lock() + # Connection healing: pymilvus 2.6 exposes no documented client-level + # reconnect knob (no retry/keepalive params on MilvusClient or + # AsyncMilvusClient — see api-reference v2.6.x). Trust the gRPC + # channel's internal handling, same as the legacy MilvusDB. If + # production drops surface a real issue, revisit with evidence + # rather than racing pymilvus's internal channel state. + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def initialize(self, embedding_dimension: int) -> None: + """Materialise the backing Milvus collection. + + Safe to call multiple times. The first caller wins; concurrent callers + block on the same lock and observe ``_loaded`` set on exit. + + Args: + embedding_dimension: Dimensionality of the dense vectors that will + be upserted. Used to size the ``vector`` field in a fresh + collection. Ignored if the collection already exists. + """ + if self._loaded: + return + async with self._load_lock: + if self._loaded: + return + self._embedding_dimension = embedding_dimension + await asyncio.to_thread(self._ensure_loaded) + self._loaded = True + + def _ensure_loaded(self) -> None: + """Create-if-absent + load the configured collection. + + Synchronous because the Milvus 2.6 admin/lifecycle endpoints + (``has_collection``, ``create_collection``, ``load_collection``, + ``alter_collection_properties``, ``describe_collection``) have no + async equivalents. + """ + try: + if self._client.has_collection(self._collection_name): + self._check_schema_version() + else: + schema = self._create_schema() + index_params = self._create_index() + try: + self._client.create_collection( + collection_name=self._collection_name, + schema=schema, + consistency_level="Strong", + index_params=index_params, + enable_dynamic_field=True, + ) + except MilvusException as e: + raise VDBCreateOrLoadCollectionError( + f"Failed to create collection `{self._collection_name}`: {e!s}", + collection_name=self._collection_name, + operation="create_collection", + ) from e + self._store_schema_version() + + try: + self._client.load_collection(self._collection_name) + except MilvusException as e: + raise VDBCreateOrLoadCollectionError( + f"Failed to load collection `{self._collection_name}`: {e!s}", + collection_name=self._collection_name, + operation="load_collection", + ) from e + + except VDBCreateOrLoadCollectionError: + raise + except VDBSchemaMigrationRequiredError: + raise + except Exception as e: + raise UnexpectedVDBError( + f"Unexpected error preparing collection `{self._collection_name}`: {e!s}", + collection_name=self._collection_name, + ) from e + + # ------------------------------------------------------------------ + # Schema / index + # ------------------------------------------------------------------ + + def _create_schema(self): + """Build the OpenRAG hybrid schema. + + Fields: auto-id ``_id`` (INT64 PK), ``text`` (VARCHAR + analyzer), + ``partition`` (VARCHAR, partition_key), ``file_id`` (VARCHAR), + ``vector`` (FLOAT_VECTOR, dim from :meth:`initialize`), one + TIMESTAMPTZ per field in :data:`INDEXED_TIME_FIELDS`, and — when + ``hybrid_search`` is on — ``sparse`` (SPARSE_FLOAT_VECTOR) wired to a + native :class:`Function` of type :data:`FunctionType.BM25` over + ``text``. + """ + if self._embedding_dimension is None: + raise VDBCreateOrLoadCollectionError( + "embedding_dimension must be set before building the schema; " + "call MilvusVectorStore.initialize(dim) first.", + collection_name=self._collection_name, + operation="create_schema", + ) + + schema = self._client.create_schema(enable_dynamic_field=True) + schema.add_field(field_name="_id", datatype=DataType.INT64, is_primary=True, auto_id=True) + schema.add_field( + field_name="text", + datatype=DataType.VARCHAR, + enable_analyzer=True, + enable_match=True, + max_length=MAX_LENGTH, + analyzer_params=analyzer_params, + ) + schema.add_field( + field_name="partition", + datatype=DataType.VARCHAR, + max_length=MAX_LENGTH, + is_partition_key=True, + ) + schema.add_field( + field_name="file_id", + datatype=DataType.VARCHAR, + max_length=MAX_LENGTH, + ) + schema.add_field( + field_name="vector", + datatype=DataType.FLOAT_VECTOR, + dim=self._embedding_dimension, + ) + + for time_field in INDEXED_TIME_FIELDS: + schema.add_field(field_name=time_field, datatype=DataType.TIMESTAMPTZ, nullable=True) + + if self._hybrid: + schema.add_field( + field_name="sparse", + datatype=DataType.SPARSE_FLOAT_VECTOR, + index_type="SPARSE_INVERTED_INDEX", + ) + schema.add_function( + Function( + name="text_bm25_emb", + function_type=FunctionType.BM25, + input_field_names=["text"], + output_field_names=["sparse"], + ) + ) + + return schema + + def _create_index(self): + """Build index params: HNSW/COSINE on ``vector``, inverted on scalars, + STL_SORT on every :data:`INDEXED_TIME_FIELDS` entry, and — only when + ``hybrid_search`` is enabled — SPARSE_INVERTED_INDEX/BM25 on + ``sparse`` (k1=1.2, b=0.75) to mirror the schema gating in + :meth:`_create_schema`. + """ + index_params = self._client.prepare_index_params() + index_params.add_index( + field_name="file_id", + index_type="INVERTED", + index_name="file_id_idx", + ) + index_params.add_index( + field_name="partition", + index_type="INVERTED", + index_name="partition_idx", + ) + index_params.add_index( + field_name="vector", + index_type="HNSW", + metric_type="COSINE", + index_params={"M": 128, "efConstruction": 256, "metric_type": "COSINE"}, + ) + if self._hybrid: + index_params.add_index( + field_name="sparse", + index_name="sparse_idx", + index_type="SPARSE_INVERTED_INDEX", + index_params={ + "metric_type": "BM25", + "inverted_index_algo": "DAAT_MAXSCORE", + "bm25_k1": 1.2, + "bm25_b": 0.75, + }, + ) + for time_field in INDEXED_TIME_FIELDS: + index_params.add_index( + field_name=time_field, + index_type="STL_SORT", + index_name=f"{time_field}_idx", + ) + return index_params + + # ------------------------------------------------------------------ + # Schema versioning + # ------------------------------------------------------------------ + + def _store_schema_version(self) -> None: + """Persist the configured schema version as a Milvus collection property.""" + self._client.alter_collection_properties( + collection_name=self._collection_name, + properties={SCHEMA_VERSION_PROPERTY_KEY: str(self._config.schema_version)}, + ) + + def _check_schema_version(self) -> None: + """Compare stored vs. configured schema version; raise on mismatch. + + Missing or unparseable values default to ``0`` so existing + pre-versioning collections always trigger an explicit migration step + rather than silently working on a stale schema. + """ + expected_version = self._config.schema_version + desc = self._client.describe_collection(self._collection_name) + raw = desc.get("properties", {}).get(SCHEMA_VERSION_PROPERTY_KEY) + try: + stored_version = int(raw) if raw is not None else 0 + except (ValueError, TypeError): + stored_version = 0 + + if stored_version != expected_version: + raise VDBSchemaMigrationRequiredError( + f"Collection `{self._collection_name}` is at schema version " + f"{stored_version} but the application requires version " + f"{expected_version}. Please perform the migration script.", + collection_name=self._collection_name, + stored_version=stored_version, + expected_version=expected_version, + ) + + # ------------------------------------------------------------------ + # Collection-arg discipline + # ------------------------------------------------------------------ + # + # The :class:`VectorStore` ABC carries a ``collection`` argument on most + # methods; in Milvus terminology a *collection* is the top-level data + # container (one per store, set by config) while a *partition* is a row + # tag implemented via ``partition_key``. This store services exactly one + # Milvus collection, so the ABC's ``collection`` arg either: + # + # * equals ``self._collection_name`` -> accepted, no-op. + # * equals the ABC default ``"default"`` -> treated as "use mine". + # * anything else -> :class:`ValueError`. + # + # Partition row-tagging lives exclusively in ``filters['partition']`` (or + # in ``Chunk.partition`` on the write path). + + _COLLECTION_DEFAULT_SENTINEL = "default" + + def _resolve_collection(self, collection: str) -> str: + if collection in (self._collection_name, self._COLLECTION_DEFAULT_SENTINEL): + return self._collection_name + raise ValueError( + f"MilvusVectorStore is bound to collection `{self._collection_name}`; " + f"got `{collection}`. One store services exactly one Milvus collection — " + "partitions go in filters, not in the `collection` argument." + ) + + # ------------------------------------------------------------------ + # Filter-expression construction + # ------------------------------------------------------------------ + + #: Filter keys with dedicated semantics, pulled out before the generic + #: ``key == value`` loop runs. ``partition`` is the partition_key row + #: tag; ``expr`` is a raw-expression escape hatch. + _SPECIAL_FILTER_KEYS = frozenset({"partition", "expr"}) + + #: Partition values that mean "do not filter by partition". + _PARTITION_WILDCARDS = frozenset({"all"}) + + # Whitespace-stripped, lowercased raw expressions that match every row. + # ``delete_by_filter`` rejects these so callers don't accidentally wipe + # the collection through ``filters={"expr": "1==1"}`` — explicit drops + # must go through :meth:`drop_collection`. + _TAUTOLOGICAL_EXPRS = frozenset({"true", "1==1"}) + + @staticmethod + def _format_value(value: Any) -> str: + """Render a scalar as a Milvus filter literal. + + Strings are double-quoted with ``\\`` and ``"`` escaped; bools are + rendered lower-case; ints / floats pass through unquoted. + """ + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, float)): + return str(value) + s = str(value).replace("\\", "\\\\").replace('"', '\\"') + return f'"{s}"' + + def _build_filter_expr(self, filters: dict[str, Any] | None) -> str: + """Translate a filter dict into a Milvus boolean expression. + + Rules: + * ``filters['partition']`` builds the partition_key clause. + Value ``"all"`` (or ``["all"]``) skips the clause. List/tuple + values become ``partition in [...]``. Mixing a wildcard with + explicit partitions in the same list raises ``ValueError`` — + that combination is rejected rather than silently widened to + every partition. + * ``filters['expr']`` is appended verbatim as an escape hatch + for callers that need operators the dict form cannot express. + * Any other key with a scalar value becomes ``key == ``. + * Any other key with a list/tuple value becomes ``key in [...]``. + Empty list/tuple short-circuits to ``"false"`` (matches no row). + + Workspace-id resolution, role checks, and other PG concerns are + upstream concerns — they resolve to ``file_id`` lists before reaching + this store. + """ + filters = dict(filters or {}) + parts: list[str] = [] + + partition = filters.pop("partition", None) + if isinstance(partition, (list, tuple)): + has_wildcard = any(p in self._PARTITION_WILDCARDS for p in partition) + if has_wildcard and len(partition) > 1: + raise ValueError("`partition` cannot mix wildcard with explicit values.") + if not has_wildcard and partition: + quoted = ", ".join(self._format_value(p) for p in partition) + parts.append(f"partition in [{quoted}]") + elif partition is not None and partition not in self._PARTITION_WILDCARDS: + parts.append(f"partition == {self._format_value(partition)}") + + raw_expr = filters.pop("expr", None) + + for key, value in filters.items(): + if key in self._SPECIAL_FILTER_KEYS: + continue # already handled above + if isinstance(value, (list, tuple)): + if not value: + return "false" # empty IN list — match nothing + quoted = ", ".join(self._format_value(v) for v in value) + parts.append(f"{key} in [{quoted}]") + else: + parts.append(f"{key} == {self._format_value(value)}") + + if raw_expr: + parts.append(str(raw_expr)) + + return " and ".join(parts) + + # ------------------------------------------------------------------ + # Sync paginated query helper (Milvus 2.6 query_iterator is sync-only) + # ------------------------------------------------------------------ + + def _iter_query( + self, + expr: str, + output_fields: list[str], + batch_size: int = 16_000, + ) -> list[dict[str, Any]]: + """Drain a Milvus 2.6 ``query_iterator`` into a list. + + Synchronous; call via :func:`asyncio.to_thread` from async methods. + """ + iterator = self._client.query_iterator( + collection_name=self._collection_name, + filter=expr, + batch_size=batch_size, + output_fields=output_fields, + ) + out: list[dict[str, Any]] = [] + try: + while True: + batch = iterator.next() + if not batch: + break + out.extend(batch) + finally: + iterator.close() + return out + + # ------------------------------------------------------------------ + # ID round-trip (Chunk.id: str <--> Milvus _id: INT64 auto_id PK) + # ------------------------------------------------------------------ + + @staticmethod + def _str_id_to_milvus(id_str: str) -> int | None: + """Coerce a ``Chunk.id`` string to a Milvus INT64 ``_id``. + + Returns ``None`` for non-numeric IDs (e.g. fresh UUIDs that have not + been round-tripped through Milvus yet) so callers can skip them + rather than crash a batch delete. + """ + try: + return int(id_str) + except (TypeError, ValueError): + return None + + @staticmethod + def _milvus_id_to_str(id_int: int) -> str: + """Convert a Milvus ``_id`` (INT64) back to the domain string form.""" + return str(id_int) + + # ------------------------------------------------------------------ + # Entity construction + # ------------------------------------------------------------------ + + @staticmethod + def _gen_chunk_order_metadata(n: int) -> list[dict[str, int | None]]: + """Generate prev/section/next IDs for a batch of ``n`` chunks. + + Uses a monotonic nanosecond base so IDs are unique across batches. + Preserves the legacy MilvusDB ordering so existing + surrounding-chunk hydration keeps working. + """ + base_ts = int(time.time_ns()) + ids = [base_ts + i for i in range(n)] + return [ + { + "prev_section_id": ids[i - 1] if i > 0 else None, + "section_id": ids[i], + "next_section_id": ids[i + 1] if i < n - 1 else None, + } + for i in range(n) + ] + + @staticmethod + def _chunk_to_entity( + chunk: Chunk, + *, + indexed_at: str, + order: dict[str, int | None], + ) -> dict[str, Any]: + """Build the Milvus insert payload for one chunk. + + Layering: start from the free-form ``chunk.metadata`` dict, then + overwrite with the typed Chunk fields so the strict domain model + always wins over caller-supplied metadata keys with the same name. + ``_id`` is intentionally omitted — Milvus assigns it via ``auto_id``. + """ + entity: dict[str, Any] = dict(chunk.metadata) + entity.update( + { + "text": chunk.text, + "vector": chunk.embedding, + "partition": chunk.partition, + "file_id": chunk.document_id, + "chunk_type": chunk.chunk_type.value, + "page": chunk.page_number, + "indexed_at": indexed_at, + **order, + } + ) + # Optional typed fields only emitted when set, to avoid stamping + # nulls into the dynamic schema. + for field, value in ( + ("chunk_index", chunk.chunk_index), + ("token_count", chunk.token_count), + ("header", chunk.header), + ("context", chunk.context), + ("content", chunk.content), + ): + if value is not None: + entity[field] = value + return entity + + # ------------------------------------------------------------------ + # VectorStore ABC — writes + # ------------------------------------------------------------------ + + async def upsert(self, chunks: list[Chunk], collection: str = "default") -> int: + """Insert pre-embedded chunks into the backing Milvus collection. + + ``chunk.partition`` is authoritative — the ``collection`` argument is + accepted for ABC compatibility but does not override per-chunk + partition values. Every chunk MUST carry a populated ``embedding``; + embedding is an upstream pipeline concern, not a store concern. + """ + self._resolve_collection(collection) + if not chunks: + return 0 + + missing = [c.id for c in chunks if c.embedding is None] + if missing: + raise VDBInsertError( + f"upsert received {len(missing)} chunk(s) with no embedding; embed before calling the vector store.", + collection_name=self._collection_name, + ) + + indexed_at = datetime.now(UTC).isoformat() + order_metadata = self._gen_chunk_order_metadata(len(chunks)) + entities = [ + self._chunk_to_entity(c, indexed_at=indexed_at, order=o) + for c, o in zip(chunks, order_metadata, strict=True) + ] + + try: + result = await self._async_client.insert( + collection_name=self._collection_name, + data=entities, + ) + except MilvusException as e: + raise VDBInsertError( + f"Milvus insert failed: {e!s}", + collection_name=self._collection_name, + ) from e + except Exception as e: + raise UnexpectedVDBError( + f"Unexpected error during Milvus insert: {e!s}", + collection_name=self._collection_name, + ) from e + + # Milvus 2.6 returns {"insert_count": N, "ids": [...], "cost": ...}. + # Fall back to len(entities) if the server omits insert_count. + return int(result.get("insert_count", len(entities))) if isinstance(result, dict) else len(entities) + + def _parse_search_response(self, response: Any) -> list[dict[str, Any]]: + """Normalise a Milvus 2.6 search/hybrid_search response to raw dicts. + + Each record has ``id`` (stringified for :class:`Chunk` round-trip), + ``score`` (distance for dense, fused RRF score for hybrid), and the + entity's output fields except ``vector``. + """ + if not response: + return [] + out: list[dict[str, Any]] = [] + for hit in response[0]: + entity = hit.get("entity", {}) if isinstance(hit, dict) else {} + record = {k: v for k, v in entity.items() if k not in _SEARCH_RESULT_DROPPED_KEYS} + record["id"] = self._milvus_id_to_str(hit.get("id")) + record["score"] = hit.get("distance") + out.append(record) + return out + + @contextmanager + def _search_errors(self, kind: str) -> Iterator[None]: + """Map Milvus failures from a search call to the VDB error taxonomy. + + Wraps the ``await`` site so :meth:`search` and :meth:`hybrid_search` + don't each repeat the same two-arm ``MilvusException`` / + ``Exception`` translation. ``kind`` names the operation for the + message (``"dense search"`` / ``"hybrid search"``). + """ + try: + yield + except MilvusException as e: + raise VDBSearchError( + f"Milvus {kind} failed: {e!s}", + collection_name=self._collection_name, + ) from e + except Exception as e: + raise UnexpectedVDBError( + f"Unexpected error during Milvus {kind}: {e!s}", + collection_name=self._collection_name, + ) from e + + def _dense_search_params(self, similarity_threshold: float | None) -> dict[str, Any]: + """Build the dense COSINE search params, optionally range-filtered. + + Returns a fresh dict each call so the frozen + :data:`DEFAULT_DENSE_SEARCH_PARAMS` module default is never mutated. + When ``similarity_threshold`` is set, Milvus range search keeps only + hits whose COSINE similarity falls in + ``(similarity_threshold, COSINE_RANGE_FILTER_MAX]`` — the same + ``radius`` / ``range_filter`` pair the legacy MilvusDB used. ``None`` + leaves it an unbounded top-k search. + """ + params = dict(DEFAULT_DENSE_SEARCH_PARAMS["params"]) + if similarity_threshold is not None: + params["radius"] = similarity_threshold + params["range_filter"] = COSINE_RANGE_FILTER_MAX + return {"metric_type": DEFAULT_DENSE_SEARCH_PARAMS["metric_type"], "params": params} + + async def search( + self, + embedding: list[float], + query_text: str | None = None, + top_k: int = 10, + collection: str = "default", + filters: dict[str, Any] | None = None, + similarity_threshold: float | None = None, + ) -> list[dict[str, Any]]: + """Similarity search — single entry point for dense and hybrid. + + Hybrid is a collection-build property, not a caller choice: the + store dispatches to :meth:`_hybrid_search` when ``config.hybrid_search`` + was on (the backing collection then has a ``sparse`` BM25 field) and + to :meth:`_dense_search` otherwise. ``query_text`` is only consumed on + the hybrid path — Milvus's server-side BM25 ``Function`` generates the + sparse vector from it; the dense path ignores it. + + Returns raw dicts (``id``, ``score``, plus entity fields except + ``vector``). ``similarity_threshold`` (when set) is the range-search + ``radius`` floor on the dense leg; see :meth:`_dense_search_params`. + """ + if self._hybrid: + return await self._hybrid_search(embedding, query_text, top_k, collection, filters, similarity_threshold) + return await self._dense_search(embedding, top_k, collection, filters, similarity_threshold) + + async def _dense_search( + self, + embedding: list[float], + top_k: int, + collection: str, + filters: dict[str, Any] | None, + similarity_threshold: float | None, + ) -> list[dict[str, Any]]: + """Dense ANN search on the ``vector`` field. + + Uses HNSW with COSINE distance and ``ef=64`` — same tuning as the + legacy MilvusDB. + """ + self._resolve_collection(collection) + expr = self._build_filter_expr(filters) + + with self._search_errors("dense search"): + response = await self._async_client.search( + collection_name=self._collection_name, + data=[embedding], + anns_field="vector", + search_params=self._dense_search_params(similarity_threshold), + limit=top_k, + filter=expr, + output_fields=["*"], + ) + + return self._parse_search_response(response) + + async def _hybrid_search( + self, + embedding: list[float], + query_text: str | None, + top_k: int, + collection: str, + filters: dict[str, Any] | None, + similarity_threshold: float | None, + ) -> list[dict[str, Any]]: + """Dense + Milvus-native BM25 sparse, fused via ``RRFRanker``. + + Only reached when the backing collection was built with + ``hybrid_search=True`` (it then has the ``sparse`` field). The + ``query_text`` is required here — Milvus's server-side + ``Function(FunctionType.BM25)`` generates the sparse vector from it, + so a missing query would silently drop the lexical signal. + + ``similarity_threshold`` (when set) range-filters the dense leg only; + the BM25 leg has no comparable distance metric, matching the legacy + MilvusDB behaviour. + + Raises: + VDBSearchError: ``query_text`` is ``None`` — the BM25 leg has no + input. + """ + self._resolve_collection(collection) + if query_text is None: + raise VDBSearchError( + f"hybrid search on collection `{self._collection_name}` requires " + "query_text for the server-side BM25 leg; got None.", + collection_name=self._collection_name, + ) + expr = self._build_filter_expr(filters) + + dense_req = AnnSearchRequest( + data=[embedding], + anns_field="vector", + param=self._dense_search_params(similarity_threshold), + limit=top_k, + expr=expr, + ) + sparse_req = AnnSearchRequest( + data=[query_text], + anns_field="sparse", + param=DEFAULT_BM25_SEARCH_PARAMS, + limit=top_k, + expr=expr, + ) + + with self._search_errors("hybrid search"): + response = await self._async_client.hybrid_search( + collection_name=self._collection_name, + reqs=[dense_req, sparse_req], + ranker=RRFRanker(RRF_K), + limit=top_k, + output_fields=["*"], + ) + + return self._parse_search_response(response) + + async def delete(self, ids: list[str], collection: str = "default") -> int: + """Delete chunks by Milvus ``_id``. + + ``Chunk.id`` is a string while the Milvus primary key is INT64. + Non-numeric IDs are silently dropped (they cannot exist in Milvus by + construction) so a partially-fresh batch doesn't fail the whole call. + The ``collection`` argument is accepted for ABC compatibility; the + Milvus delete is scoped to the backing collection regardless, and + rows are uniquely keyed by ``_id``. + """ + self._resolve_collection(collection) + if not ids: + return 0 + + numeric_ids = [n for n in (self._str_id_to_milvus(i) for i in ids) if n is not None] + if not numeric_ids: + return 0 + + try: + result = await self._async_client.delete( + collection_name=self._collection_name, + ids=numeric_ids, + ) + except MilvusException as e: + raise VDBDeleteError( + f"Milvus delete failed: {e!s}", + collection_name=self._collection_name, + ) from e + except Exception as e: + raise UnexpectedVDBError( + f"Unexpected error during Milvus delete: {e!s}", + collection_name=self._collection_name, + ) from e + + return int(result.get("delete_count", 0)) if isinstance(result, dict) else 0 + + async def upsert_entities(self, entities: list[dict[str, Any]], collection: str = "default") -> int: + """Upsert raw Milvus entities that already include vector data. + + This is intentionally narrower than the VectorStore port: it supports + file metadata patch/copy paths where re-embedding would be wrong and + the existing Milvus rows already carry the vectors to preserve. + """ + self._resolve_collection(collection) + if not entities: + return 0 + + try: + result = await self._async_client.upsert( + collection_name=self._collection_name, + data=entities, + ) + except MilvusException as e: + raise VDBInsertError( + f"Milvus raw entity upsert failed: {e!s}", + collection_name=self._collection_name, + ) from e + except Exception as e: + raise UnexpectedVDBError( + f"Unexpected error during Milvus raw entity upsert: {e!s}", + collection_name=self._collection_name, + ) from e + + return int(result.get("upsert_count", len(entities))) if isinstance(result, dict) else len(entities) + + async def insert_entities(self, entities: list[dict[str, Any]], collection: str = "default") -> int: + """Insert raw Milvus entities that already include vector data.""" + self._resolve_collection(collection) + if not entities: + return 0 + + try: + result = await self._async_client.insert( + collection_name=self._collection_name, + data=entities, + ) + except MilvusException as e: + raise VDBInsertError( + f"Milvus raw entity insert failed: {e!s}", + collection_name=self._collection_name, + ) from e + except Exception as e: + raise UnexpectedVDBError( + f"Unexpected error during Milvus raw entity insert: {e!s}", + collection_name=self._collection_name, + ) from e + + return int(result.get("insert_count", len(entities))) if isinstance(result, dict) else len(entities) + + async def ensure_collection(self, name: str, dimension: int, **kwargs: Any) -> None: + """Public entry point for materialising the backing collection. + + Thin wrapper over :meth:`initialize`: validates ``name`` against the + bound collection (so a future per-tenant store factory cannot + accidentally cross-wire one tenant's collection name into another's + store) and forwards ``dimension``. Idempotent. + + Raises: + ValueError: ``name`` is neither ``self._collection_name`` nor + the ABC sentinel ``"default"``. + ValueError: the store is already initialized with a different + embedding dimension — re-initialising would invalidate the + index, so callers must drop and re-create explicitly. + """ + self._resolve_collection(name) + if self._loaded and self._embedding_dimension != dimension: + raise ValueError( + f"MilvusVectorStore already initialised at " + f"dimension={self._embedding_dimension}; " + f"got ensure_collection(dimension={dimension}). " + "Drop the collection before re-sizing." + ) + await self.initialize(dimension) + + async def drop_collection(self, name: str) -> None: + """Destructive: drop the entire backing Milvus collection. + + For administrative / test fixture use only. To remove rows for a + specific partition or any filterable subset, call + :meth:`delete_by_filter` instead — that is the surface the 7C shim + uses for partition-level deletion. + """ + self._resolve_collection(name) + try: + await asyncio.to_thread(self._client.drop_collection, self._collection_name) + except MilvusException as e: + raise VDBDeleteError( + f"Failed to drop collection `{self._collection_name}`: {e!s}", + collection_name=self._collection_name, + ) from e + except Exception as e: + raise UnexpectedVDBError( + f"Unexpected error dropping collection `{self._collection_name}`: {e!s}", + collection_name=self._collection_name, + ) from e + self._loaded = False + self._embedding_dimension = None + + # ------------------------------------------------------------------ + # Milvus-specific (not on the VectorStore ABC) + # ------------------------------------------------------------------ + + async def delete_by_filter(self, filters: dict[str, Any]) -> int: + """Delete every row whose entity matches the filter expression. + + Used by callers that want to remove a partition's worth of rows + without first paginating all chunk IDs (e.g. the legacy + ``delete_partition`` flow). Guarded so an accidental empty / + wildcard filter does NOT nuke the entire collection — explicit + drop is :meth:`drop_collection`. + + Raises: + ValueError: ``filters`` builds an empty Milvus expression + (no clauses, or only a wildcard partition), or resolves to + a tautological raw expression such as ``"1==1"``/``"true"`` + that would wipe the collection. + """ + expr = self._build_filter_expr(filters) + normalized = "".join(expr.lower().split()) if expr else "" + if not expr or normalized in self._TAUTOLOGICAL_EXPRS: + raise ValueError( + "delete_by_filter requires a non-empty, non-tautological " + "filter expression. To delete every row, call " + "drop_collection() explicitly." + ) + try: + result = await self._async_client.delete( + collection_name=self._collection_name, + filter=expr, + ) + except MilvusException as e: + raise VDBDeleteError( + f"Milvus delete-by-filter failed (expr=`{expr}`): {e!s}", + collection_name=self._collection_name, + ) from e + except Exception as e: + raise UnexpectedVDBError( + f"Unexpected error during Milvus delete-by-filter: {e!s}", + collection_name=self._collection_name, + ) from e + + return int(result.get("delete_count", 0)) if isinstance(result, dict) else 0 + + async def collection_exists(self, name: str) -> bool: + """Report whether the Milvus collection exists on the server. + + Accepts ``self._collection_name`` or the ABC default ``"default"``; + any other name falsifies (we don't query other collections — this + store services exactly one). + """ + if name not in (self._collection_name, self._COLLECTION_DEFAULT_SENTINEL): + return False + return await asyncio.to_thread(self._client.has_collection, self._collection_name) + + async def query_ids_by_filter( + self, + collection: str, + filters: dict[str, Any], + ) -> list[str]: + """Return ``Chunk.id`` strings for every row matching ``filters``. + + Uses Milvus 2.6 ``query_iterator`` under the hood so result-set size + is bounded only by Milvus pagination, not by a server-side + ``limit``. The returned IDs are the INT64 ``_id`` values stringified + for round-trip with :class:`Chunk`. + """ + self._resolve_collection(collection) + expr = self._build_filter_expr(filters) + rows = await asyncio.to_thread(self._iter_query, expr, ["_id"]) + return [self._milvus_id_to_str(r["_id"]) for r in rows if "_id" in r] + + async def query_chunks_by_filter( + self, + collection: str, + filters: dict[str, Any], + output_fields: list[str] | None = None, + ) -> list[dict[str, Any]]: + """Return full row data for every chunk matching ``filters``. + + ``output_fields`` defaults to ``["*"]``. Milvus 2.6 quirk: ``"*"`` + does NOT include the dense vector — callers that need the vector + must pass ``output_fields=["*", "vector"]`` explicitly. + """ + self._resolve_collection(collection) + expr = self._build_filter_expr(filters) + fields = output_fields or ["*"] + return await asyncio.to_thread(self._iter_query, expr, fields) diff --git a/openrag/services/storage/postgres_store.py b/openrag/services/storage/postgres_store.py new file mode 100644 index 000000000..ed90d73a1 --- /dev/null +++ b/openrag/services/storage/postgres_store.py @@ -0,0 +1,231 @@ +"""PostgresStore — concrete :class:`~openrag.core.ports.catalog_store.CatalogStore`. + +Composes the connection manager (Phase 7A.1) with every repository +implementation (Phase 7A.2) into a single aggregate root that orchestrators +and the Phase 7C shim consume through the ``CatalogStore`` ABC. + +The store owns the lifecycle: + +* :meth:`initialize` opens the asyncpg pool, then runs Alembic migrations to + ``head``. The order matters — the legacy ORM's ``Base.metadata.create_all`` + used to fast-forward the schema before migrations ran, which is why every + Alembic revision is idempotent (see ``CLAUDE.md`` "Alembic Migration + Idempotency"). The new store keeps that contract. +* :meth:`shutdown` closes the pool. Repositories share the pool via a + ``pool_getter`` callable, so the pool can be reinitialised in tests without + rebuilding the repos. + +The optional :pyattr:`pool` property is an escape hatch for cross-repo +transactional work — Phase 8 orchestrators will reach for it to wrap multiple +repo writes in a single :func:`asyncpg.Pool.acquire` + ``conn.transaction()`` +context. It is not part of the ABC contract; clients that only need a single +repository call should never touch it. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from core.ports.catalog_store import CatalogStore +from services.persistence.audit_log_repo import PgAuditLogRepository +from services.persistence.chunk_repo import PgChunkRepository +from services.persistence.connection import ConnectionManager +from services.persistence.conversation_repo import PgConversationRepository +from services.persistence.document_repo import PgDocumentRepository +from services.persistence.entity_repo import PgEntityRepository +from services.persistence.idempotency_repo import PgIdempotencyRepository +from services.persistence.job_repo import PgJobRepository +from services.persistence.model_endpoint_repo import PgModelEndpointRepository +from services.persistence.oidc_session_repo import PgOIDCSessionRepository +from services.persistence.partition_membership_repo import PgPartitionMembershipRepository +from services.persistence.partition_repo import PgPartitionRepository +from services.persistence.preset_repo import PgPresetRepository +from services.persistence.prompt_repo import PgPromptRepository +from services.persistence.topic_tag_repo import PgTopicTagRepository +from services.persistence.user_repo import PgUserRepository +from services.persistence.workspace_repo import PgWorkspaceRepository + +if TYPE_CHECKING: + import asyncpg + from core.config.infrastructure import RDBConfig + from core.ports.audit_log_repo import AuditLogRepository + from core.ports.chunk_repo import ChunkRepository + from core.ports.conversation_repo import ConversationRepository + from core.ports.document_repo import DocumentRepository + from core.ports.entity_repo import EntityRepository + from core.ports.idempotency_repo import IdempotencyRepository + from core.ports.job_repo import JobRepository + from core.ports.model_endpoint_repo import ModelEndpointRepository + from core.ports.oidc_session_repo import OIDCSessionRepository + from core.ports.partition_membership_repo import PartitionMembershipRepository + from core.ports.partition_repo import PartitionRepository + from core.ports.preset_repo import PresetRepository + from core.ports.prompt_repo import PromptRepository + from core.ports.topic_tag_repo import TopicTagRepository + from core.ports.user_repo import UserRepository + from core.ports.workspace_repo import WorkspaceRepository + + +class PostgresStore(CatalogStore): + """asyncpg-backed :class:`CatalogStore` composing all repository ports.""" + + def __init__(self, config: RDBConfig, *, run_migrations: bool = True) -> None: + self._conn = ConnectionManager(config) + self._run_migrations = run_migrations + self._initialized = False + + # Repositories take a pool_getter callable instead of a pool reference + # so they always see the live pool even if ConnectionManager is + # reinitialised between tests. + pool_getter = self._pool_getter + + self._document_repo = PgDocumentRepository(pool_getter) + self._user_repo = PgUserRepository(pool_getter) + self._partition_repo = PgPartitionRepository(pool_getter) + self._membership_repo = PgPartitionMembershipRepository(pool_getter) + self._oidc_session_repo = PgOIDCSessionRepository(pool_getter) + self._workspace_repo = PgWorkspaceRepository(pool_getter) + + # Stubs — every method raises StubRepositoryError until the matching + # table exists. Listed in the post-refactoring roadmap. + self._job_repo = PgJobRepository(pool_getter) + self._chunk_repo = PgChunkRepository(pool_getter) + self._prompt_repo = PgPromptRepository(pool_getter) + self._conversation_repo = PgConversationRepository(pool_getter) + self._audit_log_repo = PgAuditLogRepository(pool_getter) + self._idempotency_repo = PgIdempotencyRepository(pool_getter) + self._entity_repo = PgEntityRepository(pool_getter) + self._topic_tag_repo = PgTopicTagRepository(pool_getter) + self._model_endpoint_repo = PgModelEndpointRepository(pool_getter) + self._preset_repo = PgPresetRepository(pool_getter) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def initialize(self) -> None: + """Open the asyncpg pool, then upgrade the schema to ``head``. + + Idempotent: ``get_vectordb()`` re-invokes the actor's ``initialize`` + on every request (it is a FastAPI ``Depends``), and the documented + contract is that re-running on a hot actor is a no-op. Without the + ``_initialized`` guard each request would re-run a full Alembic + ``command.upgrade``, which under concurrent load starves the Postgres + pool and surfaces as 500s. + + Order matters: the pool must exist before Alembic runs because the + legacy ``PartitionFileManager`` bootstrapped tables synchronously via + ``Base.metadata.create_all`` *before* migrations. The Phase 7 + migrations therefore guard every DDL with an inspector check, which + keeps re-runs safe regardless of pool state. + """ + if self._initialized: + return + await self._conn.initialize() + if self._run_migrations: + await self._conn.run_migrations() + self._initialized = True + + async def shutdown(self) -> None: + await self._conn.shutdown() + self._initialized = False + + # ------------------------------------------------------------------ + # Connection access (escape hatch for Phase 8 orchestrators) + # ------------------------------------------------------------------ + + @property + def pool(self) -> asyncpg.Pool: + """Raw asyncpg pool for cross-repo transactional work. + + Phase 8 orchestrators need to wrap multi-repo writes in a single + transaction (e.g. delete-document + delete-chunks). This property + exposes the pool *only* to that caller — most code paths should + never touch it. + """ + return self._conn.pool + + # ------------------------------------------------------------------ + # Real repos + # ------------------------------------------------------------------ + + @property + def document_repo(self) -> DocumentRepository: + return self._document_repo + + @property + def user_repo(self) -> UserRepository: + return self._user_repo + + @property + def partition_repo(self) -> PartitionRepository: + return self._partition_repo + + @property + def membership_repo(self) -> PartitionMembershipRepository: + return self._membership_repo + + @property + def oidc_session_repo(self) -> OIDCSessionRepository: + return self._oidc_session_repo + + @property + def workspace_repo(self) -> WorkspaceRepository: + return self._workspace_repo + + # ------------------------------------------------------------------ + # Stub repos — methods raise StubRepositoryError until the matching + # tables and orchestrators are added in the post-refactoring roadmap. + # ------------------------------------------------------------------ + + @property + def job_repo(self) -> JobRepository: + return self._job_repo + + @property + def chunk_repo(self) -> ChunkRepository: + return self._chunk_repo + + @property + def prompt_repo(self) -> PromptRepository: + return self._prompt_repo + + @property + def conversation_repo(self) -> ConversationRepository: + return self._conversation_repo + + @property + def audit_log_repo(self) -> AuditLogRepository: + return self._audit_log_repo + + @property + def idempotency_repo(self) -> IdempotencyRepository: + return self._idempotency_repo + + @property + def entity_repo(self) -> EntityRepository: + return self._entity_repo + + @property + def topic_tag_repo(self) -> TopicTagRepository: + return self._topic_tag_repo + + @property + def model_endpoint_repo(self) -> ModelEndpointRepository: + return self._model_endpoint_repo + + @property + def preset_repo(self) -> PresetRepository: + return self._preset_repo + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _pool_getter(self) -> asyncpg.Pool: + # Resolved at call time so repos can keep working after a + # shutdown()/initialize() cycle in tests. + return self._conn.pool + + +__all__ = ["PostgresStore"] diff --git a/openrag/services/storage/test_milvus_store.py b/openrag/services/storage/test_milvus_store.py new file mode 100644 index 000000000..35286c13e --- /dev/null +++ b/openrag/services/storage/test_milvus_store.py @@ -0,0 +1,429 @@ +"""Unit tests for the pure-logic surface of :class:`MilvusVectorStore`. + +These tests instantiate the store with both Milvus clients mocked out, so they +exercise filter-expression construction, ID coercion, entity layering, and the +``collection`` argument discipline without touching a live Milvus. + +Integration tests that round-trip through a real Milvus 2.6 container live in +:mod:`test_milvus_store_integration` and are gated by the ``integration`` +pytest marker. +""" + +from __future__ import annotations + +import re +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from openrag.core.config.infrastructure import VectorDBConfig +from openrag.core.models.chunk import Chunk, ChunkType +from openrag.core.utils.exceptions import VDBSearchError +from openrag.services.storage.milvus_store import MilvusVectorStore + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def vdb_config() -> VectorDBConfig: + """A config bound to a non-default collection name so ``_resolve_collection`` + has a real value to validate against (a ``vdb_test`` default would collide + with the ABC sentinel in some assertions). + """ + return VectorDBConfig( + host="milvus-test", + port=19530, + collection_name="test_collection", + hybrid_search=True, + schema_version=1, + ) + + +@pytest.fixture +def store(vdb_config: VectorDBConfig, monkeypatch: pytest.MonkeyPatch) -> MilvusVectorStore: + """A ``MilvusVectorStore`` with both pymilvus clients mocked. + + Use this for pure-logic tests. Methods that drive the client (``upsert``, + ``search``, ...) will hit the mocks; assert on mock calls if you must. + + We monkeypatch the symbols directly inside the loaded module rather than + using ``unittest.mock.patch`` because the project's pythonpath setup + (``pythonpath = ./openrag`` plus the ``openrag`` package itself on + sys.path) lets the same source file get registered under two different + module names depending on the import form, which makes string-based + patch paths fragile. + """ + import openrag.services.storage.milvus_store as _store_mod + + monkeypatch.setattr(_store_mod, "MilvusClient", MagicMock()) + monkeypatch.setattr(_store_mod, "AsyncMilvusClient", MagicMock()) + return MilvusVectorStore(vdb_config) + + +# --------------------------------------------------------------------------- +# _format_value +# --------------------------------------------------------------------------- + + +class TestFormatValue: + def test_int_renders_unquoted(self) -> None: + assert MilvusVectorStore._format_value(42) == "42" + + def test_float_renders_unquoted(self) -> None: + assert MilvusVectorStore._format_value(3.14) == "3.14" + + def test_true_renders_lowercase(self) -> None: + assert MilvusVectorStore._format_value(True) == "true" + + def test_false_renders_lowercase(self) -> None: + # bool is an int subclass — make sure we hit the bool branch first. + assert MilvusVectorStore._format_value(False) == "false" + + def test_string_is_double_quoted(self) -> None: + assert MilvusVectorStore._format_value("alice") == '"alice"' + + def test_string_escapes_double_quotes(self) -> None: + assert MilvusVectorStore._format_value('a"b') == '"a\\"b"' + + def test_string_escapes_backslashes(self) -> None: + assert MilvusVectorStore._format_value("a\\b") == '"a\\\\b"' + + def test_string_escape_order(self) -> None: + # Backslash must be escaped before the quote so we don't double-escape + # the quote's preceding backslash. + assert MilvusVectorStore._format_value('a\\"b') == '"a\\\\\\"b"' + + +# --------------------------------------------------------------------------- +# _build_filter_expr +# --------------------------------------------------------------------------- + + +class TestBuildFilterExpr: + def test_none_yields_empty(self, store: MilvusVectorStore) -> None: + assert store._build_filter_expr(None) == "" + + def test_empty_dict_yields_empty(self, store: MilvusVectorStore) -> None: + assert store._build_filter_expr({}) == "" + + def test_scalar_partition(self, store: MilvusVectorStore) -> None: + assert store._build_filter_expr({"partition": "p1"}) == 'partition == "p1"' + + def test_list_partition(self, store: MilvusVectorStore) -> None: + expr = store._build_filter_expr({"partition": ["p1", "p2"]}) + assert expr == 'partition in ["p1", "p2"]' + + def test_partition_wildcard_is_skipped(self, store: MilvusVectorStore) -> None: + # 'all' is the documented wildcard — should produce no partition clause. + assert store._build_filter_expr({"partition": "all"}) == "" + + def test_partition_wildcard_alone_in_list_is_skipped(self, store: MilvusVectorStore) -> None: + # Wildcard on its own in a list is still a wildcard — no partition clause. + assert store._build_filter_expr({"partition": ["all"]}) == "" + + def test_partition_wildcard_mixed_with_explicit_raises(self, store: MilvusVectorStore) -> None: + # Mixing the wildcard with explicit partitions would silently widen the + # query/delete scope to every partition. Reject rather than absorb. + with pytest.raises(ValueError, match="cannot mix wildcard"): + store._build_filter_expr({"partition": ["all", "p1"]}) + + def test_empty_partition_list_is_skipped(self, store: MilvusVectorStore) -> None: + # No partitions means no partition clause (not match-nothing). + assert store._build_filter_expr({"partition": []}) == "" + + def test_scalar_field(self, store: MilvusVectorStore) -> None: + assert store._build_filter_expr({"file_id": "abc"}) == 'file_id == "abc"' + + def test_list_field_becomes_in(self, store: MilvusVectorStore) -> None: + expr = store._build_filter_expr({"file_id": ["a", "b"]}) + assert expr == 'file_id in ["a", "b"]' + + def test_empty_list_field_matches_nothing(self, store: MilvusVectorStore) -> None: + # An empty IN list cannot be expressed in Milvus, so short-circuit to + # the explicit no-match literal — callers get an empty result set + # instead of a syntax error. + assert store._build_filter_expr({"file_id": []}) == "false" + + def test_raw_expr_appended(self, store: MilvusVectorStore) -> None: + expr = store._build_filter_expr({"expr": "created_at > ISO '2025-01-01'"}) + assert expr == "created_at > ISO '2025-01-01'" + + def test_raw_expr_combined_with_partition(self, store: MilvusVectorStore) -> None: + expr = store._build_filter_expr({"partition": "p1", "expr": "page > 5"}) + assert expr == 'partition == "p1" and page > 5' + + def test_partition_and_field_joined_with_and(self, store: MilvusVectorStore) -> None: + expr = store._build_filter_expr({"partition": "p1", "file_id": "f1"}) + assert expr == 'partition == "p1" and file_id == "f1"' + + def test_int_value_passes_through(self, store: MilvusVectorStore) -> None: + assert store._build_filter_expr({"page": 7}) == "page == 7" + + +# --------------------------------------------------------------------------- +# _resolve_collection +# --------------------------------------------------------------------------- + + +class TestResolveCollection: + def test_bound_name_passes(self, store: MilvusVectorStore) -> None: + assert store._resolve_collection("test_collection") == "test_collection" + + def test_default_sentinel_passes(self, store: MilvusVectorStore) -> None: + # The ABC default 'default' resolves to the bound collection — without + # this, every ABC-typed caller that omits the kwarg would crash. + assert store._resolve_collection("default") == "test_collection" + + def test_other_name_raises(self, store: MilvusVectorStore) -> None: + with pytest.raises(ValueError, match=re.escape("test_collection")): + store._resolve_collection("some_other_collection") + + def test_error_mentions_partition_guidance(self, store: MilvusVectorStore) -> None: + # The error must tell callers where partitions actually go, otherwise + # the failure looks like a generic bad-arg and people retry with the + # partition name as the collection. + with pytest.raises(ValueError, match="partitions go in filters"): + store._resolve_collection("bad-name") + + +# --------------------------------------------------------------------------- +# ID round-trip +# --------------------------------------------------------------------------- + + +class TestIdRoundTrip: + def test_numeric_string_coerces(self) -> None: + assert MilvusVectorStore._str_id_to_milvus("12345") == 12345 + + def test_non_numeric_returns_none(self) -> None: + # UUIDs (e.g. freshly-built Chunks before insert) must not crash a batch + # delete — the contract is "silently skip", asserted here. + assert MilvusVectorStore._str_id_to_milvus("not-a-number") is None + + def test_empty_string_returns_none(self) -> None: + assert MilvusVectorStore._str_id_to_milvus("") is None + + def test_int_to_string_roundtrip(self) -> None: + assert MilvusVectorStore._milvus_id_to_str(12345) == "12345" + + def test_full_roundtrip(self) -> None: + original = 9876543210 + as_str = MilvusVectorStore._milvus_id_to_str(original) + back = MilvusVectorStore._str_id_to_milvus(as_str) + assert back == original + + +# --------------------------------------------------------------------------- +# _gen_chunk_order_metadata +# --------------------------------------------------------------------------- + + +class TestChunkOrderMetadata: + def test_zero_chunks(self) -> None: + assert MilvusVectorStore._gen_chunk_order_metadata(0) == [] + + def test_single_chunk_has_no_neighbours(self) -> None: + out = MilvusVectorStore._gen_chunk_order_metadata(1) + assert len(out) == 1 + assert out[0]["prev_section_id"] is None + assert out[0]["next_section_id"] is None + assert isinstance(out[0]["section_id"], int) + + def test_three_chunks_form_linked_list(self) -> None: + out = MilvusVectorStore._gen_chunk_order_metadata(3) + # Mid chunk points to both neighbours. + assert out[1]["prev_section_id"] == out[0]["section_id"] + assert out[1]["next_section_id"] == out[2]["section_id"] + # Edges have one None each. + assert out[0]["prev_section_id"] is None + assert out[2]["next_section_id"] is None + # Section IDs are monotonically increasing within a batch. + assert out[0]["section_id"] < out[1]["section_id"] < out[2]["section_id"] + + +# --------------------------------------------------------------------------- +# _chunk_to_entity +# --------------------------------------------------------------------------- + + +def _make_chunk(**overrides: Any) -> Chunk: + defaults: dict[str, Any] = { + "text": "hello", + "document_id": "doc-1", + "partition": "p1", + "embedding": [0.1, 0.2, 0.3], + "chunk_type": ChunkType.TEXT, + "metadata": {"author": "alice"}, + } + defaults.update(overrides) + return Chunk(**defaults) + + +class TestChunkToEntity: + @staticmethod + def _entity(**overrides: Any) -> dict[str, Any]: + chunk = _make_chunk(**overrides) + order = {"prev_section_id": 1, "section_id": 2, "next_section_id": 3} + return MilvusVectorStore._chunk_to_entity( + chunk, + indexed_at="2026-01-01T00:00:00+00:00", + order=order, + ) + + def test_typed_fields_present(self) -> None: + entity = self._entity() + assert entity["text"] == "hello" + assert entity["partition"] == "p1" + assert entity["file_id"] == "doc-1" + assert entity["vector"] == [0.1, 0.2, 0.3] + assert entity["chunk_type"] == "text" + + def test_indexed_at_stamped(self) -> None: + assert self._entity()["indexed_at"] == "2026-01-01T00:00:00+00:00" + + def test_order_metadata_merged(self) -> None: + entity = self._entity() + assert entity["prev_section_id"] == 1 + assert entity["section_id"] == 2 + assert entity["next_section_id"] == 3 + + def test_metadata_passthrough(self) -> None: + # Arbitrary metadata keys flow into the entity by design (dynamic schema). + assert self._entity()["author"] == "alice" + + def test_typed_fields_win_over_metadata(self) -> None: + # If caller-supplied metadata collides with a typed field, the typed + # value wins — strict domain model > free-form dict. + entity = self._entity(metadata={"partition": "WRONG", "file_id": "WRONG"}) + assert entity["partition"] == "p1" + assert entity["file_id"] == "doc-1" + + def test_none_optional_fields_are_omitted(self) -> None: + # token_count/header/context/content default to None; they must not + # be stamped into the dynamic schema as nulls. + entity = self._entity() + for absent in ("token_count", "header", "context", "content"): + assert absent not in entity, f"{absent} should be omitted when None" + + def test_set_optional_fields_present(self) -> None: + entity = self._entity(token_count=42, header="H1", context="ctx", content="C") + assert entity["token_count"] == 42 + assert entity["header"] == "H1" + assert entity["context"] == "ctx" + assert entity["content"] == "C" + + def test_id_is_not_in_entity(self) -> None: + # Milvus assigns _id via auto_id=True; including it in the payload + # would be rejected on insert. + entity = self._entity() + assert "_id" not in entity + + +# --------------------------------------------------------------------------- +# Surface-level ABC-vs-bound-collection enforcement +# --------------------------------------------------------------------------- + + +class TestCollectionArgDiscipline: + @pytest.mark.asyncio + async def test_upsert_rejects_foreign_collection(self, store: MilvusVectorStore) -> None: + with pytest.raises(ValueError, match="test_collection"): + await store.upsert([_make_chunk()], collection="some-other-name") + + @pytest.mark.asyncio + async def test_search_rejects_foreign_collection(self, store: MilvusVectorStore) -> None: + with pytest.raises(ValueError): + await store.search([0.1, 0.2], collection="some-other-name") + + @pytest.mark.asyncio + async def test_delete_rejects_foreign_collection(self, store: MilvusVectorStore) -> None: + with pytest.raises(ValueError): + await store.delete(["1"], collection="some-other-name") + + @pytest.mark.asyncio + async def test_drop_rejects_foreign_collection(self, store: MilvusVectorStore) -> None: + with pytest.raises(ValueError): + await store.drop_collection("some-other-name") + + @pytest.mark.asyncio + async def test_collection_exists_returns_false_for_foreign(self, store: MilvusVectorStore) -> None: + # Falsifies rather than raising — `collection_exists` is asked + # questions about names it does not own and answers "no, not here". + assert await store.collection_exists("some-other-name") is False + + @pytest.mark.asyncio + async def test_upsert_empty_list_is_noop(self, store: MilvusVectorStore) -> None: + # No client calls should happen, and the return must be 0. + result = await store.upsert([]) + assert result == 0 + store._async_client.insert.assert_not_called() # type: ignore[attr-defined] + + @pytest.mark.asyncio + async def test_delete_empty_list_is_noop(self, store: MilvusVectorStore) -> None: + result = await store.delete([]) + assert result == 0 + store._async_client.delete.assert_not_called() # type: ignore[attr-defined] + + @pytest.mark.asyncio + async def test_delete_by_filter_empty_filter_raises(self, store: MilvusVectorStore) -> None: + # Guards against an accidental wildcard wiping the whole collection. + with pytest.raises(ValueError, match="drop_collection"): + await store.delete_by_filter({}) + + @pytest.mark.asyncio + async def test_delete_by_filter_partition_wildcard_raises(self, store: MilvusVectorStore) -> None: + # 'all' produces an empty expression — same guard must fire. + with pytest.raises(ValueError, match="drop_collection"): + await store.delete_by_filter({"partition": "all"}) + + @pytest.mark.asyncio + @pytest.mark.parametrize("tautology", ["1==1", "1 == 1", "true", "TRUE", " True "]) + async def test_delete_by_filter_tautological_expr_raises(self, store: MilvusVectorStore, tautology: str) -> None: + # Raw `expr` tautologies bypass the dict-form guards but would still + # delete every row — the safety contract must reject them too. + with pytest.raises(ValueError, match="drop_collection"): + await store.delete_by_filter({"expr": tautology}) + + +# --------------------------------------------------------------------------- +# Hybrid dispatch +# --------------------------------------------------------------------------- + + +class TestHybridDispatch: + @pytest.mark.asyncio + async def test_hybrid_disabled_store_routes_to_dense( + self, + vdb_config: VectorDBConfig, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """``search()`` on a store built with ``hybrid_search=False`` must take + the dense path — the collection has no ``sparse`` field, so the BM25 + leg must never be reached. + """ + import openrag.services.storage.milvus_store as _store_mod + + monkeypatch.setattr(_store_mod, "MilvusClient", MagicMock()) + monkeypatch.setattr(_store_mod, "AsyncMilvusClient", MagicMock()) + cfg = vdb_config.model_copy(update={"hybrid_search": False}) + store = MilvusVectorStore(cfg) + store._async_client.search = AsyncMock(return_value=[]) + store._async_client.hybrid_search = AsyncMock(return_value=[]) + + result = await store.search([0.1, 0.2], collection="default") + + assert result == [] + store._async_client.search.assert_awaited_once() + store._async_client.hybrid_search.assert_not_called() + + @pytest.mark.asyncio + async def test_hybrid_store_requires_query_text(self, store: MilvusVectorStore) -> None: + """The ``store`` fixture is hybrid-enabled; its BM25 leg has no input + when ``query_text`` is omitted, so ``search()`` must refuse rather + than silently drop the lexical signal. + """ + with pytest.raises(VDBSearchError, match="query_text"): + await store.search([0.1, 0.2], collection="default") diff --git a/openrag/services/storage/test_vector_store_searcher.py b/openrag/services/storage/test_vector_store_searcher.py new file mode 100644 index 000000000..01d9f92aa --- /dev/null +++ b/openrag/services/storage/test_vector_store_searcher.py @@ -0,0 +1,250 @@ +"""Unit tests for :class:`VectorStoreSearcher`. + +All I/O (VectorStore, Embedder, DocumentRepository) is mocked so tests run +without a live Milvus, vLLM, or database. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from openrag.services.storage.vector_store_searcher import VectorStoreSearcher, _dict_to_chunk + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_EMBED_VEC = [0.1] * 8 + + +def _make_row(id_: str, text: str = "hello", partition: str = "p1", **extra) -> dict: + return {"id": id_, "text": text, "partition": partition, "file_id": "f1", **extra} + + +def _make_searcher( + search_results=None, + filter_results=None, + file_ids_by_rel=None, + ancestor_ids=None, +) -> tuple[VectorStoreSearcher, MagicMock, MagicMock, MagicMock]: + store = MagicMock() + store.search = AsyncMock(return_value=search_results or []) + store.query_chunks_by_filter = AsyncMock(return_value=filter_results or []) + + embedder = MagicMock() + embedder.embed = AsyncMock(return_value=[_EMBED_VEC]) + + doc_repo = MagicMock() + doc_repo.get_file_ids_by_relationship = AsyncMock(return_value=file_ids_by_rel or []) + doc_repo.get_ancestor_file_ids = AsyncMock(return_value=ancestor_ids or []) + + searcher = VectorStoreSearcher( + vector_store=store, + embedder=embedder, + document_repo=doc_repo, + collection="test_col", + ) + return searcher, store, embedder, doc_repo + + +# --------------------------------------------------------------------------- +# _dict_to_chunk +# --------------------------------------------------------------------------- + + +def test_dict_to_chunk_uses_id_field(): + row = {"id": "abc", "text": "t", "partition": "p", "file_id": "f"} + c = _dict_to_chunk(row) + assert c.id == "abc" + + +def test_dict_to_chunk_uses_underscore_id_as_fallback(): + row = {"_id": 42, "text": "t", "partition": "p", "file_id": "f"} + c = _dict_to_chunk(row) + assert c.id == "42" + + +def test_dict_to_chunk_metadata_excludes_reserved_keys(): + row = {"id": "x", "text": "t", "partition": "p", "file_id": "f", "score": 0.9, "extra_key": "val"} + c = _dict_to_chunk(row) + assert "score" not in c.metadata + assert "extra_key" in c.metadata + + +# --------------------------------------------------------------------------- +# search() +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_search_embeds_query_and_calls_store(): + searcher, store, embedder, _ = _make_searcher( + search_results=[_make_row("1")], + ) + await searcher.search(query="hello", partition=["p1"], top_k=5, with_surrounding_chunks=False) + + embedder.embed.assert_awaited_once_with(["hello"]) + store.search.assert_awaited_once() + call_kwargs = store.search.call_args.kwargs + assert call_kwargs["embedding"] == _EMBED_VEC + assert call_kwargs["query_text"] == "hello" + assert call_kwargs["top_k"] == 5 + assert call_kwargs["filters"]["partition"] == ["p1"] + + +@pytest.mark.asyncio +async def test_search_passes_filter_expr(): + searcher, store, *_ = _make_searcher(search_results=[]) + await searcher.search( + query="q", + partition=["p1"], + top_k=3, + filter="file_id == 'x'", + with_surrounding_chunks=False, + ) + call_kwargs = store.search.call_args.kwargs + assert call_kwargs["filters"]["expr"] == "file_id == 'x'" + + +@pytest.mark.asyncio +async def test_search_with_surrounding_chunks_deduplicates(): + main_row = _make_row("1", prev_section_id="s0", next_section_id="s2") + surrounding_rows = [_make_row("0"), _make_row("1")] # "1" is a duplicate + searcher, store, _, _ = _make_searcher( + search_results=[main_row], + filter_results=surrounding_rows, + ) + chunks = await searcher.search(query="q", partition=["p1"], top_k=5) + + ids = [c.id for c in chunks] + assert ids.count("1") == 1 + assert "0" in ids + + +@pytest.mark.asyncio +async def test_search_skips_surrounding_when_disabled(): + searcher, store, _, _ = _make_searcher(search_results=[_make_row("1")]) + await searcher.search(query="q", partition=["p1"], top_k=5, with_surrounding_chunks=False) + store.query_chunks_by_filter.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_search_skips_surrounding_when_no_results(): + searcher, store, _, _ = _make_searcher(search_results=[]) + await searcher.search(query="q", partition=["p1"], top_k=5, with_surrounding_chunks=True) + store.query_chunks_by_filter.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# multi_query_search() +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_multi_query_search_embeds_all_queries(): + embedder = MagicMock() + embedder.embed = AsyncMock(return_value=[_EMBED_VEC, _EMBED_VEC]) + + store = MagicMock() + store.search = AsyncMock(return_value=[]) + store.query_chunks_by_filter = AsyncMock(return_value=[]) + + searcher = VectorStoreSearcher( + vector_store=store, + embedder=embedder, + document_repo=MagicMock(), + collection="col", + ) + await searcher.multi_query_search( + queries=["q1", "q2"], partition=["p1"], top_k_per_query=3, with_surrounding_chunks=False + ) + embedder.embed.assert_awaited_once_with(["q1", "q2"]) + + +@pytest.mark.asyncio +async def test_multi_query_search_deduplicates_across_queries(): + embedder = MagicMock() + embedder.embed = AsyncMock(return_value=[_EMBED_VEC, _EMBED_VEC]) + + # Both queries return the same chunk "1" plus a unique one each + store = MagicMock() + store.search = AsyncMock( + side_effect=[ + [_make_row("1"), _make_row("2")], + [_make_row("1"), _make_row("3")], + ] + ) + store.query_chunks_by_filter = AsyncMock(return_value=[]) + + searcher = VectorStoreSearcher( + vector_store=store, + embedder=embedder, + document_repo=MagicMock(), + collection="col", + ) + chunks = await searcher.multi_query_search( + queries=["q1", "q2"], partition=["p1"], top_k_per_query=5, with_surrounding_chunks=False + ) + ids = [c.id for c in chunks] + assert ids.count("1") == 1 + assert sorted(ids) == ["1", "2", "3"] + + +# --------------------------------------------------------------------------- +# get_related_chunks() +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_related_chunks_returns_empty_when_no_file_ids(): + searcher, store, _, doc_repo = _make_searcher(file_ids_by_rel=[]) + result = await searcher.get_related_chunks(partition="p1", relationship_id="r1", limit=10) + assert result == [] + store.query_chunks_by_filter.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_related_chunks_queries_store_with_file_ids(): + rows = [_make_row("1"), _make_row("2"), _make_row("3")] + searcher, store, _, doc_repo = _make_searcher( + file_ids_by_rel=["f1", "f2"], + filter_results=rows, + ) + chunks = await searcher.get_related_chunks(partition="p1", relationship_id="r1", limit=2) + assert len(chunks) == 2 + doc_repo.get_file_ids_by_relationship.assert_awaited_once_with(partition="p1", relationship_id="r1") + call_args = store.query_chunks_by_filter.call_args + assert call_args.args[1]["file_id"] == ["f1", "f2"] + + +# --------------------------------------------------------------------------- +# get_ancestor_chunks() +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_ancestor_chunks_returns_empty_when_no_ancestors(): + searcher, store, _, _ = _make_searcher(ancestor_ids=[]) + result = await searcher.get_ancestor_chunks(partition="p1", file_id="f1", limit=5) + assert result == [] + store.query_chunks_by_filter.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_ancestor_chunks_applies_limit(): + rows = [_make_row(str(i)) for i in range(10)] + searcher, store, _, doc_repo = _make_searcher( + ancestor_ids=["a1", "a2"], + filter_results=rows, + ) + chunks = await searcher.get_ancestor_chunks(partition="p1", file_id="f1", limit=4) + assert len(chunks) == 4 + + +@pytest.mark.asyncio +async def test_get_ancestor_chunks_passes_max_depth(): + searcher, _, _, doc_repo = _make_searcher(ancestor_ids=[]) + await searcher.get_ancestor_chunks(partition="p1", file_id="f1", limit=5, max_ancestor_depth=2) + doc_repo.get_ancestor_file_ids.assert_awaited_once_with(partition="p1", file_id="f1", max_ancestor_depth=2) diff --git a/openrag/services/storage/vector_store_searcher.py b/openrag/services/storage/vector_store_searcher.py new file mode 100644 index 000000000..5b7d3ca70 --- /dev/null +++ b/openrag/services/storage/vector_store_searcher.py @@ -0,0 +1,181 @@ +"""``RetrievalSearcher`` backed directly by ``VectorStore`` + ``Embedder``. + +Replaces ``MilvusRayShim`` — embeds queries in-process and calls the +``VectorStore`` without routing through a Ray actor. +""" + +from __future__ import annotations + +import asyncio +import uuid +from typing import Any + +from core.embeddings import Embedder +from core.models.chunk import Chunk, _coerce_chunk_type +from core.ports.document_repo import DocumentRepository +from core.retrieval.searcher import RetrievalSearcher +from core.vector_stores import VectorStore + + +def _dict_to_chunk(row: dict[str, Any]) -> Chunk: + """Convert a VectorStore result dict to a domain Chunk. + + ``search()`` returns ``"id"`` (string already stringified by the store); + ``query_chunks_by_filter()`` returns ``"_id"`` (raw Milvus INT64). + """ + raw_id = row.get("id") or row.get("_id") + chunk_id = str(raw_id) if raw_id is not None else str(uuid.uuid4()) + skip = {"text", "vector", "_id", "id", "score", "file_id", "partition", "page", "chunk_type"} + metadata = {k: v for k, v in row.items() if k not in skip} + return Chunk( + id=chunk_id, + document_id=row.get("file_id", ""), + text=row.get("text", ""), + partition=row.get("partition", "default"), + page_number=row.get("page"), + chunk_type=_coerce_chunk_type(row.get("chunk_type", "text")), + metadata=metadata, + ) + + +class VectorStoreSearcher(RetrievalSearcher): + """``RetrievalSearcher`` that uses ``VectorStore`` + ``Embedder`` directly. + + This replaces the transitional ``MilvusRayShim`` used during Phase 8. + Queries are embedded in-process; surrounding / related / ancestor chunk + lookups go through ``VectorStore.query_chunks_by_filter``. + """ + + def __init__( + self, + vector_store: VectorStore, + embedder: Embedder, + document_repo: DocumentRepository, + collection: str, + ) -> None: + self._store = vector_store + self._embedder = embedder + self._document_repo = document_repo + self._collection = collection + + async def search( + self, + query: str, + partition: list[str], + top_k: int, + filter: str | None = None, + filter_params: dict | None = None, + similarity_threshold: float = 0.0, + with_surrounding_chunks: bool = True, + ) -> list[Chunk]: + (embedding,) = await self._embedder.embed([query]) + filters: dict[str, Any] = {"partition": partition} + if filter: + filters["expr"] = filter + results = await self._store.search( + embedding=embedding, + query_text=query, + collection=self._collection, + filters=filters, + top_k=top_k, + similarity_threshold=similarity_threshold or None, + ) + chunks = [_dict_to_chunk(r) for r in results] + if with_surrounding_chunks and chunks: + surrounding = await self._fetch_surrounding(chunks) + seen = {c.id for c in chunks} + chunks.extend(c for c in surrounding if c.id not in seen) + return chunks + + async def multi_query_search( + self, + queries: list[str], + partition: list[str], + top_k_per_query: int, + filter: str | None = None, + filter_params: dict | None = None, + similarity_threshold: float = 0.0, + with_surrounding_chunks: bool = True, + ) -> list[Chunk]: + embeddings = await self._embedder.embed(queries) + filters: dict[str, Any] = {"partition": partition} + if filter: + filters["expr"] = filter + per_query = await asyncio.gather( + *[ + self._store.search( + embedding=emb, + query_text=q, + collection=self._collection, + filters=filters, + top_k=top_k_per_query, + similarity_threshold=similarity_threshold or None, + ) + for emb, q in zip(embeddings, queries) + ] + ) + seen_ids: set[str] = set() + chunks: list[Chunk] = [] + for results in per_query: + for r in results: + c = _dict_to_chunk(r) + if c.id not in seen_ids: + seen_ids.add(c.id) + chunks.append(c) + if with_surrounding_chunks and chunks: + surrounding = await self._fetch_surrounding(chunks) + chunks.extend(c for c in surrounding if c.id not in seen_ids) + return chunks + + async def get_related_chunks( + self, + partition: str, + relationship_id: str, + limit: int, + ) -> list[Chunk]: + file_ids = await self._document_repo.get_file_ids_by_relationship( + partition=partition, relationship_id=relationship_id + ) + if not file_ids: + return [] + rows = await self._store.query_chunks_by_filter( + self._collection, + {"partition": partition, "file_id": file_ids}, + ) + return [_dict_to_chunk(r) for r in rows[:limit]] + + async def get_ancestor_chunks( + self, + partition: str, + file_id: str, + limit: int, + max_ancestor_depth: int | None = None, + ) -> list[Chunk]: + ancestor_ids = await self._document_repo.get_ancestor_file_ids( + partition=partition, file_id=file_id, max_ancestor_depth=max_ancestor_depth + ) + if not ancestor_ids: + return [] + rows = await self._store.query_chunks_by_filter( + self._collection, + {"partition": partition, "file_id": ancestor_ids}, + ) + return [_dict_to_chunk(r) for r in rows[:limit]] + + async def _fetch_surrounding(self, chunks: list[Chunk]) -> list[Chunk]: + section_ids = [ + sid + for c in chunks + for sid in (c.metadata.get("prev_section_id"), c.metadata.get("next_section_id")) + if sid is not None + ] + if not section_ids: + return [] + rows = await self._store.query_chunks_by_filter( + self._collection, + {"section_id": section_ids}, + ) + return [_dict_to_chunk(r) for r in rows] + + +__all__ = ["VectorStoreSearcher"] diff --git a/openrag/services/workers/__init__.py b/openrag/services/workers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/services/workers/batch_ingest.py b/openrag/services/workers/batch_ingest.py new file mode 100644 index 000000000..7a18d0c16 --- /dev/null +++ b/openrag/services/workers/batch_ingest.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import asyncio +from collections.abc import MutableMapping, Sequence +from typing import Any + +from services.workers.pipeline_builder import IndexingPipeline +from utils.logger import get_logger + +_logger = get_logger().bind(component="batch_ingest") + + +async def ingest_batch( + pipeline: IndexingPipeline, + rows: Sequence[MutableMapping[str, Any]], + *, + concurrency: int | None = None, +) -> list[MutableMapping[str, Any]]: + """Run each row through *pipeline*, capturing per-row exceptions. + + A failure on one row does not abort the others. Each failed row will have + ``row["stage"]`` set to the failing stage name and ``row["error"]`` set to + the exception message — exactly as the individual stages do. + + Args: + pipeline: The assembled indexing pipeline to run each row through. + rows: Rows to process. Each row is a mutable mapping passed directly to + ``pipeline.run()``. + concurrency: Maximum number of rows processed concurrently. ``None`` + means all rows are dispatched at once. Must be >= 1 when provided. + """ + if concurrency is not None and concurrency < 1: + raise ValueError(f"concurrency must be >= 1, got {concurrency}") + + sem = asyncio.Semaphore(concurrency) if concurrency is not None else None + + async def _run_one(row: MutableMapping[str, Any]) -> MutableMapping[str, Any]: + if sem is not None: + async with sem: + await _pipeline_run_catching(pipeline, row) + else: + await _pipeline_run_catching(pipeline, row) + return row + + return list(await asyncio.gather(*(_run_one(row) for row in rows))) + + +async def _pipeline_run_catching( + pipeline: IndexingPipeline, + row: MutableMapping[str, Any], +) -> None: + try: + await pipeline.run(row) + except Exception as exc: + row.setdefault("stage", "pipeline_failed") + row.setdefault("error", str(exc)) + _logger.bind(stage=row.get("stage"), partition=row.get("partition")).warning( + "Unhandled exception escaping pipeline stage", + error=str(exc), + ) + + +__all__ = ["ingest_batch"] diff --git a/openrag/utils/dependencies.py b/openrag/services/workers/bootstrap.py similarity index 75% rename from openrag/utils/dependencies.py rename to openrag/services/workers/bootstrap.py index ee21cd004..3e689bff3 100644 --- a/openrag/utils/dependencies.py +++ b/openrag/services/workers/bootstrap.py @@ -1,17 +1,35 @@ +"""Ray-actor bootstrap for the worker pool. + +Imported at startup (explicitly by ``main.py``) to create the long-lived +detached actors the request path looks up by name: + +* TaskStateManager — shared task-state actor +* DocSerializer — loader dispatcher +* MarkerPool / DoclingPool / WhisperPool / WhisperActor — GPU parsers +* llmSemaphore / vlmSemaphore / audioSemaphore — cluster-wide rate limiters + +This is the **only** non-startup module outside ``services/workers/`` that +imports Ray — the phase-9 plan permits Ray in ``services/workers/`` plus +the ``main.py`` startup sequence. Before phase 9 the same logic lived in +``utils/dependencies.py``; that location violated the "Ray only in +workers + startup" rule, so the bootstrap moved here. + +``actor_creation_map`` is read by ``routers/actors.py`` to restart a +named actor on-demand from the admin endpoint. +""" + from functools import wraps import ray -from components.indexer.indexer import Indexer, TaskStateManager from components.indexer.loaders.audio import WhisperActor, WhisperPool from components.indexer.loaders.pdf_loaders.docling2 import DoclingPool from components.indexer.loaders.pdf_loaders.marker import MarkerPool from components.indexer.loaders.serializer import DocSerializer -from components.indexer.vectordb.vectordb import ConnectorFactory -from components.utils import DistributedSemaphoreActor from config import load_config +from services.inference.distributed_semaphore import DistributedSemaphoreActor +from services.workers.task_state import TaskStateManager from utils.logger import get_logger -# load config config = load_config() logger = get_logger() @@ -54,15 +72,6 @@ def get_marker_pool(): return get_or_create_actor("MarkerPool", MarkerPool, lifetime="detached") -def get_indexer(): - return get_or_create_actor("Indexer", Indexer, lifetime="detached") - - -def get_vectordb(): - vectordb_cls = ConnectorFactory().get_vectordb_cls() - return get_or_create_actor("Vectordb", vectordb_cls, lifetime="detached") - - def init_audio_actor(): use_whisper_lang_detector = config.loader.transcriber.use_whisper_lang_detector file_loaders = config.loader.file_loaders @@ -110,5 +119,3 @@ def init_audio_semaphore(): task_state_manager = get_task_state_manager() serializer = get_serializer() -vectordb = get_vectordb() -indexer = get_indexer() diff --git a/openrag/services/workers/dispatcher.py b/openrag/services/workers/dispatcher.py new file mode 100644 index 000000000..69f1224de --- /dev/null +++ b/openrag/services/workers/dispatcher.py @@ -0,0 +1,246 @@ +from __future__ import annotations + +import uuid +from typing import Any + +from core.indexing.dispatcher import IndexingDispatcher + +DEFAULT_TIMEOUT = 60.0 + + +class WorkerDispatcher(IndexingDispatcher): + """Dispatcher that routes new indexing jobs through ``IndexerPool``. + + File mutation paths use the storage ports directly so the API no longer + depends on the legacy ``Indexer`` actor being present. + """ + + _FILE_METADATA_EXCLUDED_KEYS = frozenset( + { + "_id", + "id", + "text", + "vector", + "page", + "section_id", + "prev_section_id", + "next_section_id", + } + ) + + def __init__( + self, + *, + pool: Any, + task_state_manager: Any, + vector_store: Any, + document_repo: Any, + workspace_repo: Any, + collection: str, + timeout: float = DEFAULT_TIMEOUT, + ) -> None: + self._pool = pool + self._tsm = task_state_manager + self._vector_store = vector_store + self._document_repo = document_repo + self._workspace_repo = workspace_repo + self._collection = collection + self._timeout = timeout + + async def _call(self, future: Any, task_description: str) -> Any: + from services.workers.ray_utils import call_ray_actor_with_timeout + + return await call_ray_actor_with_timeout( + future=future, + timeout=self._timeout, + task_description=task_description, + ) + + async def dispatch_indexing( + self, + *, + path: str, + metadata: dict, + partition: str, + user: dict | None, + workspace_ids: list[str] | None, + replace: bool, + ) -> str: + task_id = uuid.uuid4().hex + + await self._call( + self._tsm.set_state.remote(task_id, "QUEUED"), + task_description=f"set_state({task_id})", + ) + + user_metadata = {key: value for key, value in metadata.items() if key not in {"file_id", "source"}} + await self._call( + self._tsm.set_details.remote( + task_id, + file_id=metadata.get("file_id"), + partition=partition, + metadata=user_metadata, + user_id=user.get("id") if user else None, + ), + task_description=f"set_details({task_id})", + ) + + task = self._pool.process_file.remote( + task_id=task_id, + path=path, + metadata=metadata, + partition=partition, + user=user, + workspace_ids=workspace_ids, + replace=replace, + ) + + await self._call( + self._tsm.set_object_ref.remote(task_id, {"ref": task}), + task_description=f"set_object_ref({task_id})", + ) + return task_id + + async def delete_file(self, file_id: str, partition: str) -> None: + ids = await self._vector_store.query_ids_by_filter( + self._collection, + {"partition": partition, "file_id": file_id}, + ) + if ids: + await self._vector_store.delete(ids, self._collection) + await self._workspace_repo.remove_file_from_all_workspaces(file_id, partition) + await self._document_repo.remove_file_from_partition(file_id=file_id, partition=partition) + + async def update_file_metadata( + self, + file_id: str, + metadata: dict, + partition: str, + user: dict | None, + ) -> None: + rows = await self._vector_store.query_chunks_by_filter( + self._collection, + {"partition": partition, "file_id": file_id}, + output_fields=["*", "vector"], + ) + if not rows: + return + + entities = [] + for row in rows: + entity = dict(row) + entity.update(metadata) + entities.append(entity) + + await self._upsert_entities(entities) + + file_metadata = self._file_metadata_from_chunk(rows[0]) + file_metadata.update(metadata) + await self._document_repo.update_file_metadata_in_db(file_id, partition, file_metadata) + + async def copy_file( + self, + file_id: str, + metadata: dict, + partition: str, + user: dict | None, + ) -> None: + rows = await self._vector_store.query_chunks_by_filter( + self._collection, + {"partition": partition, "file_id": file_id}, + output_fields=["*", "vector"], + ) + if not rows: + return + + entities = [] + for row in rows: + entity = dict(row) + entity.pop("_id", None) + entity.update(metadata) + entities.append(entity) + + await self._insert_entities(entities) + + target_file_id = metadata.get("file_id", file_id) + target_partition = metadata.get("partition", partition) + file_metadata = self._file_metadata_from_chunk(rows[0]) + file_metadata.update(metadata) + await self._document_repo.add_file_to_partition( + file_id=target_file_id, + partition=target_partition, + file_metadata=file_metadata, + user_id=user.get("id") if user else None, + relationship_id=file_metadata.get("relationship_id"), + parent_id=file_metadata.get("parent_id"), + ) + + async def _upsert_entities(self, entities: list[dict[str, Any]]) -> None: + upsert_entities = getattr(self._vector_store, "upsert_entities", None) + if upsert_entities is None: + raise TypeError("vector_store must expose upsert_entities for file metadata mutations") + await upsert_entities(entities, self._collection) + + async def _insert_entities(self, entities: list[dict[str, Any]]) -> None: + insert_entities = getattr(self._vector_store, "insert_entities", None) + if insert_entities is None: + raise TypeError("vector_store must expose insert_entities for file copy mutations") + await insert_entities(entities, self._collection) + + def _file_metadata_from_chunk(self, chunk: dict[str, Any]) -> dict[str, Any]: + return {k: v for k, v in chunk.items() if k not in self._FILE_METADATA_EXCLUDED_KEYS} + + async def get_task_state(self, task_id: str) -> str | None: + return await self._call( + self._tsm.get_state.remote(task_id), + task_description=f"get_state({task_id})", + ) + + async def get_task_error(self, task_id: str) -> str | None: + return await self._call( + self._tsm.get_error.remote(task_id), + task_description=f"get_error({task_id})", + ) + + async def cancel_task(self, task_id: str) -> bool: + import ray + + obj_ref = await self._call( + self._tsm.get_object_ref.remote(task_id), + task_description=f"get_object_ref({task_id})", + ) + if obj_ref is None: + return False + + ray.cancel(obj_ref["ref"], recursive=True) + await self._call( + self._tsm.set_state.remote(task_id, "CANCELLED"), + task_description=f"set_state({task_id})", + ) + return True + + +def from_ray_namespace( + namespace: str = "openrag", + timeout: float = DEFAULT_TIMEOUT, + *, + vector_store: Any, + document_repo: Any, + workspace_repo: Any, + collection: str, +) -> WorkerDispatcher: + import ray + from services.workers.indexer_pool import build_indexer_pool + + return WorkerDispatcher( + pool=build_indexer_pool(namespace=namespace), + task_state_manager=ray.get_actor("TaskStateManager", namespace=namespace), + vector_store=vector_store, + document_repo=document_repo, + workspace_repo=workspace_repo, + collection=collection, + timeout=timeout, + ) + + +__all__ = ["WorkerDispatcher", "from_ray_namespace"] diff --git a/openrag/services/workers/indexer_actor.py b/openrag/services/workers/indexer_actor.py new file mode 100644 index 000000000..a619b571f --- /dev/null +++ b/openrag/services/workers/indexer_actor.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import traceback +from pathlib import Path +from typing import Any + +from core.models.document import Document +from services.workers.pipeline_builder import IndexingPipeline + + +class IndexerWorker: + """Pure-Python core of the thin indexer actor. + + ``@ray.remote`` is not applied here so the class is directly + instantiable in tests. The production Ray actor wraps this class + (or applies ``@ray.remote`` at startup). + + State transitions reported to *task_state_manager* are compatible + with the existing queue-monitoring states: + + ``SERIALIZING`` — processing has started (parse + chunk + embed + store) + ``COMPLETED`` — pipeline finished successfully + ``FAILED`` — pipeline raised; set via ``set_failed_if_not_cancelled`` + + Callers are responsible for setting ``QUEUED`` *before* dispatching + the task, and for storing the object ref via ``set_object_ref``. + """ + + def __init__( + self, + pipeline: IndexingPipeline, + task_state_manager: Any, + document_repo: Any = None, + ) -> None: + self._pipeline = pipeline + self._tsm = task_state_manager + self._document_repo = document_repo + + async def process_file( + self, + *, + task_id: str, + path: str, + metadata: dict[str, Any], + partition: str, + user: dict[str, Any] | None = None, + workspace_ids: list[str] | None = None, + replace: bool = False, + ) -> dict[str, Any]: + """Run one file through the indexing pipeline. + + Returns a plain dict ``{"stored_count": int, "stage": "stored"}`` + on success. On failure, state is set to FAILED and the exception + is re-raised so the Ray task is marked as errored. + """ + await self._tsm.set_state.remote(task_id, "SERIALIZING") + try: + document = _load_document(path, metadata, partition) + row: dict[str, Any] = { + "document": document, + "partition": partition, + "filename": Path(path).name, + "language": metadata.get("language", "en"), + "replace": replace, + "user": user, + "workspace_ids": workspace_ids, + } + await self._pipeline.run(row) + if self._document_repo is not None: + await _write_catalog_record( + doc_repo=self._document_repo, + metadata=metadata, + partition=partition, + user=user, + replace=replace, + ) + await self._tsm.set_state.remote(task_id, "COMPLETED") + return {"stored_count": row.get("stored_count", 0), "stage": row.get("stage", "")} + except Exception: + tb = traceback.format_exc() + await self._tsm.set_failed_if_not_cancelled.remote(task_id, tb) + raise + + +async def _write_catalog_record( + *, + doc_repo: Any, + metadata: dict[str, Any], + partition: str, + user: dict[str, Any] | None, + replace: bool, +) -> None: + file_id = metadata.get("file_id", "") + file_metadata = {key: value for key, value in metadata.items() if key != "page"} + if replace: + await doc_repo.update_file_in_partition( + file_id=file_id, + partition=partition, + file_metadata=file_metadata, + relationship_id=metadata.get("relationship_id"), + parent_id=metadata.get("parent_id"), + ) + return + + await doc_repo.add_file_to_partition( + file_id=file_id, + partition=partition, + file_metadata=file_metadata, + user_id=user.get("id") if user else None, + relationship_id=metadata.get("relationship_id"), + parent_id=metadata.get("parent_id"), + ) + + +def _load_document(path: str, metadata: dict[str, Any], partition: str) -> Document: + p = Path(path) + return Document( + filename=metadata.get("file_id") or p.name, + raw_bytes=p.read_bytes(), + content_type=Document.detect_content_type(p.name), + partition=partition, + metadata=metadata, + ) + + +__all__ = ["IndexerWorker"] diff --git a/openrag/services/workers/indexer_pool.py b/openrag/services/workers/indexer_pool.py new file mode 100644 index 000000000..54d74e72b --- /dev/null +++ b/openrag/services/workers/indexer_pool.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import asyncio +from typing import Any + +import ray +from services.workers.indexer_actor import IndexerWorker + + +@ray.remote +class IndexerPool: + """Thin Ray actor wrapper around ``IndexerWorker``.""" + + def __init__(self) -> None: + import services.inference.vllm_client # noqa: F401 + from config import load_config + from core.embeddings import embedder_registry + from services.storage.milvus_store import MilvusVectorStore + from services.storage.postgres_store import PostgresStore + from services.workers.parsers.doc_serializer_bridge import DocSerializerBridgeParser + from services.workers.pipeline_builder import build_indexing_pipeline + + cfg = load_config() + + parser = DocSerializerBridgeParser(config=cfg) + chunker = _build_chunker(cfg) + + embed_cfg = cfg.embedder + embedder = embedder_registry.create( + "vllm", + endpoint=embed_cfg.base_url, + model_name=embed_cfg.model_name, + api_key=embed_cfg.api_key, + max_model_len=embed_cfg.max_model_len, + ) + self._vector_store = MilvusVectorStore(cfg.vectordb) + task_state_manager = ray.get_actor("TaskStateManager", namespace="openrag") + pipeline = build_indexing_pipeline( + parser=parser, + chunker=chunker, + embedder=embedder, + vector_store=self._vector_store, + ) + rdb_cfg = cfg.rdb.model_copy(update={"database": f"partitions_for_collection_{cfg.vectordb.collection_name}"}) + self._catalog_store = PostgresStore(rdb_cfg, run_migrations=False) + self._catalog_initialized = False + self._worker = IndexerWorker( + pipeline=pipeline, + task_state_manager=task_state_manager, + document_repo=self._catalog_store.document_repo, + ) + + async def _ensure_catalog(self) -> None: + if not self._catalog_initialized: + await self._catalog_store.initialize() + self._catalog_initialized = True + + async def process_file( + self, + *, + task_id: str, + path: str, + metadata: dict[str, Any], + partition: str, + user: dict[str, Any] | None = None, + workspace_ids: list[str] | None = None, + replace: bool = False, + ) -> dict[str, Any]: + await self._ensure_catalog() + result = await self._worker.process_file( + task_id=task_id, + path=path, + metadata=metadata, + partition=partition, + user=user, + workspace_ids=workspace_ids, + replace=replace, + ) + file_id = metadata.get("file_id", "") + if workspace_ids and not replace and file_id: + try: + await asyncio.gather( + *( + self._catalog_store.workspace_repo.add_files_to_workspace(workspace_id, [file_id]) + for workspace_id in workspace_ids + ) + ) + except Exception: + pass + return result + + +def build_indexer_pool(namespace: str = "openrag") -> Any: + return IndexerPool.options( # type: ignore[attr-defined] + name="IndexerPool", + namespace=namespace, + get_if_exists=True, + ).remote() + + +def _build_chunker(cfg: Any) -> Any: + from components.indexer.chunker.chunker import ChunkerFactory + + legacy_chunker = ChunkerFactory.create_chunker(cfg) + if hasattr(legacy_chunker, "chunk"): + return legacy_chunker + core_chunker = getattr(legacy_chunker, "_core_splitter", None) + if core_chunker is None or not hasattr(core_chunker, "chunk"): + raise TypeError("Configured chunker does not expose a chunk(document, partition) method") + return core_chunker + + +__all__ = ["IndexerPool", "build_indexer_pool"] diff --git a/openrag/services/workers/parsers/__init__.py b/openrag/services/workers/parsers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/services/workers/parsers/doc_serializer.py b/openrag/services/workers/parsers/doc_serializer.py new file mode 100644 index 000000000..6aca97d18 --- /dev/null +++ b/openrag/services/workers/parsers/doc_serializer.py @@ -0,0 +1,92 @@ +"""DocSerializer Ray actor. + +Moved from ``components/indexer/loaders/serializer.py``; the old module +re-exports this class for backward compatibility. +""" + +from __future__ import annotations + +import gc +from pathlib import Path + +import ray +import torch +from components.indexer.loaders import get_loader_classes +from config import load_config +from langchain_core.documents.base import Document + +config = load_config() + +if torch.cuda.is_available(): + NUM_GPUS = config.ray.num_gpus +else: + NUM_GPUS = 0 + +DICT_MIMETYPES = config.loader.mimetypes.to_dict() + + +@ray.remote(max_restarts=5) +class DocSerializer: + def __init__(self, data_dir=None, **kwargs) -> None: + from config import load_config + from utils.logger import get_logger + + self.logger = get_logger() + self.config = load_config() + self.data_dir = data_dir + self.kwargs = kwargs + self.kwargs["config"] = self.config + self.save_markdown = self.config.loader.save_markdown + + self.loader_classes = get_loader_classes(config=self.config) + self.logger.info("DocSerializer initialized.") + + async def serialize_document( + self, + task_id: str, + path: str | Path, + metadata: dict | None = None, + ) -> Document: + metadata = metadata or {} + log = self.logger.bind( + file_id=metadata.get("file_id"), + partition=metadata.get("partition"), + task_id=task_id, + ) + task_state_manager = ray.get_actor("TaskStateManager", namespace="openrag") + await task_state_manager.set_state.remote(task_id, "SERIALIZING") + + log.info("Starting document serialization") + + p = Path(path) + file_ext = p.suffix.lower() + mimetype = metadata.get("mimetype", None) + if mimetype is None: + loader_cls = self.loader_classes.get(file_ext) + else: + loader_cls = self.loader_classes.get(DICT_MIMETYPES.get(mimetype)) + + if loader_cls is None: + log.warning(f"No loader available for {p.name}") + raise ValueError(f"No loader available for file type {file_ext}.") + + log.debug(f"Loading document: {p.name} with loader {loader_cls.__name__}") + loader = loader_cls(**self.kwargs) + + try: + doc: Document = await loader.aload_document( + file_path=path, metadata=metadata, save_markdown=self.save_markdown + ) + del loader + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + log.info("Document serialized successfully") + return doc + except Exception as e: + log.exception("Failed to serialize document", error=str(e)) + raise + + +__all__ = ["DocSerializer"] diff --git a/openrag/services/workers/parsers/doc_serializer_adapter.py b/openrag/services/workers/parsers/doc_serializer_adapter.py new file mode 100644 index 000000000..7842ff38d --- /dev/null +++ b/openrag/services/workers/parsers/doc_serializer_adapter.py @@ -0,0 +1,37 @@ +"""FileSerializer adapter over the DocSerializer Ray actor. + +Replaces ``services/storage/serializer_ray_shim.py`` (Phase 9E). The adapter +lives in the workers layer because it wraps a worker Ray actor; the storage +layer no longer references Ray directly. +""" + +from __future__ import annotations + +from core.indexing.serializer import FileSerializer + +_FALLBACK_TASK_ID = "tools-extract" + + +class DocSerializerAdapter(FileSerializer): + """Implements FileSerializer by delegating to the DocSerializer Ray actor.""" + + async def serialize(self, path: str, metadata: dict) -> str: + import ray + from config import load_config + from services.workers.ray_utils import call_ray_actor_with_timeout + + cfg = load_config() + timeout = cfg.ray.indexer.serialize_timeout + task_id = ray.get_runtime_context().get_task_id() or _FALLBACK_TASK_ID + serializer = ray.get_actor("DocSerializer", namespace="openrag") + doc = await call_ray_actor_with_timeout( + future=serializer.serialize_document.remote(task_id, path, metadata=metadata or {}), + timeout=timeout, + task_description=f"Serialization task {task_id}", + ) + return doc.page_content + + +def from_ray_namespace() -> DocSerializerAdapter: + """Build the adapter. Convenience for the composition root.""" + return DocSerializerAdapter() diff --git a/openrag/services/workers/parsers/doc_serializer_bridge.py b/openrag/services/workers/parsers/doc_serializer_bridge.py new file mode 100644 index 000000000..2c2f556b7 --- /dev/null +++ b/openrag/services/workers/parsers/doc_serializer_bridge.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import tempfile +from pathlib import Path +from typing import Any + +from core.indexing.parsers.document_parser import DocumentParser +from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock + + +class DocSerializerBridgeParser(DocumentParser): + """Transitional parser backed by the legacy loader registry.""" + + def __init__(self, config: Any) -> None: + from components.indexer.loaders import get_loader_classes + + self._config = config + self._loader_classes = get_loader_classes(config=config) + self._save_markdown = getattr(config.loader, "save_markdown", False) + + def supported_types(self) -> list[str]: + return [doc_type.value for doc_type in DocumentType] + + async def parse(self, document: Document) -> ProcessedDocument: + suffix = _suffix_from_document(document) + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as handle: + handle.write(document.raw_bytes or b"") + temp_path = handle.name + try: + return await self._load_via_legacy(temp_path, document) + finally: + Path(temp_path).unlink(missing_ok=True) + + async def _load_via_legacy(self, path: str, document: Document) -> ProcessedDocument: + metadata = dict(document.metadata or {}) + loader_cls = self._loader_for(path, metadata) + if loader_cls is None: + raise ValueError(f"No loader registered for file extension {Path(path).suffix.lower()!r}") + + loader = loader_cls(config=self._config) + lang_doc = await loader.aload_document( + file_path=path, + metadata=metadata, + save_markdown=self._save_markdown, + ) + + return ProcessedDocument( + document_id=document.filename or "unknown", + text_blocks=[TextBlock(text=lang_doc.page_content or "")], + metadata=lang_doc.metadata or {}, + ) + + def _loader_for(self, path: str, metadata: dict[str, Any]) -> Any | None: + mimetype = metadata.get("mimetype") + if mimetype: + try: + from services.workers.parsers.doc_serializer import DICT_MIMETYPES + + loader_cls = self._loader_classes.get(DICT_MIMETYPES.get(mimetype)) + except Exception: + loader_cls = None + if loader_cls is not None: + return loader_cls + + return self._loader_classes.get(Path(path).suffix.lower()) + + +def _suffix_from_document(document: Document) -> str: + source = (document.metadata or {}).get("source") + if source: + suffix = Path(str(source)).suffix + if suffix: + return suffix + if document.filename: + suffix = Path(document.filename).suffix + if suffix: + return suffix + return f".{document.content_type.value}" if document.content_type else "" + + +__all__ = ["DocSerializerBridgeParser"] diff --git a/openrag/services/workers/parsers/docling_workers.py b/openrag/services/workers/parsers/docling_workers.py new file mode 100644 index 000000000..7c0a9ef99 --- /dev/null +++ b/openrag/services/workers/parsers/docling_workers.py @@ -0,0 +1,189 @@ +"""Docling-backed PDF Ray actors and BasePooledParser facade. + +``DoclingWorker`` and ``DoclingPool`` are Ray actors. +``DoclingLoader`` is a :class:`~core.indexing.parsers.document_parser.BasePooledParser` +that wraps the pool so the core pipeline can call ``parse()`` uniformly. + +The old module ``components/indexer/loaders/pdf_loaders/docling2.py`` re-exports +``DoclingWorker`` and ``DoclingPool`` for legacy import paths. +""" + +from __future__ import annotations + +import asyncio + +import ray +import torch +from config import load_config +from core.indexing.image_preprocessor import pil_to_png_bytes +from core.indexing.parsers.document_parser import BasePooledParser +from core.models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend +from docling.datamodel.base_models import InputFormat +from docling.datamodel.document import ConversionResult +from docling.datamodel.pipeline_options import ( + AcceleratorDevice, + AcceleratorOptions, + PdfPipelineOptions, + TableFormerMode, + TableStructureOptions, +) +from docling.document_converter import DocumentConverter, PdfFormatOption +from utils.logger import get_logger + +from ..ray_utils import call_ray_actor_with_timeout, retry_with_backoff, with_timeout + +logger = get_logger() +config = load_config() + +if torch.cuda.is_available(): + DOCLING_NUM_GPUS = config.loader.docling_num_gpus +else: + DOCLING_NUM_GPUS = 0 + +DOCLING_MAX_TASKS_PER_WORKER = config.loader.docling_max_tasks_per_worker + + +@ray.remote(num_gpus=DOCLING_NUM_GPUS) +class DoclingWorker: + def __init__(self): + img_scale = 2 + pipeline_options = PdfPipelineOptions( + do_ocr=True, + do_table_structure=True, + generate_picture_images=True, + images_scale=img_scale, + ) + pipeline_options.table_structure_options = TableStructureOptions( + do_cell_matching=True, mode=TableFormerMode.ACCURATE + ) + pipeline_options.accelerator_options = AcceleratorOptions(device=AcceleratorDevice.AUTO) + self.converter = DocumentConverter( + format_options={ + InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options, backend=PyPdfiumDocumentBackend) + } + ) + + async def convert(self, file_path: str) -> ConversionResult: + with torch.no_grad(): + return await asyncio.to_thread(self.converter.convert, str(file_path)) + + +@ray.remote +class DoclingPool: + def __init__(self): + self.logger = get_logger() + self.config = load_config() + self.pool_size = self.config.loader.docling_pool_size + + self.actors = [DoclingWorker.remote() for _ in range(self.pool_size)] + self._queue: asyncio.Queue[ray.actor.ActorHandle] = asyncio.Queue() + + for _ in range(DOCLING_MAX_TASKS_PER_WORKER): + for actor in self.actors: + self._queue.put_nowait(actor) + + total_slots = self.pool_size * DOCLING_MAX_TASKS_PER_WORKER + self.logger.info( + f"Docling pool: {self.pool_size} actors × {DOCLING_MAX_TASKS_PER_WORKER} slots = " + f"{total_slots} PDF concurrency" + ) + + async def process_pdf(self, file_path: str) -> ConversionResult: + timeout = self.config.loader.docling_timeout + + async def attempt(i: int) -> ConversionResult: + actor: DoclingWorker = await self._queue.get() + try: + return await call_ray_actor_with_timeout( + actor.convert.remote(file_path), + timeout=timeout, + task_description=f"DoclingPool PDF ({file_path})", + ) + finally: + await self._queue.put(actor) + + return await retry_with_backoff( + attempt, + max_retries=self.config.loader.docling_max_task_retry, + base_delay=self.config.loader.docling_retry_base_delay, + task_description=f"DoclingPool PDF ({file_path})", + ) + + +class DoclingLoader(BasePooledParser): + """``BasePooledParser`` facade over the Docling Ray pool. + + Materialises ``Document.raw_bytes`` to a temporary file, dispatches to the + named ``DoclingPool`` actor, and converts the ``ConversionResult`` into a + ``ProcessedDocument`` with one ``TextBlock`` per page and one ``ImageBlock`` + per picture. Image captioning is left to the downstream caption stage. + """ + + def __init__(self) -> None: + self.worker: DoclingPool = ray.get_actor("DoclingPool", namespace="openrag") + + def supported_types(self) -> list[str]: + return [DocumentType.PDF.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + async with document.as_temporary_file() as path: + result: ConversionResult = await self._dispatch(str(path)) + + text_blocks = self._build_text_blocks(result) + image_blocks = self._build_image_blocks(result) + + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=image_blocks, + metadata=dict(document.metadata), + page_count=len(result.pages), + ) + + @with_timeout( + seconds=config.loader.docling_timeout, + description="DoclingLoader PDF loading ({file_path})", + ) + async def _dispatch(self, file_path: str) -> ConversionResult: + return self.worker.process_pdf.remote(file_path) + + @staticmethod + def _build_text_blocks(result: ConversionResult) -> list[TextBlock]: + blocks: list[TextBlock] = [] + n_pages = len(result.pages) + for i in range(1, n_pages + 1): + text = result.document.export_to_markdown(page_no=i).strip() + blocks.append(TextBlock(text=text, page_number=i)) + return blocks + + @staticmethod + def _build_image_blocks(result: ConversionResult) -> list[ImageBlock]: + blocks: list[ImageBlock] = [] + for idx, picture in enumerate(result.document.pictures): + try: + pil_image = picture.image.pil_image + if pil_image is None: + continue + png_bytes = pil_to_png_bytes(pil_image) + except Exception as exc: + logger.warning(f"Failed to encode Docling picture {idx}: {exc}") + continue + ref = f"![](docling_img_{idx})" + blocks.append( + ImageBlock( + image_bytes=png_bytes, + mime_type="image/png", + metadata={"markdown_ref": ref}, + ) + ) + return blocks + + +__all__ = ["DoclingLoader", "DoclingPool", "DoclingWorker"] diff --git a/openrag/services/workers/parsers/marker_workers.py b/openrag/services/workers/parsers/marker_workers.py new file mode 100644 index 000000000..3a495ca64 --- /dev/null +++ b/openrag/services/workers/parsers/marker_workers.py @@ -0,0 +1,455 @@ +import asyncio +import gc +import re +import time + +import pypdfium2 +import ray +import torch +from config import load_config +from core.indexing.image_preprocessor import pil_to_png_bytes +from core.indexing.parsers.document_parser import BasePooledParser +from core.models.document import ( + Document, + DocumentType, + ImageBlock, + ProcessedDocument, + TextBlock, +) +from marker.converters.pdf import PdfConverter +from utils.logger import get_logger + +from ..ray_utils import with_retry, with_timeout + +logger = get_logger() +config = load_config() + +if torch.cuda.is_available(): + MARKER_NUM_GPUS = config.loader.marker_num_gpus +else: # On CPU + MARKER_NUM_GPUS = 0 + + +@ray.remote(num_gpus=MARKER_NUM_GPUS, max_restarts=5) +class MarkerWorker: + def __init__(self): + import os + + from config import load_config + from utils.logger import get_logger + + self.logger = get_logger() + self.config = load_config() + self.page_sep = "[PAGE_SEP]" + + self._workers = self.config.loader.marker_max_processes + + self.converter_config = { + "output_format": "markdown", + "paginate_output": True, + "page_separator": self.page_sep, + "pdftext_workers": self.config.loader.marker_pdftext_workers, + "disable_multiprocessing": False, + } + os.environ["RAY_ADDRESS"] = "auto" + + self.executor = None + self.init_resources() + + def init_resources(self): + from marker.models import create_model_dict + + self.model_dict = create_model_dict() + for v in self.model_dict.values(): + if hasattr(v.model, "share_memory"): + v.model.share_memory() + + self.setup_mp() + + def setup_mp(self): + """Initialize ProcessPoolExecutor for PDF processing. + + We use ProcessPoolExecutor instead of multiprocessing.Pool because: + - Ray actors run as daemon processes + - Pool workers are daemonic by default and cannot spawn children + - The pdftext library (used by Marker) internally spawns processes + - ProcessPoolExecutor workers are non-daemon, allowing nested process creation + """ + from concurrent.futures import ProcessPoolExecutor + + import torch.multiprocessing as mp + + if self.executor: + self.logger.warning("Resetting ProcessPoolExecutor") + self.executor.shutdown(wait=False, cancel_futures=True) + self.executor = None + + # Ensure spawn method for CUDA compatibility + try: + if mp.get_start_method(allow_none=True) != "spawn": + mp.set_start_method("spawn", force=True) + except RuntimeError: + self.logger.warning("Process start method already set, using existing method") + + self.logger.info(f"Initializing MarkerWorker with {self._workers} workers") + self.executor = ProcessPoolExecutor( + max_workers=self._workers, + initializer=self._worker_init, + initargs=(self.model_dict,), + mp_context=mp.get_context("spawn"), + max_tasks_per_child=self.config.loader.marker_max_tasks_per_child, + ) + self.logger.info("MarkerWorker initialized with ProcessPoolExecutor") + + @staticmethod + def _worker_init(model_dict): + global worker_model_dict + worker_model_dict = model_dict + logger.debug("Worker initialized with model dictionary") + + @staticmethod + def _process_pdf(file_path, config): + global worker_model_dict + + page_range = config.get("page_range") + if page_range is not None: + label = f"[p{page_range[0]}-{page_range[-1]}]" + else: + label = "(all pages)" + + try: + logger.debug("Processing PDF", path=file_path, label=label) + converter = PdfConverter( + artifact_dict=worker_model_dict, + config=config, + ) + render = converter(file_path) + return render + except Exception as e: + logger.exception("Error processing PDF", path=file_path, label=label, error=str(e)) + raise + finally: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + async def process_pdf(self, file_path: str, page_range: list[int] | None = None): + from concurrent.futures import TimeoutError as FuturesTimeoutError + + converter_config = self.converter_config.copy() + if page_range is not None: + converter_config["page_range"] = page_range + + loop = asyncio.get_event_loop() + timeout = self.config.loader.marker_timeout + + def run_with_timeout(): + future = self.executor.submit(self._process_pdf, file_path, converter_config) + try: + result = future.result(timeout=timeout) + return result + except FuturesTimeoutError: + self.logger.exception("MarkerWorker child process timed out", path=file_path) + raise + except Exception: + self.logger.exception("Error processing with MarkerWorker", path=file_path) + raise + + result = await loop.run_in_executor(None, run_with_timeout) + return result.markdown, result.images + + def is_pool_broken(self): + # ProcessPoolExecutor auto-replaces dead/finished workers on next + # submit(), so counting live processes is unreliable and unnecessary. + # Only a None or shut-down executor requires reinitialization. + return self.executor is None or bool(getattr(self.executor, "_broken", False)) + + def __del__(self): + """Clean up ProcessPoolExecutor on actor destruction""" + if self.executor: + try: + self.executor.shutdown(wait=False, cancel_futures=True) + except Exception: + pass # Best effort cleanup + + +@ray.remote(max_restarts=5) +class MarkerPool: + def __init__(self): + from config import load_config + from utils.logger import get_logger + + self.logger = get_logger() + self.config = load_config() + self.max_processes = self.config.loader.marker_max_processes + self.pool_size = self.config.loader.marker_pool_size + self.actors = [MarkerWorker.remote() for _ in range(self.pool_size)] + self._queue: asyncio.Queue[ray.actor.ActorHandle] = asyncio.Queue() + + for _ in range(self.max_processes): + for actor in self.actors: + self._queue.put_nowait(actor) + + self.logger.info( + f"Marker pool: {self.pool_size} actors × {self.max_processes} slots = " + f"{self.pool_size * self.max_processes} PDF concurrency" + ) + + @staticmethod + def _get_page_count(file_path: str) -> int: + pdf = pypdfium2.PdfDocument(file_path) + try: + return len(pdf) + finally: + pdf.close() + + @staticmethod + def _create_chunks(page_count: int, chunk_size: int) -> list[tuple[list[int], str]]: + if page_count <= chunk_size: + return [(list(range(page_count)), f"({page_count}p)")] + chunks = [] + for start in range(0, page_count, chunk_size): + end = min(start + chunk_size, page_count) + page_range = list(range(start, end)) + label = f"[p{start}-{end - 1}]" + chunks.append((page_range, label)) + return chunks + + @with_timeout( + seconds=config.loader.marker_timeout, + description="MarkerWorker pool health check", + ) + async def _check_pool_broken(self, worker): + return worker.is_pool_broken.remote() + + @with_timeout( + seconds=config.loader.marker_timeout, + description="MarkerWorker pool reset", + ) + async def _reset_worker_pool(self, worker): + return worker.setup_mp.remote() + + async def ensure_worker_pool_healthy(self, worker): + if await self._check_pool_broken(worker): + self.logger.warning("Worker ProcessPoolExecutor is broken. Reinitializing pool...") + await self._reset_worker_pool(worker) + + @with_timeout( + seconds=config.loader.marker_timeout, + description="MarkerPool PDF {label} ({file_path})", + ) + async def _run_chunk(self, worker, file_path: str, page_range: list[int] | None, label: str): + return worker.process_pdf.remote(file_path, page_range=page_range) + + @with_retry( + max_retries=config.loader.marker_max_task_retry, + base_delay=config.loader.marker_retry_base_delay, + description="MarkerPool PDF {label} ({file_path})", + ) + async def _process_chunk(self, file_path: str, page_range: list[int] | None, label: str): + """Acquire a worker slot, process a PDF chunk, and release the slot. + + A fresh worker is acquired per attempt so a flaky worker can be + sidestepped and ``ensure_worker_pool_healthy`` re-runs each time. + Retries are handled by ``@with_retry``. + """ + worker = await self._queue.get() + try: + self.logger.info(f"MarkerWorker allocated for {label}") + await self.ensure_worker_pool_healthy(worker) + return await self._run_chunk(worker, file_path, page_range, label) + finally: + await self._queue.put(worker) + self.logger.debug(f"MarkerWorker returned to pool for {label}") + + async def process_pdf(self, file_path: str): + chunk_size = self.config.loader.marker_chunk_size + + if chunk_size <= 0: + return await self._process_chunk(file_path, page_range=None, label="(all pages)") + + page_count = self._get_page_count(file_path) + chunks = self._create_chunks(page_count, chunk_size) + + if len(chunks) == 1: + page_range, label = chunks[0] + return await self._process_chunk(file_path, page_range=None, label=label) + + self.logger.info( + f"Splitting {page_count}-page PDF into {len(chunks)} chunks of ~{chunk_size} pages for parallel processing" + ) + + tasks = [asyncio.create_task(self._process_chunk(file_path, page_range, label)) for page_range, label in chunks] + try: + results = await asyncio.gather(*tasks) + except Exception: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + + # Reassemble: concatenate markdown in order, merge image dicts + all_markdown = [] + all_images = {} + for markdown, images in results: + all_markdown.append(markdown) + all_images.update(images) + + combined_markdown = "\n\n".join(all_markdown) + return combined_markdown, all_images + + +_MARKER_KEY_PAGE_RE = re.compile(r"_page_(\d+)_") + + +def _marker_key_to_page(key: str) -> int | None: + """Extract the 1-indexed page number from a Marker image key. + + Marker emits keys like ``_page_0_Picture_1.jpeg`` (0-indexed). We + return ``N + 1`` so callers see 1-indexed pages aligned with the + ``[PAGE_N]`` markers produced by the post-processing step. + Returns ``None`` if the key doesn't match the expected pattern. + """ + match = _MARKER_KEY_PAGE_RE.search(key) + if match is None: + return None + try: + return int(match.group(1)) + 1 + except (TypeError, ValueError): + return None + + +class MarkerLoader(BasePooledParser): + """Public ``BasePooledParser`` facade for the Marker Ray pool. + + Holds a handle to the named ``MarkerPool`` Ray actor and dispatches + each ``parse()`` call to it. Marker requires a file path on disk, so + ``Document.raw_bytes`` is materialized to a temporary file (via + ``Document.as_temporary_file``) before handoff. + + Output: one ``TextBlock`` per page (1-indexed ``page_number``) plus + one ``ImageBlock`` per Marker image. Each ``ImageBlock`` carries the + ``![](key)`` markdown ref in ``metadata['markdown_ref']`` so a + downstream caption stage can substitute the wrapped caption back + into the markdown by string match. Captioning is not done here — + see :class:`ImageBlock` for the parser→caption contract. + """ + + PAGE_SEP = "[PAGE_SEP]" + _PAGE_MARKER_RE = re.compile(r"\{(\d+)\}" + re.escape(PAGE_SEP)) + + def __init__(self) -> None: + self.worker = ray.get_actor("MarkerPool", namespace="openrag") + + def supported_types(self) -> list[str]: + return [DocumentType.PDF.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + async with document.as_temporary_file() as path: + markdown, images = await self._dispatch(str(path)) + + pages = self._split_pages(markdown) + image_blocks = self._build_image_blocks(images) + text_blocks = [TextBlock(text=text, page_number=page) for page, text in pages] + + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=image_blocks, + metadata=dict(document.metadata), + page_count=pages[-1][0] if pages else 0, + ) + + # ----- helpers ----- + + @with_timeout( + seconds=config.loader.marker_timeout, + description="MarkerLoader PDF loading ({file_path})", + ) + async def _convert_pdf(self, file_path: str): + return self.worker.process_pdf.remote(file_path) + + async def _dispatch(self, file_path: str) -> tuple[str, dict]: + start = time.time() + try: + markdown, images = await self._convert_pdf(file_path) + if not markdown: + raise RuntimeError(f"Conversion failed for {file_path}") + duration = time.time() - start + logger.info(f"Processed {file_path} in {duration:.2f}s") + return markdown, images or {} + except Exception: + logger.exception("Error in MarkerLoader.parse", path=file_path) + raise + + @staticmethod + def _build_image_blocks(images: dict) -> list[ImageBlock]: + """Convert Marker's ``{key: PIL_image}`` dict into ``ImageBlock``s. + + Each block records the ``![](key)`` markdown ref in + ``metadata['markdown_ref']`` so a downstream caption stage can + substitute the wrapped caption back into the text. The page + number is parsed from Marker's key format + (``_page_{N}_Picture_{i}.{ext}``) and stored 1-indexed to match + the ``[PAGE_N]`` markers in the post-processed markdown. + """ + blocks: list[ImageBlock] = [] + for key, pil_image in images.items(): + try: + png_bytes = pil_to_png_bytes(pil_image) + except Exception as exc: + logger.warning(f"Failed to encode Marker image {key}: {exc}") + continue + blocks.append( + ImageBlock( + image_bytes=png_bytes, + page_number=_marker_key_to_page(str(key)), + mime_type="image/png", + metadata={"markdown_ref": f"![]({key})", "marker_key": str(key)}, + ) + ) + return blocks + + @classmethod + def _split_pages(cls, markdown: str) -> list[tuple[int, str]]: + """Clean Marker output and split it into ``[(page_number, text), …]``. + + Marker emits ``{1}[PAGE_SEP]{2}[PAGE_SEP]…``. We + drop the leading ``[PAGE_SEP]`` segment (Marker prefixes one), + strip ``
``, then split on each ``{N}[PAGE_SEP]`` marker — + the captured ``N`` is the 1-indexed page that just ended. + + Blank pages are preserved (text=``""``) so ``page_number`` and + ``page_count`` reflect the source document, not just the + non-empty subset. Trailing text after the last marker (rare) is + assigned to ``last_page + 1``. Markdown with no markers collapses + to a single page-1 entry. + """ + if markdown is None: + return [] + if cls.PAGE_SEP in markdown: + markdown = markdown.split(cls.PAGE_SEP, 1)[1] + markdown = markdown.replace("
", "") + + pairs: list[tuple[int, str]] = [] + cursor = 0 + last_page = 0 + for match in cls._PAGE_MARKER_RE.finditer(markdown): + page = int(match.group(1)) + text = markdown[cursor : match.start()].strip() + pairs.append((page, text)) + cursor = match.end() + last_page = page + tail = markdown[cursor:].strip() + if tail: + pairs.append((last_page + 1, tail)) + elif not pairs and markdown.strip(): + pairs.append((1, markdown.strip())) + return pairs diff --git a/openrag/services/workers/parsers/whisper_workers.py b/openrag/services/workers/parsers/whisper_workers.py new file mode 100644 index 000000000..1893d88bd --- /dev/null +++ b/openrag/services/workers/parsers/whisper_workers.py @@ -0,0 +1,201 @@ +import asyncio +from pathlib import Path + +import ray +import torch +from config import load_config +from core.indexing.parsers.document_parser import BasePooledParser +from core.models.document import ( + Document, + DocumentType, + ProcessedDocument, + TextBlock, +) +from faster_whisper import WhisperModel +from utils.logger import get_logger + +from ..ray_utils import with_retry, with_timeout + +logger = get_logger() +config = load_config() + + +if torch.cuda.is_available(): + WHISPER_NUM_GPUS = config.loader.local_whisper.whisper_num_gpus +else: # On CPU + WHISPER_NUM_GPUS = 0 + +WHISPER_CONCURRENCY_PER_WORKER = config.loader.local_whisper.whisper_concurrency_per_worker + + +# Duration of the audio sample used for language detection +LANG_DETECT_SAMPLE_MS = 30_000 # 30 s + + +@ray.remote( + num_gpus=WHISPER_NUM_GPUS, max_restarts=5, max_concurrency=WHISPER_CONCURRENCY_PER_WORKER +) # Ensure each worker processes one file at a time +class WhisperActor: + def __init__(self): + import torch + from config import load_config + from utils.logger import get_logger + + self.logger = get_logger() + self.config = load_config() + + device = "cuda" if torch.cuda.is_available() else "cpu" + compute_type = "float16" if device == "cuda" else "int8" + model_name = self.config.loader.local_whisper.model + + self.logger.info("Loading Whisper model", model_name=model_name, device=device, compute_type=compute_type) + self.model = WhisperModel(model_name, device=device, compute_type=compute_type) + self.logger.info("Whisper model loaded successfully", model_name=model_name, device=device) + + async def transcribe(self, wav_path: str | Path) -> str: + self.logger.info("Transcribing audio file", file_path=Path(wav_path).name) + + def _transcribe_sync() -> str: + segments, _ = self.model.transcribe(str(wav_path)) + return "".join(segment.text for segment in segments) + + return await asyncio.to_thread(_transcribe_sync) + + async def detect_language( + self, + file_path: str | Path, + fallback_language: str = "en", + sample_ms: int = LANG_DETECT_SAMPLE_MS, + ) -> str: + import tempfile + + from pydub import AudioSegment + + file_path = Path(file_path) + self.logger.info("Detecting language for audio file", file_path=file_path.name) + fd, tmp_path = tempfile.mkstemp(prefix=f"{file_path.stem}_langdetect_", suffix=".wav") + try: + + def _prepare_and_detect() -> str: + import os + + os.close(fd) + sound = AudioSegment.from_file(file_path) + sample = sound[:sample_ms] + sample.export(tmp_path, format="wav") + _, info = self.model.transcribe(tmp_path, beam_size=1, max_new_tokens=1) + return info.language + + return await asyncio.to_thread(_prepare_and_detect) + except Exception: + self.logger.exception("Error detecting language", file_path=file_path.name) + return fallback_language + finally: + Path(tmp_path).unlink(missing_ok=True) + + +@ray.remote +class WhisperPool: + """Ray-actor pool of ``WhisperActor``s. Internal — the public + ``BasePooledParser`` face is ``LocalWhisperLoader``. + """ + + def __init__(self): + from utils.logger import get_logger + + self.logger = get_logger() + + n_workers = config.loader.local_whisper.whisper_n_workers + self.logger.info(f"Starting WhisperPool with {n_workers} workers") + self.workers = [WhisperActor.remote() for _ in range(n_workers)] + self._pending = [0] * n_workers + + @with_timeout( + seconds=config.loader.local_whisper.whisper_timeout, + description="WhisperPool transcribe ({path})", + ) + async def _transcribe_chunk(self, idx: int, path): + return self.workers[idx].transcribe.remote(path) + + @with_retry( + max_retries=config.loader.local_whisper.whisper_max_task_retry, + base_delay=config.loader.local_whisper.whisper_retry_base_delay, + description="WhisperPool transcribe ({path})", + ) + async def transcribe(self, path): + idx = min(range(len(self._pending)), key=lambda j: self._pending[j]) + self._pending[idx] += 1 + try: + return await self._transcribe_chunk(idx, path) + finally: + self._pending[idx] -= 1 + + +async def detect_language_via_actor( + file_path: str | Path, + *, + fallback_language: str = "en", + timeout: float | None = None, +) -> str | None: + """Detect the spoken language of ``file_path`` via the singleton WhisperActor. + + Wraps the ``.remote()`` call so non-worker code (the OpenAI audio + loader's optional language detector) can stay Ray-free. Returns + ``None`` on failure so callers can fall back to a default behaviour. + """ + from ..ray_utils import call_ray_actor_with_timeout + + try: + actor = WhisperActor.options(name="WhisperActor", namespace="openrag", get_if_exists=True).remote() + except Exception: + logger.exception("Error getting WhisperActor") + return None + + effective_timeout = timeout if timeout is not None else config.loader.local_whisper.whisper_timeout + try: + return await call_ray_actor_with_timeout( + actor.detect_language.remote(file_path, fallback_language), + timeout=effective_timeout, + task_description=f"WhisperActor detect_language ({Path(file_path).name})", + ) + except Exception: + logger.exception("Language detection failed", file_path=str(file_path)) + return None + + +class LocalWhisperLoader(BasePooledParser): + """Public ``BasePooledParser`` facade for the local-Whisper Ray pool. + + Holds a handle to the named ``WhisperPool`` Ray actor and dispatches + each ``parse()`` call to it. Whisper requires a file path on disk, + so ``Document.raw_bytes`` is written to a NamedTemporaryFile before + handoff. + """ + + def __init__(self): + self.whisper_actor: WhisperPool = ray.get_actor("WhisperPool", namespace="openrag") + + def supported_types(self) -> list[str]: + return [DocumentType.AUDIO.value, DocumentType.VIDEO.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + async with document.as_temporary_file() as path: + try: + text = await self.whisper_actor.transcribe.remote(str(path)) + except Exception as e: + logger.error("Error transcribing audio", error=str(e)) + raise + + text_blocks = [TextBlock(text=text, page_number=1)] if text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + metadata=dict(document.metadata), + page_count=1 if text else 0, + ) diff --git a/openrag/services/workers/pipeline_builder.py b/openrag/services/workers/pipeline_builder.py new file mode 100644 index 000000000..1bbb41d6a --- /dev/null +++ b/openrag/services/workers/pipeline_builder.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from collections.abc import MutableMapping +from dataclasses import dataclass +from typing import Any + +from core.chunking.chunking_strategy import ChunkingStrategy +from core.embeddings.embedder import Embedder +from core.indexing.contextualize import ChunkContextualizer +from core.indexing.parsers.document_parser import DocumentParser +from core.vector_stores.vector_store import VectorStore +from core.vlm.vlm import VLM +from services.workers.stages.caption import caption_stage +from services.workers.stages.chunk import chunk_stage +from services.workers.stages.contextualize import contextualize_stage +from services.workers.stages.embed import embed_stage +from services.workers.stages.parse import parse_stage +from services.workers.stages.store import store_stage + + +@dataclass(slots=True, frozen=True) +class PipelineTimeouts: + """Per-stage timeout configuration for an indexing pipeline row.""" + + parse: float | None = None + caption: float | None = None + caption_per_image: float = 0.0 + chunk: float | None = None + contextualize: float | None = None + contextualize_per_chunk: float = 0.0 + embed: float | None = None + embed_per_chunk: float = 0.0 + store: float | None = None + store_per_chunk: float = 0.0 + + +@dataclass(slots=True, frozen=True) +class IndexingPipeline: + """Sequential indexing pipeline assembled from worker stage functions.""" + + parser: DocumentParser + chunker: ChunkingStrategy + embedder: Embedder + vector_store: VectorStore + vlm: VLM | None = None + contextualizer: ChunkContextualizer | None = None + timeouts: PipelineTimeouts = PipelineTimeouts() + + async def run(self, row: MutableMapping[str, Any]) -> MutableMapping[str, Any]: + """Run a single row through parse, optional enrichments, embed, and store.""" + + await parse_stage(row, self.parser, timeout=self.timeouts.parse) + if self.vlm is not None: + await caption_stage( + row, + self.vlm, + timeout=self.timeouts.caption, + per_image_timeout=self.timeouts.caption_per_image, + ) + await chunk_stage(row, self.chunker, timeout=self.timeouts.chunk) + if self.contextualizer is not None: + await contextualize_stage( + row, + self.contextualizer, + timeout=self.timeouts.contextualize, + per_chunk_timeout=self.timeouts.contextualize_per_chunk, + ) + await embed_stage( + row, + self.embedder, + timeout=self.timeouts.embed, + per_chunk_timeout=self.timeouts.embed_per_chunk, + ) + await store_stage( + row, + self.vector_store, + timeout=self.timeouts.store, + per_chunk_timeout=self.timeouts.store_per_chunk, + ) + return row + + +def build_indexing_pipeline( + *, + parser: DocumentParser, + chunker: ChunkingStrategy, + embedder: Embedder, + vector_store: VectorStore, + vlm: VLM | None = None, + contextualizer: ChunkContextualizer | None = None, + timeouts: PipelineTimeouts | None = None, +) -> IndexingPipeline: + """Build the default sequential indexing pipeline.""" + + return IndexingPipeline( + parser=parser, + chunker=chunker, + embedder=embedder, + vector_store=vector_store, + vlm=vlm, + contextualizer=contextualizer, + timeouts=timeouts or PipelineTimeouts(), + ) + + +__all__ = ["IndexingPipeline", "PipelineTimeouts", "build_indexing_pipeline"] diff --git a/openrag/services/workers/ray_utils.py b/openrag/services/workers/ray_utils.py new file mode 100644 index 000000000..5061bc244 --- /dev/null +++ b/openrag/services/workers/ray_utils.py @@ -0,0 +1,229 @@ +"""Ray-actor concurrency helpers. + +Two pairs of utilities, each in function and decorator form: + +- timeout: ``call_ray_actor_with_timeout`` / ``@with_timeout``. Awaits a + ``ray.ObjectRef`` with proper cancel-on-timeout semantics. +- retry: ``retry_with_backoff`` / ``@with_retry``. Exponential backoff + + jitter; ``CancelledError`` is never retried. + +Use the decorator form when params are static (or pulled from a +module-level config); use the function form when params are dynamic per +call. Cancellation paths are translated into a predictable shape: + +- caller-side timeout → ``ray.cancel(future)`` then re-raise ``TimeoutError`` +- caller-side ``asyncio.CancelledError`` → ``ray.cancel(future)`` then re-raise +- worker-side ``TaskCancelledError`` → re-raise as-is +- worker-side ``RayTaskError`` → re-raise as ``RuntimeError`` (cause preserved) +""" + +from __future__ import annotations + +import asyncio +import functools +import inspect +import random +from collections.abc import Callable +from typing import Any + +import ray +from ray.exceptions import RayTaskError, TaskCancelledError +from utils.logger import get_logger + +logger = get_logger() + +__all__ = [ + "call_ray_actor_with_timeout", + "retry_with_backoff", + "with_retry", + "with_timeout", +] + + +def _resolve_description(template: str, fn: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]) -> str: + """Format ``template`` with the wrapped call's bound arguments. + + A description like ``"PDF parse ({file_path})"`` gets ``{file_path}`` + substituted with the value passed to ``fn`` for that parameter. If + ``template`` contains no ``{`` it is returned unchanged — no inspect + cost in the hot path for plain-string descriptions. + + ``KeyError`` from a missing placeholder is caught and the raw + template is returned, so a typo in a placeholder name degrades to a + log-line oddity rather than a runtime crash on the wrapped call. + """ + if "{" not in template: + return template + try: + bound = inspect.signature(fn).bind(*args, **kwargs).arguments + return template.format(**bound) + except (KeyError, TypeError) as exc: + logger.warning(f"description template missing placeholder for {exc}") + return template + + +# --------------------------------------------------------------------------- +# Timeout +# --------------------------------------------------------------------------- + + +async def call_ray_actor_with_timeout( + future: ray.ObjectRef, + timeout: float, + task_description: str = "Ray task", +) -> Any: + """Await a Ray ``ObjectRef`` with a timeout, propagating cancellation. + + Raises: + TimeoutError: If the task exceeds ``timeout``. + asyncio.CancelledError: If the calling coroutine is cancelled. + TaskCancelledError: If the Ray task was cancelled by the worker. + RuntimeError: If the Ray task failed (original exception chained). + """ + try: + result = await asyncio.wait_for(asyncio.gather(future), timeout=timeout) + return result[0] + + except TimeoutError: + logger.warning(f"{task_description} timed out, cancelling Ray task") + ray.cancel(future, recursive=True) + raise + + except asyncio.CancelledError: + logger.warning(f"{task_description} cancelled, cancelling Ray task") + ray.cancel(future, recursive=True) + raise + + except TaskCancelledError: + logger.warning(f"{task_description} Ray task was cancelled") + raise + + except RayTaskError as e: + raise RuntimeError(f"{task_description} failed") from e + + +def with_timeout( + *, + seconds: float, + description: str = "Ray task", +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Decorator: wrap an async function returning a ``ray.ObjectRef``. + + The wrapped function is called normally; its return value (an + ``ObjectRef``) is then awaited via ``call_ray_actor_with_timeout``. + + ``description`` may embed any of the wrapped function's parameter + names as ``str.format``-style placeholders; they are substituted + with the per-call argument values for log lines. + + Example:: + + @with_timeout( + seconds=30.0, + description="caption_image ({path})", + ) + async def caption(self, path): + return self.actor.caption.remote(path) + """ + + def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: + @functools.wraps(fn) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + desc = _resolve_description(description, fn, args, kwargs) + future = fn(*args, **kwargs) + if asyncio.iscoroutine(future): + future = await future + return await call_ray_actor_with_timeout(future, seconds, desc) + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------- +# Retry +# --------------------------------------------------------------------------- + + +async def retry_with_backoff( + attempt_fn: Callable[[int], Any], + max_retries: int, + base_delay: float, + task_description: str = "task", + jitter: bool = True, +) -> Any: + """Run ``attempt_fn(attempt_index)`` with exponential backoff. + + Backoff is ``base_delay * 2**attempt`` seconds, plus uniform jitter + in ``[0, base_delay)`` when ``jitter=True``. ``attempt_fn`` is an + async callable; it owns acquire/release of any per-attempt resources + so a flaky resource can be sidestepped on retry. + """ + last_exc: Exception | None = None + for attempt in range(max_retries + 1): + try: + return await attempt_fn(attempt) + except (asyncio.CancelledError, TaskCancelledError): + raise + except Exception as e: + last_exc = e + if attempt >= max_retries: + logger.error(f"{task_description} failed after {attempt + 1} attempts: {e}") + raise + delay = base_delay * (2**attempt) + if jitter: + delay += random.uniform(0, base_delay) + logger.warning( + f"{task_description} failed (attempt {attempt + 1}/{max_retries + 1}): {e}. Retrying in {delay:.1f}s..." + ) + await asyncio.sleep(delay) + + raise last_exc # unreachable + + +def with_retry( + *, + max_retries: int, + base_delay: float, + description: str = "task", + jitter: bool = True, +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Decorator: retry an async function with exponential backoff + jitter. + + Each invocation counts as one attempt. ``CancelledError`` is never + retried. + + ``description`` may embed any of the wrapped function's parameter + names as ``str.format``-style placeholders; they are substituted + with the per-call argument values for log lines. + + Example:: + + @with_retry( + max_retries=3, + base_delay=0.5, + description="transcribe ({path})", + ) + async def transcribe(self, path): + return await self.actor.transcribe.remote(path) + """ + + def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: + @functools.wraps(fn) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + desc = _resolve_description(description, fn, args, kwargs) + + async def attempt(_i: int) -> Any: + return await fn(*args, **kwargs) + + return await retry_with_backoff( + attempt, + max_retries=max_retries, + base_delay=base_delay, + task_description=desc, + jitter=jitter, + ) + + return wrapper + + return decorator diff --git a/openrag/services/workers/result_aggregation.py b/openrag/services/workers/result_aggregation.py new file mode 100644 index 000000000..110fa82bf --- /dev/null +++ b/openrag/services/workers/result_aggregation.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from collections.abc import MutableMapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +_SUCCESS_STAGE = "stored" + + +@dataclass(frozen=True) +class RowFailure: + """Stage and error message for a single failed row.""" + + stage: str + error: str + + +@dataclass(frozen=True) +class BatchIngestSummary: + """Aggregated result of a batch ingestion run.""" + + total: int + succeeded: int + failed: int + stored_count: int + failures: tuple[RowFailure, ...] = field(default_factory=tuple) + + @property + def success_rate(self) -> float: + return self.succeeded / self.total if self.total > 0 else 0.0 + + +def aggregate_batch_results( + rows: Sequence[MutableMapping[str, Any]], +) -> BatchIngestSummary: + """Summarise processed rows from :func:`ingest_batch`. + + A row is counted as succeeded when ``row["stage"] == "stored"``. + All other rows are counted as failed regardless of whether an exception + was raised. + """ + succeeded = 0 + stored_count = 0 + failures: list[RowFailure] = [] + + for row in rows: + stage = row.get("stage", "") + if stage == _SUCCESS_STAGE: + succeeded += 1 + try: + stored_count += int(row.get("stored_count", 0)) + except (TypeError, ValueError): + pass + else: + failures.append( + RowFailure( + stage=str(stage), + error=str(row.get("error", "")), + ) + ) + + return BatchIngestSummary( + total=len(rows), + succeeded=succeeded, + failed=len(failures), + stored_count=stored_count, + failures=tuple(failures), + ) + + +__all__ = ["BatchIngestSummary", "RowFailure", "aggregate_batch_results"] diff --git a/openrag/services/workers/stages/__init__.py b/openrag/services/workers/stages/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/services/workers/stages/_common.py b/openrag/services/workers/stages/_common.py new file mode 100644 index 000000000..1c4192a4e --- /dev/null +++ b/openrag/services/workers/stages/_common.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, MutableMapping +from typing import Any + +_CREDENTIAL_KEYS = frozenset({"credentials", "credential", "api_key", "token", "secret", "password"}) + + +def scrub_credentials(row: MutableMapping[str, Any]) -> None: + for key in _CREDENTIAL_KEYS: + row.pop(key, None) + + +def stage_timeout(base_timeout: float | None, item_count: int, *, per_item_timeout: float = 0.0) -> float | None: + if base_timeout is None: + return None + return base_timeout + max(0, item_count) * per_item_timeout + + +async def run_with_optional_timeout[T]( + operation: Callable[[], Awaitable[T]], + timeout: float | None, +) -> T: + if timeout is None: + return await operation() + return await asyncio.wait_for(operation(), timeout=timeout) diff --git a/openrag/services/workers/stages/caption.py b/openrag/services/workers/stages/caption.py new file mode 100644 index 000000000..2dfc01d30 --- /dev/null +++ b/openrag/services/workers/stages/caption.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from collections.abc import MutableMapping +from typing import Any + +from core.models.document import ImageBlock, ProcessedDocument, TextBlock +from core.prompts.vlm_prompt_builder import wrap_caption +from core.vlm.vlm import VLM +from services.workers.stages._common import run_with_optional_timeout, scrub_credentials, stage_timeout + + +async def caption_stage( + row: MutableMapping[str, Any], + vlm: VLM, + *, + timeout: float | None = None, + per_image_timeout: float = 0.0, +) -> MutableMapping[str, Any]: + """Caption images in ``row["processed_document"]`` and mutate the row.""" + + try: + processed_document = row.get("processed_document") + if not isinstance(processed_document, ProcessedDocument): + raise ValueError("caption_stage row must contain a ProcessedDocument under 'processed_document'") + + prompt = row.get("caption_prompt") + if prompt is not None: + prompt = str(prompt) + + effective_timeout = stage_timeout(timeout, len(processed_document.images), per_item_timeout=per_image_timeout) + row["processed_document"] = await run_with_optional_timeout( + lambda: _caption_document(processed_document, vlm, prompt), + effective_timeout, + ) + row["stage"] = "captioned" + row.pop("error", None) + return row + except Exception as exc: + row["stage"] = "caption_failed" + row["error"] = str(exc) + raise + finally: + scrub_credentials(row) + + +async def _caption_document( + processed_document: ProcessedDocument, + vlm: VLM, + prompt: str | None, +) -> ProcessedDocument: + """Return a copy of ``processed_document`` with captions materialized.""" + text_blocks = list(processed_document.text_blocks) + captioned_images: list[ImageBlock] = [] + + for image in processed_document.images: + caption = await vlm.caption_image(image.image_bytes, prompt=prompt) + wrapped = wrap_caption(caption) + captioned_image = image.model_copy(update={"caption": caption}) + captioned_images.append(captioned_image) + if not _replace_markdown_ref(text_blocks, image, wrapped): + text_blocks.append(TextBlock(text=wrapped, page_number=image.page_number)) + + return processed_document.model_copy( + update={ + "text_blocks": text_blocks, + "images": captioned_images, + } + ) + + +def _replace_markdown_ref(text_blocks: list[TextBlock], image: ImageBlock, wrapped_caption: str) -> bool: + """Replace an image markdown placeholder in text blocks when one exists.""" + markdown_ref = image.metadata.get("markdown_ref") + if not markdown_ref: + return False + + replaced = False + for index, block in enumerate(text_blocks): + if markdown_ref not in block.text: + continue + text_blocks[index] = block.model_copy(update={"text": block.text.replace(markdown_ref, wrapped_caption)}) + replaced = True + return replaced diff --git a/openrag/services/workers/stages/chunk.py b/openrag/services/workers/stages/chunk.py new file mode 100644 index 000000000..661c79b83 --- /dev/null +++ b/openrag/services/workers/stages/chunk.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import asyncio +from collections.abc import MutableMapping +from typing import Any + +from core.chunking.chunking_strategy import ChunkingStrategy +from core.models.chunk import Chunk +from core.models.document import ProcessedDocument +from services.workers.stages._common import run_with_optional_timeout, scrub_credentials + + +async def chunk_stage( + row: MutableMapping[str, Any], + chunker: ChunkingStrategy, + *, + timeout: float | None = None, +) -> MutableMapping[str, Any]: + """Chunk ``row["processed_document"]`` and mutate the row with chunks.""" + + try: + processed_document = row.get("processed_document") + if not isinstance(processed_document, ProcessedDocument): + raise ValueError("chunk_stage row must contain a ProcessedDocument under 'processed_document'") + + partition = str(row.get("partition") or "default") + chunks = await _chunk_with_timeout(chunker, processed_document, partition, timeout) + row["chunks"] = chunks + row["stage"] = "chunked" + row.pop("error", None) + return row + except Exception as exc: + row["stage"] = "chunk_failed" + row["error"] = str(exc) + raise + finally: + scrub_credentials(row) + + +async def _chunk_with_timeout( + chunker: ChunkingStrategy, + processed_document: ProcessedDocument, + partition: str, + timeout: float | None, +) -> list[Chunk]: + """Run the synchronous chunker without blocking the event loop.""" + + async def run() -> list[Chunk]: + return await asyncio.to_thread(chunker.chunk, processed_document, partition) + + return await run_with_optional_timeout(run, timeout) diff --git a/openrag/services/workers/stages/contextualize.py b/openrag/services/workers/stages/contextualize.py new file mode 100644 index 000000000..6df3c7625 --- /dev/null +++ b/openrag/services/workers/stages/contextualize.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from collections.abc import MutableMapping +from typing import Any + +from core.indexing.contextualize import ChunkContextualizer +from core.models.chunk import Chunk +from services.workers.stages._common import run_with_optional_timeout, scrub_credentials, stage_timeout + + +async def contextualize_stage( + row: MutableMapping[str, Any], + contextualizer: ChunkContextualizer, + *, + timeout: float | None = None, + per_chunk_timeout: float = 0.0, +) -> MutableMapping[str, Any]: + """Contextualize ``row["chunks"]`` in place.""" + + try: + chunks = row.get("chunks") + if not _is_chunk_list(chunks): + raise ValueError("contextualize_stage row must contain a list[Chunk] under 'chunks'") + + filename = str(row.get("filename") or "") + language = str(row.get("language") or row.get("lang") or "en") + effective_timeout = stage_timeout(timeout, len(chunks), per_item_timeout=per_chunk_timeout) + row["chunks"] = await run_with_optional_timeout( + lambda: contextualizer.contextualize(chunks, filename=filename, lang=language), + effective_timeout, + ) + row["stage"] = "contextualized" + row.pop("error", None) + return row + except Exception as exc: + row["stage"] = "contextualize_failed" + row["error"] = str(exc) + raise + finally: + scrub_credentials(row) + + +def _is_chunk_list(value: Any) -> bool: + """Return whether ``value`` is a concrete list of domain chunks.""" + return isinstance(value, list) and all(isinstance(chunk, Chunk) for chunk in value) diff --git a/openrag/services/workers/stages/embed.py b/openrag/services/workers/stages/embed.py new file mode 100644 index 000000000..28954f128 --- /dev/null +++ b/openrag/services/workers/stages/embed.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +from collections.abc import MutableMapping +from typing import Any + +from core.embeddings.embedder import Embedder +from core.models.chunk import Chunk +from services.workers.stages._common import run_with_optional_timeout, scrub_credentials, stage_timeout + + +async def embed_stage( + row: MutableMapping[str, Any], + embedder: Embedder, + *, + timeout: float | None = None, + per_chunk_timeout: float = 0.0, +) -> MutableMapping[str, Any]: + """Embed ``row["chunks"]`` and replace them with embedded copies.""" + + try: + chunks = row.get("chunks") + if not _is_chunk_list(chunks): + raise ValueError("embed_stage row must contain a list[Chunk] under 'chunks'") + + effective_timeout = stage_timeout(timeout, len(chunks), per_item_timeout=per_chunk_timeout) + texts = [chunk.text for chunk in chunks] + vectors = await run_with_optional_timeout(lambda: embedder.embed(texts), effective_timeout) + if len(vectors) != len(chunks): + raise ValueError("embedder returned a different number of vectors than chunks") + row["chunks"] = [chunk.with_embedding(vector) for chunk, vector in zip(chunks, vectors, strict=True)] + row["stage"] = "embedded" + row.pop("error", None) + return row + except Exception as exc: + row["stage"] = "embed_failed" + row["error"] = str(exc) + raise + finally: + scrub_credentials(row) + + +def _is_chunk_list(value: Any) -> bool: + """Return whether ``value`` is a concrete list of domain chunks.""" + return isinstance(value, list) and all(isinstance(chunk, Chunk) for chunk in value) diff --git a/openrag/services/workers/stages/parse.py b/openrag/services/workers/stages/parse.py new file mode 100644 index 000000000..29f225aed --- /dev/null +++ b/openrag/services/workers/stages/parse.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from collections.abc import MutableMapping +from typing import Any + +from core.indexing.parsers.document_parser import DocumentParser +from core.models.document import Document, ProcessedDocument +from services.workers.stages._common import run_with_optional_timeout, scrub_credentials + + +async def parse_stage( + row: MutableMapping[str, Any], + parser: DocumentParser, + *, + timeout: float | None = None, +) -> MutableMapping[str, Any]: + """Parse ``row["document"]`` and mutate the row with the stage result.""" + + try: + document = row.get("document") + if not isinstance(document, Document): + raise ValueError("parse_stage row must contain a Document under 'document'") + + processed = await _parse_with_timeout(parser, document, timeout) + row["processed_document"] = processed + row["stage"] = "parsed" + row.pop("error", None) + return row + except Exception as exc: + row["stage"] = "parse_failed" + row["error"] = str(exc) + raise + finally: + scrub_credentials(row) + + +async def _parse_with_timeout( + parser: DocumentParser, + document: Document, + timeout: float | None, +) -> ProcessedDocument: + return await run_with_optional_timeout(lambda: parser.parse(document), timeout) diff --git a/openrag/services/workers/stages/store.py b/openrag/services/workers/stages/store.py new file mode 100644 index 000000000..3f6bd50fc --- /dev/null +++ b/openrag/services/workers/stages/store.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +from collections.abc import MutableMapping +from typing import Any + +from core.models.chunk import Chunk +from core.vector_stores.vector_store import VectorStore +from services.workers.stages._common import run_with_optional_timeout, scrub_credentials, stage_timeout + + +async def store_stage( + row: MutableMapping[str, Any], + vector_store: VectorStore, + *, + timeout: float | None = None, + per_chunk_timeout: float = 0.0, +) -> MutableMapping[str, Any]: + """Upsert ``row["chunks"]`` into the configured vector collection. + + Tenant routing stays on each chunk's ``partition`` field. The vector + store collection argument remains the configured backend collection. + """ + + try: + chunks = row.get("chunks") + if not _is_chunk_list(chunks): + raise ValueError("store_stage row must contain a list[Chunk] under 'chunks'") + if chunks: + embedding = chunks[0].embedding + if embedding is None: + raise ValueError("store_stage received chunks without embeddings") + await vector_store.ensure_collection("default", len(embedding)) + + effective_timeout = stage_timeout(timeout, len(chunks), per_item_timeout=per_chunk_timeout) + row["stored_count"] = await run_with_optional_timeout( + lambda: vector_store.upsert(chunks), + effective_timeout, + ) + row["stage"] = "stored" + row.pop("error", None) + return row + except Exception as exc: + row["stage"] = "store_failed" + row["error"] = str(exc) + raise + finally: + scrub_credentials(row) + + +def _is_chunk_list(value: Any) -> bool: + """Return whether ``value`` is a concrete list of domain chunks.""" + return isinstance(value, list) and all(isinstance(chunk, Chunk) for chunk in value) diff --git a/openrag/services/workers/stages/test_parse.py b/openrag/services/workers/stages/test_parse.py new file mode 100644 index 000000000..85b77e000 --- /dev/null +++ b/openrag/services/workers/stages/test_parse.py @@ -0,0 +1,73 @@ +import pytest +from core.indexing.parsers.document_parser import DocumentParser +from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock +from services.workers.stages.parse import parse_stage + + +class FakeParser(DocumentParser): + def __init__(self, output: ProcessedDocument | None = None, error: Exception | None = None) -> None: + self.output = output + self.error = error + self.seen_documents: list[Document] = [] + + async def parse(self, document: Document) -> ProcessedDocument: + self.seen_documents.append(document) + if self.error is not None: + raise self.error + assert self.output is not None + return self.output + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + +@pytest.mark.asyncio +async def test_parse_stage_mutates_row_with_processed_document_and_scrubs_credentials(): + document = Document(id="doc-1", filename="note.txt", content_type=DocumentType.TEXT, text="hello") + processed = ProcessedDocument( + document_id="doc-1", + text_blocks=[TextBlock(text="hello", page_number=1)], + metadata={"file_id": "file-1"}, + page_count=1, + ) + parser = FakeParser(output=processed) + row = { + "document": document, + "credentials": {"api_key": "secret"}, + "token": "secret-token", + } + + result = await parse_stage(row, parser) + + assert result is row + assert parser.seen_documents == [document] + assert row["processed_document"] == processed + assert row["stage"] == "parsed" + assert "error" not in row + assert "credentials" not in row + assert "token" not in row + + +@pytest.mark.asyncio +async def test_parse_stage_marks_error_and_scrubs_credentials_when_parser_fails(): + document = Document(id="doc-1", filename="note.txt", content_type=DocumentType.TEXT, text="hello") + row = {"document": document, "api_key": "secret"} + + with pytest.raises(ValueError, match="parse failed"): + await parse_stage(row, FakeParser(error=ValueError("parse failed"))) + + assert row["stage"] == "parse_failed" + assert row["error"] == "parse failed" + assert "api_key" not in row + + +@pytest.mark.asyncio +async def test_parse_stage_requires_document_in_row(): + row = {"api_key": "secret"} + + with pytest.raises(ValueError, match="document"): + await parse_stage(row, FakeParser()) + + assert row["stage"] == "parse_failed" + assert row["error"] == "parse_stage row must contain a Document under 'document'" + assert "api_key" not in row diff --git a/openrag/services/workers/stages/test_pipeline_stages.py b/openrag/services/workers/stages/test_pipeline_stages.py new file mode 100644 index 000000000..7887a0d22 --- /dev/null +++ b/openrag/services/workers/stages/test_pipeline_stages.py @@ -0,0 +1,300 @@ +import re + +import pytest +from core.chunking.chunking_strategy import ChunkingStrategy +from core.embeddings.embedder import Embedder +from core.indexing.contextualize import ChunkContextualizer +from core.models.chunk import Chunk +from core.models.document import ImageBlock, ProcessedDocument, TextBlock +from core.prompts.vlm_prompt_builder import wrap_caption +from core.vector_stores.vector_store import VectorStore +from core.vlm.vlm import VLM +from services.workers.stages.caption import caption_stage +from services.workers.stages.chunk import chunk_stage +from services.workers.stages.contextualize import contextualize_stage +from services.workers.stages.embed import embed_stage +from services.workers.stages.store import store_stage + + +class FakeChunker(ChunkingStrategy): + def __init__(self, chunks: list[Chunk], error: Exception | None = None) -> None: + self.chunks = chunks + self.error = error + self.calls: list[tuple[ProcessedDocument, str]] = [] + + def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]: + self.calls.append((document, partition)) + if self.error is not None: + raise self.error + return self.chunks + + +class FakeContextualizer(ChunkContextualizer): + def __init__(self, chunks: list[Chunk], error: Exception | None = None) -> None: + self.chunks = chunks + self.error = error + self.calls: list[tuple[list[Chunk], str, str]] = [] + + async def contextualize(self, chunks, *, filename: str = "", lang: str = "en") -> list[Chunk]: + self.calls.append((list(chunks), filename, lang)) + if self.error is not None: + raise self.error + return self.chunks + + +class FakeEmbedder(Embedder): + def __init__(self, vectors: list[list[float]], error: Exception | None = None) -> None: + self.vectors = vectors + self.error = error + self.text_batches: list[list[str]] = [] + + async def embed(self, texts: list[str]) -> list[list[float]]: + self.text_batches.append(texts) + if self.error is not None: + raise self.error + return self.vectors + + async def embed_single(self, text: str) -> list[float]: + return (await self.embed([text]))[0] + + @property + def dimension(self) -> int: + return 2 + + +class FakeVLM(VLM): + def __init__(self, captions: list[str], error: Exception | None = None) -> None: + self.captions = captions + self.error = error + self.calls: list[tuple[bytes, str | None]] = [] + + async def caption_image(self, image_bytes: bytes, prompt: str | None = None) -> str: + self.calls.append((image_bytes, prompt)) + if self.error is not None: + raise self.error + return self.captions[len(self.calls) - 1] + + async def caption_images_batch(self, images: list[bytes], prompt: str | None = None) -> list[str]: + return [await self.caption_image(image, prompt=prompt) for image in images] + + +class FakeVectorStore(VectorStore): + def __init__(self, count: int, error: Exception | None = None) -> None: + self.count = count + self.error = error + self.calls: list[tuple[list[Chunk], str]] = [] + self.ensure_calls: list[tuple[str, int]] = [] + + async def upsert(self, chunks: list[Chunk], collection: str = "default") -> int: + self.calls.append((chunks, collection)) + if self.error is not None: + raise self.error + return self.count + + async def search( + self, embedding, query_text=None, top_k=10, collection="default", filters=None, similarity_threshold=None + ): + return [] + + async def delete(self, ids: list[str], collection: str = "default") -> int: + return 0 + + async def ensure_collection(self, name: str, dimension: int, **kwargs) -> None: + self.ensure_calls.append((name, dimension)) + return None + + async def drop_collection(self, name: str) -> None: + return None + + async def collection_exists(self, name: str) -> bool: + return True + + async def query_ids_by_filter(self, collection: str, filters: dict) -> list[str]: + return [] + + async def query_chunks_by_filter(self, collection: str, filters: dict, output_fields=None) -> list[dict]: + return [] + + +@pytest.mark.asyncio +async def test_caption_stage_replaces_markdown_refs_and_scrubs_credentials(): + markdown_ref = "![](image-1)" + processed = ProcessedDocument( + document_id="doc-1", + text_blocks=[TextBlock(text=f"before {markdown_ref} after", page_number=2)], + images=[ImageBlock(image_bytes=b"img", page_number=2, metadata={"markdown_ref": markdown_ref})], + ) + row = {"processed_document": processed, "caption_prompt": "Describe", "token": "secret"} + + await caption_stage(row, FakeVLM(["a chart"])) + + assert row["processed_document"].images[0].caption == "a chart" + assert row["processed_document"].text_blocks[0].text == f"before {wrap_caption('a chart')} after" + assert row["stage"] == "captioned" + assert "token" not in row + + +@pytest.mark.asyncio +async def test_caption_stage_appends_standalone_image_captions(): + processed = ProcessedDocument( + document_id="doc-1", + text_blocks=[TextBlock(text="body", page_number=1)], + images=[ImageBlock(image_bytes=b"img", page_number=3)], + ) + row = {"processed_document": processed} + + await caption_stage(row, FakeVLM(["a diagram"])) + + assert row["processed_document"].text_blocks[-1] == TextBlock(text=wrap_caption("a diagram"), page_number=3) + assert row["stage"] == "captioned" + + +@pytest.mark.asyncio +async def test_chunk_stage_mutates_row_and_scrubs_credentials(): + processed = ProcessedDocument(document_id="doc-1", text_blocks=[TextBlock(text="hello")]) + chunks = [Chunk(id="c1", text="hello", partition="p1")] + row = {"processed_document": processed, "partition": "p1", "api_key": "secret"} + + result = await chunk_stage(row, FakeChunker(chunks)) + + assert result is row + assert row["chunks"] == chunks + assert row["stage"] == "chunked" + assert "api_key" not in row + + +@pytest.mark.asyncio +async def test_contextualize_stage_uses_filename_and_language(): + chunks = [Chunk(id="c1", text="hello")] + contextualized = [Chunk(id="c1", text="ctx hello", context="ctx")] + contextualizer = FakeContextualizer(contextualized) + row = {"chunks": chunks, "filename": "note.md", "language": "fr", "token": "secret"} + + await contextualize_stage(row, contextualizer) + + assert contextualizer.calls == [(chunks, "note.md", "fr")] + assert row["chunks"] == contextualized + assert row["stage"] == "contextualized" + assert "token" not in row + + +@pytest.mark.asyncio +async def test_embed_stage_attaches_vectors_by_chunk_order(): + chunks = [Chunk(id="c1", text="alpha"), Chunk(id="c2", text="beta")] + embedder = FakeEmbedder([[1.0, 0.0], [0.0, 1.0]]) + row = {"chunks": chunks, "secret": "value"} + + await embed_stage(row, embedder) + + assert embedder.text_batches == [["alpha", "beta"]] + assert [chunk.embedding for chunk in row["chunks"]] == [[1.0, 0.0], [0.0, 1.0]] + assert row["stage"] == "embedded" + assert "secret" not in row + + +@pytest.mark.asyncio +async def test_store_stage_upserts_to_default_collection_with_chunk_partitions(): + chunks = [Chunk(id="c1", text="alpha", embedding=[1.0])] + store = FakeVectorStore(count=1) + row = {"chunks": chunks, "partition": "tenant-a", "credentials": {"token": "secret"}} + + await store_stage(row, store) + + assert store.ensure_calls == [("default", 1)] + assert store.calls == [(chunks, "default")] + assert row["stored_count"] == 1 + assert row["stage"] == "stored" + assert "credentials" not in row + + +@pytest.mark.asyncio +async def test_store_stage_rejects_chunks_without_embeddings(): + chunks = [Chunk(id="c1", text="alpha")] + store = FakeVectorStore(count=0) + row = {"chunks": chunks} + + with pytest.raises(ValueError, match="without embeddings"): + await store_stage(row, store) + + assert store.calls == [] + assert row["stage"] == "store_failed" + + +@pytest.mark.asyncio +async def test_stage_marks_error_and_scrubs_credentials_on_failure(): + row = {"chunks": [Chunk(id="c1", text="alpha")], "password": "secret"} + + with pytest.raises(RuntimeError, match="embed failed"): + await embed_stage(row, FakeEmbedder([], error=RuntimeError("embed failed"))) + + assert row["stage"] == "embed_failed" + assert row["error"] == "embed failed" + assert "password" not in row + + +@pytest.mark.asyncio +async def test_caption_stage_marks_error_and_scrubs_credentials_on_failure(): + processed = ProcessedDocument( + document_id="doc-1", + images=[ImageBlock(image_bytes=b"img", page_number=1)], + ) + row = {"processed_document": processed, "api_key": "secret"} + + with pytest.raises(RuntimeError, match="caption failed"): + await caption_stage(row, FakeVLM([], error=RuntimeError("caption failed"))) + + assert row["stage"] == "caption_failed" + assert row["error"] == "caption failed" + assert "api_key" not in row + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stage_fn", "dependency", "expected_stage", "expected_error"), + [ + ( + caption_stage, + FakeVLM(["unused"]), + "caption_failed", + "caption_stage row must contain a ProcessedDocument under 'processed_document'", + ), + ( + chunk_stage, + FakeChunker([]), + "chunk_failed", + "chunk_stage row must contain a ProcessedDocument under 'processed_document'", + ), + ( + contextualize_stage, + FakeContextualizer([]), + "contextualize_failed", + "contextualize_stage row must contain a list[Chunk] under 'chunks'", + ), + ( + embed_stage, + FakeEmbedder([]), + "embed_failed", + "embed_stage row must contain a list[Chunk] under 'chunks'", + ), + ( + store_stage, + FakeVectorStore(count=0), + "store_failed", + "store_stage row must contain a list[Chunk] under 'chunks'", + ), + ], +) +async def test_stages_mark_error_and_scrub_credentials_on_invalid_input( + stage_fn, + dependency, + expected_stage: str, + expected_error: str, +): + row = {"token": "secret"} + + with pytest.raises(ValueError, match=re.escape(expected_error)): + await stage_fn(row, dependency) + + assert row["stage"] == expected_stage + assert row["error"] == expected_error + assert "token" not in row diff --git a/openrag/services/workers/task_state.py b/openrag/services/workers/task_state.py new file mode 100644 index 000000000..9f2955a44 --- /dev/null +++ b/openrag/services/workers/task_state.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from typing import Any + +import ray + +try: + from config import load_config as _load_config + + _cfg = _load_config() + _POOL_SIZE: int = _cfg.ray.pool_size + _MAX_TASKS_PER_WORKER: int = _cfg.ray.max_tasks_per_worker +except (ImportError, AttributeError) as _cfg_err: + import logging as _logging + + _logging.getLogger(__name__).warning( + "Could not load ray config for TaskStateManager pool info: %s — using defaults", _cfg_err + ) + _POOL_SIZE = 1 + _MAX_TASKS_PER_WORKER = 1 + + +@dataclass +class TaskInfo: + state: str | None = None + error: str | None = None + details: dict[str, Any] = field(default_factory=dict) + object_ref: ray.ObjectRef | None = None + + +@ray.remote(concurrency_groups={"set": 1000, "get": 1000, "queue_info": 1000}) +class TaskStateManager: + def __init__(self) -> None: + self.tasks: dict[str, TaskInfo] = {} + self.user_index: dict[int, set[str]] = {} + self.lock = asyncio.Lock() + + async def _ensure_task(self, task_id: str) -> TaskInfo: + if task_id not in self.tasks: + self.tasks[task_id] = TaskInfo() + return self.tasks[task_id] + + @ray.method(concurrency_group="set") + async def set_state(self, task_id: str, state: str) -> None: + async with self.lock: + info = await self._ensure_task(task_id) + info.state = state + + @ray.method(concurrency_group="set") + async def set_error(self, task_id: str, tb_str: str) -> None: + async with self.lock: + info = await self._ensure_task(task_id) + info.error = tb_str + + @ray.method(concurrency_group="set") + async def set_failed_if_not_cancelled(self, task_id: str, tb_str: str) -> bool: + """Atomically set state to FAILED and record the traceback, unless already CANCELLED.""" + async with self.lock: + info = self.tasks.get(task_id) + if info is None or info.state == "CANCELLED": + return False + info.state = "FAILED" + info.error = tb_str + return True + + @ray.method(concurrency_group="set") + async def set_details( + self, + task_id: str, + *, + file_id: str, + partition: int, + metadata: dict, + user_id: int, + ) -> None: + async with self.lock: + info = await self._ensure_task(task_id) + info.details = { + "file_id": file_id, + "partition": partition, + "metadata": metadata, + "user_id": user_id, + } + self.user_index.setdefault(user_id, set()).add(task_id) + + @ray.method(concurrency_group="set") + async def set_object_ref(self, task_id: str, object_ref: ray.ObjectRef) -> None: + async with self.lock: + info = await self._ensure_task(task_id) + info.object_ref = object_ref + + @ray.method(concurrency_group="get") + async def get_state(self, task_id: str) -> str | None: + async with self.lock: + info = self.tasks.get(task_id) + return info.state if info else None + + @ray.method(concurrency_group="get") + async def get_error(self, task_id: str) -> str | None: + async with self.lock: + info = self.tasks.get(task_id) + return info.error if info else None + + @ray.method(concurrency_group="get") + async def get_details(self, task_id: str) -> dict | None: + async with self.lock: + info = self.tasks.get(task_id) + return info.details if info else None + + @ray.method(concurrency_group="get") + async def get_object_ref(self, task_id: str) -> ray.ObjectRef | None: + async with self.lock: + info = self.tasks.get(task_id) + return info.object_ref if info else None + + @ray.method(concurrency_group="queue_info") + async def get_all_states(self) -> dict[str, str | None]: + async with self.lock: + return {tid: info.state for tid, info in self.tasks.items()} + + @ray.method(concurrency_group="queue_info") + async def get_all_info(self) -> dict[str, dict]: + async with self.lock: + return { + task_id: { + "state": info.state, + "error": info.error, + "details": info.details, + } + for task_id, info in self.tasks.items() + } + + @ray.method(concurrency_group="queue_info") + async def get_all_user_info(self, user_id: int) -> dict[str, dict]: + async with self.lock: + task_ids = self.user_index.get(user_id, set()) + return { + tid: { + "state": self.tasks[tid].state, + "error": self.tasks[tid].error, + "details": self.tasks[tid].details, + } + for tid in task_ids + if tid in self.tasks + } + + @ray.method(concurrency_group="queue_info") + async def get_pool_info(self) -> dict[str, int]: + return { + "pool_size": _POOL_SIZE, + "max_tasks_per_worker": _MAX_TASKS_PER_WORKER, + "total_capacity": _POOL_SIZE * _MAX_TASKS_PER_WORKER, + } + + @ray.method(concurrency_group="queue_info") + async def get_user_pending_task_count(self, user_id: int) -> int: + async with self.lock: + task_ids = self.user_index.get(user_id, set()) + pending_states = {"QUEUED", "SERIALIZING", "CHUNKING", "INSERTING"} + return sum(1 for tid in task_ids if (info := self.tasks.get(tid)) and info.state in pending_states) + + +__all__ = ["TaskInfo", "TaskStateManager"] diff --git a/openrag/services/workers/test_batch_ingest.py b/openrag/services/workers/test_batch_ingest.py new file mode 100644 index 000000000..bcf23cc7d --- /dev/null +++ b/openrag/services/workers/test_batch_ingest.py @@ -0,0 +1,388 @@ +from __future__ import annotations + +import asyncio + +import pytest +from core.models.chunk import Chunk +from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock +from services.workers.batch_ingest import ingest_batch +from services.workers.pipeline_builder import build_indexing_pipeline +from services.workers.result_aggregation import aggregate_batch_results + +# --------------------------------------------------------------------------- +# Fakes (minimal — no ABC inheritance needed for these tests) +# --------------------------------------------------------------------------- + + +class FakeParser: + def __init__(self, processed: ProcessedDocument) -> None: + self.processed = processed + + async def parse(self, document: Document) -> ProcessedDocument: + return self.processed + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + +class _FailParser: + def __init__(self, error: Exception) -> None: + self.error = error + + async def parse(self, document: Document) -> ProcessedDocument: + raise self.error + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + +class FakeChunker: + def __init__(self, chunks: list[Chunk]) -> None: + self.chunks = chunks + + def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]: + return self.chunks + + +class FakeEmbedder: + def __init__(self, vectors: list[list[float]]) -> None: + self.vectors = vectors + + async def embed(self, texts: list[str]) -> list[list[float]]: + return self.vectors[: len(texts)] + + +class FakeVectorStore: + def __init__(self) -> None: + self.calls: list[tuple] = [] + self.ensure_calls: list[tuple[str, int]] = [] + + async def upsert(self, chunks: list[Chunk], collection: str = "default") -> int: + self.calls.append((chunks, collection)) + return len(chunks) + + async def ensure_collection(self, name: str, dimension: int, **kwargs) -> None: + self.ensure_calls.append((name, dimension)) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_doc(filename: str = "doc.txt") -> Document: + return Document(filename=filename, text="hello", partition="p") + + +def _make_processed(doc: Document) -> ProcessedDocument: + return ProcessedDocument(document_id=doc.id, text_blocks=[TextBlock(text="hello")]) + + +def _make_chunk(doc: Document) -> Chunk: + return Chunk(id=doc.id, text="hello", partition="p") + + +def _make_pipeline(doc: Document, *, fail: bool = False) -> tuple: + processed = _make_processed(doc) + chunk = _make_chunk(doc) + parser = _FailParser(RuntimeError("parse error")) if fail else FakeParser(processed) + chunker = FakeChunker([chunk]) + embedder = FakeEmbedder([[1.0]]) + vector_store = FakeVectorStore() + pipeline = build_indexing_pipeline( + parser=parser, + chunker=chunker, + embedder=embedder, + vector_store=vector_store, + ) + return pipeline, vector_store + + +# --------------------------------------------------------------------------- +# Tests — ingest_batch +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_ingest_batch_all_succeed(): + docs = [_make_doc(f"doc{i}.txt") for i in range(3)] + pipeline, store = _make_pipeline(docs[0]) + rows = [{"document": doc, "partition": "p"} for doc in docs] + + result = await ingest_batch(pipeline, rows) + + assert len(result) == 3 + assert all(r["stage"] == "stored" for r in result) + assert all(r["stored_count"] == 1 for r in result) + + +@pytest.mark.asyncio +async def test_ingest_batch_partial_failure_does_not_abort_others(): + docs = [_make_doc(f"doc{i}.txt") for i in range(3)] + processed = _make_processed(docs[0]) + chunk = _make_chunk(docs[0]) + + # Parser fails on the *second* call only + call_count = 0 + + class SelectiveParser: + async def parse(self, document: Document) -> ProcessedDocument: + nonlocal call_count + call_count += 1 + if call_count == 2: + raise RuntimeError("second parse fails") + return processed + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + pipeline = build_indexing_pipeline( + parser=SelectiveParser(), + chunker=FakeChunker([chunk]), + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + ) + rows = [{"document": doc, "partition": "p"} for doc in docs] + + result = await ingest_batch(pipeline, rows) + + stages = [r["stage"] for r in result] + assert stages.count("stored") == 2 + assert stages.count("parse_failed") == 1 + failed_row = next(r for r in result if r["stage"] == "parse_failed") + assert failed_row["error"] == "second parse fails" + + +@pytest.mark.asyncio +async def test_ingest_batch_all_fail(): + docs = [_make_doc(f"doc{i}.txt") for i in range(2)] + pipeline, _ = _make_pipeline(docs[0], fail=True) + rows = [{"document": doc, "partition": "p"} for doc in docs] + + result = await ingest_batch(pipeline, rows) + + assert all(r["stage"] == "parse_failed" for r in result) + + +@pytest.mark.asyncio +async def test_ingest_batch_returns_same_row_objects(): + doc = _make_doc() + pipeline, _ = _make_pipeline(doc) + row: dict = {"document": doc, "partition": "p"} + + result = await ingest_batch(pipeline, [row]) + + assert result[0] is row + + +@pytest.mark.asyncio +async def test_ingest_batch_concurrency_cap_limits_parallelism(): + """At most *concurrency* rows run simultaneously.""" + active: list[int] = [] + peak: list[int] = [] + + class SlowParser: + async def parse(self, document: Document) -> ProcessedDocument: + active.append(1) + peak.append(len(active)) + await asyncio.sleep(0) + active.pop() + return ProcessedDocument(document_id=document.id, text_blocks=[TextBlock(text="x")]) + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + doc = _make_doc() + chunk = _make_chunk(doc) + pipeline = build_indexing_pipeline( + parser=SlowParser(), + chunker=FakeChunker([chunk]), + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + ) + rows = [{"document": _make_doc(f"d{i}.txt"), "partition": "p"} for i in range(5)] + + await ingest_batch(pipeline, rows, concurrency=2) + + assert max(peak) <= 2 + + +@pytest.mark.asyncio +async def test_ingest_batch_empty_input_returns_empty_list(): + pipeline, _ = _make_pipeline(_make_doc()) + result = await ingest_batch(pipeline, []) + assert result == [] + + +@pytest.mark.asyncio +async def test_ingest_batch_concurrency_zero_raises(): + pipeline, _ = _make_pipeline(_make_doc()) + with pytest.raises(ValueError, match="concurrency must be >= 1"): + await ingest_batch(pipeline, [], concurrency=0) + + +@pytest.mark.asyncio +async def test_ingest_batch_concurrency_negative_raises(): + pipeline, _ = _make_pipeline(_make_doc()) + with pytest.raises(ValueError, match="concurrency must be >= 1"): + await ingest_batch(pipeline, [], concurrency=-1) + + +# --------------------------------------------------------------------------- +# Tests — aggregate_batch_results +# --------------------------------------------------------------------------- + + +def test_aggregate_all_succeeded(): + rows = [ + {"stage": "stored", "stored_count": 3}, + {"stage": "stored", "stored_count": 2}, + ] + summary = aggregate_batch_results(rows) + + assert summary.total == 2 + assert summary.succeeded == 2 + assert summary.failed == 0 + assert summary.stored_count == 5 + assert summary.failures == () + assert summary.success_rate == 1.0 + + +def test_aggregate_mixed_results(): + rows = [ + {"stage": "stored", "stored_count": 4}, + {"stage": "embed_failed", "error": "timeout"}, + {"stage": "chunk_failed", "error": "empty doc"}, + ] + summary = aggregate_batch_results(rows) + + assert summary.total == 3 + assert summary.succeeded == 1 + assert summary.failed == 2 + assert summary.stored_count == 4 + assert len(summary.failures) == 2 + assert {f.stage for f in summary.failures} == {"embed_failed", "chunk_failed"} + + +@pytest.mark.asyncio +async def test_e2e_multiple_documents_pipeline_and_summary(): + docs = [ + Document(filename="ok-1.txt", text="alpha", partition="tenant-a"), + Document(filename="bad.txt", text="bad", partition="tenant-a"), + Document(filename="ok-2.txt", text="beta", partition="tenant-b"), + ] + + class TextParser: + async def parse(self, document: Document) -> ProcessedDocument: + if document.filename == "bad.txt": + raise RuntimeError("unsupported document") + return ProcessedDocument( + document_id=document.id, + text_blocks=[TextBlock(text=document.text or "")], + ) + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + class TextChunker: + def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]: + return [ + Chunk( + id=f"{document.document_id}-chunk", + document_id=document.document_id, + text=document.text_blocks[0].text, + partition=partition, + ) + ] + + store = FakeVectorStore() + pipeline = build_indexing_pipeline( + parser=TextParser(), + chunker=TextChunker(), + embedder=FakeEmbedder([[1.0]]), + vector_store=store, + ) + rows = [{"document": doc, "partition": doc.partition} for doc in docs] + + processed_rows = await ingest_batch(pipeline, rows, concurrency=2) + summary = aggregate_batch_results(processed_rows) + + assert [row["stage"] for row in processed_rows] == ["stored", "parse_failed", "stored"] + assert summary.total == 3 + assert summary.succeeded == 2 + assert summary.failed == 1 + assert summary.stored_count == 2 + assert summary.success_rate == pytest.approx(2 / 3) + assert summary.failures[0].stage == "parse_failed" + assert summary.failures[0].error == "unsupported document" + assert store.calls[0][1] == "default" + assert store.calls[1][1] == "default" + assert processed_rows[0]["chunks"][0].partition == "tenant-a" + assert processed_rows[2]["chunks"][0].partition == "tenant-b" + + +def test_aggregate_all_failed(): + rows = [ + {"stage": "parse_failed", "error": "bad file"}, + {"stage": "parse_failed", "error": "another bad file"}, + ] + summary = aggregate_batch_results(rows) + + assert summary.succeeded == 0 + assert summary.failed == 2 + assert summary.stored_count == 0 + assert summary.success_rate == 0.0 + + +def test_aggregate_empty_input(): + summary = aggregate_batch_results([]) + + assert summary.total == 0 + assert summary.succeeded == 0 + assert summary.failed == 0 + assert summary.stored_count == 0 + assert summary.success_rate == 0.0 + + +def test_aggregate_missing_stage_counts_as_failure(): + rows = [{"stored_count": 1}] # no "stage" key + summary = aggregate_batch_results(rows) + + assert summary.failed == 1 + assert summary.succeeded == 0 + + +@pytest.mark.asyncio +async def test_ingest_then_aggregate_round_trip(): + docs = [_make_doc(f"doc{i}.txt") for i in range(4)] + processed = _make_processed(docs[0]) + + call_count = 0 + + class SelectiveParser: + async def parse(self, document: Document) -> ProcessedDocument: + nonlocal call_count + call_count += 1 + if call_count == 3: + raise RuntimeError("oops") + return processed + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + chunk = _make_chunk(docs[0]) + pipeline = build_indexing_pipeline( + parser=SelectiveParser(), + chunker=FakeChunker([chunk]), + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + ) + rows = [{"document": doc, "partition": "p"} for doc in docs] + + result = await ingest_batch(pipeline, rows) + summary = aggregate_batch_results(result) + + assert summary.total == 4 + assert summary.succeeded == 3 + assert summary.failed == 1 + assert summary.stored_count == 3 diff --git a/openrag/services/workers/test_dispatcher.py b/openrag/services/workers/test_dispatcher.py new file mode 100644 index 000000000..e13d1ff54 --- /dev/null +++ b/openrag/services/workers/test_dispatcher.py @@ -0,0 +1,231 @@ +from __future__ import annotations + +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +def _remote_mock(return_value: Any = None) -> MagicMock: + method = MagicMock() + method.remote = AsyncMock(return_value=return_value) + return method + + +def _pool_with_ref(ref: object) -> MagicMock: + pool = MagicMock() + pool.process_file = MagicMock() + pool.process_file.remote = MagicMock(return_value=ref) + return pool + + +def _vector_store() -> MagicMock: + store = MagicMock() + store.query_ids_by_filter = AsyncMock(return_value=["1", "2"]) + store.query_chunks_by_filter = AsyncMock( + return_value=[ + { + "_id": 1, + "text": "hello", + "vector": [0.1, 0.2], + "file_id": "file-1", + "partition": "tenant-a", + "page": 1, + "section_id": 11, + "title": "old", + } + ] + ) + store.delete = AsyncMock() + store.upsert_entities = AsyncMock() + store.insert_entities = AsyncMock() + return store + + +def _document_repo() -> MagicMock: + repo = MagicMock() + repo.remove_file_from_partition = AsyncMock() + repo.update_file_metadata_in_db = AsyncMock(return_value=True) + repo.add_file_to_partition = AsyncMock(return_value=True) + return repo + + +def _workspace_repo() -> MagicMock: + repo = MagicMock() + repo.remove_file_from_all_workspaces = AsyncMock() + return repo + + +def _task_state_manager() -> MagicMock: + tsm = MagicMock() + tsm.set_state = _remote_mock() + tsm.set_details = _remote_mock() + tsm.set_object_ref = _remote_mock() + tsm.get_state = _remote_mock("SERIALIZING") + tsm.get_error = _remote_mock("traceback") + tsm.get_object_ref = _remote_mock({"ref": object()}) + return tsm + + +def test_from_ray_namespace_does_not_require_legacy_indexer_actor() -> None: + from services.workers.dispatcher import WorkerDispatcher, from_ray_namespace + + tsm = _task_state_manager() + pool = _pool_with_ref(object()) + + def fake_get_actor(name: str, namespace: str): + assert namespace == "openrag" + if name == "TaskStateManager": + return tsm + raise AssertionError(f"unexpected eager actor lookup: {name}") + + with ( + patch("ray.get_actor", side_effect=fake_get_actor), + patch("services.workers.indexer_pool.build_indexer_pool", return_value=pool), + ): + dispatcher = from_ray_namespace( + vector_store=_vector_store(), + document_repo=_document_repo(), + workspace_repo=_workspace_repo(), + collection="default", + ) + + assert isinstance(dispatcher, WorkerDispatcher) + + +@pytest.mark.asyncio +async def test_dispatch_indexing_queues_worker_pool_task_and_records_ref() -> None: + from services.workers.dispatcher import WorkerDispatcher + + ref = object() + pool = _pool_with_ref(ref) + tsm = _task_state_manager() + dispatcher = WorkerDispatcher( + pool=pool, + task_state_manager=tsm, + vector_store=_vector_store(), + document_repo=_document_repo(), + workspace_repo=_workspace_repo(), + collection="default", + ) + + with patch("services.workers.dispatcher.uuid") as mock_uuid: + mock_uuid.uuid4.return_value.hex = "task-1" + task_id = await dispatcher.dispatch_indexing( + path="/data/report.txt", + metadata={"file_id": "file-1", "source": "/data/report.txt", "filename": "report.txt"}, + partition="tenant-a", + user={"id": 42}, + workspace_ids=["ws-1"], + replace=True, + ) + + assert task_id == "task-1" + tsm.set_state.remote.assert_called_once_with("task-1", "QUEUED") + tsm.set_details.remote.assert_called_once_with( + "task-1", + file_id="file-1", + partition="tenant-a", + metadata={"filename": "report.txt"}, + user_id=42, + ) + pool.process_file.remote.assert_called_once_with( + task_id="task-1", + path="/data/report.txt", + metadata={"file_id": "file-1", "source": "/data/report.txt", "filename": "report.txt"}, + partition="tenant-a", + user={"id": 42}, + workspace_ids=["ws-1"], + replace=True, + ) + tsm.set_object_ref.remote.assert_called_once_with("task-1", {"ref": ref}) + + +@pytest.mark.asyncio +async def test_worker_dispatcher_mutates_files_without_legacy_indexer() -> None: + from services.workers.dispatcher import WorkerDispatcher + + vector_store = _vector_store() + document_repo = _document_repo() + workspace_repo = _workspace_repo() + dispatcher = WorkerDispatcher( + pool=_pool_with_ref(object()), + task_state_manager=_task_state_manager(), + vector_store=vector_store, + document_repo=document_repo, + workspace_repo=workspace_repo, + collection="default", + ) + + await dispatcher.delete_file("file-1", "tenant-a") + await dispatcher.update_file_metadata("file-1", {"title": "new"}, "tenant-a", user={"id": 7}) + await dispatcher.copy_file("file-1", {"file_id": "copy-1", "partition": "tenant-b"}, "tenant-b", user=None) + + vector_store.query_ids_by_filter.assert_called_once_with("default", {"partition": "tenant-a", "file_id": "file-1"}) + vector_store.delete.assert_called_once_with(["1", "2"], "default") + workspace_repo.remove_file_from_all_workspaces.assert_called_once_with("file-1", "tenant-a") + document_repo.remove_file_from_partition.assert_called_once_with(file_id="file-1", partition="tenant-a") + document_repo.update_file_metadata_in_db.assert_called_once_with( + "file-1", + "tenant-a", + {"file_id": "file-1", "partition": "tenant-a", "title": "new"}, + ) + document_repo.add_file_to_partition.assert_called_once_with( + file_id="copy-1", + partition="tenant-b", + file_metadata={"file_id": "copy-1", "partition": "tenant-b", "title": "old"}, + user_id=None, + relationship_id=None, + parent_id=None, + ) + vector_store.upsert_entities.assert_awaited_once() + vector_store.insert_entities.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_cancel_task_uses_stored_pool_object_ref() -> None: + from services.workers.dispatcher import WorkerDispatcher + + ref = object() + tsm = _task_state_manager() + tsm.get_object_ref.remote = AsyncMock(return_value={"ref": ref}) + tsm.get_state.remote = AsyncMock(return_value="SERIALIZING") + dispatcher = WorkerDispatcher( + pool=_pool_with_ref(object()), + task_state_manager=tsm, + vector_store=_vector_store(), + document_repo=_document_repo(), + workspace_repo=_workspace_repo(), + collection="default", + ) + + with patch("ray.cancel") as cancel: + result = await dispatcher.cancel_task("task-1") + + assert result is True + cancel.assert_called_once_with(ref, recursive=True) + tsm.set_state.remote.assert_called_once_with("task-1", "CANCELLED") + + +@pytest.mark.asyncio +async def test_cancel_task_marks_cancelled_even_if_worker_finished_first() -> None: + from services.workers.dispatcher import WorkerDispatcher + + ref = object() + tsm = _task_state_manager() + tsm.get_object_ref.remote = AsyncMock(return_value={"ref": ref}) + tsm.get_state.remote = AsyncMock(return_value="COMPLETED") + dispatcher = WorkerDispatcher( + pool=_pool_with_ref(object()), + task_state_manager=tsm, + vector_store=_vector_store(), + document_repo=_document_repo(), + workspace_repo=_workspace_repo(), + collection="default", + ) + + with patch("ray.cancel"): + result = await dispatcher.cancel_task("task-1") + + assert result is True + tsm.set_state.remote.assert_called_once_with("task-1", "CANCELLED") diff --git a/openrag/services/workers/test_indexer_pool.py b/openrag/services/workers/test_indexer_pool.py new file mode 100644 index 000000000..127d95acc --- /dev/null +++ b/openrag/services/workers/test_indexer_pool.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +import pytest + + +class _NativeChunker: + def chunk(self, document, partition: str = "default"): + return [] + + +class _LegacyChunker: + def __init__(self) -> None: + self._core_splitter = _NativeChunker() + + +class _BrokenLegacyChunker: + pass + + +def test_build_chunker_returns_native_chunker(monkeypatch: pytest.MonkeyPatch) -> None: + from components.indexer.chunker.chunker import ChunkerFactory + from services.workers.indexer_pool import _build_chunker + + native = _NativeChunker() + monkeypatch.setattr(ChunkerFactory, "create_chunker", staticmethod(lambda _cfg: native)) + + assert _build_chunker(object()) is native + + +def test_build_chunker_unwraps_legacy_core_splitter(monkeypatch: pytest.MonkeyPatch) -> None: + from components.indexer.chunker.chunker import ChunkerFactory + from services.workers.indexer_pool import _build_chunker + + legacy = _LegacyChunker() + monkeypatch.setattr(ChunkerFactory, "create_chunker", staticmethod(lambda _cfg: legacy)) + + assert _build_chunker(object()) is legacy._core_splitter + + +def test_build_chunker_rejects_invalid_legacy_chunker(monkeypatch: pytest.MonkeyPatch) -> None: + from components.indexer.chunker.chunker import ChunkerFactory + from services.workers.indexer_pool import _build_chunker + + monkeypatch.setattr(ChunkerFactory, "create_chunker", staticmethod(lambda _cfg: _BrokenLegacyChunker())) + + with pytest.raises(TypeError, match="chunk"): + _build_chunker(object()) diff --git a/openrag/services/workers/test_indexer_worker.py b/openrag/services/workers/test_indexer_worker.py new file mode 100644 index 000000000..2785e6a4f --- /dev/null +++ b/openrag/services/workers/test_indexer_worker.py @@ -0,0 +1,327 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +from core.models.chunk import Chunk +from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock +from services.workers.indexer_actor import IndexerWorker, _load_document +from services.workers.pipeline_builder import build_indexing_pipeline + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +class FakeParser: + def __init__(self, processed: ProcessedDocument) -> None: + self.processed = processed + self.calls: list[Document] = [] + + async def parse(self, document: Document) -> ProcessedDocument: + self.calls.append(document) + return self.processed + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + +class FakeChunker: + def __init__(self, chunks: list[Chunk]) -> None: + self.chunks = chunks + + def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]: + return self.chunks + + +class FakeEmbedder: + def __init__(self) -> None: + self.calls: list[list[str]] = [] + + async def embed(self, texts: list[str]) -> list[list[float]]: + self.calls.append(texts) + return [[1.0] for _ in texts] + + +class FakeVectorStore: + def __init__(self) -> None: + self.calls: list[tuple] = [] + self.ensure_calls: list[tuple[str, int]] = [] + + async def upsert(self, chunks: list[Chunk], collection: str = "default") -> int: + self.calls.append((chunks, collection)) + return len(chunks) + + async def ensure_collection(self, name: str, dimension: int, **kwargs: Any) -> None: + self.ensure_calls.append((name, dimension)) + + +def _fake_tsm() -> MagicMock: + """Task-state-manager mock whose .remote() methods return awaitables.""" + tsm = MagicMock() + tsm.set_state = MagicMock() + tsm.set_state.remote = AsyncMock(return_value=None) + tsm.set_failed_if_not_cancelled = MagicMock() + tsm.set_failed_if_not_cancelled.remote = AsyncMock(return_value=True) + return tsm + + +def _make_pipeline(processed: ProcessedDocument, chunks: list[Chunk]) -> Any: + return build_indexing_pipeline( + parser=FakeParser(processed), + chunker=FakeChunker(chunks), + embedder=FakeEmbedder(), + vector_store=FakeVectorStore(), + ) + + +class FakeDocumentRepo: + def __init__(self) -> None: + self.add_calls: list[dict[str, Any]] = [] + self.update_calls: list[dict[str, Any]] = [] + + async def add_file_to_partition(self, **kwargs: Any) -> bool: + self.add_calls.append(kwargs) + return True + + async def update_file_in_partition(self, **kwargs: Any) -> bool: + self.update_calls.append(kwargs) + return True + + +# --------------------------------------------------------------------------- +# Tests — _load_document helper +# --------------------------------------------------------------------------- + + +def test_load_document_reads_bytes_and_detects_type(tmp_path: Path) -> None: + p = tmp_path / "report.pdf" + p.write_bytes(b"%PDF-1.4") + doc = _load_document(str(p), {"file_id": "fid-1"}, "tenant-a") + + assert doc.raw_bytes == b"%PDF-1.4" + assert doc.content_type == DocumentType.PDF + assert doc.partition == "tenant-a" + assert doc.filename == "fid-1" + + +def test_load_document_falls_back_to_filename_when_no_file_id(tmp_path: Path) -> None: + p = tmp_path / "note.txt" + p.write_bytes(b"hi") + doc = _load_document(str(p), {}, "p") + + assert doc.filename == "note.txt" + + +# --------------------------------------------------------------------------- +# Tests — IndexerWorker.process_file +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_process_file_success_sets_state_and_returns_count(tmp_path: Path) -> None: + path = tmp_path / "doc.txt" + path.write_bytes(b"content") + processed = ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="content")]) + chunks = [Chunk(id="c1", text="content", partition="p")] + pipeline = _make_pipeline(processed, chunks) + tsm = _fake_tsm() + + worker = IndexerWorker(pipeline=pipeline, task_state_manager=tsm) + result = await worker.process_file( + task_id="t1", + path=str(path), + metadata={"file_id": "f1"}, + partition="p", + ) + + assert result["stored_count"] == 1 + assert result["stage"] == "stored" + state_calls = [call.args for call in tsm.set_state.remote.call_args_list] + assert ("t1", "SERIALIZING") in state_calls + assert ("t1", "COMPLETED") in state_calls + tsm.set_failed_if_not_cancelled.remote.assert_not_called() + + +@pytest.mark.asyncio +async def test_process_file_pipeline_failure_sets_failed_and_reraises(tmp_path: Path) -> None: + path = tmp_path / "bad.txt" + path.write_bytes(b"x") + + class BrokenParser: + async def parse(self, document: Document) -> ProcessedDocument: + raise RuntimeError("parser exploded") + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + pipeline = build_indexing_pipeline( + parser=BrokenParser(), + chunker=FakeChunker([]), + embedder=FakeEmbedder(), + vector_store=FakeVectorStore(), + ) + tsm = _fake_tsm() + worker = IndexerWorker(pipeline=pipeline, task_state_manager=tsm) + + with pytest.raises(RuntimeError, match="parser exploded"): + await worker.process_file( + task_id="t2", + path=str(path), + metadata={}, + partition="p", + ) + + tsm.set_state.remote.assert_called_once_with("t2", "SERIALIZING") + tsm.set_failed_if_not_cancelled.remote.assert_called_once() + call_args = tsm.set_failed_if_not_cancelled.remote.call_args + assert call_args.args[0] == "t2" + assert "parser exploded" in call_args.args[1] + + +@pytest.mark.asyncio +async def test_process_file_missing_path_raises_and_sets_failed() -> None: + processed = ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="x")]) + pipeline = _make_pipeline(processed, [Chunk(id="c1", text="x")]) + tsm = _fake_tsm() + worker = IndexerWorker(pipeline=pipeline, task_state_manager=tsm) + + with pytest.raises(FileNotFoundError): + await worker.process_file( + task_id="t3", + path="/nonexistent/file.txt", + metadata={}, + partition="p", + ) + + tsm.set_failed_if_not_cancelled.remote.assert_called_once() + + +@pytest.mark.asyncio +async def test_process_file_passes_partition_and_filename_to_row(tmp_path: Path) -> None: + path = tmp_path / "note.txt" + path.write_bytes(b"hello") + + seen_partitions: list[str] = [] + + class TrackingChunker: + def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]: + seen_partitions.append(partition) + return [Chunk(id="c1", text="hello", partition=partition)] + + pipeline = build_indexing_pipeline( + parser=FakeParser(ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="hello")])), + chunker=TrackingChunker(), + embedder=FakeEmbedder(), + vector_store=FakeVectorStore(), + ) + tsm = _fake_tsm() + worker = IndexerWorker(pipeline=pipeline, task_state_manager=tsm) + await worker.process_file( + task_id="t4", + path=str(path), + metadata={"file_id": "fid"}, + partition="tenant-b", + ) + + assert seen_partitions == ["tenant-b"] + + +@pytest.mark.asyncio +async def test_process_file_creates_catalog_record_after_successful_pipeline(tmp_path: Path) -> None: + path = tmp_path / "doc.txt" + path.write_bytes(b"content") + processed = ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="content")]) + chunks = [Chunk(id="c1", text="content", partition="p")] + repo = FakeDocumentRepo() + worker = IndexerWorker( + pipeline=_make_pipeline(processed, chunks), + task_state_manager=_fake_tsm(), + document_repo=repo, + ) + + await worker.process_file( + task_id="t-new", + path=str(path), + metadata={"file_id": "f1", "relationship_id": "rel", "parent_id": "parent"}, + partition="p", + user={"id": 42}, + ) + + assert repo.add_calls == [ + { + "file_id": "f1", + "partition": "p", + "file_metadata": {"file_id": "f1", "relationship_id": "rel", "parent_id": "parent"}, + "user_id": 42, + "relationship_id": "rel", + "parent_id": "parent", + } + ] + assert repo.update_calls == [] + + +@pytest.mark.asyncio +async def test_process_file_updates_catalog_record_on_replace(tmp_path: Path) -> None: + path = tmp_path / "doc.txt" + path.write_bytes(b"content") + processed = ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="content")]) + chunks = [Chunk(id="c1", text="content", partition="p")] + repo = FakeDocumentRepo() + worker = IndexerWorker( + pipeline=_make_pipeline(processed, chunks), + task_state_manager=_fake_tsm(), + document_repo=repo, + ) + + await worker.process_file( + task_id="t-replace", + path=str(path), + metadata={"file_id": "f1"}, + partition="p", + replace=True, + ) + + assert repo.update_calls == [ + { + "file_id": "f1", + "partition": "p", + "file_metadata": {"file_id": "f1"}, + "relationship_id": None, + "parent_id": None, + } + ] + assert repo.add_calls == [] + + +@pytest.mark.asyncio +async def test_process_file_catalog_failure_sets_failed_state(tmp_path: Path) -> None: + path = tmp_path / "doc.txt" + path.write_bytes(b"content") + processed = ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="content")]) + chunks = [Chunk(id="c1", text="content", partition="p")] + tsm = _fake_tsm() + + class BrokenRepo: + async def add_file_to_partition(self, **kwargs: Any) -> bool: + raise RuntimeError("pg down") + + worker = IndexerWorker( + pipeline=_make_pipeline(processed, chunks), + task_state_manager=tsm, + document_repo=BrokenRepo(), + ) + + with pytest.raises(RuntimeError, match="pg down"): + await worker.process_file( + task_id="t-fail", + path=str(path), + metadata={"file_id": "f1"}, + partition="p", + ) + + tsm.set_failed_if_not_cancelled.remote.assert_called_once() + completed_calls = [call for call in tsm.set_state.remote.call_args_list if call.args == ("t-fail", "COMPLETED")] + assert completed_calls == [] diff --git a/openrag/services/workers/test_pipeline_builder.py b/openrag/services/workers/test_pipeline_builder.py new file mode 100644 index 000000000..c6366b0a9 --- /dev/null +++ b/openrag/services/workers/test_pipeline_builder.py @@ -0,0 +1,107 @@ +import pytest +from core.models.chunk import Chunk +from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock +from services.workers.pipeline_builder import build_indexing_pipeline + + +class FakeParser: + def __init__(self, processed: ProcessedDocument) -> None: + self.processed = processed + self.calls: list[Document] = [] + + async def parse(self, document: Document) -> ProcessedDocument: + self.calls.append(document) + return self.processed + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + +class FakeChunker: + def __init__(self, chunks: list[Chunk], error: Exception | None = None) -> None: + self.chunks = chunks + self.error = error + self.calls: list[tuple[ProcessedDocument, str]] = [] + + def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]: + self.calls.append((document, partition)) + if self.error is not None: + raise self.error + return self.chunks + + +class FakeEmbedder: + def __init__(self, vectors: list[list[float]]) -> None: + self.vectors = vectors + self.calls: list[list[str]] = [] + + async def embed(self, texts: list[str]) -> list[list[float]]: + self.calls.append(texts) + return self.vectors + + +class FakeVectorStore: + def __init__(self) -> None: + self.calls: list[tuple[list[Chunk], str]] = [] + self.ensure_calls: list[tuple[str, int]] = [] + + async def upsert(self, chunks: list[Chunk], collection: str = "default") -> int: + self.calls.append((chunks, collection)) + return len(chunks) + + async def ensure_collection(self, name: str, dimension: int, **kwargs) -> None: + self.ensure_calls.append((name, dimension)) + + +@pytest.mark.asyncio +async def test_pipeline_runs_required_stages_in_order_and_keeps_row_object(): + document = Document(filename="note.txt", text="hello", partition="tenant-a") + processed = ProcessedDocument(document_id=document.id, text_blocks=[TextBlock(text="hello")]) + chunks = [Chunk(id="c1", text="hello", partition="tenant-a")] + parser = FakeParser(processed) + chunker = FakeChunker(chunks) + embedder = FakeEmbedder([[1.0, 0.0]]) + vector_store = FakeVectorStore() + pipeline = build_indexing_pipeline( + parser=parser, + chunker=chunker, + embedder=embedder, + vector_store=vector_store, + ) + row = {"document": document, "partition": "tenant-a", "token": "secret"} + + result = await pipeline.run(row) + + assert result is row + assert parser.calls == [document] + assert chunker.calls == [(processed, "tenant-a")] + assert embedder.calls == [["hello"]] + assert vector_store.ensure_calls == [("default", 2)] + assert vector_store.calls == [(row["chunks"], "default")] + assert row["stage"] == "stored" + assert row["stored_count"] == 1 + assert row["chunks"][0].embedding == [1.0, 0.0] + assert "token" not in row + + +@pytest.mark.asyncio +async def test_pipeline_stops_before_later_stages_when_a_stage_fails(): + document = Document(filename="note.txt", text="hello", partition="tenant-a") + processed = ProcessedDocument(document_id=document.id, text_blocks=[TextBlock(text="hello")]) + chunker = FakeChunker([], error=RuntimeError("chunk failed")) + vector_store = FakeVectorStore() + pipeline = build_indexing_pipeline( + parser=FakeParser(processed), + chunker=chunker, + embedder=FakeEmbedder([]), + vector_store=vector_store, + ) + row = {"document": document, "password": "secret"} + + with pytest.raises(RuntimeError, match="chunk failed"): + await pipeline.run(row) + + assert row["stage"] == "chunk_failed" + assert row["error"] == "chunk failed" + assert vector_store.calls == [] + assert "password" not in row diff --git a/openrag/test_token_validation.py b/openrag/test_token_validation.py index e3fcade1d..06eab5661 100644 --- a/openrag/test_token_validation.py +++ b/openrag/test_token_validation.py @@ -5,25 +5,17 @@ import resolution, causing a circular import. """ -import sys -from types import ModuleType -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest # Prevent Ray from scanning the working directory (which may contain -# permission-restricted folders like db/) and stub RagPipeline to -# avoid heavy initialization during test collection. +# permission-restricted folders like db/). import ray # noqa: E402 if not ray.is_initialized(): ray.init(runtime_env={"working_dir": None}, ignore_reinit_error=True) -if "components.pipeline" not in sys.modules: - _stub = ModuleType("components.pipeline") - _stub.RagPipeline = MagicMock() - sys.modules["components.pipeline"] = _stub - from models.openai import OpenAIChatCompletionRequest, OpenAICompletionRequest # noqa: E402 from routers.openai import validate_tokens_limit # noqa: E402 diff --git a/openrag/utils/exceptions/__init__.py b/openrag/utils/exceptions/__init__.py index 9b5ed21c9..7ed1e2a97 100644 --- a/openrag/utils/exceptions/__init__.py +++ b/openrag/utils/exceptions/__init__.py @@ -1 +1,3 @@ -from .base import * +# Re-export from canonical location for backward compatibility. +# New code should import from `core.utils.exceptions` directly. +from core.utils.exceptions import * # noqa: F401,F403 diff --git a/openrag/utils/exceptions/base.py b/openrag/utils/exceptions/base.py index 32a3be8fd..c0fed1294 100644 --- a/openrag/utils/exceptions/base.py +++ b/openrag/utils/exceptions/base.py @@ -1,39 +1,7 @@ -from fastapi import status - - -class OpenRAGError(Exception): - """Base class for all OpenRAG exceptions.""" - - def __init__( - self, - message: str, - code: str, - status_code: int = status.HTTP_500_INTERNAL_SERVER_ERROR, - **kwargs, - ): - self.message = message - self.code = code - self.status_code = status_code - self.extra = kwargs or {} - super().__init__(f"{self.code}: {self.message}") - - def to_dict(self) -> dict: - return { - "detail": f"[{self.code}]: {self.message}", - "extra": self.extra, - } - - -# Subclass exceptions for specific error types -class EmbeddingError(OpenRAGError): - """Base exception for all embedding-related errors.""" - - def __init__(self, message, code, status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, **kwargs): - super().__init__(message, code, status_code, **kwargs) - - -class VDBError(OpenRAGError): - """Base exception for all vector database-related errors.""" - - def __init__(self, message, code, status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, **kwargs): - super().__init__(message, code, status_code, **kwargs) +# Re-export from canonical location for backward compatibility. +# New code should import from `core.utils.exceptions` directly. +from core.utils.exceptions import ( # noqa: F401 + EmbeddingError, + OpenRAGError, + VDBError, +) diff --git a/openrag/utils/exceptions/embeddings.py b/openrag/utils/exceptions/embeddings.py index 7f8b808a7..84be79cef 100644 --- a/openrag/utils/exceptions/embeddings.py +++ b/openrag/utils/exceptions/embeddings.py @@ -1,32 +1,7 @@ -from .base import EmbeddingError - - -class EmbeddingAPIError(EmbeddingError): - """Raised when there's an API error with the embedding provider.""" - - def __init__(self, message: str, **kwargs): - super().__init__( - message=message, - code="EMBEDDING_API_ERROR", - status_code=500, - **kwargs, - ) - - -class EmbeddingResponseError(EmbeddingError): - """Raised when the response from the embedding provider is invalid or unexpected.""" - - def __init__(self, message: str, **kwargs): - super().__init__(message=message, code="EMBEDDING_RESPONSE_ERROR", status_code=422, **kwargs) - - -class UnexpectedEmbeddingError(EmbeddingError): - """Raised for unexpected errors in embedding operations.""" - - def __init__(self, message: str, **kwargs): - super().__init__( - message=message, - code="EMBEDDING_UNEXPECTED_ERROR", - status_code=500, - **kwargs, - ) +# Re-export from canonical location for backward compatibility. +# New code should import from `core.utils.exceptions` directly. +from core.utils.exceptions import ( # noqa: F401 + EmbeddingAPIError, + EmbeddingResponseError, + UnexpectedEmbeddingError, +) diff --git a/openrag/utils/exceptions/vectordb.py b/openrag/utils/exceptions/vectordb.py index e9bf7cbc1..54dd12df1 100644 --- a/openrag/utils/exceptions/vectordb.py +++ b/openrag/utils/exceptions/vectordb.py @@ -1,130 +1,17 @@ -from .base import VDBError - - -class VDBConnectionError(VDBError): - """Raised when connection to vector database fails.""" - - def __init__(self, message: str, **kwargs): - super().__init__( - message=message, - code="VDB_CONNECTION_ERROR", - status_code=503, - **kwargs, - ) - - -class VDBCreateOrLoadCollectionError(VDBError): - """Raised when there's an issue with collection operations.""" - - def __init__(self, message: str, **kwargs): - super().__init__(message=message, code="VDB_COLLECTION_ERROR", status_code=422, **kwargs) - - -class VDBInsertError(VDBError): - """Raised when data insertion fails.""" - - def __init__(self, message: str, status_code: int = 422, **kwargs): - super().__init__(message=message, code="VDB_INSERT_ERROR", status_code=status_code, **kwargs) - - -class VDBFileIDAlreadyExistsError(VDBError): - """Raised when a file already exists in the vector database.""" - - def __init__(self, message: str, **kwargs): - super().__init__(message=message, code="VDB_FILE_ALREADY_EXISTS", status_code=409, **kwargs) - - -class VDBDeleteError(VDBError): - """Raised when data deletion fails.""" - - def __init__( - self, - message: str, - status_code=422, - **kwargs, - ): - super().__init__(message=message, code="VDB_DELETE_ERROR", status_code=status_code, **kwargs) - - -class VDBSearchError(VDBError): - """Raised when vector search fails.""" - - def __init__(self, message: str, **kwargs): - super().__init__( - message=message, - code="VDB_SEARCH_ERROR", - status_code=422, - **kwargs, - ) - - -class VDBPartitionNotFound(VDBError): - """Raised when a partition is not found in the vector database.""" - - def __init__(self, message: str, **kwargs): - super().__init__( - message=message, - code="VDB_PARTITION_NOT_FOUND", - status_code=404, - **kwargs, - ) - - -class VDBFileNotFoundError(VDBError): - """Raised when a file is not found in the vector database.""" - - def __init__(self, message: str, **kwargs): - super().__init__( - message=message, - code="VDB_FILE_NOT_FOUND", - status_code=404, - **kwargs, - ) - - -class VDBUserNotFound(VDBError): - """Raised when a user is not found in the vector database.""" - - def __init__(self, message: str, **kwargs): - super().__init__( - message=message, - code="VDB_USER_NOT_FOUND", - status_code=404, - **kwargs, - ) - - -class VDBMembershipNotFound(VDBError): - """Raised when a partition membership is not found in the vector database.""" - - def __init__(self, message: str, **kwargs): - super().__init__( - message=message, - code="VDB_MEMBERSHIP_NOT_FOUND", - status_code=404, - **kwargs, - ) - - -class VDBSchemaMigrationRequiredError(VDBError): - """Raised when the collection schema version does not match the expected version.""" - - def __init__(self, message: str, **kwargs): - super().__init__( - message=message, - code="VDB_SCHEMA_MIGRATION_REQUIRED", - status_code=503, - **kwargs, - ) - - -class UnexpectedVDBError(VDBError): - """Raised for unexpected errors in vector database operations.""" - - def __init__(self, message: str, **kwargs): - super().__init__( - message=message, - code="VDB_UNEXPECTED_ERROR", - status_code=500, - **kwargs, - ) +# Re-export from canonical location for backward compatibility. +# New code should import from `core.utils.exceptions` directly. +from core.utils.exceptions import ( # noqa: F401 + UnexpectedVDBError, + VDBConnectionError, + VDBCreateOrLoadCollectionError, + VDBDeleteError, + VDBError, + VDBFileIDAlreadyExistsError, + VDBFileNotFoundError, + VDBInsertError, + VDBMembershipNotFound, + VDBPartitionNotFound, + VDBSchemaMigrationRequiredError, + VDBSearchError, + VDBUserNotFound, +) diff --git a/openrag/utils/external_resource_errors.py b/openrag/utils/external_resource_errors.py index 3752a0e3f..49f930a26 100644 --- a/openrag/utils/external_resource_errors.py +++ b/openrag/utils/external_resource_errors.py @@ -1,65 +1,13 @@ -""" -Utilities for detecting external resource access errors. +"""Re-export from canonical location for backwards compatibility.""" -When VLM models fetch external image URLs, HTTP errors (403, 404, etc.) from -remote servers get wrapped in InternalServerError, which is misleading. -This module detects such errors for better logging. -""" - -import re - -# HTTP error codes indicating external resource issues -EXTERNAL_ERROR_CODES = frozenset( - { - # 4xx client errors - "400", - "401", - "403", - "404", - "405", - "408", - "410", - "429", - "451", - # 5xx gateway errors - "502", - "503", - "504", - } +from core.utils.external_errors import ( + EXTERNAL_ERROR_CODES, + EXTERNAL_ERROR_INDICATORS, + is_external_resource_error, ) -# Error type indicators for external fetch failures -EXTERNAL_ERROR_INDICATORS = ( - "ClientResponseError", - "HTTPError", - "ConnectionError", - "TimeoutError", - "SSLError", -) - - -def is_external_resource_error(error: Exception) -> tuple[bool, str, str]: - """ - Check if an error is caused by an external resource access issue. - - Returns: - (is_external_error, status_code, url) - status_code and url are empty - strings if not detected. - """ - error_str = str(error) - - # Find HTTP 4xx/5xx status code (first one that's in our allowed set) - status_code = "" - for match in re.finditer(r"\b([45]\d{2})\b", error_str): - if match.group(1) in EXTERNAL_ERROR_CODES: - status_code = match.group(1) - break - - # Extract URL - url_match = re.search(r"https?://[^\s'\"\)>]+", error_str) - url = url_match.group(0) if url_match else "" - - # Check for error type indicators - has_indicator = any(ind in error_str for ind in EXTERNAL_ERROR_INDICATORS) - - return bool(status_code) or has_indicator, status_code, url +__all__ = [ + "EXTERNAL_ERROR_CODES", + "EXTERNAL_ERROR_INDICATORS", + "is_external_resource_error", +] diff --git a/pyproject.toml b/pyproject.toml index 2a9ba5f2d..caa14330a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,8 @@ dependencies = [ "authlib>=1.3", "itsdangerous>=2.2", "cryptography>=42", + "tenacity>=8.2.0", + "aiobreaker>=1.2.0", ] [dependency-groups] diff --git a/scripts/check_layer_imports.py b/scripts/check_layer_imports.py new file mode 100644 index 000000000..e11159803 --- /dev/null +++ b/scripts/check_layer_imports.py @@ -0,0 +1,136 @@ +"""Layer import guard for the hexagonal refactoring. + +Walks every .py file under the four layer roots and flags imports that +violate the dependency rule: + + api -> di, core (NOT services directly) + di -> core, services (free across boundaries) + services -> core (NOT api, NOT di) + core -> (nothing in openrag) (pure domain) + +Only files inside openrag/{core,services,api,di}/ are checked. Legacy +paths (openrag/components/, openrag/routers/, openrag/models/, ...) are +ignored until they are migrated. + +Usage: + python scripts/check_layer_imports.py + +Exit code 0 on pass, 1 on any violation. Prints one line per violation: + path/to/file.py:LINE core -> services (openrag.services.foo) +""" + +from __future__ import annotations + +import ast +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +OPENRAG = REPO_ROOT / "openrag" + +LAYERS = ("core", "services", "api", "di") + +FORBIDDEN: dict[str, set[str]] = { + "core": {"services", "api", "di"}, + "services": {"api", "di"}, + "api": {"services"}, + "di": set(), +} + + +def layer_of(module: str) -> str | None: + """Return the layer name if `module` is openrag.[...] or [...], else None.""" + parts = module.split(".") + if len(parts) >= 2 and parts[0] == "openrag" and parts[1] in LAYERS: + return parts[1] + if parts[0] in LAYERS: + return parts[0] + return None + + +def file_layer(path: Path) -> str | None: + """Return the layer the file belongs to, or None if outside the four roots.""" + try: + rel = path.relative_to(OPENRAG) + except ValueError: + return None + top = rel.parts[0] if rel.parts else "" + return top if top in LAYERS else None + + +def iter_imports(tree: ast.AST): + """Yield (lineno, dotted_module) for every Import/ImportFrom node.""" + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + yield node.lineno, alias.name + elif isinstance(node, ast.ImportFrom): + if node.level: + # relative import — resolve against the file's package later + yield node.lineno, ("__relative__", node.level, node.module or "") + else: + yield node.lineno, node.module or "" + + +def resolve_relative(file_path: Path, level: int, module: str) -> str: + """Turn `from ..foo import bar` into an absolute dotted module.""" + rel = file_path.relative_to(REPO_ROOT).with_suffix("") + parts = list(rel.parts) + # __init__.py sits one level shallower + if parts[-1] == "__init__": + parts.pop() + # `level` dots = climb that many packages + anchor = parts[:-level] if level <= len(parts) else [] + if module: + anchor.extend(module.split(".")) + return ".".join(anchor) + + +def check_file(path: Path) -> list[str]: + src_layer = file_layer(path) + if src_layer is None: + return [] + try: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + except SyntaxError as exc: + return [f"{path}:{exc.lineno} syntax error in file: {exc.msg}"] + + violations: list[str] = [] + forbidden = FORBIDDEN[src_layer] + for lineno, entry in iter_imports(tree): + if isinstance(entry, tuple) and entry and entry[0] == "__relative__": + _, level, module = entry + dotted = resolve_relative(path, level, module) + else: + dotted = entry + tgt_layer = layer_of(dotted) + if tgt_layer is None or tgt_layer == src_layer: + continue + if tgt_layer in forbidden: + rel = path.relative_to(REPO_ROOT) + violations.append(f"{rel}:{lineno} {src_layer} -> {tgt_layer} ({dotted})") + return violations + + +def main() -> int: + if not OPENRAG.is_dir(): + print(f"error: {OPENRAG} not found", file=sys.stderr) + return 2 + + all_violations: list[str] = [] + for path in sorted(OPENRAG.rglob("*.py")): + all_violations.extend(check_file(path)) + + if all_violations: + print("layer import violations:") + for v in all_violations: + print(f" {v}") + print(f"\n{len(all_violations)} violation(s)") + return 1 + + print("layer import guard: OK") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/postgres-init/01-chainlit.sql b/scripts/postgres-init/01-chainlit.sql new file mode 100644 index 000000000..e72738cca --- /dev/null +++ b/scripts/postgres-init/01-chainlit.sql @@ -0,0 +1,64 @@ +CREATE DATABASE chainlit; + +\connect chainlit + +CREATE EXTENSION IF NOT EXISTS pgcrypto; + +CREATE TABLE IF NOT EXISTS "User" ( + id uuid DEFAULT gen_random_uuid() PRIMARY KEY, + identifier text NOT NULL UNIQUE, + metadata text DEFAULT '{}'::text NOT NULL, + "createdAt" timestamptz DEFAULT now() NOT NULL, + "updatedAt" timestamptz DEFAULT now() NOT NULL +); + +CREATE TABLE IF NOT EXISTS "Thread" ( + id uuid DEFAULT gen_random_uuid() PRIMARY KEY, + name text, + "userId" uuid REFERENCES "User"(id) ON DELETE SET NULL, + tags text[], + metadata text DEFAULT '{}'::text NOT NULL, + "createdAt" timestamptz DEFAULT now() NOT NULL, + "deletedAt" timestamptz +); + +CREATE TABLE IF NOT EXISTS "Step" ( + id uuid DEFAULT gen_random_uuid() PRIMARY KEY, + "threadId" uuid REFERENCES "Thread"(id) ON DELETE CASCADE, + "parentId" uuid REFERENCES "Step"(id) ON DELETE SET NULL, + input jsonb DEFAULT '{}'::jsonb, + metadata jsonb DEFAULT '{}'::jsonb NOT NULL, + name text, + output jsonb DEFAULT '{}'::jsonb NOT NULL, + type text NOT NULL, + "startTime" timestamptz, + "endTime" timestamptz, + "showInput" text, + "isError" boolean, + "createdAt" timestamptz DEFAULT now() +); + +CREATE TABLE IF NOT EXISTS "Feedback" ( + id text PRIMARY KEY, + "stepId" uuid NOT NULL REFERENCES "Step"(id) ON DELETE CASCADE, + name text DEFAULT 'user_feedback'::text NOT NULL, + value double precision NOT NULL, + comment text +); + +CREATE TABLE IF NOT EXISTS "Element" ( + id text PRIMARY KEY, + "threadId" uuid REFERENCES "Thread"(id) ON DELETE CASCADE, + "stepId" uuid REFERENCES "Step"(id) ON DELETE CASCADE, + metadata jsonb DEFAULT '{}'::jsonb NOT NULL, + mime text, + name text NOT NULL, + "objectKey" text, + url text, + "chainlitKey" text, + display text, + size text, + language text, + page integer, + props jsonb DEFAULT '{}'::jsonb NOT NULL +); diff --git a/test_copilot.html b/test_copilot.html deleted file mode 100644 index 4aa5f9a60..000000000 --- a/test_copilot.html +++ /dev/null @@ -1,18 +0,0 @@ - - - - - - - - - \ No newline at end of file diff --git a/tests/api_tests/test_oidc_lifecycle.py b/tests/api_tests/test_oidc_lifecycle.py index f6cf24f09..d01705860 100644 --- a/tests/api_tests/test_oidc_lifecycle.py +++ b/tests/api_tests/test_oidc_lifecycle.py @@ -206,13 +206,15 @@ def _impl_create_oidc_session(self, **kwargs): return row def _impl_get_oidc_session_by_token(self, session_token_plain: str): - sid_key = self._sessions_by_token.get(session_token_plain) - if sid_key is None: - return None - row = self._sessions[sid_key] - if row.get("revoked_at"): - return None - return row + # Faithful to production: PartitionFileManager.get_oidc_session_by_token + # hashes the plaintext and matches on session_token_hash. Post-Phase-8 + # the row is written by AuthService via the repo adapter (hash only), + # so the legacy plaintext index no longer applies. + token_hash = hash_session_token(session_token_plain) + for row in self._sessions.values(): + if row.get("session_token_hash") == token_hash and not row.get("revoked_at"): + return row + return None def _impl_get_user(self, user_id: int): return self._users_by_id.get(user_id) @@ -236,50 +238,195 @@ def _impl_revoke_oidc_sessions_by_sid(self, sid: str) -> int: # --------------------------------------------------------------------------- _stub_vdb = _StubVectorDB() +_stub_task_state_manager = types.SimpleNamespace( + get_user_pending_task_count=_RayMethodStub("get_user_pending_task_count", lambda *a, **kw: 0, []) +) def _install_stubs(): stub = types.ModuleType("utils.dependencies") stub.get_vectordb = lambda: _stub_vdb - stub.get_task_state_manager = lambda: None + stub.get_task_state_manager = lambda: _stub_task_state_manager stub.get_serializer = lambda: None stub.get_indexer = lambda: None stub.get_marker_pool = lambda: None sys.modules["utils.dependencies"] = stub + def _logger(): + logger = types.SimpleNamespace( + debug=lambda *args, **kwargs: None, + info=lambda *args, **kwargs: None, + warning=lambda *args, **kwargs: None, + error=lambda *args, **kwargs: None, + exception=lambda *args, **kwargs: None, + ) + logger.bind = lambda *args, **kwargs: logger + return logger + + logger_stub = types.ModuleType("utils.logger") + logger_stub.escape_markup = lambda s: s.replace("\\", "\\\\").replace("<", "\\<").replace(">", "\\>") + logger_stub.mask_email = ( + lambda email: f"{email.partition('@')[0][0]}***@{email.partition('@')[2]}" + if isinstance(email, str) and "@" in email and email.partition("@")[0] + else "***" + ) + logger_stub.get_logger = _logger + sys.modules["utils.logger"] = logger_stub + openai_stub = types.ModuleType("openai") + openai_stub.AsyncOpenAI = object + openai_stub.APITimeoutError = TimeoutError + openai_stub.APIConnectionError = ConnectionError + openai_stub.APIError = Exception + sys.modules.setdefault("openai", openai_stub) + _install_stubs() -# Reload auth deps + routers after stub installation -from components.auth import deps as _auth_deps # noqa: E402 +# Reload routers after stub installation. Post-Phase-8 the auth/users routers +# resolve AuthService/UserService from the DI providers, so we import the +# provider symbols here (di.providers is not popped, so these are the same +# function objects the routers close over) to key dependency_overrides. +from components.auth import OIDCClient # noqa: E402 +from components.auth.session_tokens import hash_session_token # noqa: E402 +from core.config.auth import OIDCConfig # noqa: E402 +from core.models.user import OIDCSession, User # noqa: E402 +from di.providers import get_auth_service, get_user_service # noqa: E402 +from services.orchestrators.auth_service import AuthService # noqa: E402 +from services.orchestrators.user_service import UserService # noqa: E402 sys.modules.pop("routers.auth", None) sys.modules.pop("routers.users", None) _auth_router_mod = importlib.import_module("routers.auth") _users_router_mod = importlib.import_module("routers.users") + +# --------------------------------------------------------------------------- +# Phase-8 repository-port adapters over the shared _StubVectorDB state. +# AuthService/UserService take repo ports, not the Ray vdb actor; these +# wrap the same dicts the auth middleware reads through _StubVectorDB so a +# session created by AuthService is visible to the middleware and vice-versa. +# --------------------------------------------------------------------------- + + +class _StubUserRepo: + def __init__(self, vdb: _StubVectorDB): + self._vdb = vdb + + @staticmethod + def _to_user(d: dict | None) -> User | None: + if d is None: + return None + return User( + id=d["id"], + display_name=d.get("display_name"), + external_user_id=d.get("external_user_id"), + email=d.get("email"), + is_admin=d.get("is_admin", False), + ) + + async def get_user_by_external_id(self, external_id: str) -> User | None: + return self._to_user(self._vdb._users_by_sub.get(external_id)) + + async def get_user(self, user_id: int) -> User | None: + return self._to_user(self._vdb._users_by_id.get(user_id)) + + async def create_user(self, user: User) -> User: + new_id = max(self._vdb._users_by_id, default=0) + 1 + user.id = new_id + rec = { + "id": new_id, + "email": user.email, + "external_user_id": user.external_user_id, + "is_admin": user.is_admin, + "display_name": user.display_name, + } + self._vdb._users_by_id[new_id] = rec + if user.external_user_id: + self._vdb._users_by_sub[user.external_user_id] = rec + return user + + async def update_user(self, user_id: int, **fields) -> User | None: + rec = self._vdb._users_by_id.get(user_id) + if rec is None: + return None + for k, v in fields.items(): + rec[k] = v + return self._to_user(rec) + + +class _StubOIDCSessionRepo: + def __init__(self, vdb: _StubVectorDB): + self._vdb = vdb + + async def create_session(self, session: OIDCSession) -> OIDCSession: + sid_key = self._vdb._next_session_id + self._vdb._next_session_id += 1 + session.id = sid_key + self._vdb._sessions[sid_key] = { + "id": sid_key, + "user_id": session.user_id, + "sub": session.sub, + "sid": session.sid, + "session_token_hash": session.session_token_hash, + "id_token_encrypted": session.id_token_encrypted, + "access_token_encrypted": session.access_token_encrypted, + "refresh_token_encrypted": session.refresh_token_encrypted, + "access_token_expires_at": session.access_token_expires_at, + "session_expires_at": session.session_expires_at, + "created_at": session.created_at, + "last_refresh_at": None, + "revoked_at": None, + } + return session + + async def get_by_token_hash(self, token_hash: str) -> OIDCSession | None: + for row in self._vdb._sessions.values(): + if row.get("session_token_hash") == token_hash and not row.get("revoked_at"): + return OIDCSession(**{k: row.get(k) for k in OIDCSession.model_fields}) + return None + + async def revoke_session(self, session_id: int) -> bool: + self._vdb._impl_revoke_oidc_session_by_id(session_id) + return True + + async def revoke_by_sid(self, sid: str) -> int: + return self._vdb._impl_revoke_oidc_sessions_by_sid(sid) + + +class _StubMembershipRepo: + async def list_user_partitions(self, user_id: int) -> list: + return [] + + +class _StubJobService: + async def get_user_pending_task_count(self, user_id) -> int: + return 0 + + # --------------------------------------------------------------------------- # Build the composite app (auth + users) # --------------------------------------------------------------------------- def _make_app(router) -> tuple[FastAPI, TestClient]: - """Build a minimal FastAPI app combining auth + users routers, with a - respx MockRouter injected into the OIDCClient singleton. respx >= 0.22 - exposes MockRouter + httpx.MockTransport(router.handler).""" + """Build a minimal FastAPI app combining the auth + users routers. + + Post-Phase-8 the routers pull AuthService/UserService from the DI + providers, so rather than patching a client singleton we build the real + services over the shared _StubVectorDB state, inject a respx-mocked + OIDCClient, and override get_auth_service / get_user_service. respx >= + 0.22 exposes MockRouter + httpx.MockTransport(router.handler).""" app = FastAPI() # Install the AuthMiddleware (from components.auth.middleware) from components.auth.middleware import AuthMiddleware # noqa: E402 - app.add_middleware(AuthMiddleware, get_vectordb=lambda: _stub_vdb) + app.add_middleware(AuthMiddleware, get_auth_service=lambda _request: auth_service) app.include_router(_auth_router_mod.router) app.include_router(_users_router_mod.router, prefix="/users") - # Override OIDCClient singleton with our mocked http transport. - _auth_deps.reset_oidc_client() - _auth_deps._client = _auth_router_mod.OIDCClient( + oidc_client = OIDCClient( issuer=ISSUER, client_id=CLIENT_ID, client_secret=CLIENT_SECRET, @@ -287,6 +434,38 @@ def _make_app(router) -> tuple[FastAPI, TestClient]: scopes=SCOPES, http_client=httpx.AsyncClient(transport=httpx.MockTransport(router.handler)), ) + cfg = OIDCConfig( + enabled=True, + issuer_url=ISSUER, + client_id=CLIENT_ID, + client_secret=CLIENT_SECRET, + redirect_uri=REDIRECT_URI, + scopes=SCOPES, + token_encryption_key=_FERNET_KEY, + claim_source="id_token", + claim_mapping="", + post_logout_redirect_uri="/", + auto_provision_login=False, + ) + user_repo = _StubUserRepo(_stub_vdb) + membership_repo = _StubMembershipRepo() + auth_service = AuthService( + user_repo=user_repo, + oidc_session_repo=_StubOIDCSessionRepo(_stub_vdb), + membership_repo=membership_repo, + oidc_client=oidc_client, + config=cfg, + ) + user_service = UserService( + user_repo=user_repo, + auth_service=auth_service, + default_file_quota=10, + partition_service=object(), + membership_repo=membership_repo, + job_service=_StubJobService(), + ) + app.dependency_overrides[get_auth_service] = lambda: auth_service + app.dependency_overrides[get_user_service] = lambda: user_service client = TestClient(app, raise_server_exceptions=True) return app, client @@ -317,7 +496,6 @@ def test_full_oidc_lifecycle(monkeypatch): monkeypatch.setenv("OIDC_POST_LOGOUT_REDIRECT_URI", "/") monkeypatch.delenv("OIDC_CLAIM_MAPPING", raising=False) monkeypatch.delenv("AUTH_TOKEN", raising=False) - _auth_deps.reset_oidc_client() # ── Pre-seed alice with the exact sub that the IdP mock will return ──────── ALICE_SUB = "alice-sub" diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 000000000..fb8931b41 --- /dev/null +++ b/tests/integration/__init__.py @@ -0,0 +1,2 @@ +# Phase 7F integration tests — exercise the asyncpg repos and PostgresStore +# composite against a real Postgres. Tests auto-skip when no DSN is reachable. diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py new file mode 100644 index 000000000..b90d0e57d --- /dev/null +++ b/tests/integration/conftest.py @@ -0,0 +1,164 @@ +"""Shared fixtures for the Phase 7F persistence-layer integration tests. + +The whole suite auto-skips when a Postgres instance is not reachable. We try +the explicit ``POSTGRES_TEST_DSN`` env var first, then fall back to the local +docker-compose ``rdb`` container (the dev DB the rest of the project assumes +is up). One ephemeral test database is created at session start, Alembic +migrations run once, and individual tests share the same +:class:`PostgresStore`; an autouse fixture truncates user-modifiable tables +between tests so each one starts clean without paying for a full +drop/migrate cycle per case. + +The integration tests live under ``tests/integration/`` rather than +alongside the repository code because they need a real database — pytest's +``testpaths`` keeps them out of the default unit run; invoke explicitly with +``uv run pytest tests/integration``. +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import asyncpg +import pytest +import pytest_asyncio + +# tests/integration is outside the pytest.ini `pythonpath = ./openrag` +# search path, so we have to add it ourselves before importing project +# modules. Mirrors how scripts/migrations/env.py wires its path. +_OPENRAG = Path(__file__).resolve().parents[2] / "openrag" +if str(_OPENRAG) not in sys.path: + sys.path.insert(0, str(_OPENRAG)) + +from core.config.infrastructure import RDBConfig # noqa: E402 +from services.storage.postgres_store import PostgresStore # noqa: E402 + +_DEFAULT_ADMIN_DSN = "postgresql://root:root_password@172.21.0.4:5432/postgres" +_TEST_DB_NAME = "openrag_phase7_test" + + +def _admin_dsn() -> str: + return os.environ.get("POSTGRES_TEST_ADMIN_DSN", _DEFAULT_ADMIN_DSN) + + +def _admin_dsn_parts() -> dict[str, str | int]: + """Decompose the admin DSN into the bits :class:`RDBConfig` needs.""" + import urllib.parse + + p = urllib.parse.urlparse(_admin_dsn()) + return { + "host": p.hostname or "localhost", + "port": p.port or 5432, + "user": p.username or "postgres", + "password": p.password or "", + } + + +async def _connect_admin() -> asyncpg.Connection | None: + """Open an admin connection or return None when the server is unreachable.""" + try: + return await asyncpg.connect(_admin_dsn(), timeout=5) + except (OSError, asyncpg.PostgresError): + return None + + +async def _drop_test_database(conn: asyncpg.Connection) -> None: + # Kick any stale sessions off the test DB before dropping it. ``DROP + # DATABASE`` fails on active connections — important when a previous + # crashed session left a pool open. + await conn.execute( + """ + SELECT pg_terminate_backend(pid) + FROM pg_stat_activity + WHERE datname = $1 AND pid <> pg_backend_pid() + """, + _TEST_DB_NAME, + ) + await conn.execute(f'DROP DATABASE IF EXISTS "{_TEST_DB_NAME}"') + + +@pytest_asyncio.fixture(scope="session", loop_scope="session") +async def _ensured_test_database() -> str: + """Create the test database from scratch; drop it at session teardown.""" + admin = await _connect_admin() + if admin is None: + pytest.skip( + f"Postgres unreachable at {_admin_dsn()}; set POSTGRES_TEST_ADMIN_DSN or start the rdb container.", + ) + try: + await _drop_test_database(admin) + await admin.execute(f'CREATE DATABASE "{_TEST_DB_NAME}"') + finally: + await admin.close() + + yield _TEST_DB_NAME + + admin = await _connect_admin() + if admin is None: + return + try: + await _drop_test_database(admin) + finally: + await admin.close() + + +@pytest.fixture(scope="session") +def test_rdb_config(_ensured_test_database: str) -> RDBConfig: + """A :class:`RDBConfig` pointed at the ephemeral test database. + + Exposed as its own fixture so per-test lifecycle assertions can build + their own short-lived :class:`PostgresStore` without disturbing the + session-scoped pool. + """ + parts = _admin_dsn_parts() + return RDBConfig( + host=str(parts["host"]), + port=int(parts["port"]), + user=str(parts["user"]), + password=str(parts["password"]), + database=_ensured_test_database, + pool_min_size=1, + pool_max_size=4, + ) + + +@pytest_asyncio.fixture(scope="session", loop_scope="session") +async def postgres_store(test_rdb_config: RDBConfig) -> PostgresStore: + """A real :class:`PostgresStore` against the freshly-created test DB. + + Migrations run once per session via :meth:`PostgresStore.initialize`. + """ + store = PostgresStore(test_rdb_config, run_migrations=True) + await store.initialize() + try: + yield store + finally: + await store.shutdown() + + +_TRUNCATE_SQL = """ +TRUNCATE TABLE + oidc_sessions, + workspace_files, + workspaces, + partition_memberships, + files, + partitions, + users +RESTART IDENTITY CASCADE +""" + + +@pytest_asyncio.fixture(autouse=True, loop_scope="session") +async def _clean_db(postgres_store: PostgresStore): + """Wipe user-modifiable tables before each test. + + ``alembic_version`` is left alone — migrations only run once per session. + Identities are restarted so primary keys reset to 1, giving each test a + deterministic ``users.id`` / ``files.id`` starting point. + """ + async with postgres_store.pool.acquire() as conn: + await conn.execute(_TRUNCATE_SQL) + yield diff --git a/tests/integration/docker-compose.yaml b/tests/integration/docker-compose.yaml new file mode 100644 index 000000000..0caa6907d --- /dev/null +++ b/tests/integration/docker-compose.yaml @@ -0,0 +1,49 @@ +services: + etcd: + image: quay.io/coreos/etcd:v3.5.25 + environment: + - ETCD_AUTO_COMPACTION_MODE=revision + - ETCD_AUTO_COMPACTION_RETENTION=1000 + - ETCD_QUOTA_BACKEND_BYTES=4294967296 + - ETCD_SNAPSHOT_COUNT=50000 + command: etcd -advertise-client-urls=http://etcd:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd + healthcheck: + test: ["CMD", "etcdctl", "endpoint", "health"] + interval: 5s + timeout: 10s + retries: 5 + + minio: + image: minio/minio:RELEASE.2024-12-18T13-15-44Z + environment: + MINIO_ACCESS_KEY: minioadmin + MINIO_SECRET_KEY: minioadmin + command: minio server /minio_data --console-address ":9001" + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"] + interval: 5s + timeout: 10s + retries: 5 + + milvus: + image: milvusdb/milvus:v2.6.11 + command: ["milvus", "run", "standalone"] + security_opt: + - seccomp:unconfined + environment: + ETCD_ENDPOINTS: etcd:2379 + MINIO_ADDRESS: minio:9000 + ports: + - "19530:19530" + - "9091:9091" + depends_on: + etcd: + condition: service_healthy + minio: + condition: service_healthy + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"] + interval: 5s + timeout: 10s + retries: 10 + start_period: 30s diff --git a/tests/integration/test_document_repo.py b/tests/integration/test_document_repo.py new file mode 100644 index 000000000..2fc15c57f --- /dev/null +++ b/tests/integration/test_document_repo.py @@ -0,0 +1,126 @@ +"""Phase 7F — PgDocumentRepository against a real Postgres.""" + +from __future__ import annotations + +import pytest +from core.models.catalog import DocumentRecord, DocumentStatus +from services.storage.postgres_store import PostgresStore + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +async def _seed_partition(store: PostgresStore, name: str = "p") -> str: + """``files`` has an FK to ``partitions``; the partition row must exist first.""" + await store.partition_repo.create_partition(name) + return name + + +def _doc(file_id: str, partition: str = "p", **extra) -> DocumentRecord: + return DocumentRecord( + id=file_id, + file_id=file_id, + partition=partition, + filename=f"{file_id}.pdf", + **extra, + ) + + +class TestCreateGetDelete: + async def test_create_then_get(self, postgres_store: PostgresStore): + partition = await _seed_partition(postgres_store) + await postgres_store.document_repo.create_document(_doc("f1", partition)) + fetched = await postgres_store.document_repo.get_document("f1") + assert fetched is not None + assert fetched.file_id == "f1" + assert fetched.partition == partition + assert fetched.filename == "f1.pdf" + + async def test_get_missing_returns_none(self, postgres_store: PostgresStore): + assert await postgres_store.document_repo.get_document("nope") is None + + async def test_delete_returns_true_on_success(self, postgres_store: PostgresStore): + partition = await _seed_partition(postgres_store) + await postgres_store.document_repo.create_document(_doc("f2", partition)) + assert await postgres_store.document_repo.delete_document("f2") is True + assert await postgres_store.document_repo.get_document("f2") is None + + async def test_delete_missing_returns_false(self, postgres_store: PostgresStore): + assert await postgres_store.document_repo.delete_document("ghost") is False + + +class TestListFilter: + async def test_list_by_partition(self, postgres_store: PostgresStore): + await _seed_partition(postgres_store, "alpha") + await _seed_partition(postgres_store, "beta") + repo = postgres_store.document_repo + await repo.create_document(_doc("a1", "alpha")) + await repo.create_document(_doc("a2", "alpha")) + await repo.create_document(_doc("b1", "beta")) + + only_alpha = await repo.list_documents(partition="alpha") + assert {d.file_id for d in only_alpha} == {"a1", "a2"} + + async def test_list_by_partition_list(self, postgres_store: PostgresStore): + await _seed_partition(postgres_store, "alpha") + await _seed_partition(postgres_store, "beta") + repo = postgres_store.document_repo + await repo.create_document(_doc("a1", "alpha")) + await repo.create_document(_doc("b1", "beta")) + both = await repo.list_documents(partition=["alpha", "beta"]) + assert {d.file_id for d in both} == {"a1", "b1"} + + async def test_count_documents(self, postgres_store: PostgresStore): + partition = await _seed_partition(postgres_store) + repo = postgres_store.document_repo + assert await repo.count_documents(partition=partition) == 0 + await repo.create_document(_doc("c1", partition)) + await repo.create_document(_doc("c2", partition)) + assert await repo.count_documents(partition=partition) == 2 + + async def test_file_exists_in_partition(self, postgres_store: PostgresStore): + partition = await _seed_partition(postgres_store) + repo = postgres_store.document_repo + assert await repo.file_exists_in_partition("e1", partition) is False + await repo.create_document(_doc("e1", partition)) + assert await repo.file_exists_in_partition("e1", partition) is True + + +class TestUpdate: + async def test_update_status_folds_into_metadata( + self, + postgres_store: PostgresStore, + ): + partition = await _seed_partition(postgres_store) + repo = postgres_store.document_repo + await repo.create_document(_doc("u1", partition)) + + updated = await repo.update_document("u1", status=DocumentStatus.COMPLETED) + assert updated is not None + assert updated.status == DocumentStatus.COMPLETED + + async def test_update_metadata_merges(self, postgres_store: PostgresStore): + partition = await _seed_partition(postgres_store) + repo = postgres_store.document_repo + await repo.create_document( + _doc("u2", partition, metadata={"a": 1, "b": 2}), + ) + updated = await repo.update_document("u2", metadata={"b": 99, "c": 3}) + assert updated is not None + # filename / status / error_message live in their own DocumentRecord + # fields after the row → domain conversion lifts them out of the JSON. + assert updated.filename == "u2.pdf" + assert updated.metadata == {"a": 1, "b": 99, "c": 3} + + async def test_update_missing_returns_none(self, postgres_store: PostgresStore): + assert await postgres_store.document_repo.update_document("nope") is None + + +class TestDeleteByPartition: + async def test_returns_deletion_count(self, postgres_store: PostgresStore): + partition = await _seed_partition(postgres_store, "trash") + repo = postgres_store.document_repo + await repo.create_document(_doc("d1", partition)) + await repo.create_document(_doc("d2", partition)) + deleted = await repo.delete_documents_by_partition(partition) + assert deleted == 2 + assert await repo.count_documents(partition=partition) == 0 diff --git a/tests/integration/test_milvus_store_integration.py b/tests/integration/test_milvus_store_integration.py new file mode 100644 index 000000000..1aca88778 --- /dev/null +++ b/tests/integration/test_milvus_store_integration.py @@ -0,0 +1,368 @@ +"""End-to-end integration tests for :class:`MilvusVectorStore`. + +These tests round-trip through a real Milvus 2.6 instance: they create a +fresh collection per test, exercise the public surface, and drop the +collection on teardown. They are gated by the ``integration`` pytest marker +and auto-skip when the configured Milvus host is not reachable. + +Run locally against the dev compose stack: + + docker compose up -d milvus + uv run pytest tests/integration/test_milvus_store_integration.py -m integration + +Lives under ``tests/integration/`` per the Phase 13C target test layout +(``tests/{unit,integration,load}``) — see +``docs/refactoring/REFACTORING_STRATEGY_v1.md``. Pure-logic tests (filter +expressions, ID coercion, ABC discipline) stay colocated at +``openrag/services/storage/test_milvus_store.py`` until the Phase 13C sweep +relocates them under ``tests/unit/``. +""" + +from __future__ import annotations + +import os +import socket +import uuid +from collections.abc import Iterator + +import pytest + +from openrag.core.config.infrastructure import VectorDBConfig +from openrag.core.models.chunk import Chunk, ChunkType +from openrag.services.storage.milvus_store import MilvusVectorStore + +pytestmark = pytest.mark.integration + + +# --------------------------------------------------------------------------- +# Reachability gate — keeps the suite green when Milvus isn't running +# --------------------------------------------------------------------------- + + +def _milvus_reachable(host: str, port: int, timeout: float = 1.0) -> bool: + try: + with socket.create_connection((host, port), timeout=timeout): + return True + except OSError: + return False + + +# Embedding dimension is intentionally tiny — smaller = faster index build, +# and the schema cares about *having* a dimension, not the specific value. +_EMBEDDING_DIM = 4 + + +def _embedding(seed: float) -> list[float]: + """Build a small deterministic vector. Same seed = same vector.""" + return [seed, seed + 0.1, seed + 0.2, seed + 0.3] + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def milvus_host_port() -> tuple[str, int]: + """Resolve the Milvus endpoint, preferring test-specific env overrides. + + Defaults to ``localhost:19530`` because the test runs on the host, not + inside the docker network where the service is named ``milvus``. + + ``VDB_HOST`` / ``MILVUS_HOST`` are NOT honoured here because pymilvus + auto-loads the project's ``.env`` at import time (see ``pymilvus.settings``), + which would inject the docker-network hostname ``milvus`` into a host-side + test run. The dedicated ``OPENRAG_TEST_VDB_HOST`` env keeps the runtime + config and the test config independent. + """ + host = os.getenv("OPENRAG_TEST_VDB_HOST", "localhost") + port = int(os.getenv("OPENRAG_TEST_VDB_PORT", "19530")) + return host, port + + +@pytest.fixture(scope="module") +def _live_milvus(milvus_host_port: tuple[str, int]) -> None: + host, port = milvus_host_port + if not _milvus_reachable(host, port): + pytest.skip(f"Milvus not reachable at {host}:{port} — skipping integration tests") + + +@pytest.fixture +def collection_name() -> str: + """A throwaway collection name per test, so parallel runs don't collide.""" + # Milvus collection names are alphanumeric/underscore; uuid hex fits. + return f"itest_{uuid.uuid4().hex[:12]}" + + +@pytest.fixture +def hybrid_config( + milvus_host_port: tuple[str, int], + collection_name: str, +) -> VectorDBConfig: + host, port = milvus_host_port + return VectorDBConfig( + host=host, + port=port, + collection_name=collection_name, + hybrid_search=True, + schema_version=1, + ) + + +@pytest.fixture +def dense_only_config( + milvus_host_port: tuple[str, int], + collection_name: str, +) -> VectorDBConfig: + host, port = milvus_host_port + return VectorDBConfig( + host=host, + port=port, + collection_name=collection_name, + hybrid_search=False, + schema_version=1, + ) + + +@pytest.fixture +def hybrid_store( + _live_milvus: None, + hybrid_config: VectorDBConfig, +) -> Iterator[MilvusVectorStore]: + """A real hybrid-enabled store wired to a freshly-named collection. + + The collection is created lazily by ``initialize()`` and dropped after + every test so suite reruns don't accumulate orphaned collections. + """ + store = MilvusVectorStore(hybrid_config) + try: + yield store + finally: + # Best-effort teardown — collection may not exist if a test never + # initialized it (or already dropped it explicitly). + try: + if store._client.has_collection(hybrid_config.collection_name): + store._client.drop_collection(hybrid_config.collection_name) + except Exception: + pass + + +@pytest.fixture +def dense_only_store( + _live_milvus: None, + dense_only_config: VectorDBConfig, +) -> Iterator[MilvusVectorStore]: + """A real dense-only store (no ``sparse`` field) on a fresh collection. + + Mirrors :func:`hybrid_store` but with ``hybrid_search=False`` so + ``search()`` exercises the dense dispatch branch end-to-end. + """ + store = MilvusVectorStore(dense_only_config) + try: + yield store + finally: + try: + if store._client.has_collection(dense_only_config.collection_name): + store._client.drop_collection(dense_only_config.collection_name) + except Exception: + pass + + +def _chunk(text: str, partition: str, seed: float, **extra) -> Chunk: + """Build a freshly-embedded chunk with sensible defaults.""" + return Chunk( + text=text, + document_id=extra.pop("document_id", "doc-1"), + partition=partition, + embedding=_embedding(seed), + chunk_type=ChunkType.TEXT, + metadata=extra, + ) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestEndToEnd: + """Happy-path round trip: create → upsert → search → query → delete → drop.""" + + @pytest.mark.asyncio + async def test_initialize_creates_collection( + self, hybrid_store: MilvusVectorStore, hybrid_config: VectorDBConfig + ) -> None: + assert await hybrid_store.collection_exists(hybrid_config.collection_name) is False + await hybrid_store.initialize(_EMBEDDING_DIM) + assert await hybrid_store.collection_exists(hybrid_config.collection_name) is True + + @pytest.mark.asyncio + async def test_initialize_is_idempotent(self, hybrid_store: MilvusVectorStore) -> None: + await hybrid_store.initialize(_EMBEDDING_DIM) + # Second call must not raise and must not re-create. + await hybrid_store.initialize(_EMBEDDING_DIM) + assert hybrid_store._loaded is True + + @pytest.mark.asyncio + async def test_ensure_collection_rejects_dimension_change( + self, hybrid_store: MilvusVectorStore, hybrid_config: VectorDBConfig + ) -> None: + await hybrid_store.initialize(_EMBEDDING_DIM) + with pytest.raises(ValueError, match="Drop the collection before re-sizing"): + await hybrid_store.ensure_collection(hybrid_config.collection_name, _EMBEDDING_DIM + 1) + + @pytest.mark.asyncio + async def test_upsert_returns_insert_count(self, hybrid_store: MilvusVectorStore) -> None: + await hybrid_store.initialize(_EMBEDDING_DIM) + chunks = [ + _chunk("alpha doc one", "p1", 0.1), + _chunk("beta doc two", "p1", 0.2), + _chunk("gamma doc three", "p1", 0.3), + ] + n = await hybrid_store.upsert(chunks) + assert n == 3 + + @pytest.mark.asyncio + async def test_upsert_without_embedding_raises(self, hybrid_store: MilvusVectorStore) -> None: + await hybrid_store.initialize(_EMBEDDING_DIM) + bad = Chunk(text="missing", partition="p1", embedding=None) + from openrag.core.utils.exceptions import VDBInsertError + + with pytest.raises(VDBInsertError, match="no embedding"): + await hybrid_store.upsert([bad]) + + @pytest.mark.asyncio + async def test_dense_search_returns_results(self, dense_only_store: MilvusVectorStore) -> None: + await dense_only_store.initialize(_EMBEDDING_DIM) + chunks = [ + _chunk("alpha", "p1", 0.1), + _chunk("beta", "p1", 0.5), + _chunk("gamma", "p1", 0.9), + ] + await dense_only_store.upsert(chunks) + # Force the collection to flush so reads see the writes — Milvus is + # eventually consistent in default mode but our config sets Strong + # consistency so the search below should see everything. + hits = await dense_only_store.search(_embedding(0.1), top_k=10) + assert len(hits) >= 1 + for hit in hits: + assert "id" in hit + assert "score" in hit + assert "vector" not in hit, "raw vector must be stripped from results" + + @pytest.mark.asyncio + async def test_search_with_partition_filter(self, dense_only_store: MilvusVectorStore) -> None: + await dense_only_store.initialize(_EMBEDDING_DIM) + await dense_only_store.upsert( + [ + _chunk("a", "p1", 0.1), + _chunk("b", "p1", 0.2), + _chunk("c", "p2", 0.3), + ] + ) + hits = await dense_only_store.search(_embedding(0.1), top_k=10, filters={"partition": "p1"}) + assert len(hits) >= 1 + for hit in hits: + assert hit["partition"] == "p1" + + +class TestHybridSearch: + @pytest.mark.asyncio + async def test_hybrid_search_returns_fused_results(self, hybrid_store: MilvusVectorStore) -> None: + await hybrid_store.initialize(_EMBEDDING_DIM) + await hybrid_store.upsert( + [ + _chunk("milvus vector database", "p1", 0.1), + _chunk("postgres relational database", "p1", 0.5), + _chunk("redis key value store", "p1", 0.9), + ] + ) + hits = await hybrid_store.search( + _embedding(0.1), + query_text="milvus database", + top_k=5, + ) + assert len(hits) >= 1 + # RRF fusion still returns the same shape — id, score, entity fields. + for hit in hits: + assert "id" in hit + assert "score" in hit + assert "text" in hit + + +class TestDeleteByFilter: + @pytest.mark.asyncio + async def test_delete_by_partition(self, hybrid_store: MilvusVectorStore) -> None: + await hybrid_store.initialize(_EMBEDDING_DIM) + await hybrid_store.upsert( + [ + _chunk("a", "p1", 0.1), + _chunk("b", "p2", 0.2), + ] + ) + deleted = await hybrid_store.delete_by_filter({"partition": "p1"}) + # We don't assert an exact count — Milvus returns delete_count, but + # the integration's value is that the call succeeds and p1 vanishes. + assert deleted >= 0 + remaining_p1 = await hybrid_store.query_ids_by_filter(hybrid_store._collection_name, {"partition": "p1"}) + assert remaining_p1 == [] + + @pytest.mark.asyncio + async def test_delete_by_filter_with_wildcard_partition_raises(self, hybrid_store: MilvusVectorStore) -> None: + await hybrid_store.initialize(_EMBEDDING_DIM) + with pytest.raises(ValueError, match="drop_collection"): + await hybrid_store.delete_by_filter({"partition": "all"}) + + +class TestQueryByFilter: + @pytest.mark.asyncio + async def test_query_ids_returns_string_ids(self, hybrid_store: MilvusVectorStore) -> None: + await hybrid_store.initialize(_EMBEDDING_DIM) + await hybrid_store.upsert([_chunk("only", "p1", 0.1)]) + ids = await hybrid_store.query_ids_by_filter(hybrid_store._collection_name, {"partition": "p1"}) + assert ids, "expected at least one row matching partition=p1" + for chunk_id in ids: + assert isinstance(chunk_id, str) + assert chunk_id.isdigit(), f"Milvus _id round-trip lost INT64 form: {chunk_id}" + + @pytest.mark.asyncio + async def test_query_chunks_returns_full_records(self, hybrid_store: MilvusVectorStore) -> None: + await hybrid_store.initialize(_EMBEDDING_DIM) + await hybrid_store.upsert([_chunk("only", "p1", 0.1)]) + rows = await hybrid_store.query_chunks_by_filter(hybrid_store._collection_name, {"partition": "p1"}) + assert rows + assert rows[0]["partition"] == "p1" + assert rows[0]["text"] == "only" + # NOTE: ``_iter_query`` does NOT strip the vector field, unlike the + # search path (which filters via ``_SEARCH_RESULT_DROPPED_KEYS``). + # ``query_chunks_by_filter`` therefore leaks the dense vector when + # called with the default ``["*"]`` output_fields — asymmetric with + # ``search()`` and contradicts the method docstring. Tracked as a + # follow-up; this test documents current behaviour so the next change + # is intentional. + + +class TestDropAndDelete: + @pytest.mark.asyncio + async def test_drop_collection_lets_initialize_recreate( + self, hybrid_store: MilvusVectorStore, hybrid_config: VectorDBConfig + ) -> None: + await hybrid_store.initialize(_EMBEDDING_DIM) + await hybrid_store.drop_collection(hybrid_config.collection_name) + assert await hybrid_store.collection_exists(hybrid_config.collection_name) is False + # After drop, the store is allowed to re-initialize from scratch — + # otherwise per-tenant lifecycles would need a new instance just to + # rebuild the collection. + await hybrid_store.initialize(_EMBEDDING_DIM) + assert await hybrid_store.collection_exists(hybrid_config.collection_name) is True + + @pytest.mark.asyncio + async def test_delete_by_id_removes_rows(self, hybrid_store: MilvusVectorStore) -> None: + await hybrid_store.initialize(_EMBEDDING_DIM) + await hybrid_store.upsert([_chunk("to-delete", "p1", 0.1)]) + ids = await hybrid_store.query_ids_by_filter(hybrid_store._collection_name, {"partition": "p1"}) + assert ids, "expected the upsert to land at least one row" + deleted = await hybrid_store.delete(ids) + assert deleted >= 0 + remaining = await hybrid_store.query_ids_by_filter(hybrid_store._collection_name, {"partition": "p1"}) + assert remaining == [] diff --git a/tests/integration/test_oidc_session_repo.py b/tests/integration/test_oidc_session_repo.py new file mode 100644 index 000000000..06016fbe7 --- /dev/null +++ b/tests/integration/test_oidc_session_repo.py @@ -0,0 +1,135 @@ +"""Phase 7F — PgOIDCSessionRepository against a real Postgres.""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta + +import pytest +from core.models.user import OIDCSession, User +from services.storage.postgres_store import PostgresStore + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +async def _seed_user(store: PostgresStore) -> User: + return await store.user_repo.create_user( + User(display_name="OIDC user", external_user_id="kc-oidc-sub"), + ) + + +def _session( + user_id: int, + token_hash: str = "deadbeef", + *, + sub: str = "kc-oidc-sub", + sid: str | None = "sess-1", + revoked: bool = False, + expires_in: timedelta = timedelta(hours=1), +) -> OIDCSession: + # The ``oidc_sessions`` columns are TIMESTAMP WITHOUT TIME ZONE — the + # OIDCSession defaults are tz-aware UTC. Production callers strip tzinfo + # before insert; mirroring that here keeps the repo a thin pass-through + # (the timezone-column mismatch is a separate carryover from the legacy + # schema and not in scope for Phase 7F). + now = datetime.now(UTC).replace(tzinfo=None) + revoked_at = now if revoked else None + return OIDCSession( + session_token_hash=token_hash, + user_id=user_id, + sub=sub, + sid=sid, + id_token_encrypted=b"id-token-bytes", + access_token_encrypted=b"access-token-bytes", + refresh_token_encrypted=b"refresh-token-bytes", + access_token_expires_at=now + expires_in, + session_expires_at=now + expires_in, + created_at=now, + revoked_at=revoked_at, + ) + + +class TestCreateGet: + async def test_create_returns_assigned_id(self, postgres_store: PostgresStore): + user = await _seed_user(postgres_store) + created = await postgres_store.oidc_session_repo.create_session( + _session(user.id), + ) + assert created.id > 0 + # encrypted byte payloads round-trip verbatim — the repo doesn't + # try to crypt them, that lives in the auth service. + assert created.id_token_encrypted == b"id-token-bytes" + assert created.refresh_token_encrypted == b"refresh-token-bytes" + + async def test_get_by_token_hash(self, postgres_store: PostgresStore): + user = await _seed_user(postgres_store) + await postgres_store.oidc_session_repo.create_session( + _session(user.id, token_hash="hash-abc"), + ) + fetched = await postgres_store.oidc_session_repo.get_by_token_hash("hash-abc") + assert fetched is not None + assert fetched.user_id == user.id + + async def test_revoked_session_hidden_from_lookup( + self, + postgres_store: PostgresStore, + ): + user = await _seed_user(postgres_store) + await postgres_store.oidc_session_repo.create_session( + _session(user.id, token_hash="hash-revoked", revoked=True), + ) + assert await postgres_store.oidc_session_repo.get_by_token_hash("hash-revoked") is None + + async def test_expired_session_hidden_from_lookup( + self, + postgres_store: PostgresStore, + ): + user = await _seed_user(postgres_store) + await postgres_store.oidc_session_repo.create_session( + _session(user.id, token_hash="hash-expired", expires_in=timedelta(seconds=-60)), + ) + assert await postgres_store.oidc_session_repo.get_by_token_hash("hash-expired") is None + + +class TestRevoke: + async def test_revoke_by_sid_marks_all(self, postgres_store: PostgresStore): + user = await _seed_user(postgres_store) + await postgres_store.oidc_session_repo.create_session( + _session(user.id, token_hash="t1", sid="shared-sid"), + ) + await postgres_store.oidc_session_repo.create_session( + _session(user.id, token_hash="t2", sid="shared-sid"), + ) + count = await postgres_store.oidc_session_repo.revoke_by_sid("shared-sid") + assert count == 2 + assert await postgres_store.oidc_session_repo.get_by_token_hash("t1") is None + + async def test_revoke_by_sid_missing_returns_zero( + self, + postgres_store: PostgresStore, + ): + assert await postgres_store.oidc_session_repo.revoke_by_sid("never") == 0 + + async def test_revoke_by_user(self, postgres_store: PostgresStore): + user = await _seed_user(postgres_store) + await postgres_store.oidc_session_repo.create_session( + _session(user.id, token_hash="u-t1"), + ) + revoked = await postgres_store.oidc_session_repo.revoke_by_user(user.id) + assert revoked == 1 + assert await postgres_store.oidc_session_repo.get_by_token_hash("u-t1") is None + + +class TestExpiry: + async def test_delete_expired_only_removes_long_dead( + self, + postgres_store: PostgresStore, + ): + user = await _seed_user(postgres_store) + # Repo keeps a 7-day grace window after expiry — see the + # implementation note in oidc_session_repo.py. A row that just + # expired must NOT be deleted yet. + await postgres_store.oidc_session_repo.create_session( + _session(user.id, token_hash="recently-expired", expires_in=timedelta(seconds=-60)), + ) + removed = await postgres_store.oidc_session_repo.delete_expired() + assert removed == 0 diff --git a/tests/integration/test_partition_membership_repo.py b/tests/integration/test_partition_membership_repo.py new file mode 100644 index 000000000..157288bba --- /dev/null +++ b/tests/integration/test_partition_membership_repo.py @@ -0,0 +1,61 @@ +"""Phase 7A.2 — PgPartitionMembershipRepository against a real Postgres. + +Split out of ``test_user_repo.py`` when partition memberships moved off +``PgUserRepository`` into their own repo (one-repo-per-entity, 7A.2). +""" + +from __future__ import annotations + +import pytest +from core.models.user import PartitionRole, User, UserPartition +from services.storage.postgres_store import PostgresStore + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +def _user(**overrides) -> User: + defaults = { + "display_name": "Alice", + "email": "alice@example.com", + "is_admin": False, + } + defaults.update(overrides) + return User(**defaults) + + +class TestPartitionMemberships: + async def test_assign_then_list(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user()) + await postgres_store.partition_repo.create_partition("docs") + await postgres_store.membership_repo.assign_partition( + UserPartition(user_id=user.id, partition="docs", role=PartitionRole.OWNER), + ) + memberships = await postgres_store.membership_repo.list_user_partitions(user.id) + assert len(memberships) == 1 + assert memberships[0].partition == "docs" + assert memberships[0].role == PartitionRole.OWNER + + async def test_assign_is_idempotent_and_updates_role( + self, + postgres_store: PostgresStore, + ): + user = await postgres_store.user_repo.create_user(_user()) + await postgres_store.partition_repo.create_partition("docs") + await postgres_store.membership_repo.assign_partition( + UserPartition(user_id=user.id, partition="docs", role=PartitionRole.VIEWER), + ) + await postgres_store.membership_repo.assign_partition( + UserPartition(user_id=user.id, partition="docs", role=PartitionRole.OWNER), + ) + memberships = await postgres_store.membership_repo.list_user_partitions(user.id) + assert len(memberships) == 1 + assert memberships[0].role == PartitionRole.OWNER + + async def test_remove_partition(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user()) + await postgres_store.partition_repo.create_partition("docs") + await postgres_store.membership_repo.assign_partition( + UserPartition(user_id=user.id, partition="docs"), + ) + assert await postgres_store.membership_repo.remove_partition(user.id, "docs") is True + assert await postgres_store.membership_repo.list_user_partitions(user.id) == [] diff --git a/tests/integration/test_partition_repo.py b/tests/integration/test_partition_repo.py new file mode 100644 index 000000000..8594e502e --- /dev/null +++ b/tests/integration/test_partition_repo.py @@ -0,0 +1,98 @@ +"""Phase 7F — PgPartitionRepository against a real Postgres.""" + +from __future__ import annotations + +import pytest +from services.storage.postgres_store import PostgresStore + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +class TestCreateList: + async def test_create_then_get(self, postgres_store: PostgresStore): + repo = postgres_store.partition_repo + created = await repo.create_partition("alpha") + assert created["partition"] == "alpha" + # ``created_at`` comes from the DB default. + assert created.get("created_at") + + async def test_list_returns_all_known(self, postgres_store: PostgresStore): + repo = postgres_store.partition_repo + await repo.create_partition("p1") + await repo.create_partition("p2") + names = {row["partition"] for row in await repo.list_partitions()} + assert {"p1", "p2"} <= names + + async def test_create_is_idempotent_per_name(self, postgres_store: PostgresStore): + repo = postgres_store.partition_repo + await repo.create_partition("dup") + # The legacy method swallows the conflict and returns the existing row + # rather than raising. Orchestrators rely on this for "ensure exists". + await repo.create_partition("dup") + assert await repo.partition_exists("dup") is True + assert len([r for r in await repo.list_partitions() if r["partition"] == "dup"]) == 1 + + +class TestExistsCounts: + async def test_partition_exists_returns_false_for_missing( + self, + postgres_store: PostgresStore, + ): + repo = postgres_store.partition_repo + assert await repo.partition_exists("never-created") is False + + async def test_total_file_count_starts_at_zero(self, postgres_store: PostgresStore): + repo = postgres_store.partition_repo + assert await repo.get_total_file_count() == 0 + + +class TestDelete: + async def test_delete_removes_the_partition_row( + self, + postgres_store: PostgresStore, + ): + repo = postgres_store.partition_repo + await repo.create_partition("doomed") + assert await repo.partition_exists("doomed") is True + removed = await repo.delete_partition("doomed") + assert removed is True + assert await repo.partition_exists("doomed") is False + + async def test_delete_missing_returns_false(self, postgres_store: PostgresStore): + repo = postgres_store.partition_repo + assert await repo.delete_partition("ghost") is False + + async def test_delete_cascades_files_and_decrements_uploader_count( + self, + postgres_store: PostgresStore, + ): + """Regression: ``files.partition_name`` has no DB-level CASCADE, so the + repo must delete file rows itself before dropping the partition. Also + verifies the per-uploader ``file_count`` decrement. + """ + partition_repo = postgres_store.partition_repo + document_repo = postgres_store.document_repo + user_repo = postgres_store.user_repo + + uploader = await user_repo.create_legacy_user(display_name="Uploader") + uploader_id = uploader["id"] + + await partition_repo.create_partition("cascade-me") + await document_repo.add_file_to_partition( + file_id="f1", + partition="cascade-me", + user_id=uploader_id, + ) + await document_repo.add_file_to_partition( + file_id="f2", + partition="cascade-me", + user_id=uploader_id, + ) + assert await partition_repo.get_partition_file_count("cascade-me") == 2 + + assert await partition_repo.delete_partition("cascade-me") is True + + assert await partition_repo.partition_exists("cascade-me") is False + assert await partition_repo.get_partition_file_count("cascade-me") == 0 + refreshed = await user_repo.get_user_dict_by_id(uploader_id) + assert refreshed["file_count"] == 0 diff --git a/tests/integration/test_postgres_store.py b/tests/integration/test_postgres_store.py new file mode 100644 index 000000000..26ff08c8c --- /dev/null +++ b/tests/integration/test_postgres_store.py @@ -0,0 +1,117 @@ +"""Phase 7F — PostgresStore composite against a real Postgres.""" + +from __future__ import annotations + +import pytest +from core.config.infrastructure import RDBConfig +from core.ports.catalog_store import CatalogStore +from services.persistence.audit_log_repo import PgAuditLogRepository +from services.persistence.chunk_repo import PgChunkRepository +from services.persistence.conversation_repo import PgConversationRepository +from services.persistence.document_repo import PgDocumentRepository +from services.persistence.entity_repo import PgEntityRepository +from services.persistence.idempotency_repo import PgIdempotencyRepository +from services.persistence.job_repo import PgJobRepository +from services.persistence.model_endpoint_repo import PgModelEndpointRepository +from services.persistence.oidc_session_repo import PgOIDCSessionRepository +from services.persistence.partition_repo import PgPartitionRepository +from services.persistence.preset_repo import PgPresetRepository +from services.persistence.prompt_repo import PgPromptRepository +from services.persistence.topic_tag_repo import PgTopicTagRepository +from services.persistence.user_repo import PgUserRepository +from services.persistence.workspace_repo import PgWorkspaceRepository +from services.storage.postgres_store import PostgresStore + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +_EXPECTED_REPO_TYPES = { + "document_repo": PgDocumentRepository, + "user_repo": PgUserRepository, + "partition_repo": PgPartitionRepository, + "oidc_session_repo": PgOIDCSessionRepository, + "workspace_repo": PgWorkspaceRepository, + "job_repo": PgJobRepository, + "chunk_repo": PgChunkRepository, + "prompt_repo": PgPromptRepository, + "conversation_repo": PgConversationRepository, + "audit_log_repo": PgAuditLogRepository, + "idempotency_repo": PgIdempotencyRepository, + "entity_repo": PgEntityRepository, + "topic_tag_repo": PgTopicTagRepository, + "model_endpoint_repo": PgModelEndpointRepository, + "preset_repo": PgPresetRepository, +} + + +class TestComposite: + async def test_satisfies_catalog_store_abc(self, postgres_store: PostgresStore): + assert isinstance(postgres_store, CatalogStore) + + async def test_pool_open_after_initialize(self, postgres_store: PostgresStore): + # ``pool`` raises if initialize() never ran. The session-scoped + # fixture initialises so this must succeed. + assert postgres_store.pool is not None + + @pytest.mark.parametrize("name", sorted(_EXPECTED_REPO_TYPES)) + async def test_all_fifteen_repos_exposed( + self, + postgres_store: PostgresStore, + name: str, + ): + repo = getattr(postgres_store, name) + assert isinstance(repo, _EXPECTED_REPO_TYPES[name]) + + async def test_repo_properties_are_idempotent(self, postgres_store: PostgresStore): + # Each property must return the exact same instance every call — + # orchestrators cache repos and rely on identity. + first = postgres_store.document_repo + assert postgres_store.document_repo is first + + +class TestMigrationIdempotency: + async def test_run_migrations_twice_is_safe(self, postgres_store: PostgresStore): + # The session fixture already ran migrations once. Re-running must + # be a no-op — every phase-7 Alembic revision is supposed to guard + # its DDL with an inspector check (see CLAUDE.md "Alembic Migration + # Idempotency"). + await postgres_store._conn.run_migrations() # noqa: SLF001 + + +class TestLifecycle: + """Initialise/shutdown a fresh store without disturbing the session pool.""" + + async def test_initialize_then_shutdown(self, test_rdb_config: RDBConfig): + # The session store already migrated the DB so we can skip migrations + # here and just exercise the pool lifecycle. + store = PostgresStore(test_rdb_config, run_migrations=False) + await store.initialize() + try: + assert store.pool is not None + finally: + await store.shutdown() + with pytest.raises(RuntimeError, match="initialize"): + _ = store.pool + + async def test_initialize_is_idempotent_no_remigrate(self, test_rdb_config: RDBConfig): + # ``get_vectordb()`` re-invokes the actor's initialize on every + # request (it is a FastAPI ``Depends``). Re-running must NOT re-run + # Alembic — a per-request migration storm starves the PG pool and + # surfaces as 500s under load (regression for the api-tests flake). + store = PostgresStore(test_rdb_config) + calls = 0 + original = store._conn.run_migrations # noqa: SLF001 + + async def counting_run_migrations() -> None: + nonlocal calls + calls += 1 + await original() + + store._conn.run_migrations = counting_run_migrations # type: ignore[method-assign] # noqa: SLF001 + try: + await store.initialize() + await store.initialize() + await store.initialize() + assert calls == 1 + finally: + await store.shutdown() diff --git a/tests/integration/test_stores.py b/tests/integration/test_stores.py new file mode 100644 index 000000000..6b1ae06bc --- /dev/null +++ b/tests/integration/test_stores.py @@ -0,0 +1,40 @@ +"""Phase 7F — cross-store full-cycle integration (Phase 7C handoff). + +Phase 7B's :class:`MilvusVectorStore` is wired through DI now, but a real +cross-store cycle (``create partition → upsert pre-embedded chunks → search +→ delete``) still needs: + +* a Milvus instance reachable from the test runner (Person B's + ``test_milvus_store_integration.py`` already covers that piece in + isolation), and +* a fixture that builds *both* stores against the same Milvus collection + + Postgres database — the existing ``postgres_store`` fixture in + ``conftest.py`` doesn't yet hand out a Milvus store. + +The combined fixture lands as part of Phase 7C (shim) so the assertion +matches the legacy ``MilvusDB`` cross-store flow byte-for-byte. Until then +this test stays ``xfail(strict=True)`` so the day the fixture is added and +the body filled in, the unintended pass trips a clear failure. +""" + +from __future__ import annotations + +import pytest + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +@pytest.mark.xfail( + reason="cross-store fixture lands in Phase 7C — see REFACTORING_DECISION_LOG.md", + strict=True, +) +async def test_cross_store_full_cycle(): + # The shape the eventual test will take: + # 1. ``postgres_store.partition_repo`` creates a partition row. + # 2. ``vector_store.upsert`` inserts pre-embedded chunks tagged with + # that partition. + # 3. ``vector_store.search`` round-trips the embedding and returns the + # ids/text. + # 4. ``postgres_store.partition_repo.delete_partition`` cascades the + # catalog rows; ``vector_store.delete`` clears the Milvus side. + raise NotImplementedError("cross-store fixture lands in Phase 7C") diff --git a/tests/integration/test_user_repo.py b/tests/integration/test_user_repo.py new file mode 100644 index 000000000..7da6139d3 --- /dev/null +++ b/tests/integration/test_user_repo.py @@ -0,0 +1,151 @@ +"""Phase 7F — PgUserRepository against a real Postgres.""" + +from __future__ import annotations + +import pytest +from core.models.user import User +from services.persistence.user_repo import _hash_token +from services.storage.postgres_store import PostgresStore + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +def _user(**overrides) -> User: + defaults = { + "display_name": "Alice", + "email": "alice@example.com", + "is_admin": False, + } + defaults.update(overrides) + return User(**defaults) + + +class TestCreateGet: + async def test_create_returns_assigned_id(self, postgres_store: PostgresStore): + created = await postgres_store.user_repo.create_user(_user()) + assert created.id > 0 + assert created.display_name == "Alice" + assert created.email == "alice@example.com" + + async def test_get_by_id(self, postgres_store: PostgresStore): + created = await postgres_store.user_repo.create_user(_user(display_name="Bob")) + fetched = await postgres_store.user_repo.get_user(created.id) + assert fetched is not None + assert fetched.display_name == "Bob" + + async def test_get_missing_returns_none(self, postgres_store: PostgresStore): + assert await postgres_store.user_repo.get_user(9999) is None + + async def test_email_lowercased_on_insert(self, postgres_store: PostgresStore): + created = await postgres_store.user_repo.create_user( + _user(email="MIXED@Example.COM"), + ) + assert created.email == "mixed@example.com" + + async def test_get_by_email_case_insensitive(self, postgres_store: PostgresStore): + await postgres_store.user_repo.create_user( + _user(email="carol@example.com"), + ) + # Lookup uppercases the input — the repo normalises. + fetched = await postgres_store.user_repo.get_user_by_email("CAROL@Example.com") + assert fetched is not None + assert fetched.email == "carol@example.com" + + async def test_get_by_external_id(self, postgres_store: PostgresStore): + await postgres_store.user_repo.create_user( + _user(external_user_id="kc-alice-uuid"), + ) + fetched = await postgres_store.user_repo.get_user_by_external_id("kc-alice-uuid") + assert fetched is not None + assert fetched.external_user_id == "kc-alice-uuid" + + +class TestLegacyTokenFlow: + async def test_create_legacy_user_returns_plaintext_token( + self, + postgres_store: PostgresStore, + ): + result = await postgres_store.user_repo.create_legacy_user( + display_name="Tokened", + external_user_id=None, + email=None, + is_admin=False, + file_quota=None, + ) + assert result["token"].startswith("or-") + # ``"or-"`` (3 chars) + ``secrets.token_hex(16)`` (32 hex chars) = 35. + assert len(result["token"]) == 35 + + async def test_get_user_by_token_hash_roundtrip( + self, + postgres_store: PostgresStore, + ): + created = await postgres_store.user_repo.create_legacy_user( + display_name="Tokened2", + external_user_id=None, + email=None, + is_admin=False, + file_quota=None, + ) + looked_up = await postgres_store.user_repo.get_user_by_token( + _hash_token(created["token"]), + ) + assert looked_up is not None + assert looked_up.id == created["id"] + + async def test_regenerate_token_invalidates_old( + self, + postgres_store: PostgresStore, + ): + created = await postgres_store.user_repo.create_legacy_user( + display_name="Tokened3", + external_user_id=None, + email=None, + is_admin=False, + file_quota=None, + ) + new = await postgres_store.user_repo.regenerate_user_token(created["id"]) + assert new is not None + assert new["token"] != created["token"] + # old hash no longer resolves + assert ( + await postgres_store.user_repo.get_user_by_token( + _hash_token(created["token"]), + ) + is None + ) + + +class TestUpdateDelete: + async def test_update_user_fields(self, postgres_store: PostgresStore): + created = await postgres_store.user_repo.create_user(_user()) + updated = await postgres_store.user_repo.update_user( + created.id, + display_name="Renamed", + ) + assert updated is not None + assert updated.display_name == "Renamed" + + async def test_unknown_field_is_ignored(self, postgres_store: PostgresStore): + created = await postgres_store.user_repo.create_user(_user()) + # ``password_hash`` is in the domain model but not the schema — + # the repo silently ignores it instead of failing. + updated = await postgres_store.user_repo.update_user( + created.id, + password_hash="ignored", + ) + assert updated is not None + assert updated.display_name == "Alice" + + async def test_delete_returns_true(self, postgres_store: PostgresStore): + created = await postgres_store.user_repo.create_user(_user()) + assert await postgres_store.user_repo.delete_user(created.id) is True + assert await postgres_store.user_repo.get_user(created.id) is None + + async def test_count_users(self, postgres_store: PostgresStore): + assert await postgres_store.user_repo.count_users() == 0 + await postgres_store.user_repo.create_user(_user(display_name="A")) + await postgres_store.user_repo.create_user( + _user(display_name="B", email="b@example.com"), + ) + assert await postgres_store.user_repo.count_users() == 2 diff --git a/tests/integration/test_workspace_repo.py b/tests/integration/test_workspace_repo.py new file mode 100644 index 000000000..d52507365 --- /dev/null +++ b/tests/integration/test_workspace_repo.py @@ -0,0 +1,139 @@ +"""Phase 7F — PgWorkspaceRepository against a real Postgres.""" + +from __future__ import annotations + +import pytest +from core.models.catalog import DocumentRecord +from core.models.workspace import Workspace +from services.storage.postgres_store import PostgresStore + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +async def _seed_partition_and_files( + store: PostgresStore, + partition: str = "ws-p", + file_ids: tuple[str, ...] = ("f1", "f2", "f3"), +) -> str: + await store.partition_repo.create_partition(partition) + for fid in file_ids: + await store.document_repo.create_document( + DocumentRecord(id=fid, file_id=fid, partition=partition, filename=f"{fid}.pdf"), + ) + return partition + + +def _workspace(workspace_id: str = "ws1", partition: str = "ws-p", **extra) -> Workspace: + return Workspace(workspace_id=workspace_id, partition=partition, **extra) + + +class TestCreateGetList: + async def test_create_then_get(self, postgres_store: PostgresStore): + await _seed_partition_and_files(postgres_store) + await postgres_store.workspace_repo.create_workspace( + _workspace("ws1", display_name="My workspace"), + ) + fetched = await postgres_store.workspace_repo.get_workspace("ws1") + assert fetched is not None + assert fetched.workspace_id == "ws1" + assert fetched.display_name == "My workspace" + + async def test_list_filters_by_partition(self, postgres_store: PostgresStore): + await _seed_partition_and_files(postgres_store, partition="a") + await _seed_partition_and_files(postgres_store, partition="b", file_ids=("b1",)) + repo = postgres_store.workspace_repo + await repo.create_workspace(_workspace("ws-a", partition="a")) + await repo.create_workspace(_workspace("ws-b", partition="b")) + only_a = await repo.list_workspaces("a") + assert {w.workspace_id for w in only_a} == {"ws-a"} + + +class TestFileMembership: + async def test_add_then_list_workspace_files(self, postgres_store: PostgresStore): + await _seed_partition_and_files(postgres_store) + repo = postgres_store.workspace_repo + await repo.create_workspace(_workspace("ws1")) + missing = await repo.add_files_to_workspace("ws1", ["f1", "f2"]) + assert missing == [] + files = await repo.list_workspace_files("ws1") + assert set(files) == {"f1", "f2"} + + async def test_add_reports_unknown_file_ids(self, postgres_store: PostgresStore): + await _seed_partition_and_files(postgres_store) + repo = postgres_store.workspace_repo + await repo.create_workspace(_workspace("ws1")) + missing = await repo.add_files_to_workspace("ws1", ["f1", "ghost", "f2"]) + assert missing == ["ghost"] + assert set(await repo.list_workspace_files("ws1")) == {"f1", "f2"} + + async def test_add_is_idempotent(self, postgres_store: PostgresStore): + await _seed_partition_and_files(postgres_store) + repo = postgres_store.workspace_repo + await repo.create_workspace(_workspace("ws1")) + await repo.add_files_to_workspace("ws1", ["f1"]) + await repo.add_files_to_workspace("ws1", ["f1"]) + assert await repo.list_workspace_files("ws1") == ["f1"] + + async def test_remove_file_from_workspace(self, postgres_store: PostgresStore): + await _seed_partition_and_files(postgres_store) + repo = postgres_store.workspace_repo + await repo.create_workspace(_workspace("ws1")) + await repo.add_files_to_workspace("ws1", ["f1", "f2"]) + assert await repo.remove_file_from_workspace("ws1", "f1") is True + assert await repo.list_workspace_files("ws1") == ["f2"] + + async def test_get_file_workspaces_is_partition_scoped( + self, + postgres_store: PostgresStore, + ): + await _seed_partition_and_files(postgres_store, partition="a") + await _seed_partition_and_files(postgres_store, partition="b", file_ids=("f1",)) + repo = postgres_store.workspace_repo + await repo.create_workspace(_workspace("ws-a", partition="a")) + await repo.create_workspace(_workspace("ws-b", partition="b")) + await repo.add_files_to_workspace("ws-a", ["f1"]) + await repo.add_files_to_workspace("ws-b", ["f1"]) + # ``f1`` exists in both partitions as distinct ``files`` rows; + # the lookup must only return the workspace in partition "a". + in_a = await repo.get_file_workspaces("f1", "a") + assert in_a == ["ws-a"] + + +class TestDeleteWorkspace: + async def test_returns_orphan_file_ids(self, postgres_store: PostgresStore): + await _seed_partition_and_files(postgres_store) + repo = postgres_store.workspace_repo + await repo.create_workspace(_workspace("ws1")) + await repo.add_files_to_workspace("ws1", ["f1", "f2"]) + orphans = await repo.delete_workspace("ws1") + # Both files were only in ws1 — both come back as orphans. + assert set(orphans) == {"f1", "f2"} + assert await repo.get_workspace("ws1") is None + + async def test_files_shared_with_other_workspaces_are_not_orphaned( + self, + postgres_store: PostgresStore, + ): + await _seed_partition_and_files(postgres_store) + repo = postgres_store.workspace_repo + await repo.create_workspace(_workspace("ws1")) + await repo.create_workspace(_workspace("ws2")) + await repo.add_files_to_workspace("ws1", ["f1", "f2"]) + await repo.add_files_to_workspace("ws2", ["f1"]) # f1 shared + orphans = await repo.delete_workspace("ws1") + # f1 is still in ws2 so it's not orphaned. f2 only lived in ws1. + assert set(orphans) == {"f2"} + + async def test_remove_file_from_all_workspaces( + self, + postgres_store: PostgresStore, + ): + await _seed_partition_and_files(postgres_store) + repo = postgres_store.workspace_repo + await repo.create_workspace(_workspace("ws1")) + await repo.create_workspace(_workspace("ws2")) + await repo.add_files_to_workspace("ws1", ["f1"]) + await repo.add_files_to_workspace("ws2", ["f1"]) + await repo.remove_file_from_all_workspaces("f1", "ws-p") + assert await repo.list_workspace_files("ws1") == [] + assert await repo.list_workspace_files("ws2") == [] diff --git a/tests/test_vectordb.py b/tests/test_vectordb.py deleted file mode 100644 index 39d19c48a..000000000 --- a/tests/test_vectordb.py +++ /dev/null @@ -1,24 +0,0 @@ -from openrag.components.indexer.vectordb import MilvusDB - - -class TestMilvusDB: - def test_complex_templating(self): - partition = ["bob.localhost"] - filter = {"file_id": "a3b2c1", "custom_param": 314} - expr, params = MilvusDB._build_expr_template_and_params(partition, filter) - assert expr == "partition in {partition} and file_id == {file_id} and custom_param == {custom_param}" - - # Note how parameter values are not converted to str - assert params == { - "partition": ["bob.localhost"], - "file_id": "a3b2c1", - "custom_param": 314, - } - - def test_templating_no_filter(self): - # If there is no filter, the search is run on every document - partition = ["all"] - filter = {} - expr, params = MilvusDB._build_expr_template_and_params(partition, filter) - assert expr == "" - assert params == {} diff --git a/uv.lock b/uv.lock index 03bc35121..3f088049d 100644 --- a/uv.lock +++ b/uv.lock @@ -6,6 +6,15 @@ resolution-markers = [ "python_full_version < '3.13'", ] +[[package]] +name = "aiobreaker" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7f/eb/749ef48d3227fd62d500ff01fcd451f10111e00d822c200eb51782ba076a/aiobreaker-1.2.0.tar.gz", hash = "sha256:217a9cfa12e520bb2dd1934bace281d1d7deb8d7630dd183a6295fd22e323ce7", size = 15947, upload-time = "2021-05-17T11:58:05.278Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ff/c8/4cd4b2834012ffc71ae3fd69187f08a17f01f3937527b6b5e077f4f5d0db/aiobreaker-1.2.0-py3-none-any.whl", hash = "sha256:f275decad78bdd161715afeee67e5dde7967de54c836648b44f4eea1b5e41d60", size = 20700, upload-time = "2021-05-17T11:58:04.192Z" }, +] + [[package]] name = "aiofile" version = "3.9.0" @@ -2635,6 +2644,7 @@ name = "openrag" version = "1.1.11" source = { editable = "." } dependencies = [ + { name = "aiobreaker" }, { name = "aiopath" }, { name = "alembic" }, { name = "asyncpg" }, @@ -2683,6 +2693,7 @@ dependencies = [ { name = "ruff" }, { name = "spire-doc" }, { name = "sqlalchemy-utils" }, + { name = "tenacity" }, { name = "torch" }, { name = "umap-learn" }, ] @@ -2699,6 +2710,7 @@ lint = [ [package.metadata] requires-dist = [ + { name = "aiobreaker", specifier = ">=1.2.0" }, { name = "aiopath", specifier = ">=0.7.7" }, { name = "alembic", specifier = ">=1.17.0" }, { name = "asyncpg", specifier = ">=0.30.0" }, @@ -2747,6 +2759,7 @@ requires-dist = [ { name = "ruff", specifier = ">=0.14.1" }, { name = "spire-doc", specifier = ">=13.1.0" }, { name = "sqlalchemy-utils" }, + { name = "tenacity", specifier = ">=8.2.0" }, { name = "torch", specifier = ">=2.4.1" }, { name = "umap-learn", specifier = ">=0.5.9.post2" }, ]