refactor(phase 14): Model Endpoints Registry and Per-Partition Presets - #444
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds Phase 14 admin APIs and persistence: model endpoint and preset registries, partition config endpoints, new Pydantic config/models, DB schema + migration, DI getters and container wiring, per-row indexation pipeline plumbing, and broad unit tests. ChangesPhase 14 Admin APIs: Model Endpoints, Presets, and Partition Configuration
Sequence Diagram(s)sequenceDiagram
participant Client
participant APIApp
participant ModelEndpointsRouter
participant DIGetter as get_model_endpoint_service
participant ModelEndpointService
Client->>APIApp: POST /model-endpoints (CreateModelEndpointRequest)
APIApp->>ModelEndpointsRouter: route to create handler
ModelEndpointsRouter->>DIGetter: resolve service from container
DIGetter->>ModelEndpointService: create_model_endpoint(payload)
ModelEndpointService-->>ModelEndpointsRouter: ModelEndpointResponse
ModelEndpointsRouter-->>APIApp: 201 Created
APIApp-->>Client: 201 Created
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
✨ Finishing Touches🧪 Generate unit tests (beta)
|
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
openrag/api/schemas/admin/preset_schemas.py (1)
52-56: ⚡ Quick winUse field name from validator context for clearer error messages.
Line 56 passes the hardcoded string
"name"to_reject_explicit_null, which works correctly for this validator since it only applies to thenamefield. However, this pattern is brittle—if someone later adds more fields to the validator decorator, the error message will remain hardcoded as"name cannot be null". Usinginfo.field_namemakes the validator more maintainable and consistent with the pattern used elsewhere.♻️ Proposed refactor to use field name from context
- `@field_validator`("name") + `@field_validator`("name", mode="before") `@classmethod` - def validate_name(cls, value: str | None) -> str | None: + def validate_name(cls, value: str | None, info) -> str | None: """Normalize the optional replacement name and reject null.""" - return _normalize_name(_reject_explicit_null("name", value)) + return _normalize_name(_reject_explicit_null(info.field_name, value))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/api/schemas/admin/preset_schemas.py` around lines 52 - 56, The validator validate_name uses a hardcoded "name" when calling _reject_explicit_null which is brittle; change the validator to accept the pydantic FieldValidationInfo (e.g., add an info parameter) and pass info.field_name into _reject_explicit_null so the error message uses the actual field name (update validate_name signature to include info and call _normalize_name(_reject_explicit_null(info.field_name, value))). Ensure imports/types for the info param match the project’s pydantic version and keep the `@field_validator`("name") decorator as-is.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@openrag/api/schemas/admin/partition_schemas.py`:
- Around line 58-62: The validator validate_non_empty_name currently passes the
hardcoded string "name" to _reject_explicit_null, causing incorrect error
messages; update validate_non_empty_name to accept the validator info (e.g.,
change signature to def validate_non_empty_name(cls, value: str | None, info) ->
str | None) and call _reject_explicit_null(info.field_name, value) before
passing the result into _normalize_name so the error message identifies the
actual field (embedder, indexation_preset, or retrieval_preset).
---
Nitpick comments:
In `@openrag/api/schemas/admin/preset_schemas.py`:
- Around line 52-56: The validator validate_name uses a hardcoded "name" when
calling _reject_explicit_null which is brittle; change the validator to accept
the pydantic FieldValidationInfo (e.g., add an info parameter) and pass
info.field_name into _reject_explicit_null so the error message uses the actual
field name (update validate_name signature to include info and call
_normalize_name(_reject_explicit_null(info.field_name, value))). Ensure
imports/types for the info param match the project’s pydantic version and keep
the `@field_validator`("name") decorator as-is.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: f7a5e4be-26cd-41c4-a7d4-741aa2584645
📒 Files selected for processing (15)
openrag/api/main.pyopenrag/api/routers/admin/model_endpoints.pyopenrag/api/routers/admin/partitions.pyopenrag/api/routers/admin/presets.pyopenrag/api/schemas/admin/__init__.pyopenrag/api/schemas/admin/model_endpoint_schemas.pyopenrag/api/schemas/admin/partition_schemas.pyopenrag/api/schemas/admin/preset_schemas.pyopenrag/di/providers.pytests/unit/api/routers/admin/test_phase14_admin_routers.pytests/unit/api/routers/admin/test_phase14_partition_routes.pytests/unit/api/routers/test_router_imports.pytests/unit/api/schemas/admin/test_phase14_schemas.pytests/unit/api/test_main_proxy_headers.pytests/unit/di/test_container.py
e596cac to
81fe53c
Compare
81fe53c to
c8f8b61
Compare
- `core/config/model_endpoints.py`: ModelEndpointConfig, ModelsConfig (frozen ConfigMixin with mutable dicts), ModelEndpointRow - `core/ports/model_endpoint_repo.py`: redesign port with 7 explicit methods (create, get, list_all, update, rename, delete, set_default); drop upsert - `services/persistence/model_endpoint_repo.py`: update stub to match new port - `tests/unit/core/config/test_model_endpoints.py`: frozen-field + mutable-dict invariant tests Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- `core/config/indexation_pipeline.py`: IndexationPipelineConfig (chunking, VLM, contextualization, entity extraction, topic tagging) - `core/config/retrieval_pipeline.py`: RetrievalPipelineConfig (flat per-partition, named endpoint refs) - `core/models/preset.py`: PresetRow, PartitionRow, PartitionConfig (runtime-resolved) - `core/config/presets.py`: PresetsConfig (frozen ConfigMixin with mutable dicts) - `core/config/root.py`: add presets + partitions fields to Settings - `tests/unit/core/config/test_pipeline_configs.py`: frozen-field + mutable-dict invariant tests Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (4)
openrag/core/config/model_endpoints.py (2)
10-10: ⚡ Quick winSwitch this relative import to a root-based absolute import.
Line 10 should follow the same import style used across this PR (
core.*from repo root).Proposed change
-from .base import ConfigMixin +from core.config.base import ConfigMixinAs per coding guidelines:
**/*.py: Use absolute imports from theopenrag/directory as the Python path root, avoiding relative imports across packages.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/core/config/model_endpoints.py` at line 10, Replace the relative import in model_endpoints.py with a root-based absolute import: change the statement that currently imports ConfigMixin via ".base" to import it from "core.config.base" so the file uses the same absolute import style as the rest of the project (look for the import of ConfigMixin in model_endpoints.py and update it to reference core.config.base.ConfigMixin).
48-48: ⚡ Quick winHarden
ModelEndpointRowvalidation to match the repository/admin contract.
ModelEndpointRowuses unconstrainedmodel_type: strand defaultbatch_size/timeoutwithoutgt=0, unlikeModelEndpointConfigand the admin request schemas (which enforceModelEndpointTypeandgt=0). Even thoughPgModelEndpointRepositoryis currently stubbed, tighteningModelEndpointRowprevents contract drift when persistence is implemented.Proposed change
-from typing import Any +from typing import Any, Literal @@ class ModelEndpointRow(BaseModel): @@ - model_type: str # embedder | reranker | llm | vlm + model_type: Literal["embedder", "reranker", "llm", "vlm"] @@ - batch_size: int = 32 - timeout: float = 30.0 + batch_size: int = Field(default=32, gt=0) + timeout: float = Field(default=30.0, gt=0)Also applies to: 51-52
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/core/config/model_endpoints.py` at line 48, ModelEndpointRow currently declares model_type: str and unconstrained batch_size/timeout which diverges from the admin/repository contract; change ModelEndpointRow to use the ModelEndpointType enum for model_type (import ModelEndpointType) and tighten batch_size and timeout with Pydantic Field validators (gt=0) and the same defaults as ModelEndpointConfig so the row schema matches the admin request/ModelEndpointConfig contract and prevents future drift.openrag/core/config/indexation_pipeline.py (1)
12-12: ⚡ Quick winUse package-root absolute import instead of a relative import.
Line 12 should use a root-based import to match the repository import convention.
Proposed change
-from .chunking import ChunkerConfig +from core.config.chunking import ChunkerConfigAs per coding guidelines:
**/*.py: Use absolute imports from theopenrag/directory as the Python path root, avoiding relative imports across packages.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/core/config/indexation_pipeline.py` at line 12, Replace the relative import in indexation_pipeline by using a package-root absolute import: change the statement importing ChunkerConfig (currently "from .chunking import ChunkerConfig") to import from the repository root module path (e.g. "from openrag.core.config.chunking import ChunkerConfig") so it follows the project's absolute-import convention and matches other modules.tests/unit/core/config/test_model_endpoints.py (1)
10-29: ⚡ Quick winAdd contract tests for
ModelEndpointRowvalidation boundaries.This file currently validates
ModelsConfigmutability/frozen behavior only. Please add negative tests for invalidmodel_typeand non-positivebatch_size/timeoutso the new row contract is regression-protected.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/core/config/test_model_endpoints.py` around lines 10 - 29, Add tests to validate ModelEndpointRow contract boundaries by creating three negative test functions in the same test module: (1) test_model_endpoint_row_invalid_model_type: instantiate ModelEndpointRow with an invalid model_type (e.g., "not-a-model") and assert it raises ValidationError via pytest.raises; (2) test_model_endpoint_row_non_positive_batch_size: create ModelEndpointRow with batch_size set to 0 and another with a negative value and assert each raises ValidationError; (3) test_model_endpoint_row_non_positive_timeout: create ModelEndpointRow with timeout 0 and with a negative value and assert each raises ValidationError. Use the ModelEndpointRow constructor directly and import ValidationError/pytest as in the file, ensuring each test targets the model_type, batch_size, and timeout fields respectively.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@openrag/core/config/indexation_pipeline.py`:
- Line 21: IndexationPipelineConfig currently types parsing_strategy and
contextualization_mode as plain str which allows any value; change their
annotations to typing.Literal with the supported options (e.g.,
parsing_strategy: Literal["pymupdf", "marker", "docling"] = "marker" and
contextualization_mode: Literal[...supported values...] = <current_default>) and
import Literal from typing (or typing_extensions for older Python) so Pydantic
validation will reject invalid config values; keep the existing default values
and update any tests or callers expecting str if necessary.
In `@openrag/core/config/retrieval_pipeline.py`:
- Around line 18-26: The schema currently allows any string for the retrieval
mode and an unconstrained similarity_threshold; update the model so "type" is a
Literal restricted to "single", "multiQuery", or "hyde" (import typing.Literal
and replace the current type: str) and tighten similarity_threshold to a
validated Field with bounds (ge=0, le=1) so out-of-range values are rejected;
keep existing Field constraints for top_k/top_n and leave reranker/llm as
optional strings.
In `@openrag/core/models/preset.py`:
- Around line 17-18: Add invariant validation for persisted partition/preset
fields so invalid states cannot be loaded from the DB: enforce that preset_type
(in the Preset model) is one of the allowed values ("indexation" or
"retrieval"), that numeric config fields like chat_history_depth are
non-negative, and that dimension (or any vector dimension fields) is a positive
integer; implement these checks in the model validation/constructor for Preset
and the related Partition/Config classes referenced around the same file (lines
near the current preset_type/config definitions and the blocks at the other
noted ranges) so any invalid persisted record raises a clear validation error
instead of leaking into runtime.
---
Nitpick comments:
In `@openrag/core/config/indexation_pipeline.py`:
- Line 12: Replace the relative import in indexation_pipeline by using a
package-root absolute import: change the statement importing ChunkerConfig
(currently "from .chunking import ChunkerConfig") to import from the repository
root module path (e.g. "from openrag.core.config.chunking import ChunkerConfig")
so it follows the project's absolute-import convention and matches other
modules.
In `@openrag/core/config/model_endpoints.py`:
- Line 10: Replace the relative import in model_endpoints.py with a root-based
absolute import: change the statement that currently imports ConfigMixin via
".base" to import it from "core.config.base" so the file uses the same absolute
import style as the rest of the project (look for the import of ConfigMixin in
model_endpoints.py and update it to reference core.config.base.ConfigMixin).
- Line 48: ModelEndpointRow currently declares model_type: str and unconstrained
batch_size/timeout which diverges from the admin/repository contract; change
ModelEndpointRow to use the ModelEndpointType enum for model_type (import
ModelEndpointType) and tighten batch_size and timeout with Pydantic Field
validators (gt=0) and the same defaults as ModelEndpointConfig so the row schema
matches the admin request/ModelEndpointConfig contract and prevents future
drift.
In `@tests/unit/core/config/test_model_endpoints.py`:
- Around line 10-29: Add tests to validate ModelEndpointRow contract boundaries
by creating three negative test functions in the same test module: (1)
test_model_endpoint_row_invalid_model_type: instantiate ModelEndpointRow with an
invalid model_type (e.g., "not-a-model") and assert it raises ValidationError
via pytest.raises; (2) test_model_endpoint_row_non_positive_batch_size: create
ModelEndpointRow with batch_size set to 0 and another with a negative value and
assert each raises ValidationError; (3)
test_model_endpoint_row_non_positive_timeout: create ModelEndpointRow with
timeout 0 and with a negative value and assert each raises ValidationError. Use
the ModelEndpointRow constructor directly and import ValidationError/pytest as
in the file, ensuring each test targets the model_type, batch_size, and timeout
fields respectively.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 3b346e97-366e-408c-97c8-a03236147482
📒 Files selected for processing (10)
openrag/core/config/indexation_pipeline.pyopenrag/core/config/model_endpoints.pyopenrag/core/config/presets.pyopenrag/core/config/retrieval_pipeline.pyopenrag/core/config/root.pyopenrag/core/models/preset.pyopenrag/core/ports/model_endpoint_repo.pyopenrag/services/persistence/model_endpoint_repo.pytests/unit/core/config/test_model_endpoints.pytests/unit/core/config/test_pipeline_configs.py
✅ Files skipped from review due to trivial changes (1)
- tests/unit/core/config/test_pipeline_configs.py
On low-resource machines the Milvus 2.6 gRPC handshake can exceed the default 10-second channel-ready wait, causing pymilvus to throw MilvusException (code=2, "Fail connecting to server … illegal connection params or server unavailable"). The underlying TCP port is open and the HTTP health endpoint responds OK — the channel just needs more time to complete the gRPC negotiation. - Pass `timeout=60` to both `MilvusClient` and `AsyncMilvusClient` at construction time in `MilvusVectorStore.__init__`. - Store the value in `self._timeout` so any future callers have a single source of truth. Root-cause confirmed by reproducing the error with the default timeout and a successful connect with timeout=60 against the same Milvus 2.6.11 container.
… migration - openrag/services/persistence/schema.py: add model_endpoints and pipeline_presets tables, extend partitions with 9 new columns, add indexation_config to files - openrag/services/persistence/migrations/alembic/versions/06dd2101ea3a_add_endpoints_presets_phase14.py: idempotent migration with composite PKs, JSONB defaults, check constraints Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…handler logs - Fix migration 06dd2101ea3a: `index_exists()` called with 1 arg instead of the required 2 (`table`, `index`), crashing Alembic at startup and causing all endpoints to return 500 (container degraded to None) - Add try/except in AuthMiddleware.dispatch() so a None container returns JSONResponse(503) instead of propagating AttributeError through ServerErrorMiddleware to a 500 - Add `message`, `method`, and `path` fields to both error handler log calls so failing requests are diagnosable without a full traceback Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…t/partition repos - PgModelEndpointRepository: 7 methods (create, get, list_all, update, rename, delete, set_default) with atomic transaction for set_default - PgPresetRepository: 5 methods (get, list_all, upsert, delete, count_partitions_using) with ON CONFLICT upsert and safe column dispatch - PgPartitionRepository: 3 new methods (get_partition_row, list_partition_rows, update_partition) + _row_to_full_dict helper - Extend PartitionRepository port with get_partition_row, list_partition_rows, update_partition abstract methods - Extend PresetRepository port with count_partitions_using abstract method - 18 new unit tests covering all new methods with fake pool/conn pattern Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
de28be1 to
80e28e7
Compare
- seed_defaults(): reads Settings + env-var overrides (EMBEDDER_ENDPOINT,
LLM_ENDPOINT, VLM_ENDPOINT, RERANKER_ENDPOINT and model counterparts),
inserts one is_default=True row per type when the DB is empty for that type
- load_all(): fetches all rows, rebuilds config.models.{embedder,llm,vlm,
reranker} dicts atomically via clear()+update(); injects virtual "default"
alias pointing to the is_default=True row per type
- create_model_endpoint(): validates model_type, raises 409 on duplicate
- update_model_endpoint(): updates fields and/or renames, evicts stale
client-cache entry, reloads config.models
- delete_model_endpoint(): guards last-endpoint-of-type, evicts cache
- set_default(): swaps is_default flag, evicts "default" cache alias
- validate_endpoint(): async HTTP probe of {url}/models, returns reachable +
model_found + models_served dict
- 17 unit tests covering all methods
- PresetService with seed_defaults(), load_all(), create_preset(), update_preset(), delete_preset(), get_preset(), list_presets() - 6 hardcoded default seeds: 3 indexation (default, legal, finance) and 3 retrieval (default, multiquery, hyde) - Pydantic validation via IndexationPipelineConfig / RetrievalPipelineConfig - Atomic in-place swap of config.presets.indexation / .retrieval dicts - Back-reference to PartitionService.load_partitions() on preset update - Referential integrity guard: raises 422 if partition(s) use the preset - 17 unit tests, all passing
This reverts commit f6df9ab.
Phase 14J's per-partition resolution made IndexingService.add_file raise PARTITION_NOT_FOUND when indexing into a partition absent from config.partitions. The legacy flow auto-created an unknown partition and then indexed, so this was a regression for that workflow. Restore the behaviour by ensuring the partition exists before resolving its indexation config: create it with default presets and the uploader as owner, refreshing the in-memory cache if a row already exists. This is a no-op in the legacy passthrough (no preset registry) or when no partition service is wired, so the strict check still stands in those paths. The container now injects the partition service into IndexingService.
The model-endpoint and preset create routes forwarded the request body the wrong way. POST /model-endpoints passed body.model_dump() (a dict) where the service expects a ModelEndpointRow, raising "'dict' object has no attribute 'model_type'". POST /presets passed the dict as the single positional 'name' argument, while create_preset takes (name, preset_type, config), raising a TypeError. Build a timestamped ModelEndpointRow for the endpoint route (mirroring the seed path) and unpack the body into keyword arguments for the preset route. The update/delete routes already forwarded correctly. Router tests are updated to assert the corrected service contracts.
The 14H admin router consumes update_partition_config and get_partition_config
(returning PartitionDetailResponse) via _require_service_method guards, but
PartitionService never exposed them, so PATCH /partition/{name} and
GET /partition/{name}/config returned 501 "update_partition_config is not
available".
Add get_partition_config (resolve a row into the detail payload: preset
references plus the fully resolved indexation/retrieval pipelines) and
update_partition_config (apply the update, then return that detail). Also
translate a missing preset reference on create/update from the resolver's
ConfigError (500) into a 422 ValidationError, since the preset name is user
input.
Seeding previously skipped the reranker endpoint whenever reranking was disabled, so a configured-but-disabled reranker never appeared in the model_endpoints registry — unlike the embedder, which is always catalogued from its base_url. Registration is about availability, not activation. Seed the reranker from RERANKER_ENDPOINT / reranker.base_url like the embedder, skipping only when truly unconfigured (empty base_url). Whether reranking runs stays the retrieval preset's enable_reranker kill-switch (which inherits reranker.enabled), so a disabled reranker is catalogued but unused by default and available for per-partition opt-in.
|
Tip For best results, initiate chat on the files or code changes.
|
|
@CodeRabbit review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
openrag/services/orchestrators/preset_service.py (1)
181-210: ⚖️ Poor tradeoffConsider transaction boundary for rename logic.
Lines 200-202 delete the old preset and then upsert the new one. If the repository doesn't wrap these in a transaction, a crash or concurrent request between the delete and upsert could leave the preset missing. If
PresetRepositoryimplementations are expected to handle this (e.g., via DB transaction or atomic swap), document that contract. Otherwise, consider adding a rename method to the repository interface that guarantees atomicity.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/orchestrators/preset_service.py` around lines 181 - 210, The rename path in update_preset (method update_preset) currently does a _repo.delete(name, preset_type) followed by _repo.upsert(effective_name, preset_type, effective_config), which can leave a gap if not transactional; update the fix by either (A) requiring and documenting that PresetRepository implementations provide transactional atomicity for delete+upsert, or (B) add an atomic rename method to PresetRepository (e.g., rename(old_name, new_name, preset_type, config)) and call that from update_preset instead of separate delete/upsert so the rename is guaranteed atomic; adjust update_preset to use the new repository method and keep validation/_partition_service reloads the same.openrag/api/middleware/auth.py (1)
156-160: 💤 Low valueConsider narrowing the exception scope to avoid hiding programming errors.
Catching bare
Exceptionwill suppress bugs likeAttributeErrororTypeErrorduring container wiring. If the intent is to catch DI/initialization failures, consider catching specific exceptions (e.g.,RuntimeErrorfrom_require_settings, or a customServiceUnavailableError). Alternatively, document that all_get_auth_servicefailures should degrade gracefully to 503.Also, logging
error_type=type(exc).__name__could expose internal class names to operators. If that's acceptable, this is fine; otherwise consider a more generic message.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/api/middleware/auth.py` around lines 156 - 160, The try/except around self._get_auth_service(request) is too broad and may hide programming errors; change it to catch only expected initialization/DI failures (e.g., catch RuntimeError and/or a new custom ServiceUnavailableError raised by _get_auth_service), re-raise any other exceptions, and keep returning JSONResponse(status_code=503, ...) only for those expected errors; also consider replacing logger.warning(..., error_type=type(exc).__name__) with a more generic marker (e.g., "service_unavailable") or conditionally include the class name behind a debug flag so internal class names are not exposed via logger.warning.openrag/services/persistence/preset_repo.py (1)
82-87: ⚡ Quick winConsider validating
preset_typebefore column selection.The dynamic column selection is safe from SQL injection because the column name is derived from a hardcoded ternary, not user input. However, if
preset_typeis anything other than"indexation"(including typos or future values), the method defaults to queryingretrieval_preset, which may return misleading results.Consider adding explicit validation at line 83:
if preset_type not in ("indexation", "retrieval"): raise ValueError(f"Invalid preset_type: {preset_type}") col = "indexation_preset" if preset_type == "indexation" else "retrieval_preset"This makes the contract explicit and helps catch bugs early.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/persistence/preset_repo.py` around lines 82 - 87, The method count_partitions_using currently defaults any non-"indexation" preset_type to "retrieval_preset", which can mask typos or invalid values; add an explicit validation at the start of count_partitions_using that checks preset_type in ("indexation", "retrieval") and raises a ValueError with the invalid value if not, then set col = "indexation_preset" if preset_type == "indexation" else "retrieval_preset" before calling self.pool.fetchval to query the partitions table.openrag/services/workers/parsers/marker_workers.py (1)
27-36: 💤 Low valueConsider updating the docstring to reflect the fallback behavior.
The docstring states "Return the configured Marker GPU request when Ray reports GPU capacity," but the function also includes fallback logic to
torch.cuda.is_available()when Ray resource querying fails. The implementation is correct and the exception handling is appropriate for infrastructure detection, but the docstring could be more complete.Consider updating line 28:
- """Return the configured Marker GPU request when Ray reports GPU capacity.""" + """Return the configured Marker GPU request when Ray reports GPU capacity. + + Falls back to CUDA availability check if Ray resource query fails. + """🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/workers/parsers/marker_workers.py` around lines 27 - 36, Update the _marker_num_gpus docstring to describe both behaviors: it returns the configured marker_num_gpus only if Ray reports GPU capacity via ray.cluster_resources().get("GPU", 0) > 0, otherwise 0; and when querying Ray fails (caught in the except block), it falls back to checking local CUDA availability via torch.cuda.is_available() and returns requested_gpus only if CUDA is available. Mention that a warning is logged via logger.warning when the Ray query fails.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@openrag/di/container.py`:
- Around line 186-197: ServiceContainer.initialize currently runs several
dependent seeding/loading calls (ServiceContainer.initialize ->
model_endpoint_service.seed_defaults/load_all,
preset_service.seed_defaults/load_all,
partition_service.seed_default_partition/load_partitions) without phase-specific
error handling, and openrag/api/mcp/server.py::_startup calls await
container.initialize() with no try/except; update both: add small try/except
blocks around each seeding/loading call inside ServiceContainer.initialize to
catch exceptions, log structured warnings/errors (with context which phase
failed) via the existing logger, and decide per-phase whether to continue or
re-raise (e.g., non-critical failures should log and continue, critical ones
re-raise); additionally wrap the call in openrag/api/mcp/server.py::_startup
with a try/except that logs the initialization failure and either degrades
startup similarly to openrag/api/main.py or exits gracefully—use the exact
symbol names ServiceContainer.initialize and _startup to locate the changes.
In `@openrag/services/orchestrators/retrieval_service.py`:
- Around line 173-174: The call in RetrievalService to
self._embedder_factory(partition_cfg.embedder) drops the returned embedder, so
partition-specific embedders are not used; either remove the unused call or
capture the return and build a partition-aware searcher/retriever with it.
Update RetrievalService to assign embedder =
self._embedder_factory(partition_cfg.embedder) and then instantiate or configure
a VectorStoreSearcher/retriever for that partition using that embedder (instead
of relying on the single pre-instantiated VectorStoreSearcher created in
openrag/di/container.py), or if per-partition searchers are not desired, remove
the call entirely to avoid side-effect-only invocation. Ensure you reference
self._embedder_factory, partition_cfg.embedder and VectorStoreSearcher when
making the change.
---
Nitpick comments:
In `@openrag/api/middleware/auth.py`:
- Around line 156-160: The try/except around self._get_auth_service(request) is
too broad and may hide programming errors; change it to catch only expected
initialization/DI failures (e.g., catch RuntimeError and/or a new custom
ServiceUnavailableError raised by _get_auth_service), re-raise any other
exceptions, and keep returning JSONResponse(status_code=503, ...) only for those
expected errors; also consider replacing logger.warning(...,
error_type=type(exc).__name__) with a more generic marker (e.g.,
"service_unavailable") or conditionally include the class name behind a debug
flag so internal class names are not exposed via logger.warning.
In `@openrag/services/orchestrators/preset_service.py`:
- Around line 181-210: The rename path in update_preset (method update_preset)
currently does a _repo.delete(name, preset_type) followed by
_repo.upsert(effective_name, preset_type, effective_config), which can leave a
gap if not transactional; update the fix by either (A) requiring and documenting
that PresetRepository implementations provide transactional atomicity for
delete+upsert, or (B) add an atomic rename method to PresetRepository (e.g.,
rename(old_name, new_name, preset_type, config)) and call that from
update_preset instead of separate delete/upsert so the rename is guaranteed
atomic; adjust update_preset to use the new repository method and keep
validation/_partition_service reloads the same.
In `@openrag/services/persistence/preset_repo.py`:
- Around line 82-87: The method count_partitions_using currently defaults any
non-"indexation" preset_type to "retrieval_preset", which can mask typos or
invalid values; add an explicit validation at the start of
count_partitions_using that checks preset_type in ("indexation", "retrieval")
and raises a ValueError with the invalid value if not, then set col =
"indexation_preset" if preset_type == "indexation" else "retrieval_preset"
before calling self.pool.fetchval to query the partitions table.
In `@openrag/services/workers/parsers/marker_workers.py`:
- Around line 27-36: Update the _marker_num_gpus docstring to describe both
behaviors: it returns the configured marker_num_gpus only if Ray reports GPU
capacity via ray.cluster_resources().get("GPU", 0) > 0, otherwise 0; and when
querying Ray fails (caught in the except block), it falls back to checking local
CUDA availability via torch.cuda.is_available() and returns requested_gpus only
if CUDA is available. Mention that a warning is logged via logger.warning when
the Ray query fails.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: d5497d3b-201a-4e9c-bddb-aa883f092db3
📒 Files selected for processing (50)
docs/refactoring/PHASE_14_MIGRATION.mdopenrag/api/error_handlers.pyopenrag/api/middleware/auth.pyopenrag/api/routers/admin/model_endpoints.pyopenrag/api/routers/admin/presets.pyopenrag/api/schemas/admin/preset_schemas.pyopenrag/core/config/indexation_pipeline.pyopenrag/core/config/model_endpoints.pyopenrag/core/config/retrieval_pipeline.pyopenrag/core/indexing/dispatcher.pyopenrag/core/models/catalog.pyopenrag/core/models/preset.pyopenrag/core/ports/partition_repo.pyopenrag/core/ports/preset_repo.pyopenrag/di/container.pyopenrag/services/orchestrators/indexing_service.pyopenrag/services/orchestrators/model_endpoint_service.pyopenrag/services/orchestrators/partition_service.pyopenrag/services/orchestrators/preset_service.pyopenrag/services/orchestrators/retrieval_service.pyopenrag/services/persistence/document_repo.pyopenrag/services/persistence/migrations/alembic/versions/06dd2101ea3a_add_endpoints_presets_phase14.pyopenrag/services/persistence/model_endpoint_repo.pyopenrag/services/persistence/partition_repo.pyopenrag/services/persistence/preset_repo.pyopenrag/services/persistence/schema.pyopenrag/services/storage/milvus_store.pyopenrag/services/workers/dispatcher.pyopenrag/services/workers/indexer_actor.pyopenrag/services/workers/indexer_pool.pyopenrag/services/workers/parsers/marker_workers.pyopenrag/services/workers/pipeline_builder.pyscripts/seed_presets.pytests/unit/api/routers/admin/test_phase14_admin_routers.pytests/unit/core/config/test_model_endpoints.pytests/unit/core/config/test_pipeline_configs.pytests/unit/core/models/test_preset_models.pytests/unit/di/test_container.pytests/unit/services/orchestrators/test_indexing_service.pytests/unit/services/orchestrators/test_model_endpoint_service.pytests/unit/services/orchestrators/test_partition_preset_resolution.pytests/unit/services/orchestrators/test_preset_service.pytests/unit/services/orchestrators/test_retrieval_service.pytests/unit/services/persistence/test_add_file_to_partition_user_id.pytests/unit/services/persistence/test_model_endpoint_repo.pytests/unit/services/persistence/test_preset_repo.pytests/unit/services/workers/parsers/test_marker_workers.pytests/unit/services/workers/test_dispatcher.pytests/unit/services/workers/test_indexer_worker.pytests/unit/services/workers/test_pipeline_builder.py
✅ Files skipped from review due to trivial changes (2)
- tests/unit/services/workers/parsers/test_marker_workers.py
- docs/refactoring/PHASE_14_MIGRATION.md
🚧 Files skipped from review as they are similar to previous changes (10)
- openrag/core/config/model_endpoints.py
- openrag/services/storage/milvus_store.py
- openrag/core/config/retrieval_pipeline.py
- openrag/core/config/indexation_pipeline.py
- openrag/core/models/preset.py
- openrag/api/routers/admin/presets.py
- tests/unit/api/routers/admin/test_phase14_admin_routers.py
- openrag/services/persistence/schema.py
- openrag/api/routers/admin/model_endpoints.py
- openrag/api/schemas/admin/preset_schemas.py
Merge image-captioning fix into the Phase 14 branch.
Wire per-partition retrieval searchers, make preset rename transactional, and harden startup/auth error handling.
Align the admin validate route with ModelEndpointService by resolving the registered endpoint before probing /models.
Resolve the registered endpoint in the admin route, then probe the endpoint URL through ModelEndpointService.validate_endpoint().
Keep ModelEndpointService.set_default as a command-style operation and have the admin route fetch the promoted endpoint before serializing the response.
Preserve configured endpoint API keys during default seeding and use the stored key when validating registered model endpoints.
Translate update request name fields to the service new_name parameter for model endpoint and preset rename operations.
Why this PR matters
Phase 14 moves OpenRAG away from one global inference/pipeline setup and gives operators a safer way to manage models and partition behavior. Model endpoints are now named, presets are reusable, and each partition can resolve its own indexation and retrieval pipeline without hardcoding everything in environment variables.
The main goal is operational control: swap or validate model endpoints, update shared presets, and let partitions pick the right pipeline while keeping the existing default behavior for deployments that do not opt into customization.
What changed
This PR adds the model endpoint registry, the preset registry, and the admin APIs needed to manage them. It also wires those registries into startup, indexing, retrieval, and partition configuration resolution.
A few review fixes were also added before final validation: preset renames are transactional, invalid preset types fail explicitly, retrieval now uses the partition-specific embedder through the correct searcher, and startup/auth error handling is stricter.
Validation
Automated validation is passing on PR #444:
Manual E2E on the server is the remaining check before merge.