diff --git a/docs/refactoring/PHASE_14_MIGRATION.md b/docs/refactoring/PHASE_14_MIGRATION.md new file mode 100644 index 000000000..87cc6b7ff --- /dev/null +++ b/docs/refactoring/PHASE_14_MIGRATION.md @@ -0,0 +1,140 @@ +# Phase 14 Migration — Model Endpoints & Per-Partition Presets + +> **Audience:** operators upgrading an existing OpenRAG deployment to the +> Phase 14 release. +> **TL;DR:** the upgrade is backward compatible. Existing partitions keep +> working untouched; the only new step is a one-time seed of the default model +> endpoints and presets, which the startup path (or `scripts/seed_presets.py`) +> performs automatically. + +--- + +## What changed + +Phase 14 introduces two DB-backed registries that replace the previous +"single global config" model: + +1. **Model Endpoint Registry** (`model_endpoints` table → `config.models`): + named endpoints per model type (`embedder`, `reranker`, `llm`, `vlm`). + Components are built on demand from these named entries by the DI + component factories. +2. **Per-Partition Preset System** (`pipeline_presets` table → + `config.presets`): named bundles of pipeline configuration, of two types — + `indexation` and `retrieval`. A partition references one preset of each type + by name; the retrieval preset in turn references a reranker/llm *endpoint* by + name. + +A partition row now carries `embedder`, `indexation_preset` and +`retrieval_preset` columns. At startup these names are resolved into a cached +`PartitionConfig` (`config.partitions`) used by the retrieval/indexing +pipelines. + +> **Note:** there is no "reranker preset". The reranker is a model *endpoint*. +> The retrieval preset only *references* it via the `reranker` (endpoint name) +> and `enable_reranker` (on/off) fields. + +--- + +## Backward compatibility + +The upgrade is designed to be a no-op for existing deployments: + +- **Schema:** the Phase 14B migration adds the new columns with + `server_default="default"`, so every pre-existing partition row automatically + references the `default` indexation and retrieval presets — no manual + backfill required. Migrations are idempotent (guarded by inspector existence + checks), so they are safe to re-run against an already-bootstrapped database. +- **Defaults derived from your current config:** the default endpoints and + presets are seeded *from your existing global `Settings`* (YAML + env vars). + For example the default reranker endpoint comes from `RERANKER_ENDPOINT` / + `reranker.base_url`, and the default retrieval preset inherits your + `reranker.enabled` kill-switch. So the seeded defaults reproduce your current + behavior rather than imposing new defaults. +- **Reranker availability:** if reranking is disabled (`RERANKER_ENABLED=false` + / `reranker.enabled: false`), no reranker endpoint is seeded and the default + retrieval preset is seeded with `enable_reranker=false`. Deployments without a + reachable reranker therefore keep working and do not start failing on + retrieval. + +--- + +## Migration steps + +### 1. Apply the upgrade + +Deploy the new images / pull the new code as usual. Alembic migrations (run at +app startup) add the new tables and columns idempotently. + +### 2. Seed the default endpoints and presets + +Seeding reads your current `Settings` and writes one default row per model type +plus the default indexation/retrieval presets and the `default` partition. + +Run the one-time utility **inside the application container** (it needs the +project venv and reaches the database over the compose network — Postgres is not +published to the host). From the repo root: + +```bash +# GPU deployment: service "openrag"; CPU deployment: service "openrag-cpu" +docker compose -f infra/compose/docker-compose.yaml exec openrag uv run python /app/scripts/seed_presets.py +``` + +Expected output (counts depend on what is configured in your environment): + +```text +Seeded model endpoints: 1 embedder, 1 reranker, 1 llm, 1 vlm +Seeded 1 indexation presets, 1 retrieval presets +Seeded 1 partition(s) +``` + +The script is **idempotent** — each phase skips any type/preset/partition that +already has rows, so it is safe to run more than once (e.g. after adjusting env +vars and re-seeding only the missing types). + +> If a model type has no endpoint configured (e.g. no `VLM_ENDPOINT` and an +> empty `vlm.base_url`), that type is skipped with a log line and simply has no +> seeded default — register one later through the admin API when needed. + +### 3. Verify + +Run `psql` inside the `rdb` container. The database name is +`partitions_for_collection_`, where `` is your +`vectordb.collection_name` (e.g. `partitions_for_collection_vdb`). From the repo +root: + +```bash +DB=partitions_for_collection_vdb_test # adjust to your collection name + +# Model endpoints — one row per configured type, the seeded one is_default=true +docker compose -f infra/compose/docker-compose.yaml exec rdb psql -U root -d "$DB" \ + -c "SELECT name, model_type, endpoint, is_default FROM model_endpoints;" + +# Presets — default indexation + retrieval (plus the multiquery / hyde presets) +docker compose -f infra/compose/docker-compose.yaml exec rdb psql -U root -d "$DB" \ + -c "SELECT name, preset_type FROM pipeline_presets;" + +# Partitions reference presets by name +docker compose -f infra/compose/docker-compose.yaml exec rdb psql -U root -d "$DB" \ + -c "SELECT name, embedder, indexation_preset, retrieval_preset FROM partitions;" +``` + +--- + +## After migration + +- **Change defaults / add endpoints:** use the admin API + (`/admin/model-endpoints`, `/admin/presets`) rather than editing YAML. Changes + are persisted to the DB, reloaded into `config.models` / `config.presets` + atomically, and the stale cached client is evicted so the next request builds + a fresh client. +- **Per-partition configuration:** create a partition with explicit presets, or + update an existing one: + + ```bash + curl -X POST http://localhost:8080/partition/research \ + -H "Authorization: Bearer " -H "Content-Type: application/json" \ + -d '{"indexation_preset": "default", "retrieval_preset": "hyde"}' + ``` + +- The global YAML config remains the **source of the seed**: it is read once to + populate the registries. Ongoing operational changes live in the DB. diff --git a/openrag/api/error_handlers.py b/openrag/api/error_handlers.py index 952667e4a..d5aa628ce 100644 --- a/openrag/api/error_handlers.py +++ b/openrag/api/error_handlers.py @@ -100,10 +100,14 @@ async def openrag_exception_handler(request: Request, exc: OpenRAGError) -> JSON available — additive, so consumers that ignore the field still parse the response correctly. """ + status_code = _status_for(exc) logger.error( - "OpenRAGError occurred", + "OpenRAGError", error_code=getattr(exc, "code", type(exc).__name__), - status_code=_status_for(exc), + status_code=status_code, + message=str(exc), + method=request.method, + path=request.url.path, ) body = exc.to_dict() request_id = _get_request_id(request) @@ -111,7 +115,7 @@ async def openrag_exception_handler(request: Request, exc: OpenRAGError) -> JSON # Copy so we never mutate the exception's own ``extra`` dict — # the same instance may be re-raised / logged elsewhere. body["extra"] = {**body.get("extra", {}), "request_id": request_id} - return JSONResponse(status_code=_status_for(exc), content=body) + return JSONResponse(status_code=status_code, content=body) async def unhandled_exception_handler(request: Request, exc: Exception) -> JSONResponse: @@ -121,7 +125,13 @@ async def unhandled_exception_handler(request: Request, exc: Exception) -> JSONR Robot Framework suite and the existing unit assertions still match after the move. """ - logger.exception("Unhandled exception", error_type=type(exc).__name__) + logger.exception( + "Unhandled exception", + error_type=type(exc).__name__, + message=str(exc), + method=request.method, + path=request.url.path, + ) extra: dict[str, object] = {} request_id = _get_request_id(request) if request_id is not None: diff --git a/openrag/api/main.py b/openrag/api/main.py index d683e6e07..e03279adb 100644 --- a/openrag/api/main.py +++ b/openrag/api/main.py @@ -40,8 +40,10 @@ from api.routers.admin.cluster import router as actors_router from api.routers.admin.indexing import router as indexer_router from api.routers.admin.jobs import router as queue_router +from api.routers.admin.model_endpoints import router as model_endpoints_router from api.routers.admin.monitoring import router as monitoring_router from api.routers.admin.partitions import router as partition_router +from api.routers.admin.presets import router as presets_router from api.routers.admin.tools import router as tools_router from api.routers.admin.users import router as users_router from api.routers.admin.workspaces import router as workspaces_router @@ -102,12 +104,16 @@ class Tags(Enum): + """OpenAPI tag labels used by mounted routers.""" + VDB = "VectorDB operations" INDEXER = "Indexer" SEARCH = "Semantic Search" OPENAI = "OpenAI Compatible API" EXTRACT = "Document extracts" PARTITION = "Partitions & files" + MODEL_ENDPOINTS = "Model Endpoints" + PRESETS = "Presets" QUEUE = "Queue management" ACTORS = "Ray Actors" USERS = "User management" @@ -219,6 +225,7 @@ async def lifespan(app: FastAPI): def custom_openapi(): + """Build the OpenAPI schema with global bearer authentication metadata.""" if app.openapi_schema: return app.openapi_schema openapi_schema = get_openapi( @@ -297,6 +304,7 @@ def root_redirect(): @app.get("/config", summary="Get current configuration", tags=["Configuration"], dependencies=[Depends(require_admin)]) def get_config(): + """Return the loaded application settings for admins.""" return settings @@ -310,6 +318,8 @@ def get_config(): app.include_router(extract_router, prefix="/extract", tags=[Tags.EXTRACT]) app.include_router(search_router, prefix="/search", tags=[Tags.SEARCH]) app.include_router(partition_router, prefix="/partition", tags=[Tags.PARTITION]) +app.include_router(model_endpoints_router, prefix="/model-endpoints", tags=[Tags.MODEL_ENDPOINTS]) +app.include_router(presets_router, prefix="/presets", tags=[Tags.PRESETS]) app.include_router(queue_router, prefix="/queue", tags=[Tags.QUEUE]) app.include_router(actors_router, prefix="/actors", tags=[Tags.ACTORS]) app.include_router(users_router, prefix="/users", tags=[Tags.USERS]) @@ -338,7 +348,7 @@ def get_config(): @serve.deployment(num_replicas=settings.ray.serve.num_replicas) @serve.ingress(app) class OpenRagAPI: - pass + """Ray Serve deployment wrapper for the FastAPI app.""" serve.start(http_options={"host": settings.ray.serve.host, "port": settings.ray.serve.port}) if WITH_CHAINLIT_UI: diff --git a/openrag/api/mcp/server.py b/openrag/api/mcp/server.py index 6df49989a..acd4d09f1 100644 --- a/openrag/api/mcp/server.py +++ b/openrag/api/mcp/server.py @@ -80,7 +80,12 @@ async def _startup() -> None: ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True) ensure_worker_bootstrap() container = ServiceContainer(config) - await container.initialize() + try: + await container.initialize() + except Exception: + logger.exception("MCP server container initialization failed") + await container.shutdown() + raise _container = container logger.info("MCP server container initialized") diff --git a/openrag/api/middleware/auth.py b/openrag/api/middleware/auth.py index 316d4482e..b94fa748c 100644 --- a/openrag/api/middleware/auth.py +++ b/openrag/api/middleware/auth.py @@ -153,7 +153,11 @@ async def dispatch(self, request: Request, call_next): user = None session = None - auth_service = self._get_auth_service(request) + try: + auth_service = self._get_auth_service(request) + except RuntimeError: + logger.warning("Auth service unavailable", reason="service_unavailable", path=request.url.path) + return JSONResponse(status_code=503, content={"detail": "Service unavailable"}) # --- 1) Cookie session (OIDC UI flow). Gated on oidc mode so the # legacy token-mode contract remains strictly Bearer-only — diff --git a/openrag/api/routers/admin/model_endpoints.py b/openrag/api/routers/admin/model_endpoints.py new file mode 100644 index 000000000..1a5b4839d --- /dev/null +++ b/openrag/api/routers/admin/model_endpoints.py @@ -0,0 +1,111 @@ +"""Admin routes for the Phase 14 model endpoint registry. + +The router is intentionally transport-only: auth, request validation and +response shaping live here, while endpoint persistence and validation are +delegated to the service resolved from the DI container. +""" + +from datetime import UTC, datetime + +from api.dependencies.auth import require_admin +from api.schemas.admin.model_endpoint_schemas import ( + CreateModelEndpointRequest, + ModelEndpointResponse, + ModelEndpointType, + UpdateModelEndpointRequest, + ValidateEndpointResponse, +) +from core.config.model_endpoints import ModelEndpointRow +from di.providers import get_model_endpoint_service +from fastapi import APIRouter, Depends, Response, status + +router = APIRouter(dependencies=[Depends(require_admin)]) + + +@router.post( + "/", + response_model=ModelEndpointResponse, + status_code=status.HTTP_201_CREATED, +) +async def create_model_endpoint( + body: CreateModelEndpointRequest, + service=Depends(get_model_endpoint_service), +): + """Register a named inference endpoint.""" + now = datetime.now(UTC) + row = ModelEndpointRow(**body.model_dump(), created_at=now, updated_at=now) + return await service.create_model_endpoint(row) + + +@router.get("/", response_model=list[ModelEndpointResponse]) +async def list_model_endpoints( + model_type: ModelEndpointType | None = None, + service=Depends(get_model_endpoint_service), +): + """List registered inference endpoints, optionally filtered by type.""" + return await service.list_model_endpoints(model_type=model_type) + + +@router.get("/{model_type}/{name}", response_model=ModelEndpointResponse) +async def get_model_endpoint( + model_type: ModelEndpointType, + name: str, + service=Depends(get_model_endpoint_service), +): + """Return one registered inference endpoint.""" + return await service.get_model_endpoint(name=name, model_type=model_type) + + +@router.put("/{model_type}/{name}", response_model=ModelEndpointResponse) +async def update_model_endpoint( + model_type: ModelEndpointType, + name: str, + body: UpdateModelEndpointRequest, + service=Depends(get_model_endpoint_service), +): + """Update a registered inference endpoint.""" + fields = body.model_dump(exclude_unset=True) + if "name" in fields: + fields["new_name"] = fields.pop("name") + return await service.update_model_endpoint( + name=name, + model_type=model_type, + **fields, + ) + + +@router.delete("/{model_type}/{name}", status_code=status.HTTP_204_NO_CONTENT) +async def delete_model_endpoint( + model_type: ModelEndpointType, + name: str, + service=Depends(get_model_endpoint_service), +): + """Delete a registered inference endpoint.""" + await service.delete_model_endpoint(name=name, model_type=model_type) + return Response(status_code=status.HTTP_204_NO_CONTENT) + + +@router.post("/{model_type}/{name}/set-default", response_model=ModelEndpointResponse) +async def set_default_model_endpoint( + model_type: ModelEndpointType, + name: str, + service=Depends(get_model_endpoint_service), +): + """Promote a registered endpoint to the default for its type.""" + await service.set_default(model_type=model_type, name=name) + return await service.get_model_endpoint(name=name, model_type=model_type) + + +@router.post("/{model_type}/{name}/validate", response_model=ValidateEndpointResponse) +async def validate_model_endpoint( + model_type: ModelEndpointType, + name: str, + service=Depends(get_model_endpoint_service), +): + """Probe a registered endpoint for reachability and model availability.""" + endpoint = await service.get_model_endpoint(name=name, model_type=model_type) + return await service.validate_endpoint( + url=endpoint.endpoint, + model_name=endpoint.model_name, + api_key=endpoint.extra.get("api_key"), + ) diff --git a/openrag/api/routers/admin/partitions.py b/openrag/api/routers/admin/partitions.py index c95b42c2c..7196d540d 100644 --- a/openrag/api/routers/admin/partitions.py +++ b/openrag/api/routers/admin/partitions.py @@ -17,6 +17,7 @@ require_partition_owner, require_partition_viewer, ) +from api.schemas.admin.partition_schemas import PartitionDetailResponse, UpdatePartitionRequest from core.utils.logging import get_logger from di.providers import get_partition_service from fastapi import APIRouter, Depends, Form, HTTPException, Request, Response, status @@ -29,9 +30,21 @@ def _quote_param_value(s: str) -> str: + """Percent-encode a path parameter value for URL generation.""" return quote(s, safe="") +def _require_service_method(service, method_name: str): + """Return a service method or fail clearly when a phased method is absent.""" + method = getattr(service, method_name, None) + if not callable(method): + raise HTTPException( + status_code=status.HTTP_501_NOT_IMPLEMENTED, + detail=f"{method_name} is not available.", + ) + return method + + @router.get( "/", description="""List all accessible partitions. @@ -49,6 +62,7 @@ async def list_existant_partitions( partitions=Depends(partitions_with_details), service=Depends(get_partition_service), ): + """List partitions visible to the current user.""" if len(partitions) == 1 and partitions[0]["partition"] == "all": partitions = await service.list_partitions() logger.debug("Returned list of existing partitions.", partition_count=len(partitions)) @@ -77,6 +91,7 @@ async def delete_partition( partition_owner=Depends(require_partition_owner), service=Depends(get_partition_service), ): + """Delete a partition owned by the current user.""" await service.delete_partition(partition) return Response(status_code=status.HTTP_204_NO_CONTENT) @@ -107,9 +122,11 @@ async def list_files( partition_viewer=Depends(require_partition_viewer), service=Depends(get_partition_service), ): + """List files stored in a partition.""" file_dicts = await service.list_files(partition, limit) def process_file(file_dict): + """Add a canonical file-detail link to one file row.""" return { "link": str( request.url_for( @@ -153,6 +170,7 @@ async def get_file( partition_viewer=Depends(require_partition_viewer), service=Depends(get_partition_service), ): + """Return metadata and chunk links for one file in a partition.""" if not await service.file_exists(file_id, partition): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -196,6 +214,7 @@ async def list_all_chunks( partition_viewer=Depends(require_partition_viewer), service=Depends(get_partition_service), ): + """List all chunks in a partition.""" items = await service.list_all_chunks(partition=partition, include_embedding=include_embedding) chunks = [ { @@ -232,6 +251,7 @@ async def create_partition( partition: str, service=Depends(get_partition_service), ): + """Create a new partition owned by the current user.""" if await service.partition_exists(partition): raise HTTPException( status_code=status.HTTP_409_CONFLICT, @@ -242,6 +262,69 @@ async def create_partition( return Response(status_code=status.HTTP_201_CREATED) +@router.patch( + "/{partition}", + response_model=PartitionDetailResponse, + description="""Update Phase 14 preset assignments for a partition. + +**Parameters:** +- `partition`: The partition name + +**Body:** +Accepts partition config fields such as: +- `description` +- `embedder` +- `indexation_preset` +- `retrieval_preset` +- `chat_history_depth` +- `chat_llm` + +**Permissions:** +- Requires partition owner role + +**Response:** +Returns the updated resolved partition configuration. +""", +) +async def update_partition_config( + partition: str, + body: UpdatePartitionRequest, + partition_owner=Depends(require_partition_owner), + service=Depends(get_partition_service), +): + """Update Phase 14 preset references for a partition.""" + method = _require_service_method(service, "update_partition_config") + return await method( + partition=partition, + **body.model_dump(exclude_unset=True), + ) + + +@router.get( + "/{partition}/config", + response_model=PartitionDetailResponse, + description="""Return the resolved Phase 14 pipeline config for a partition. + +**Parameters:** +- `partition`: The partition name + +**Permissions:** +- Requires partition viewer role or higher + +**Response:** +Returns partition metadata, preset references, and resolved indexation/retrieval pipeline configs. +""", +) +async def get_partition_config( + partition: str, + partition_viewer=Depends(require_partition_viewer), + service=Depends(get_partition_service), +): + """Return the resolved Phase 14 config for a partition.""" + method = _require_service_method(service, "get_partition_config") + return await method(partition=partition) + + @router.get( "/{partition}/users", description="""List all users with access to a partition. @@ -401,6 +484,7 @@ async def get_related_files( partition_viewer=Depends(require_partition_viewer), service=Depends(get_partition_service), ): + """Return files sharing a relationship identifier.""" files = await service.get_related_files(partition=partition, relationship_id=relationship_id) return JSONResponse(status_code=status.HTTP_200_OK, content={"files": files}) @@ -438,6 +522,7 @@ async def get_file_ancestors( partition_viewer=Depends(require_partition_viewer), service=Depends(get_partition_service), ): + """Return the ancestor path for one file.""" if not await service.file_exists(file_id, partition): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, diff --git a/openrag/api/routers/admin/presets.py b/openrag/api/routers/admin/presets.py new file mode 100644 index 000000000..8a3ce09cf --- /dev/null +++ b/openrag/api/routers/admin/presets.py @@ -0,0 +1,94 @@ +"""Admin routes for the Phase 14 pipeline preset registry.""" + +from api.dependencies.auth import require_admin +from api.schemas.admin.preset_schemas import ( + CreatePresetRequest, + PresetOptionsResponse, + PresetResponse, + PresetType, + UpdatePresetRequest, +) +from core.chunking import chunking_registry +from core.rerankers.registry import reranker_registry +from core.retrieval import retriever_registry +from di.providers import get_preset_service +from fastapi import APIRouter, Depends, Response, status + +router = APIRouter(dependencies=[Depends(require_admin)]) + +_DEFAULT_RERANKER_PROVIDERS = ["infinity", "openai"] + + +def _registered_or_default(registered: list[str], defaults: list[str]) -> list[str]: + """Return registry values, falling back to known defaults before DI imports providers.""" + return registered or defaults + + +@router.get("/options", response_model=PresetOptionsResponse) +async def get_preset_options(): + """Return available preset strategy choices.""" + return PresetOptionsResponse( + chunking_strategies=chunking_registry.list_registered(), + retrieval_types=retriever_registry.list_registered(), + reranker_providers=_registered_or_default( + reranker_registry.list_registered(), + _DEFAULT_RERANKER_PROVIDERS, + ), + ) + + +@router.post("/", response_model=PresetResponse, status_code=status.HTTP_201_CREATED) +async def create_preset( + body: CreatePresetRequest, + service=Depends(get_preset_service), +): + """Create a named pipeline preset.""" + return await service.create_preset(**body.model_dump()) + + +@router.get("/", response_model=list[PresetResponse]) +async def list_presets( + preset_type: PresetType | None = None, + service=Depends(get_preset_service), +): + """List pipeline presets, optionally filtered by type.""" + return await service.list_presets(preset_type=preset_type) + + +@router.get("/{preset_type}/{name}", response_model=PresetResponse) +async def get_preset( + preset_type: PresetType, + name: str, + service=Depends(get_preset_service), +): + """Return one pipeline preset.""" + return await service.get_preset(name=name, preset_type=preset_type) + + +@router.put("/{preset_type}/{name}", response_model=PresetResponse) +async def update_preset( + preset_type: PresetType, + name: str, + body: UpdatePresetRequest, + service=Depends(get_preset_service), +): + """Update a pipeline preset.""" + fields = body.model_dump(exclude_unset=True) + if "name" in fields: + fields["new_name"] = fields.pop("name") + return await service.update_preset( + name=name, + preset_type=preset_type, + **fields, + ) + + +@router.delete("/{preset_type}/{name}", status_code=status.HTTP_204_NO_CONTENT) +async def delete_preset( + preset_type: PresetType, + name: str, + service=Depends(get_preset_service), +): + """Delete a pipeline preset.""" + await service.delete_preset(name=name, preset_type=preset_type) + return Response(status_code=status.HTTP_204_NO_CONTENT) diff --git a/openrag/api/schemas/admin/__init__.py b/openrag/api/schemas/admin/__init__.py index f1395c01e..bf529cc5d 100644 --- a/openrag/api/schemas/admin/__init__.py +++ b/openrag/api/schemas/admin/__init__.py @@ -1,4 +1,19 @@ from api.schemas.admin.common import DocumentsResponse, FilesResponse, MessageResponse, TaskStatusResponse +from api.schemas.admin.model_endpoint_schemas import ( + CreateModelEndpointRequest, + ModelEndpointResponse, + ModelEndpointType, + UpdateModelEndpointRequest, + ValidateEndpointResponse, +) +from api.schemas.admin.partition_schemas import CreatePartitionRequest, PartitionDetailResponse, UpdatePartitionRequest +from api.schemas.admin.preset_schemas import ( + CreatePresetRequest, + PresetOptionsResponse, + PresetResponse, + PresetType, + UpdatePresetRequest, +) from api.schemas.admin.tools import ToolInfo from api.schemas.admin.users import UserCreate, UserPublic, UserUpdate from api.schemas.admin.workspaces import AddFilesRequest, CreateWorkspaceRequest @@ -6,12 +21,25 @@ __all__ = [ "AddFilesRequest", "CreateWorkspaceRequest", + "CreateModelEndpointRequest", + "CreatePartitionRequest", + "CreatePresetRequest", "DocumentsResponse", "FilesResponse", "MessageResponse", + "ModelEndpointResponse", + "ModelEndpointType", + "PartitionDetailResponse", + "PresetOptionsResponse", + "PresetResponse", + "PresetType", "TaskStatusResponse", "ToolInfo", + "UpdateModelEndpointRequest", + "UpdatePartitionRequest", + "UpdatePresetRequest", "UserCreate", "UserPublic", "UserUpdate", + "ValidateEndpointResponse", ] diff --git a/openrag/api/schemas/admin/model_endpoint_schemas.py b/openrag/api/schemas/admin/model_endpoint_schemas.py new file mode 100644 index 000000000..028b9c054 --- /dev/null +++ b/openrag/api/schemas/admin/model_endpoint_schemas.py @@ -0,0 +1,121 @@ +"""Admin schemas for the Phase 14 model endpoint registry.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + +ModelEndpointType = Literal["embedder", "reranker", "llm", "vlm"] + + +def _normalize_name(value: str) -> str: + """Trim a user-facing registry name and reject blank values.""" + value = value.strip() + if not value: + raise ValueError("name must be non-empty") + return value + + +def _normalize_endpoint(value: str) -> str: + """Trim an endpoint URL and reject values that normalize to empty.""" + normalized = value.strip().rstrip("/") + if not normalized: + raise ValueError("endpoint must be non-empty") + return normalized + + +class CreateModelEndpointRequest(BaseModel): + """Request body for registering a model endpoint.""" + + model_config = ConfigDict(extra="forbid") + + name: str + model_type: ModelEndpointType + endpoint: str + model_name: str | None = None + batch_size: int = Field(default=32, gt=0) + timeout: float = Field(default=30.0, gt=0) + extra: dict[str, Any] = Field(default_factory=dict) + is_default: bool = False + + @field_validator("name") + @classmethod + def validate_name(cls, value: str) -> str: + """Normalize the endpoint registry name.""" + return _normalize_name(value) + + @field_validator("endpoint") + @classmethod + def validate_endpoint(cls, value: str) -> str: + """Normalize the endpoint URL.""" + return _normalize_endpoint(value) + + +class UpdateModelEndpointRequest(BaseModel): + """Request body for updating a registered model endpoint.""" + + model_config = ConfigDict(extra="forbid") + + name: str | None = None + endpoint: str | None = None + model_name: str | None = None + batch_size: int | None = Field(default=None, gt=0) + timeout: float | None = Field(default=None, gt=0) + extra: dict[str, Any] | None = None + is_default: bool | None = None + + @field_validator("name") + @classmethod + def validate_name(cls, value: str | None) -> str | None: + """Normalize the optional replacement name.""" + return _normalize_name(value) if value is not None else None + + @field_validator("endpoint") + @classmethod + def validate_endpoint(cls, value: str | None) -> str | None: + """Normalize the optional replacement endpoint URL.""" + if value is None: + return None + return _normalize_endpoint(value) + + @model_validator(mode="after") + def require_at_least_one_update(self) -> UpdateModelEndpointRequest: + """Reject empty update payloads.""" + if not self.model_fields_set: + raise ValueError("at least one field must be provided") + return self + + +class ModelEndpointResponse(BaseModel): + """Response body for a registered model endpoint.""" + + name: str + model_type: ModelEndpointType + endpoint: str + model_name: str | None + batch_size: int + timeout: float + extra: dict[str, Any] + is_default: bool + created_at: datetime + updated_at: datetime + + +class ValidateEndpointResponse(BaseModel): + """Response body for a model endpoint validation probe.""" + + reachable: bool + model_found: bool | None = None + models_served: list[str] | None = None + detail: str | None = None + + +__all__ = [ + "CreateModelEndpointRequest", + "ModelEndpointResponse", + "ModelEndpointType", + "UpdateModelEndpointRequest", + "ValidateEndpointResponse", +] diff --git a/openrag/api/schemas/admin/partition_schemas.py b/openrag/api/schemas/admin/partition_schemas.py new file mode 100644 index 000000000..95ed20e43 --- /dev/null +++ b/openrag/api/schemas/admin/partition_schemas.py @@ -0,0 +1,102 @@ +"""Admin schemas for Phase 14 partition preset assignment.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator, model_validator + + +def _normalize_name(value: str) -> str: + """Trim a partition/preset reference and reject blank values.""" + value = value.strip() + if not value: + raise ValueError("name must be non-empty") + return value + + +def _reject_explicit_null(field_name: str, value): + """Reject explicit null while still allowing omitted optional fields.""" + if value is None: + raise ValueError(f"{field_name} cannot be null") + return value + + +class CreatePartitionRequest(BaseModel): + """Request body for creating a partition with preset references.""" + + model_config = ConfigDict(extra="forbid") + + name: str + description: str = "" + embedder: str = "default" + indexation_preset: str = "default" + retrieval_preset: str = "default" + chat_history_depth: int = Field(default=0, ge=0) + chat_llm: str | None = None + + @field_validator("name", "embedder", "indexation_preset", "retrieval_preset") + @classmethod + def validate_non_empty_name(cls, value: str) -> str: + """Normalize non-null partition and preset names.""" + return _normalize_name(value) + + +class UpdatePartitionRequest(BaseModel): + """Request body for updating partition preset assignments.""" + + model_config = ConfigDict(extra="forbid") + + description: str | None = None + embedder: str | None = None + indexation_preset: str | None = None + retrieval_preset: str | None = None + chat_history_depth: int | None = Field(default=None, ge=0) + chat_llm: str | None = None + + @field_validator("embedder", "indexation_preset", "retrieval_preset") + @classmethod + def validate_non_empty_name(cls, value: str | None, info: ValidationInfo) -> str | None: + """Normalize updated references and reject explicit null.""" + return _normalize_name(_reject_explicit_null(info.field_name, value)) + + @field_validator("description") + @classmethod + def reject_null_description(cls, value: str | None) -> str: + """Reject explicit null for description updates.""" + return _reject_explicit_null("description", value) + + @field_validator("chat_history_depth") + @classmethod + def reject_null_chat_history_depth(cls, value: int | None) -> int: + """Reject explicit null for chat history depth updates.""" + return _reject_explicit_null("chat_history_depth", value) + + @model_validator(mode="after") + def require_at_least_one_update(self) -> UpdatePartitionRequest: + """Reject empty update payloads.""" + if not self.model_fields_set: + raise ValueError("at least one field must be provided") + return self + + +class PartitionDetailResponse(BaseModel): + """Response body for a resolved partition configuration.""" + + name: str + description: str + embedder: str + indexation_preset: str + retrieval_preset: str + indexation_pipeline: dict[str, Any] + retrieval_pipeline: dict[str, Any] + dimension: int + created_at: datetime + + +__all__ = [ + "CreatePartitionRequest", + "PartitionDetailResponse", + "UpdatePartitionRequest", +] diff --git a/openrag/api/schemas/admin/preset_schemas.py b/openrag/api/schemas/admin/preset_schemas.py new file mode 100644 index 000000000..e5bfd516a --- /dev/null +++ b/openrag/api/schemas/admin/preset_schemas.py @@ -0,0 +1,96 @@ +"""Admin schemas for the Phase 14 pipeline preset registry.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator, model_validator + +PresetType = Literal["indexation", "retrieval"] + + +def _normalize_name(value: str) -> str: + """Trim a preset name and reject blank values.""" + value = value.strip() + if not value: + raise ValueError("name must be non-empty") + return value + + +def _reject_explicit_null(field_name: str, value): + """Reject explicit null while still allowing omitted optional fields.""" + if value is None: + raise ValueError(f"{field_name} cannot be null") + return value + + +class CreatePresetRequest(BaseModel): + """Request body for creating an indexation or retrieval preset.""" + + model_config = ConfigDict(extra="forbid") + + name: str + preset_type: PresetType + config: dict[str, Any] = Field(default_factory=dict) + + @field_validator("name") + @classmethod + def validate_name(cls, value: str) -> str: + """Normalize the preset name.""" + return _normalize_name(value) + + +class UpdatePresetRequest(BaseModel): + """Request body for updating an existing preset.""" + + model_config = ConfigDict(extra="forbid") + + name: str | None = None + config: dict[str, Any] | None = None + + @field_validator("name") + @classmethod + def validate_name(cls, value: str | None, info: ValidationInfo) -> str | None: + """Normalize the optional replacement name and reject null.""" + return _normalize_name(_reject_explicit_null(info.field_name, value)) + + @field_validator("config") + @classmethod + def validate_config(cls, value: dict[str, Any] | None) -> dict[str, Any]: + """Reject explicit null for preset config updates.""" + return _reject_explicit_null("config", value) + + @model_validator(mode="after") + def require_at_least_one_update(self) -> UpdatePresetRequest: + """Reject empty update payloads.""" + if not self.model_fields_set or not any(getattr(self, field) is not None for field in ("name", "config")): + raise ValueError("at least one field must be provided") + return self + + +class PresetResponse(BaseModel): + """Response body for a stored pipeline preset.""" + + name: str + preset_type: PresetType + config: dict[str, Any] + created_at: datetime + updated_at: datetime + + +class PresetOptionsResponse(BaseModel): + """Response body listing allowed preset option values.""" + + chunking_strategies: list[str] + retrieval_types: list[str] + reranker_providers: list[str] + + +__all__ = [ + "CreatePresetRequest", + "PresetOptionsResponse", + "PresetResponse", + "PresetType", + "UpdatePresetRequest", +] diff --git a/openrag/core/config/indexation_pipeline.py b/openrag/core/config/indexation_pipeline.py new file mode 100644 index 000000000..793137c5c --- /dev/null +++ b/openrag/core/config/indexation_pipeline.py @@ -0,0 +1,53 @@ +"""Per-partition indexation pipeline configuration. + +Stored as JSONB in the ``pipeline_presets`` table (preset_type='indexation'). +At runtime, PartitionService deserializes each preset row into this model +and caches the resolved PartitionConfig. +""" + +from __future__ import annotations + +from typing import Literal + +from core.config.chunking import ChunkerConfig +from pydantic import BaseModel, ConfigDict, Field + + +class IndexationPipelineConfig(BaseModel): + """Indexation pipeline settings for one partition preset.""" + + model_config = ConfigDict(extra="ignore") + + chunking: ChunkerConfig = Field(default_factory=ChunkerConfig) + parsing_strategy: Literal["pymupdf", "marker", "docling"] = "marker" + + # VLM / image captioning + vlm: str | None = None # endpoint name; None = use global default + enable_image_captioning: bool = True + + # Contextualization (LLM-generated chunk context) + enable_contextualization: bool = False + contextualization_llm: str | None = None + contextualization_mode: Literal["none", "simple", "structured"] = "none" + + # Metadata extraction + enable_metadata_extraction: bool = True + metadata_extraction_llm: str | None = None + + # Prompt name overrides (None = use active prompt for the partition) + vlm_caption_prompt_name: str | None = None + contextualization_prompt_name: str | None = None + + # Entity extraction + enable_entity_extraction: bool = True + entity_labels: list[str] = Field( + default_factory=lambda: ["person", "organization", "location", "event"], + ) + + # Topic tagging + enable_topic_tagging: bool = True + max_topic_tags: int = Field(default=7, ge=1, le=50) + topic_tagging_llm: str | None = None + + +__all__ = ["IndexationPipelineConfig"] diff --git a/openrag/core/config/model_endpoints.py b/openrag/core/config/model_endpoints.py new file mode 100644 index 000000000..90a92d474 --- /dev/null +++ b/openrag/core/config/model_endpoints.py @@ -0,0 +1,60 @@ +"""Named model endpoint registry — multi-endpoint config for embedders, LLMs, rerankers, VLMs.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Literal + +from core.config.base import ConfigMixin +from pydantic import BaseModel, Field + +ModelEndpointType = Literal["embedder", "reranker", "llm", "vlm"] + + +class ModelEndpointConfig(BaseModel): + """A single registered inference endpoint. + + ``extra`` holds implementation-specific keys: + ``{"implementation": "vllm"}`` → VLLMEmbedder + ``{"implementation": "ollama"}`` → OllamaEmbedder + ``{"implementation": "infinity"}``→ InfinityReranker + ``{"api_key": "sk-..."}`` → passed to client constructor + """ + + endpoint: str + model_name: str | None = None + batch_size: int = Field(default=32, gt=0) + timeout: float = Field(default=30.0, gt=0) + extra: dict[str, Any] = Field(default_factory=dict) + + +class ModelsConfig(ConfigMixin): + """Named endpoint dictionaries — one per model type. + + Fields are frozen (Pydantic ConfigMixin), but the dict objects they + hold are mutable. Services perform atomic-ish in-place swaps via + ``dict.clear() + dict.update()`` rather than reassigning the field. + """ + + embedder: dict[str, ModelEndpointConfig] = Field(default_factory=dict) + reranker: dict[str, ModelEndpointConfig] = Field(default_factory=dict) + llm: dict[str, ModelEndpointConfig] = Field(default_factory=dict) + vlm: dict[str, ModelEndpointConfig] = Field(default_factory=dict) + + +class ModelEndpointRow(BaseModel): + """DB representation of a model endpoint (returned by the repository).""" + + name: str + model_type: ModelEndpointType + endpoint: str + model_name: str | None = None + batch_size: int = Field(default=32, gt=0) + timeout: float = Field(default=30.0, gt=0) + extra: dict[str, Any] = Field(default_factory=dict) + is_default: bool = False + created_at: datetime + updated_at: datetime + + +__all__ = ["ModelEndpointConfig", "ModelsConfig", "ModelEndpointRow", "ModelEndpointType"] diff --git a/openrag/core/config/presets.py b/openrag/core/config/presets.py new file mode 100644 index 000000000..a8957edf6 --- /dev/null +++ b/openrag/core/config/presets.py @@ -0,0 +1,24 @@ +"""In-memory preset cache — populated from DB by PresetService.load_all().""" + +from __future__ import annotations + +from typing import Any + +from pydantic import Field + +from .base import ConfigMixin + + +class PresetsConfig(ConfigMixin): + """Named preset dicts — one dict per preset type. + + Like ModelsConfig, the fields are frozen but the dicts they hold + are mutable. PresetService.load_all() performs clear()+update() + for atomic in-place swaps. + """ + + indexation: dict[str, dict[str, Any]] = Field(default_factory=dict) + retrieval: dict[str, dict[str, Any]] = Field(default_factory=dict) + + +__all__ = ["PresetsConfig"] diff --git a/openrag/core/config/retrieval_pipeline.py b/openrag/core/config/retrieval_pipeline.py new file mode 100644 index 000000000..9030a96aa --- /dev/null +++ b/openrag/core/config/retrieval_pipeline.py @@ -0,0 +1,33 @@ +"""Per-partition retrieval pipeline configuration. + +Stored as JSONB in the ``pipeline_presets`` table (preset_type='retrieval'). +Endpoint name fields (``reranker``, ``llm``) are resolved at query time by +the component factories registered in the DI container. +""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field + + +class RetrievalPipelineConfig(BaseModel): + """Retrieval pipeline settings for one partition preset.""" + + model_config = ConfigDict(extra="ignore") + + type: Literal["single", "multiQuery", "hyde"] = "single" + reranker: str | None = None # endpoint name; None = use default + llm: str | None = None # endpoint name for multiQuery / hyde expansion + + top_k: int = Field(default=50, gt=0, le=1000) # vector candidates fetched + top_n: int = Field(default=10, gt=0, le=1000) # final results after reranking + enable_reranker: bool = True + similarity_threshold: float = Field(default=0.6, ge=0.0, le=1.0) + include_related: bool = True + include_ancestors: bool = True + rrf_k: int = Field(default=60, gt=0, le=1000) # Reciprocal Rank Fusion constant + + +__all__ = ["RetrievalPipelineConfig"] diff --git a/openrag/core/config/root.py b/openrag/core/config/root.py index c44483ace..73adf8e72 100644 --- a/openrag/core/config/root.py +++ b/openrag/core/config/root.py @@ -2,6 +2,7 @@ from __future__ import annotations +from core.models.preset import PartitionConfig from pydantic import Field from .base import ConfigMixin @@ -24,6 +25,8 @@ VerboseConfig, ) from .mcp import MCPServerConfig +from .model_endpoints import ModelsConfig +from .presets import PresetsConfig from .retrieval import ( MapReduceConfig, RAGConfig, @@ -63,3 +66,6 @@ class Settings(ConfigMixin): rag: RAGConfig = Field(default_factory=RAGConfig) websearch: WebSearchConfig = Field(default_factory=StaanWebSearchConfig) mcp: MCPServerConfig = Field(default_factory=MCPServerConfig) + models: ModelsConfig = Field(default_factory=ModelsConfig) + presets: PresetsConfig = Field(default_factory=PresetsConfig) + partitions: dict[str, PartitionConfig] = Field(default_factory=dict) diff --git a/openrag/core/indexing/dispatcher.py b/openrag/core/indexing/dispatcher.py index e8cbb7b17..03ed0f001 100644 --- a/openrag/core/indexing/dispatcher.py +++ b/openrag/core/indexing/dispatcher.py @@ -35,6 +35,8 @@ async def dispatch_indexing( user: dict | None, workspace_ids: list[str] | None, replace: bool, + indexation_config: dict | None = None, + embedder_name: str | None = None, ) -> str: """Queue an (re)indexing job, register its task state, return its id.""" ... diff --git a/openrag/core/models/catalog.py b/openrag/core/models/catalog.py index 014ed94ed..721047f26 100644 --- a/openrag/core/models/catalog.py +++ b/openrag/core/models/catalog.py @@ -36,6 +36,7 @@ class DocumentRecord(BaseModel): filename: str = "" partition: str = "default" metadata: dict[str, Any] = Field(default_factory=dict) + indexation_config: dict[str, Any] | None = None status: DocumentStatus = DocumentStatus.QUEUED error_message: str | None = None created_by: int | None = None diff --git a/openrag/core/models/preset.py b/openrag/core/models/preset.py new file mode 100644 index 000000000..7fce4d529 --- /dev/null +++ b/openrag/core/models/preset.py @@ -0,0 +1,59 @@ +"""Preset and partition domain models.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Literal + +from core.config.indexation_pipeline import IndexationPipelineConfig +from core.config.retrieval_pipeline import RetrievalPipelineConfig +from pydantic import BaseModel, Field + +PresetType = Literal["indexation", "retrieval"] + + +class PresetRow(BaseModel): + """DB representation of a pipeline preset row.""" + + name: str + preset_type: PresetType + config: dict[str, Any] + created_at: datetime + updated_at: datetime + + +class PartitionRow(BaseModel): + """DB representation of a partition with preset references.""" + + name: str + display_name: str | None = None + description: str = "" + embedder: str = "default" + indexation_preset: str = "default" + retrieval_preset: str = "default" + dimension: int = Field(default=1024, gt=0) + collection_name: str | None = None + chat_history_depth: int = Field(default=0, ge=0) + chat_llm: str | None = None + created_at: datetime + updated_at: datetime + + +class PartitionConfig(BaseModel): + """Fully resolved partition config — preset names looked up and validated. + + Built by PartitionService.resolve_partition_row() and cached in + Settings.partitions at startup and on every preset change. + """ + + name: str + description: str = "" + embedder: str = "default" + indexation: IndexationPipelineConfig + retrieval: RetrievalPipelineConfig + collection_name: str | None = None + chat_history_depth: int = Field(default=0, ge=0) + chat_llm: str | None = None + + +__all__ = ["PresetRow", "PartitionRow", "PartitionConfig", "PresetType"] diff --git a/openrag/core/ports/model_endpoint_repo.py b/openrag/core/ports/model_endpoint_repo.py index e907a0ce0..f95d09465 100644 --- a/openrag/core/ports/model_endpoint_repo.py +++ b/openrag/core/ports/model_endpoint_repo.py @@ -4,18 +4,29 @@ from abc import ABC, abstractmethod +from core.config.model_endpoints import ModelEndpointRow + class ModelEndpointRepository(ABC): - """CRUD operations for model endpoint configurations.""" + """CRUD operations for named model endpoint configurations.""" + + @abstractmethod + async def create(self, row: ModelEndpointRow) -> ModelEndpointRow: ... @abstractmethod - async def get(self, name: str, model_type: str) -> dict | None: ... + async def get(self, name: str, model_type: str) -> ModelEndpointRow | None: ... @abstractmethod - async def list_all(self, model_type: str | None = None) -> list[dict]: ... + async def list_all(self, model_type: str | None = None) -> list[ModelEndpointRow]: ... @abstractmethod - async def upsert(self, name: str, model_type: str, config: dict) -> dict: ... + async def update(self, name: str, model_type: str, **fields: object) -> ModelEndpointRow | None: ... + + @abstractmethod + async def rename(self, name: str, model_type: str, new_name: str) -> None: ... @abstractmethod async def delete(self, name: str, model_type: str) -> bool: ... + + @abstractmethod + async def set_default(self, model_type: str, name: str) -> None: ... diff --git a/openrag/core/ports/partition_repo.py b/openrag/core/ports/partition_repo.py index 00898ce76..9a479bf8b 100644 --- a/openrag/core/ports/partition_repo.py +++ b/openrag/core/ports/partition_repo.py @@ -22,3 +22,12 @@ async def delete_partition(self, name: str) -> bool: ... @abstractmethod async def partition_exists(self, name: str) -> bool: ... + + @abstractmethod + async def get_partition_row(self, name: str) -> dict | None: ... + + @abstractmethod + async def list_partition_rows(self) -> list[dict]: ... + + @abstractmethod + async def update_partition(self, name: str, **fields: object) -> dict | None: ... diff --git a/openrag/core/ports/preset_repo.py b/openrag/core/ports/preset_repo.py index 363b654fc..0516691b6 100644 --- a/openrag/core/ports/preset_repo.py +++ b/openrag/core/ports/preset_repo.py @@ -17,5 +17,11 @@ 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 rename(self, old_name: str, new_name: str, preset_type: str, config: dict) -> dict: ... + @abstractmethod async def delete(self, name: str, preset_type: str) -> bool: ... + + @abstractmethod + async def count_partitions_using(self, name: str, preset_type: str) -> int: ... diff --git a/openrag/di/container.py b/openrag/di/container.py index d9cd3ba4e..837c0c0a8 100644 --- a/openrag/di/container.py +++ b/openrag/di/container.py @@ -21,6 +21,7 @@ from __future__ import annotations import os +from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any from core.embeddings import embedder_registry @@ -29,6 +30,7 @@ from core.utils.logging import get_logger from core.vlm import vlm_registry from di.embedders import register_embedders +from di.factories import make_component_factory from di.llms import register_llms from di.repositories import create_catalog_store from di.rerankers import register_rerankers @@ -60,7 +62,9 @@ from services.orchestrators.indexing_service import IndexingService from services.orchestrators.job_service import JobService from services.orchestrators.mcp_service import MCPService + from services.orchestrators.model_endpoint_service import ModelEndpointService from services.orchestrators.partition_service import PartitionService + from services.orchestrators.preset_service import PresetService from services.orchestrators.query_service import QueryService from services.orchestrators.retrieval_service import RetrievalService from services.orchestrators.user_service import UserService @@ -96,11 +100,22 @@ def __init__(self, settings: Settings | None = None) -> None: self._settings = settings self._initialized = False self._inference_clients: list[Any] = [] + self._client_caches: list[dict[str, Any]] = [] + self._embedder_cache: dict[str, Any] = {} + self._reranker_cache: dict[str, Any] = {} + self._llm_cache: dict[str, Any] = {} + self._vlm_cache: dict[str, Any] = {} + self.embedder_factory = self._missing_named_factory("embedder") + self.reranker_factory = self._missing_named_factory("reranker") + self.llm_factory = self._missing_named_factory("llm") + self.vlm_factory = self._missing_named_factory("vlm") 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._model_endpoint_service: ModelEndpointService | None = None + self._preset_service: PresetService | None = None self._workspace_service: WorkspaceService | None = None self._retrieval_service: RetrievalService | None = None self._query_service: QueryService | None = None @@ -108,6 +123,8 @@ def __init__(self, settings: Settings | None = None) -> None: self._job_service: JobService | None = None self._conversion_service: ConversionService | None = None self._mcp_service: MCPService | None = None + if settings is not None: + self._wire_named_component_factories(settings) def _require_settings(self) -> Settings: """Settings guard for the settings-dependent service properties. @@ -121,6 +138,40 @@ def _require_settings(self) -> Settings: raise RuntimeError(_NO_SETTINGS_MESSAGE) return self._settings + def _missing_named_factory(self, _kind: str): + def factory(_name: str = "default"): + raise RuntimeError(_NO_SETTINGS_MESSAGE) + + return factory + + def _wire_named_component_factories(self, settings: Settings) -> None: + """Wire Phase 14 named inference factories from ``settings.models``.""" + models = settings.models + self.embedder_factory, self._embedder_cache = make_component_factory( + registry=embedder_registry, + config_section=models.embedder, + default_impl="vllm", + client_caches=self._client_caches, + ) + self.reranker_factory, self._reranker_cache = make_component_factory( + registry=reranker_registry, + config_section=models.reranker, + default_impl="infinity", + client_caches=self._client_caches, + ) + self.llm_factory, self._llm_cache = make_component_factory( + registry=llm_registry, + config_section=models.llm, + default_impl="vllm", + client_caches=self._client_caches, + ) + self.vlm_factory, self._vlm_cache = make_component_factory( + registry=vlm_registry, + config_section=models.vlm, + default_impl="vllm", + client_caches=self._client_caches, + ) + # ------------------------------------------------------------------ # Lifecycle # ------------------------------------------------------------------ @@ -128,13 +179,28 @@ def _require_settings(self) -> Settings: async def initialize(self) -> None: """Open the storage adapters (asyncpg pool + Alembic migrations).""" if self._catalog_store is not None: - logger.info("ServiceContainer.initialize: initializing catalog store") - await self._catalog_store.initialize() - logger.info("ServiceContainer.initialize: ensuring admin user") - await self.user_repo.ensure_admin_user(os.getenv("AUTH_TOKEN")) - logger.info("ServiceContainer.initialize: admin user ready") + await self._initialize_step("initializing catalog store", self._catalog_store.initialize) + await self._initialize_step( + "ensuring admin user", + lambda: self.user_repo.ensure_admin_user(os.getenv("AUTH_TOKEN")), + ) + await self._initialize_step("seeding model endpoints", self.model_endpoint_service.seed_defaults) + await self._initialize_step("loading model endpoints", self.model_endpoint_service.load_all) + await self._initialize_step("seeding pipeline presets", self.preset_service.seed_defaults) + await self._initialize_step("loading pipeline presets", self.preset_service.load_all) + await self._initialize_step("ensuring default partition", self.partition_service.seed_default_partition) + await self._initialize_step("loading partition configs", self.partition_service.load_partitions) self._initialized = True + async def _initialize_step(self, label: str, operation: Callable[[], Awaitable[Any]]) -> None: + """Run one startup step with consistent failure logging.""" + logger.info(f"ServiceContainer.initialize: {label}") + try: + await operation() + except Exception: + logger.exception("ServiceContainer initialization step failed", step=label) + raise + async def shutdown(self) -> None: """Close inference clients and storage adapters cleanly. @@ -142,19 +208,33 @@ async def shutdown(self) -> None: remaining clients, the database pool, or the state reset. """ try: + seen_client_ids: set[int] = set() for client in self._inference_clients: - aclose = getattr(client, "aclose", None) - if aclose is not None: - try: - await aclose() - except Exception: - logger.exception("Failed to close inference client") + await self._close_inference_client(client, seen_client_ids) + for cache in self._client_caches: + for client in list(cache.values()): + await self._close_inference_client(client, seen_client_ids) + cache.clear() if self._catalog_store is not None: await self._catalog_store.shutdown() finally: self._inference_clients.clear() self._initialized = False + async def _close_inference_client(self, client: Any, seen_client_ids: set[int]) -> None: + """Close one tracked inference client once, best-effort.""" + client_id = id(client) + if client_id in seen_client_ids: + return + seen_client_ids.add(client_id) + aclose = getattr(client, "aclose", None) + if aclose is None: + return + try: + await aclose() + except Exception: + logger.exception("Failed to close inference client") + @property def is_initialized(self) -> bool: """True once :meth:`initialize` has completed its async I/O.""" @@ -319,9 +399,42 @@ def partition_service(self) -> PartitionService: vector_store=self.vector_store, user_repo=self.user_repo, collection=settings.vectordb.collection_name, + config=settings, ) return self._partition_service + @property + def model_endpoint_service(self) -> ModelEndpointService: + """ModelEndpointService — DB-backed named model endpoint registry.""" + if self._model_endpoint_service is None: + from services.orchestrators.model_endpoint_service import ModelEndpointService + + self._model_endpoint_service = ModelEndpointService( + model_endpoint_repo=self.model_endpoint_repo, + config=self._require_settings(), + partition_service=self.partition_service, + client_caches={ + "embedder": self._embedder_cache, + "reranker": self._reranker_cache, + "llm": self._llm_cache, + "vlm": self._vlm_cache, + }, + ) + return self._model_endpoint_service + + @property + def preset_service(self) -> PresetService: + """PresetService — DB-backed pipeline preset registry.""" + if self._preset_service is None: + from services.orchestrators.preset_service import PresetService + + self._preset_service = PresetService( + preset_repo=self.preset_repo, + config=self._require_settings(), + partition_service=self.partition_service, + ) + return self._preset_service + @property def workspace_service(self) -> WorkspaceService: """WorkspaceService — lazily built, cached for the container's lifetime.""" @@ -359,6 +472,15 @@ def retrieval_service(self) -> RetrievalService: document_repo=self.document_repo, collection=settings.vectordb.collection_name, ) + + def searcher_factory(embedder_name: str): + return VectorStoreSearcher( + vector_store=self.vector_store, + embedder=self.embedder_factory(embedder_name), + document_repo=self.document_repo, + collection=settings.vectordb.collection_name, + ) + llm_cfg = settings.llm.model_dump() llm = self.create_llm( "vllm", @@ -382,6 +504,9 @@ def retrieval_service(self) -> RetrievalService: reranker=reranker, llm=llm, config=settings, + searcher_factory=searcher_factory, + reranker_factory=self.reranker_factory, + llm_factory=self.llm_factory, ) return self._retrieval_service @@ -437,6 +562,8 @@ def indexing_service(self) -> IndexingService: workspace_repo=self.workspace_repo, collection=settings.vectordb.collection_name, ), + config=settings, + partition_service=self.partition_service, ) return self._indexing_service diff --git a/openrag/di/providers.py b/openrag/di/providers.py index 67c19ba5e..866cccba7 100644 --- a/openrag/di/providers.py +++ b/openrag/di/providers.py @@ -11,7 +11,7 @@ from __future__ import annotations import threading -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from fastapi import HTTPException, Request, status @@ -128,6 +128,27 @@ def get_mcp_service(request: Request = None) -> MCPService: return _require_initialized(request).mcp_service +def _get_optional_service(container: ServiceContainer, attribute_name: str) -> Any: + """Resolve a service that can be absent until its phase lands.""" + service = getattr(container, attribute_name, None) + if service is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=f"{attribute_name} is not available.", + ) + return service + + +def get_model_endpoint_service(request: Request = None) -> Any: + """Resolve the Phase 14 model endpoint orchestrator from the active container.""" + return _get_optional_service(_require_initialized(request), "model_endpoint_service") + + +def get_preset_service(request: Request = None) -> Any: + """Resolve the Phase 14 preset orchestrator from the active container.""" + return _get_optional_service(_require_initialized(request), "preset_service") + + def get_config(request: Request = None): """Resolve application configuration from the active container.""" return _require_initialized(request).config @@ -141,7 +162,9 @@ def get_config(request: Request = None): "get_indexing_service", "get_job_service", "get_mcp_service", + "get_model_endpoint_service", "get_partition_service", + "get_preset_service", "get_query_service", "get_retrieval_service", "get_user_service", diff --git a/openrag/services/orchestrators/indexing_service.py b/openrag/services/orchestrators/indexing_service.py index 3850bd9c8..a4bf05dff 100644 --- a/openrag/services/orchestrators/indexing_service.py +++ b/openrag/services/orchestrators/indexing_service.py @@ -14,15 +14,18 @@ from __future__ import annotations from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any +from core.utils.exceptions import PartitionNotFoundError from core.utils.filename import extract_temporal_fields from core.utils.logging import get_logger if TYPE_CHECKING: + from core.config.root import Settings from core.indexing.dispatcher import IndexingDispatcher from core.ports.document_repo import DocumentRepository from core.ports.workspace_repo import WorkspaceRepository + from services.orchestrators.partition_service import PartitionService logger = get_logger() @@ -53,10 +56,14 @@ def __init__( document_repo: DocumentRepository, workspace_repo: WorkspaceRepository, dispatcher: IndexingDispatcher, + config: Settings | None = None, + partition_service: PartitionService | None = None, ) -> None: self._document_repo = document_repo self._workspace_repo = workspace_repo self._dispatcher = dispatcher + self._config = config + self._partition_service = partition_service # ------------------------------------------------------------------ # Lookups (used by the thin router for its byte-identical guards) @@ -103,6 +110,42 @@ def _build_metadata( metadata.update(extract_temporal_fields(metadata, temporal_fields=TEMPORAL_FIELDS)) return metadata + def _partition_configs(self) -> dict[str, Any]: + if self._config is None: + return {} + return getattr(self._config, "partitions", {}) or {} + + def _resolve_indexation_dispatch_config(self, partition: str) -> tuple[dict | None, str | None]: + partitions = self._partition_configs() + if not partitions: + return None, None + if partition not in partitions: + raise PartitionNotFoundError(f"Partition '{partition}' does not exist.") + partition_cfg = partitions[partition] + return partition_cfg.indexation.model_dump(mode="json"), partition_cfg.embedder + + async def _ensure_partition_exists(self, partition: str, user: dict | None) -> None: + """Auto-create the partition on first index, matching legacy behaviour. + + Indexing into an unknown partition historically created it (with the + uploader as owner) and then indexed. Phase 14J's per-partition + resolution requires the partition to be present in ``config.partitions``, + so create it here with default presets before resolving. No-op when + there is no preset registry to resolve against (legacy passthrough) or + no partition service wired. + """ + if self._partition_service is None or not self._partition_configs(): + return + if partition in self._partition_configs(): + return + if await self._partition_service.partition_exists(partition): + # Row exists but the in-memory cache is stale — refresh it. + await self._partition_service.load_partitions() + return + user_id = (user or {}).get("id") or 1 + await self._partition_service.create_partition(partition, user_id=user_id) + logger.bind(partition=partition, user_id=user_id).info("Auto-created partition on index.") + async def add_file( self, *, @@ -128,6 +171,8 @@ async def add_file( sanitized_filename=sanitized_filename, original_filename=original_filename, ) + await self._ensure_partition_exists(partition, user) + indexation_config, embedder_name = self._resolve_indexation_dispatch_config(partition) return await self._dispatcher.dispatch_indexing( path=file_path, metadata=full_metadata, @@ -135,6 +180,8 @@ async def add_file( user=user, workspace_ids=workspace_ids, replace=replace, + indexation_config=indexation_config, + embedder_name=embedder_name, ) async def delete_file(self, file_id: str, partition: str) -> None: diff --git a/openrag/services/orchestrators/model_endpoint_service.py b/openrag/services/orchestrators/model_endpoint_service.py new file mode 100644 index 000000000..64dc34a4b --- /dev/null +++ b/openrag/services/orchestrators/model_endpoint_service.py @@ -0,0 +1,319 @@ +"""ModelEndpointService — CRUD, env-driven seeding, and client-cache invalidation. + +Orchestrates ModelEndpointRepository to maintain the named endpoint registry +in the DB and in the in-memory config (Settings.models). On first boot it +seeds one default endpoint per model type from existing env/config values +so the system works without any admin interaction. +""" + +from __future__ import annotations + +import os +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any + +from core.config.model_endpoints import ModelEndpointConfig, ModelEndpointRow +from core.utils.exceptions import NotFoundError, ValidationError +from core.utils.logging import get_logger + +if TYPE_CHECKING: + from core.config.root import Settings + from core.ports.model_endpoint_repo import ModelEndpointRepository + +logger = get_logger() + +_VALID_TYPES = frozenset({"embedder", "reranker", "llm", "vlm"}) + + +def _slug(model_name: str) -> str: + """'owner/model-name' → 'model-name'; '' → 'default'.""" + return model_name.split("/")[-1] if model_name else "default" + + +def _with_api_key(extra: dict[str, Any], api_key: str | None) -> dict[str, Any]: + """Add ``api_key`` to endpoint extras when configured.""" + if api_key: + return {**extra, "api_key": api_key} + return extra + + +class ModelEndpointService: + """CRUD and lifecycle management for named model endpoints.""" + + def __init__( + self, + *, + model_endpoint_repo: ModelEndpointRepository, + config: Settings, + partition_service: Any = None, + client_caches: dict[str, dict[str, Any]] | None = None, + ) -> None: + self._repo = model_endpoint_repo + self._config = config + self._partition_service = partition_service + self._client_caches: dict[str, dict[str, Any]] = client_caches or {} + + # ------------------------------------------------------------------ + # Startup lifecycle + # ------------------------------------------------------------------ + + async def seed_defaults(self) -> None: + """Insert one default endpoint per type if the DB is empty for that type. + + Seeds are derived from existing Settings / env-var values so that + existing deployments continue working after the Phase 14 upgrade + without any admin intervention. + """ + seeds = self._build_default_seeds() + now = datetime.now(UTC) + for model_type, data in seeds.items(): + existing = await self._repo.list_all(model_type=model_type) + if existing: + continue + endpoint: str = data["endpoint"] + model_name: str = data["model_name"] + if not endpoint: + logger.info(f"No {model_type} endpoint configured — skipping seed.") + continue + row = ModelEndpointRow( + name=_slug(model_name or ""), + model_type=model_type, + endpoint=endpoint, + model_name=model_name or None, + extra=data.get("extra", {}), + is_default=True, + created_at=now, + updated_at=now, + ) + await self._repo.create(row) + logger.info(f"Seeded default {model_type} endpoint '{row.name}'.") + + def _build_default_seeds(self) -> dict[str, dict[str, Any]]: + """Build seed data from env overrides + existing Settings fallbacks.""" + s = self._config + return { + "embedder": { + "endpoint": os.getenv("EMBEDDER_ENDPOINT", s.embedder.base_url), + "model_name": os.getenv("EMBEDDING_MODEL", s.embedder.model_name), + "extra": _with_api_key( + {"implementation": "vllm"}, + os.getenv("EMBEDDER_API_KEY", s.embedder.api_key), + ), + }, + "llm": { + "endpoint": os.getenv("LLM_ENDPOINT", s.llm.base_url), + "model_name": os.getenv("LLM_MODEL", s.llm.model), + "extra": _with_api_key( + {"implementation": "vllm"}, + os.getenv("API_KEY", s.llm.api_key), + ), + }, + "vlm": { + "endpoint": os.getenv("VLM_ENDPOINT", s.vlm.base_url), + "model_name": os.getenv("VLM_MODEL", s.vlm.model), + "extra": _with_api_key( + {"implementation": "vllm"}, + os.getenv("VLM_API_KEY", s.vlm.api_key), + ), + }, + "reranker": { + # Catalog the reranker endpoint whenever it is configured, like + # the embedder — registration is about availability, not whether + # reranking is on. Activation is the retrieval preset's + # enable_reranker kill-switch, which inherits reranker.enabled, + # so a disabled reranker is seeded but unused by default and + # remains available for per-partition opt-in. + "endpoint": os.getenv("RERANKER_ENDPOINT", s.reranker.base_url), + "model_name": os.getenv("RERANKER_MODEL", s.reranker.model_name), + "extra": _with_api_key( + {"implementation": s.reranker.provider}, + os.getenv("RERANKER_API_KEY", s.reranker.api_key), + ), + }, + } + + async def load_all(self) -> None: + """Fetch all endpoints from DB, rebuild config.models dicts atomically. + + Adds a virtual 'default' alias pointing to the is_default=True row for + each model type so component factories can resolve 'default' without + knowing the actual endpoint name. + """ + rows = await self._repo.list_all() + buckets: dict[str, dict[str, ModelEndpointConfig]] = {t: {} for t in _VALID_TYPES} + default_cfgs: dict[str, ModelEndpointConfig] = {} + + for row in rows: + bucket = buckets.get(row.model_type) + if bucket is None: + continue + cfg = ModelEndpointConfig( + endpoint=row.endpoint, + model_name=row.model_name, + batch_size=row.batch_size, + timeout=row.timeout, + extra=row.extra, + ) + bucket[row.name] = cfg + if row.is_default: + default_cfgs[row.model_type] = cfg + + for model_type, default_cfg in default_cfgs.items(): + buckets[model_type]["default"] = default_cfg + + models = self._config.models + for attr in ("embedder", "reranker", "llm", "vlm"): + target: dict = getattr(models, attr) + target.clear() + target.update(buckets[attr]) + + logger.info( + "Loaded model endpoints.", + n_embedder=len(buckets["embedder"]), + n_llm=len(buckets["llm"]), + n_reranker=len(buckets["reranker"]), + n_vlm=len(buckets["vlm"]), + ) + + # ------------------------------------------------------------------ + # CRUD + # ------------------------------------------------------------------ + + async def create_model_endpoint(self, row: ModelEndpointRow) -> ModelEndpointRow: + """Register a new endpoint; raises 409 if (name, model_type) already exists.""" + if row.model_type not in _VALID_TYPES: + raise ValidationError(f"Invalid model_type '{row.model_type}'. Must be one of: {sorted(_VALID_TYPES)}") + existing = await self._repo.get(row.name, row.model_type) + if existing is not None: + raise ValidationError( + f"Endpoint '{row.name}' of type '{row.model_type}' already exists.", + status_code=409, + code="ENDPOINT_EXISTS", + ) + result = await self._repo.create(row) + await self.load_all() + return result + + async def get_model_endpoint(self, name: str, model_type: str) -> ModelEndpointRow: + """Fetch one endpoint row; raises 404 if not found.""" + row = await self._repo.get(name, model_type) + if row is None: + raise NotFoundError(f"Endpoint '{name}' of type '{model_type}' not found.") + return row + + async def list_model_endpoints(self, model_type: str | None = None) -> list[ModelEndpointRow]: + return await self._repo.list_all(model_type=model_type) + + async def update_model_endpoint(self, name: str, model_type: str, **fields: object) -> ModelEndpointRow: + """Update endpoint fields and/or rename it. + + Pass ``new_name=`` to rename. After any change the in-memory config is + reloaded and the stale cached client instance is evicted so the next + request builds a fresh client against the updated config. + """ + existing = await self._repo.get(name, model_type) + if existing is None: + raise NotFoundError(f"Endpoint '{name}' of type '{model_type}' not found.") + + new_name: str | None = fields.pop("new_name", None) # type: ignore[assignment] + + if fields: + updated = await self._repo.update(name, model_type, **fields) + else: + updated = existing + + effective_name = name + if new_name and new_name != name: + await self._repo.rename(name, model_type, new_name) + effective_name = new_name + self._invalidate_client_cache(model_type, name) + + self._invalidate_client_cache(model_type, effective_name) + await self.load_all() + return await self._repo.get(effective_name, model_type) or (updated or existing) + + async def delete_model_endpoint(self, name: str, model_type: str) -> None: + """Delete an endpoint. + + Raises 404 if not found, 422 if it is the last endpoint of its type + (would leave components with no fallback). + """ + existing = await self._repo.get(name, model_type) + if existing is None: + raise NotFoundError(f"Endpoint '{name}' of type '{model_type}' not found.") + + all_of_type = await self._repo.list_all(model_type=model_type) + if len(all_of_type) <= 1: + raise ValidationError(f"Cannot delete the last '{model_type}' endpoint. Register a replacement first.") + + await self._repo.delete(name, model_type) + self._invalidate_client_cache(model_type, name) + await self.load_all() + + async def set_default(self, model_type: str, name: str) -> None: + """Promote ``name`` to the default endpoint for ``model_type``.""" + existing = await self._repo.get(name, model_type) + if existing is None: + raise NotFoundError(f"Endpoint '{name}' of type '{model_type}' not found.") + await self._repo.set_default(model_type, name) + self._invalidate_client_cache(model_type, "default") + await self.load_all() + + # ------------------------------------------------------------------ + # Endpoint validation + # ------------------------------------------------------------------ + + async def validate_endpoint( + self, + url: str, + model_name: str | None = None, + *, + api_key: str | None = None, + ) -> dict[str, Any]: + """Probe ``{url}/models`` to verify the endpoint is reachable and serving. + + Returns a dict with ``reachable``, ``model_found``, ``models_served``, + and ``detail`` keys — suitable as a ``ValidateEndpointResponse`` payload. + """ + import httpx # local import to avoid hard dep in tests that mock it + + result: dict[str, Any] = { + "reachable": False, + "model_found": None, + "models_served": None, + "detail": None, + } + models_url = url.rstrip("/") + "/models" + headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} + try: + async with httpx.AsyncClient(timeout=5.0, headers=headers) as client: + resp = await client.get(models_url) + result["reachable"] = True + if resp.status_code == 200: + data = resp.json() + served = [m["id"] for m in data.get("data", []) if "id" in m] + result["models_served"] = served + if model_name is not None: + result["model_found"] = model_name in served + else: + result["detail"] = f"HTTP {resp.status_code}" + except httpx.ConnectError as exc: + result["detail"] = f"Connection error: {exc}" + except httpx.TimeoutException: + result["detail"] = "Connection timed out" + except Exception as exc: # noqa: BLE001 + result["detail"] = str(exc) + return result + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + def _invalidate_client_cache(self, model_type: str, name: str) -> None: + """Evict ``name`` from the component-factory cache for ``model_type``.""" + cache = self._client_caches.get(model_type) + if cache is not None: + cache.pop(name, None) + + +__all__ = ["ModelEndpointService"] diff --git a/openrag/services/orchestrators/partition_service.py b/openrag/services/orchestrators/partition_service.py index ec981d427..c4d4d8869 100644 --- a/openrag/services/orchestrators/partition_service.py +++ b/openrag/services/orchestrators/partition_service.py @@ -28,7 +28,11 @@ from typing import TYPE_CHECKING, Any import numpy as np +from core.config.indexation_pipeline import IndexationPipelineConfig +from core.config.retrieval_pipeline import RetrievalPipelineConfig +from core.models.preset import PartitionConfig from core.utils.exceptions import ( + ConfigError, NotFoundError, PartitionNotFoundError, UserNotFoundError, @@ -37,6 +41,7 @@ from core.utils.logging import get_logger if TYPE_CHECKING: + from core.config.root import Settings from core.ports.document_repo import DocumentRepository from core.ports.partition_membership_repo import PartitionMembershipRepository from core.ports.partition_repo import PartitionRepository @@ -58,6 +63,7 @@ def __init__( vector_store: VectorStore, user_repo: UserRepository, collection: str, + config: Settings | None = None, ) -> None: self._partition_repo = partition_repo self._membership_repo = membership_repo @@ -65,6 +71,7 @@ def __init__( self._vector_store = vector_store self._user_repo = user_repo self._collection = collection + self._config = config # ------------------------------------------------------------------ # Existence guards (mirror the legacy _check_* helpers, core exceptions) @@ -113,12 +120,28 @@ async def partition_exists(self, partition: str) -> bool: 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``. + async def create_partition( + self, + partition: str, + user_id: int, + *, + description: str = "", + embedder: str = "default", + indexation_preset: str = "default", + retrieval_preset: str = "default", + chat_history_depth: int = 0, + chat_llm: str | None = None, + ) -> None: + """Create a partition owned by ``user_id`` with preset references. 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. + + When ``config`` was supplied to the service, the referenced presets + are validated *before* the row is written (so a bad preset name fails + fast and atomically), the non-default config columns are persisted, + and the in-memory partition cache is re-resolved. """ if await self._partition_repo.partition_exists(name=partition): raise ValidationError( @@ -126,7 +149,29 @@ async def create_partition(self, partition: str, user_id: int) -> None: status_code=409, code="PARTITION_EXISTS", ) + + config_fields = { + "description": description, + "embedder": embedder, + "indexation_preset": indexation_preset, + "retrieval_preset": retrieval_preset, + "chat_history_depth": chat_history_depth, + "chat_llm": chat_llm, + } + + # Validate the preset references before touching the DB. + if self._config is not None: + self._validate_preset_refs({"partition": partition, **config_fields}) + await self._partition_repo.create_partition(name=partition, user_id=user_id) + + # Persist the config columns (the insert only sets server defaults) + # and re-resolve the in-memory cache. Only done in the Phase 14 flow + # where a config was supplied. + if self._config is not None: + await self._partition_repo.update_partition(partition, **config_fields) + await self.load_partitions() + logger.info(f"Partition '{partition}' created by user_id {user_id}.") async def delete_partition(self, partition: str) -> None: @@ -151,6 +196,136 @@ async def delete_partition(self, partition: str) -> None: ) logger.info("Partition successfully deleted.", partition=partition) + async def update_partition(self, partition: str, **fields: object) -> dict | None: + """Update a partition's config columns and re-resolve the cache. + + ``None`` values are ignored (so partial PATCH semantics work). When a + preset reference changes, the merged row is validated *before* the + write so an unknown preset name fails fast. + """ + await self._ensure_partition(partition) + updates = {k: v for k, v in fields.items() if v is not None} + + if self._config is not None and updates: + current = await self._partition_repo.get_partition_row(partition) + if current is None: + raise PartitionNotFoundError(f"Partition '{partition}' does not exist.") + self._validate_preset_refs({**current, **updates}) + + result = await self._partition_repo.update_partition(partition, **updates) + + if self._config is not None: + await self.load_partitions() + + logger.info("Partition updated.", partition=partition, fields=sorted(updates)) + return result + + async def get_partition_config(self, partition: str) -> dict: + """Return the resolved Phase 14 detail for a partition. + + Shapes a ``PartitionDetailResponse``: the stored preset references plus + the fully resolved indexation/retrieval pipelines. Raises 404 if the + partition does not exist. + """ + self._require_config() + row = await self._partition_repo.get_partition_row(partition) + if row is None: + raise PartitionNotFoundError(f"Partition '{partition}' does not exist.") + return self._partition_detail(row, self.resolve_partition_row(row)) + + async def update_partition_config(self, partition: str, **fields: object) -> dict: + """Update a partition's preset references and return the resolved detail.""" + await self.update_partition(partition, **fields) + return await self.get_partition_config(partition) + + def _validate_preset_refs(self, row: dict) -> None: + """Validate a row's preset references for create/update. + + The preset names come from user input, so a missing preset is a client + error: translate the resolver's ConfigError (which maps to 500) into a + 422 ValidationError. + """ + try: + self.resolve_partition_row(row) + except ConfigError as exc: + raise ValidationError(exc.message, code="PRESET_NOT_FOUND") from exc + + def _partition_detail(self, row: dict, cfg: PartitionConfig) -> dict: + """Shape a resolved row into the ``PartitionDetailResponse`` payload.""" + return { + "name": cfg.name, + "description": cfg.description, + "embedder": cfg.embedder, + "indexation_preset": row.get("indexation_preset") or "default", + "retrieval_preset": row.get("retrieval_preset") or "default", + "indexation_pipeline": cfg.indexation.model_dump(mode="json"), + "retrieval_pipeline": cfg.retrieval.model_dump(mode="json"), + "dimension": row.get("dimension"), + "created_at": row.get("created_at"), + } + + # ------------------------------------------------------------------ + # Preset resolution + in-memory cache + # ------------------------------------------------------------------ + + def resolve_partition_row(self, row: dict) -> PartitionConfig: + """Resolve a partition DB row into a fully-validated PartitionConfig. + + Looks the referenced preset names up in ``config.presets`` and builds + the Pydantic pipeline configs. Raises :class:`ConfigError` if either + referenced preset is missing — the caller decides whether that is a + startup failure or a 4xx on create/update. + """ + cfg = self._require_config() + name = row["partition"] + idx_name = row.get("indexation_preset") or "default" + ret_name = row.get("retrieval_preset") or "default" + + idx_preset = cfg.presets.indexation.get(idx_name) + if idx_preset is None: + raise ConfigError(f"Indexation preset '{idx_name}' referenced by partition '{name}' not found.") + ret_preset = cfg.presets.retrieval.get(ret_name) + if ret_preset is None: + raise ConfigError(f"Retrieval preset '{ret_name}' referenced by partition '{name}' not found.") + + return PartitionConfig( + name=name, + description=row.get("description") or "", + embedder=row.get("embedder") or "default", + indexation=IndexationPipelineConfig(**idx_preset), + retrieval=RetrievalPipelineConfig(**ret_preset), + collection_name=row.get("collection_name"), + chat_history_depth=row.get("chat_history_depth") or 0, + chat_llm=row.get("chat_llm"), + ) + + async def load_partitions(self) -> None: + """Resolve every partition row and swap the in-memory cache atomically. + + Uses clear()+update() so any reference already held to the + ``config.partitions`` dict stays valid after the swap. + """ + cfg = self._require_config() + rows = await self._partition_repo.list_partition_rows() + resolved = {row["partition"]: self.resolve_partition_row(row) for row in rows} + + cache = cfg.partitions + cache.clear() + cache.update(resolved) + logger.info("Loaded partition configs.", n_partitions=len(resolved)) + + async def seed_default_partition(self, user_id: int = 1) -> None: + """Ensure the 'default' partition exists with default presets.""" + if await self._partition_repo.partition_exists(name="default"): + return + await self._partition_repo.create_partition(name="default", user_id=user_id) + logger.info("Seeded 'default' partition.") + + def _require_config(self) -> Settings: + if self._config is None: + raise ConfigError("PartitionService was constructed without a config; preset resolution unavailable.") + return self._config + # ------------------------------------------------------------------ # File / chunk reads # ------------------------------------------------------------------ diff --git a/openrag/services/orchestrators/preset_service.py b/openrag/services/orchestrators/preset_service.py new file mode 100644 index 000000000..ceb7bf060 --- /dev/null +++ b/openrag/services/orchestrators/preset_service.py @@ -0,0 +1,247 @@ +"""PresetService — CRUD, default seeding, validation, and live re-resolution. + +Orchestrates PresetRepository to maintain the named pipeline preset registry +in the DB and in the in-memory config (Settings.presets). On first boot it +seeds 6 default presets (3 indexation + 3 retrieval) so partitions resolve +without any admin intervention. + +When a preset is updated the service calls partition_service.load_partitions() +so all affected partitions pick up the change without a restart. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from core.config.indexation_pipeline import IndexationPipelineConfig +from core.config.retrieval_pipeline import RetrievalPipelineConfig +from core.utils.exceptions import NotFoundError, ValidationError +from core.utils.logging import get_logger +from services.orchestrators.partition_service import PartitionService + +if TYPE_CHECKING: + from core.config.root import Settings + from core.ports.preset_repo import PresetRepository + +logger = get_logger() + +_VALID_PRESET_TYPES = frozenset({"indexation", "retrieval"}) + +_DEFAULT_SEEDS: dict[str, dict[str, dict[str, Any]]] = { + "indexation": { + "default": { + "chunking": {"name": "recursive_splitter", "chunk_size": 512, "chunk_overlap_rate": 0.2}, + "parsing_strategy": "marker", + "enable_image_captioning": True, + "enable_contextualization": False, + "enable_entity_extraction": True, + "enable_topic_tagging": True, + }, + "legal": { + "chunking": {"name": "recursive_splitter", "chunk_size": 1024, "chunk_overlap_rate": 0.25}, + "parsing_strategy": "marker", + "enable_image_captioning": True, + "enable_contextualization": True, + "contextualization_mode": "structured", + }, + "finance": { + "chunking": {"name": "recursive_splitter", "chunk_size": 768, "chunk_overlap_rate": 0.2}, + "parsing_strategy": "marker", + "enable_image_captioning": True, + "enable_contextualization": False, + }, + }, + # ``enable_reranker`` is intentionally omitted here — it is injected at seed + # time from the global ``reranker.enabled`` kill-switch (see _finalize_seed) + # so a deployment without a reranker (CPU-only, CI) does not force reranking + # on every partition and then fail against an unreachable reranker endpoint. + "retrieval": { + "default": { + "type": "single", + "top_k": 50, + "top_n": 10, + "similarity_threshold": 0.6, + }, + "multiquery": { + "type": "multiQuery", + "top_k": 50, + "top_n": 10, + }, + "hyde": { + "type": "hyde", + "top_k": 50, + "top_n": 10, + }, + }, +} + + +class PresetService: + """CRUD and lifecycle management for named pipeline presets.""" + + def __init__( + self, + *, + preset_repo: PresetRepository, + config: Settings, + partition_service: PartitionService | None = None, + ) -> None: + self._repo = preset_repo + self._config = config + self._partition_service = partition_service + + # ------------------------------------------------------------------ + # Startup lifecycle + # ------------------------------------------------------------------ + + async def seed_defaults(self) -> None: + """Insert default presets if the DB has no rows for that type. + + Each preset type is seeded independently — if indexation presets exist + but retrieval is empty, only retrieval gets seeded. + """ + for preset_type, presets in _DEFAULT_SEEDS.items(): + existing = await self._repo.list_all(preset_type=preset_type) + if existing: + continue + for name, config in presets.items(): + await self._repo.upsert(name, preset_type, self._finalize_seed(preset_type, config)) + logger.info(f"Seeded default {preset_type} preset '{name}'.") + + def _finalize_seed(self, preset_type: str, config: dict[str, Any]) -> dict[str, Any]: + """Overlay env-derived values onto a static default seed. + + Default retrieval presets inherit the global reranker kill-switch + (``reranker.enabled``): when no reranker is available the default + presets ship with ``enable_reranker=False`` so partition resolution + never tries to build a reranker against a missing/unreachable endpoint. + """ + if preset_type == "retrieval": + return {**config, "enable_reranker": self._config.reranker.enabled} + return config + + async def load_all(self) -> None: + """Fetch all presets from DB, rebuild config.presets dicts atomically. + + Uses clear()+update() so any existing references to the dict objects + (e.g. held by PartitionService) stay valid after the swap. + """ + rows = await self._repo.list_all() + idx_bucket: dict[str, dict[str, Any]] = {} + ret_bucket: dict[str, dict[str, Any]] = {} + + for row in rows: + if row["preset_type"] == "indexation": + idx_bucket[row["name"]] = row["config"] + elif row["preset_type"] == "retrieval": + ret_bucket[row["name"]] = row["config"] + + presets = self._config.presets + presets.indexation.clear() + presets.indexation.update(idx_bucket) + presets.retrieval.clear() + presets.retrieval.update(ret_bucket) + + logger.info( + "Loaded pipeline presets.", + n_indexation=len(idx_bucket), + n_retrieval=len(ret_bucket), + ) + + # ------------------------------------------------------------------ + # CRUD + # ------------------------------------------------------------------ + + async def create_preset(self, name: str, preset_type: str, config: dict[str, Any]) -> dict: + """Register a new preset; raises 409 if (name, preset_type) already exists.""" + if preset_type not in _VALID_PRESET_TYPES: + raise ValidationError(f"Invalid preset_type '{preset_type}'. Must be 'indexation' or 'retrieval'.") + existing = await self._repo.get(name, preset_type) + if existing is not None: + raise ValidationError( + f"Preset '{name}' of type '{preset_type}' already exists.", + status_code=409, + code="PRESET_EXISTS", + ) + self._validate_config(preset_type, config) + result = await self._repo.upsert(name, preset_type, config) + await self.load_all() + return result + + async def get_preset(self, name: str, preset_type: str) -> dict: + """Fetch one preset row; raises 404 if not found.""" + row = await self._repo.get(name, preset_type) + if row is None: + raise NotFoundError(f"Preset '{name}' of type '{preset_type}' not found.") + return row + + async def list_presets(self, preset_type: str | None = None) -> list[dict]: + return await self._repo.list_all(preset_type=preset_type) + + async def update_preset(self, name: str, preset_type: str, **fields: object) -> dict: + """Update preset config and/or rename it. + + Pass ``new_name=`` to rename. If the config changes, all partitions + referencing this preset are re-resolved immediately. + """ + existing = await self._repo.get(name, preset_type) + if existing is None: + raise NotFoundError(f"Preset '{name}' of type '{preset_type}' not found.") + + new_name: str | None = fields.pop("new_name", None) # type: ignore[assignment] + new_config: dict | None = fields.pop("config", None) # type: ignore[assignment] + + if new_config is not None: + self._validate_config(preset_type, new_config) + + effective_config = new_config if new_config is not None else existing["config"] + effective_name = name + + if new_name and new_name != name: + effective_name = new_name + result = await self._repo.rename(name, effective_name, preset_type, effective_config) + else: + result = await self._repo.upsert(effective_name, preset_type, effective_config) + await self.load_all() + + if self._partition_service is not None: + await self._partition_service.load_partitions() + + return result + + async def delete_preset(self, name: str, preset_type: str) -> None: + """Delete a preset. + + Raises 404 if not found, 422 if any partition references it. + """ + existing = await self._repo.get(name, preset_type) + if existing is None: + raise NotFoundError(f"Preset '{name}' of type '{preset_type}' not found.") + + count = await self._repo.count_partitions_using(name, preset_type) + if count > 0: + raise ValidationError( + f"Cannot delete preset '{name}': {count} partition(s) are using it. " + "Reassign them to a different preset first." + ) + + await self._repo.delete(name, preset_type) + await self.load_all() + + # ------------------------------------------------------------------ + # Validation + # ------------------------------------------------------------------ + + def _validate_config(self, preset_type: str, config: dict[str, Any]) -> None: + """Instantiate the Pydantic pipeline model to validate the config dict.""" + try: + match preset_type: + case "indexation": + IndexationPipelineConfig(**config) + case "retrieval": + RetrievalPipelineConfig(**config) + except Exception as exc: + raise ValidationError(f"Invalid {preset_type} preset config: {exc}") from exc + + +__all__ = ["PresetService"] diff --git a/openrag/services/orchestrators/retrieval_service.py b/openrag/services/orchestrators/retrieval_service.py index 5ec81db3e..e9830265a 100644 --- a/openrag/services/orchestrators/retrieval_service.py +++ b/openrag/services/orchestrators/retrieval_service.py @@ -25,7 +25,8 @@ from __future__ import annotations import asyncio -from typing import TYPE_CHECKING +from collections.abc import Callable +from typing import TYPE_CHECKING, Any from core.prompts import load_template_by_key from core.retrieval.pipeline import RetrieverPipeline @@ -36,6 +37,7 @@ _expand_with_related_chunks, ) from core.retrieval.rrf import rrf_reranking +from core.utils.exceptions import PartitionNotFoundError from core.utils.logging import get_logger if TYPE_CHECKING: @@ -60,14 +62,34 @@ def __init__( self, *, searcher: RetrievalSearcher, - reranker: Reranker | None, - llm: LLM | None, + reranker: Reranker | None = None, + llm: LLM | None = None, config: Settings, + searcher_factory: Callable[[str], RetrievalSearcher] | None = None, + reranker_factory: Callable[[str], Reranker] | None = None, + llm_factory: Callable[[str], LLM] | None = None, ) -> None: self._searcher = searcher + self._config = config + self._legacy_reranker = reranker + self._legacy_llm = llm + self._searcher_factory = searcher_factory + self._reranker_factory = reranker_factory + self._llm_factory = llm_factory + self._pipeline = self._build_legacy_pipeline(reranker=reranker, llm=llm) + + logger.debug( + "RetrievalService ready", + retriever=config.retriever.type, + reranker_enabled=config.reranker.enabled and reranker is not None, + partition_configs=len(getattr(config, "partitions", {}) or {}), + ) + + def _build_legacy_pipeline(self, *, reranker: Reranker | None, llm: LLM | None) -> RetrieverPipeline: + config = self._config rcfg = config.retriever common = { - "searcher": searcher, + "searcher": self._searcher, "top_k": rcfg.top_k, "similarity_threshold": rcfg.similarity_threshold, "with_surrounding_chunks": rcfg.with_surrounding_chunks, @@ -94,17 +116,110 @@ def __init__( else: retriever = SingleRetriever(**common) - self._pipeline = RetrieverPipeline( + return 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, + + def _build_retriever( + self, + *, + rtype: str, + common: dict[str, Any], + llm: LLM | None, + k_queries: int, + combine: bool, + ): + if rtype == "multiQuery": + return MultiQueryRetriever( + llm=llm, + multi_query_template=load_template_by_key( + self._config.paths.prompts_dir, + self._config.prompts, + "multi_query", + ), + k_queries=k_queries, + **common, + ) + if rtype == "hyde": + return HyDeRetriever( + llm=llm, + hyde_template=load_template_by_key(self._config.paths.prompts_dir, self._config.prompts, "hyde"), + combine=combine, + **common, + ) + return SingleRetriever(**common) + + def _partition_configs(self) -> dict[str, Any]: + return getattr(self._config, "partitions", {}) or {} + + def _require_partition_config(self, partition: str): + partitions = self._partition_configs() + if partition not in partitions: + raise PartitionNotFoundError(f"Partition '{partition}' does not exist.") + return partitions[partition] + + def _legacy_retriever_value(self, name: str, default: Any) -> Any: + return getattr(self._config.retriever, name, default) + + def _pipeline_for_partition(self, partition: str) -> tuple[RetrieverPipeline, int | None]: + if partition == "all" or not self._partition_configs(): + return self._pipeline, None + + partition_cfg = self._require_partition_config(partition) + pipeline_cfg = partition_cfg.retrieval + searcher = ( + self._searcher_factory(partition_cfg.embedder) if self._searcher_factory is not None else self._searcher + ) + + rtype = pipeline_cfg.type + llm = self._legacy_llm + if rtype in {"multiQuery", "hyde"} and self._llm_factory is not None: + llm = self._llm_factory(pipeline_cfg.llm or partition_cfg.chat_llm or "default") + + reranker = None + if pipeline_cfg.enable_reranker: + if self._reranker_factory is not None: + reranker = self._reranker_factory(pipeline_cfg.reranker or "default") + else: + reranker = self._legacy_reranker + + retriever = self._build_retriever( + rtype=rtype, + common={ + "searcher": searcher, + "top_k": pipeline_cfg.top_k, + "similarity_threshold": pipeline_cfg.similarity_threshold, + "with_surrounding_chunks": self._legacy_retriever_value("with_surrounding_chunks", False), + "include_related": pipeline_cfg.include_related, + "include_ancestors": pipeline_cfg.include_ancestors, + "related_limit": self._legacy_retriever_value("related_limit", 10), + "max_ancestor_depth": self._legacy_retriever_value("max_ancestor_depth", None), + }, + llm=llm, + k_queries=self._legacy_retriever_value("k_queries", 3), + combine=self._legacy_retriever_value("combine", False), + ) + pipeline = RetrieverPipeline( + retriever=retriever, + reranker=reranker, + reranker_top_k=pipeline_cfg.top_n, + allow_filterless_fallback=self._legacy_retriever_value("allow_filterless_fallback", True), ) + return pipeline, pipeline_cfg.top_n + + def _pipeline_groups_for_partitions( + self, partitions: list[str] + ) -> list[tuple[list[str], RetrieverPipeline, int | None]]: + if not partitions or "all" in partitions or not self._partition_configs(): + return [(["all"] if "all" in partitions else partitions, self._pipeline, None)] + return [ + ([partition], pipeline, default_top_k) + for partition in partitions + for pipeline, default_top_k in [self._pipeline_for_partition(partition)] + ] # ------------------------------------------------------------------ # Raw semantic search (powers routers/search.py — was indexer.asearch) @@ -164,12 +279,18 @@ async def retrieve( 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, + ranked_lists = await asyncio.gather( + *[ + pipeline.retrieve_docs( + partition=partition_group, + query=query, + top_k=top_k if top_k is not None else default_top_k, + filter_params=filter_params, + ) + for partition_group, pipeline, default_top_k in self._pipeline_groups_for_partitions(partitions) + ] ) + return ranked_lists[0] if len(ranked_lists) == 1 else self.fuse(ranked_lists, top_k=top_k) async def retrieve_multi( self, @@ -180,12 +301,18 @@ async def retrieve_multi( 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, + ranked_lists = await asyncio.gather( + *[ + pipeline.get_relevant_docs( + partition=partition_group, + search_queries=search_queries, + top_k=top_k if top_k is not None else default_top_k, + filter_params=filter_params, + ) + for partition_group, pipeline, default_top_k in self._pipeline_groups_for_partitions(partitions) + ] ) + return ranked_lists[0] if len(ranked_lists) == 1 else self.fuse(ranked_lists, top_k=top_k) async def retrieve_per_query( self, @@ -202,15 +329,7 @@ async def retrieve_per_query( 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 - ] + *[self.retrieve(partitions=partitions, query=q, top_k=top_k, filter_params=filter_params) for q in queries] ) @staticmethod diff --git a/openrag/services/persistence/document_repo.py b/openrag/services/persistence/document_repo.py index 5fc487b4f..85e8d23f3 100644 --- a/openrag/services/persistence/document_repo.py +++ b/openrag/services/persistence/document_repo.py @@ -69,12 +69,13 @@ async def create_document(self, doc: DocumentRecord) -> DocumentRecord: 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) + indexation_config, created_by, relationship_id, parent_id) + VALUES ($1, $2, $3::json, $4::jsonb, $5, $6, $7) """, file_id, doc.partition, metadata, + doc.indexation_config, doc.created_by, doc.relationship_id, doc.parent_id, @@ -157,6 +158,9 @@ async def update_document(self, document_id: str, **fields: Any) -> DocumentReco # We always rewrite file_metadata so JSON-only updates are persisted. params.append(metadata) sets.append(f"file_metadata = ${len(params)}::json") + if "indexation_config" in fields: + params.append(fields.pop("indexation_config")) + sets.append(f"indexation_config = ${len(params)}::jsonb") for column in ("relationship_id", "parent_id", "created_by"): if column in fields: @@ -266,6 +270,7 @@ async def add_file_to_partition( # noqa: PLR0913 — legacy signature pinned user_id: int | None = None, relationship_id: str | None = None, parent_id: str | None = None, + indexation_config: dict | None = None, ) -> bool: """TODO(phase-9): remove. Mirror of legacy ``add_file_to_partition``. @@ -309,12 +314,13 @@ async def add_file_to_partition( # noqa: PLR0913 — legacy signature pinned 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) + indexation_config, created_by, relationship_id, parent_id) + VALUES ($1, $2, $3::json, $4::jsonb, $5, $6, $7) """, file_id, partition, file_metadata or {}, + indexation_config, user_id, relationship_id, parent_id, @@ -386,6 +392,7 @@ async def update_file_in_partition( file_metadata: dict | None = None, relationship_id: object = _UNSET, parent_id: object = _UNSET, + indexation_config: object = _UNSET, ) -> bool: """TODO(phase-9): remove. PUT-style in-place update. @@ -404,6 +411,9 @@ async def update_file_in_partition( if parent_id is not self._UNSET: params.append(parent_id) sets.append(f"parent_id = ${len(params)}") + if indexation_config is not self._UNSET: + params.append(indexation_config) + sets.append(f"indexation_config = ${len(params)}::jsonb") if not sets: # Match legacy: report whether the row exists at all. return await self.file_exists_in_partition(file_id, partition) @@ -569,6 +579,7 @@ def _row_to_document(row: asyncpg.Record) -> DocumentRecord: created_by=row["created_by"], relationship_id=row["relationship_id"], parent_id=row["parent_id"], + indexation_config=row["indexation_config"], ) diff --git a/openrag/services/persistence/migrations/alembic/versions/06dd2101ea3a_add_endpoints_presets_phase14.py b/openrag/services/persistence/migrations/alembic/versions/06dd2101ea3a_add_endpoints_presets_phase14.py new file mode 100644 index 000000000..0ca267772 --- /dev/null +++ b/openrag/services/persistence/migrations/alembic/versions/06dd2101ea3a_add_endpoints_presets_phase14.py @@ -0,0 +1,167 @@ +"""add_endpoints_presets_phase14 + +Revision ID: 06dd2101ea3a +Revises: f5b6c918f741 +Create Date: 2026-06-04 10:24:14.741985 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from schema_helpers import column_exists, index_exists, table_exists +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision: str = "06dd2101ea3a" +down_revision: str | Sequence[str] | None = "f5b6c918f741" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + if not table_exists("model_endpoints"): + op.create_table( + "model_endpoints", + sa.Column("name", sa.String(), nullable=False), + sa.Column("model_type", sa.String(), nullable=False), + sa.Column("endpoint", sa.String(), nullable=False), + sa.Column("model_name", sa.String(), nullable=True), + sa.Column("batch_size", sa.Integer(), server_default="32", nullable=False), + sa.Column("timeout", sa.Float(), server_default="30.0", nullable=False), + sa.Column( + "extra", + postgresql.JSONB(astext_type=sa.Text()), + server_default=sa.text("'{}'::jsonb"), + nullable=False, + ), + sa.Column("is_default", sa.Boolean(), server_default="false", nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.CheckConstraint( + "model_type IN ('embedder','reranker','llm','vlm')", + name="ck_model_endpoint_type", + ), + sa.PrimaryKeyConstraint("name", "model_type"), + ) + + if not table_exists("pipeline_presets"): + op.create_table( + "pipeline_presets", + sa.Column("name", sa.String(), nullable=False), + sa.Column("preset_type", sa.String(), nullable=False), + sa.Column( + "config", + postgresql.JSONB(astext_type=sa.Text()), + nullable=False, + ), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.CheckConstraint( + "preset_type IN ('indexation','retrieval')", + name="ck_pipeline_preset_type", + ), + sa.PrimaryKeyConstraint("name", "preset_type"), + ) + + if not column_exists("files", "indexation_config"): + op.add_column( + "files", + sa.Column( + "indexation_config", + postgresql.JSONB(astext_type=sa.Text()), + nullable=True, + ), + ) + + _partition_columns = { + "description": sa.Column("description", sa.String(), server_default=sa.text("''"), nullable=False), + "embedder": sa.Column( + "embedder", + sa.String(), + server_default=sa.text("'default'"), + nullable=False, + ), + "indexation_preset": sa.Column( + "indexation_preset", + sa.String(), + server_default=sa.text("'default'"), + nullable=False, + ), + "retrieval_preset": sa.Column( + "retrieval_preset", + sa.String(), + server_default=sa.text("'default'"), + nullable=False, + ), + "dimension": sa.Column("dimension", sa.Integer(), server_default="1024", nullable=False), + "collection_name": sa.Column("collection_name", sa.String(), nullable=True), + "chat_history_depth": sa.Column("chat_history_depth", sa.Integer(), server_default="0", nullable=False), + "chat_llm": sa.Column("chat_llm", sa.String(), nullable=True), + "updated_at": sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + } + for col_name, col_def in _partition_columns.items(): + if not column_exists("partitions", col_name): + op.add_column("partitions", col_def) + + if not index_exists("workspaces", "ix_workspaces_partition_name"): + op.create_index( + op.f("ix_workspaces_partition_name"), + "workspaces", + ["partition_name"], + unique=False, + ) + + +def downgrade() -> None: + if index_exists("workspaces", "ix_workspaces_partition_name"): + op.drop_index(op.f("ix_workspaces_partition_name"), table_name="workspaces") + + for col in [ + "updated_at", + "chat_llm", + "chat_history_depth", + "collection_name", + "dimension", + "retrieval_preset", + "indexation_preset", + "embedder", + "description", + ]: + if column_exists("partitions", col): + op.drop_column("partitions", col) + + if column_exists("files", "indexation_config"): + op.drop_column("files", "indexation_config") + + if table_exists("pipeline_presets"): + op.drop_table("pipeline_presets") + + if table_exists("model_endpoints"): + op.drop_table("model_endpoints") diff --git a/openrag/services/persistence/model_endpoint_repo.py b/openrag/services/persistence/model_endpoint_repo.py index 7e73e65b4..e41b30fb8 100644 --- a/openrag/services/persistence/model_endpoint_repo.py +++ b/openrag/services/persistence/model_endpoint_repo.py @@ -1,31 +1,136 @@ -"""Stub :class:`ModelEndpointRepository`. +"""asyncpg-backed :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. +Manages the ``model_endpoints`` table — named inference endpoint +configurations for embedders, LLMs, rerankers, and VLMs. Phase 14D +replaces the earlier stub with real SQL. """ from __future__ import annotations +from collections.abc import Callable +from typing import TYPE_CHECKING + +from core.config.model_endpoints import ModelEndpointRow from core.ports.model_endpoint_repo import ModelEndpointRepository -from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented +if TYPE_CHECKING: + import asyncpg + +_ALLOWED_UPDATE_FIELDS = frozenset({"endpoint", "model_name", "batch_size", "timeout", "extra", "is_default"}) + + +class PgModelEndpointRepository(ModelEndpointRepository): + """asyncpg-backed implementation of :class:`ModelEndpointRepository`.""" + + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter + + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() -class PgModelEndpointRepository(_StubRepositoryBase, ModelEndpointRepository): - """TODO: real impl once the ``model_endpoints`` table is added.""" + @staticmethod + def _to_model(row: asyncpg.Record) -> ModelEndpointRow: + return ModelEndpointRow( + name=row["name"], + model_type=row["model_type"], + endpoint=row["endpoint"], + model_name=row["model_name"], + batch_size=row["batch_size"], + timeout=row["timeout"], + extra=row["extra"] or {}, + is_default=row["is_default"], + created_at=row["created_at"], + updated_at=row["updated_at"], + ) - async def get(self, name: str, model_type: str) -> dict | None: - raise stub_not_implemented("DB-backed model endpoints") + async def create(self, row: ModelEndpointRow) -> ModelEndpointRow: + rec = await self.pool.fetchrow( + """ + INSERT INTO model_endpoints + (name, model_type, endpoint, model_name, batch_size, timeout, extra, is_default) + VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8) + RETURNING * + """, + row.name, + row.model_type, + row.endpoint, + row.model_name, + row.batch_size, + row.timeout, + row.extra, + row.is_default, + ) + return self._to_model(rec) - async def list_all(self, model_type: str | None = None) -> list[dict]: - raise stub_not_implemented("DB-backed model endpoints") + async def get(self, name: str, model_type: str) -> ModelEndpointRow | None: + rec = await self.pool.fetchrow( + "SELECT * FROM model_endpoints WHERE name = $1 AND model_type = $2", + name, + model_type, + ) + return self._to_model(rec) if rec else None - async def upsert(self, name: str, model_type: str, config: dict) -> dict: - raise stub_not_implemented("DB-backed model endpoints") + async def list_all(self, model_type: str | None = None) -> list[ModelEndpointRow]: + if model_type is not None: + rows = await self.pool.fetch( + "SELECT * FROM model_endpoints WHERE model_type = $1 ORDER BY name", + model_type, + ) + else: + rows = await self.pool.fetch( + "SELECT * FROM model_endpoints ORDER BY model_type, name", + ) + return [self._to_model(r) for r in rows] + + async def update(self, name: str, model_type: str, **fields: object) -> ModelEndpointRow | None: + updates = {k: v for k, v in fields.items() if k in _ALLOWED_UPDATE_FIELDS} + if not updates: + return await self.get(name, model_type) + + params: list = [name, model_type] + sets: list[str] = [] + for col, val in updates.items(): + idx = len(params) + 1 + sets.append(f"{col} = ${idx}::jsonb" if col == "extra" else f"{col} = ${idx}") + params.append(val) + + rec = await self.pool.fetchrow( + f"UPDATE model_endpoints SET {', '.join(sets)}, updated_at = now() " + f"WHERE name = $1 AND model_type = $2 RETURNING *", + *params, + ) + return self._to_model(rec) if rec else None + + async def rename(self, name: str, model_type: str, new_name: str) -> None: + await self.pool.execute( + "UPDATE model_endpoints SET name = $3, updated_at = now() WHERE name = $1 AND model_type = $2", + name, + model_type, + new_name, + ) async def delete(self, name: str, model_type: str) -> bool: - raise stub_not_implemented("DB-backed model endpoints") + result = await self.pool.execute( + "DELETE FROM model_endpoints WHERE name = $1 AND model_type = $2", + name, + model_type, + ) + return result == "DELETE 1" + + async def set_default(self, model_type: str, name: str) -> None: + async with self.pool.acquire() as conn: + async with conn.transaction(): + await conn.execute( + "UPDATE model_endpoints SET is_default = false, updated_at = now() WHERE model_type = $1", + model_type, + ) + await conn.execute( + "UPDATE model_endpoints SET is_default = true, updated_at = now() " + "WHERE name = $1 AND model_type = $2", + name, + model_type, + ) __all__ = ["PgModelEndpointRepository"] diff --git a/openrag/services/persistence/partition_repo.py b/openrag/services/persistence/partition_repo.py index b1b29cbf5..dcddf198d 100644 --- a/openrag/services/persistence/partition_repo.py +++ b/openrag/services/persistence/partition_repo.py @@ -136,6 +136,51 @@ async def partition_exists(self, name: str) -> bool: name, ) + # ── Phase 14 — full config row methods ─────────────────────────── + + async def get_partition_row(self, name: str) -> dict | None: + row = await self.pool.fetchrow( + "SELECT * FROM partitions WHERE partition = $1", + name, + ) + return self._row_to_full_dict(row) if row else None + + async def list_partition_rows(self) -> list[dict]: + rows = await self.pool.fetch( + "SELECT * FROM partitions ORDER BY created_at", + ) + return [self._row_to_full_dict(r) for r in rows] + + async def update_partition(self, name: str, **fields: object) -> dict | None: + _ALLOWED = frozenset( + { + "description", + "embedder", + "indexation_preset", + "retrieval_preset", + "dimension", + "collection_name", + "chat_history_depth", + "chat_llm", + } + ) + updates = {k: v for k, v in fields.items() if k in _ALLOWED} + if not updates: + return await self.get_partition_row(name) + + params: list = [name] + sets: list[str] = [] + for col, val in updates.items(): + idx = len(params) + 1 + sets.append(f"{col} = ${idx}") + params.append(val) + + row = await self.pool.fetchrow( + f"UPDATE partitions SET {', '.join(sets)}, updated_at = now() WHERE partition = $1 RETURNING *", + *params, + ) + return self._row_to_full_dict(row) if row else None + # ── Legacy method names used by the Phase 7C shim ──────────────── async def get_partition_file_count(self, partition: str) -> int: @@ -160,5 +205,22 @@ def _row_to_dict(row: asyncpg.Record) -> dict: "created_at": created.isoformat() if created else None, } + @staticmethod + def _row_to_full_dict(row: asyncpg.Record) -> dict: + """Full partition row including all Phase 14 config columns.""" + return { + "partition": row["partition"], + "description": row["description"], + "embedder": row["embedder"], + "indexation_preset": row["indexation_preset"], + "retrieval_preset": row["retrieval_preset"], + "dimension": row["dimension"], + "collection_name": row["collection_name"], + "chat_history_depth": row["chat_history_depth"], + "chat_llm": row["chat_llm"], + "created_at": row["created_at"], + "updated_at": row["updated_at"], + } + __all__ = ["PgPartitionRepository"] diff --git a/openrag/services/persistence/preset_repo.py b/openrag/services/persistence/preset_repo.py index be2caede0..c63c0f217 100644 --- a/openrag/services/persistence/preset_repo.py +++ b/openrag/services/persistence/preset_repo.py @@ -1,31 +1,120 @@ -"""Stub :class:`PresetRepository`. +"""asyncpg-backed :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. +Manages the ``pipeline_presets`` table — named pipeline configuration +bundles for indexation and retrieval. Phase 14D replaces the earlier +stub with real SQL. """ from __future__ import annotations +from collections.abc import Callable +from typing import TYPE_CHECKING + from core.ports.preset_repo import PresetRepository -from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented +if TYPE_CHECKING: + import asyncpg + +_VALID_PRESET_COLUMNS = { + "indexation": "indexation_preset", + "retrieval": "retrieval_preset", +} + +_UPSERT_PRESET_SQL = """ + INSERT INTO pipeline_presets (name, preset_type, config) + VALUES ($1, $2, $3::jsonb) + ON CONFLICT (name, preset_type) DO UPDATE + SET config = EXCLUDED.config, updated_at = now() + RETURNING * + """ + + +class PgPresetRepository(PresetRepository): + """asyncpg-backed implementation of :class:`PresetRepository`.""" + + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter -class PgPresetRepository(_StubRepositoryBase, PresetRepository): - """TODO: real impl once the ``presets`` table is added — see REFACTORING P0 plan.""" + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() + + @staticmethod + def _row_to_dict(row: asyncpg.Record) -> dict: + return { + "name": row["name"], + "preset_type": row["preset_type"], + "config": row["config"] or {}, + "created_at": row["created_at"], + "updated_at": row["updated_at"], + } async def get(self, name: str, preset_type: str) -> dict | None: - raise stub_not_implemented("Per-partition pipeline presets") + rec = await self.pool.fetchrow( + "SELECT * FROM pipeline_presets WHERE name = $1 AND preset_type = $2", + name, + preset_type, + ) + return self._row_to_dict(rec) if rec else None async def list_all(self, preset_type: str | None = None) -> list[dict]: - raise stub_not_implemented("Per-partition pipeline presets") + if preset_type is not None: + rows = await self.pool.fetch( + "SELECT * FROM pipeline_presets WHERE preset_type = $1 ORDER BY name", + preset_type, + ) + else: + rows = await self.pool.fetch( + "SELECT * FROM pipeline_presets ORDER BY preset_type, name", + ) + return [self._row_to_dict(r) for r in rows] async def upsert(self, name: str, preset_type: str, config: dict) -> dict: - raise stub_not_implemented("Per-partition pipeline presets") + rec = await self.pool.fetchrow( + _UPSERT_PRESET_SQL, + name, + preset_type, + config, + ) + return self._row_to_dict(rec) + + async def rename(self, old_name: str, new_name: str, preset_type: str, config: dict) -> dict: + async with self.pool.acquire() as conn: + async with conn.transaction(): + await conn.execute( + "DELETE FROM pipeline_presets WHERE name = $1 AND preset_type = $2", + old_name, + preset_type, + ) + rec = await conn.fetchrow( + _UPSERT_PRESET_SQL, + new_name, + preset_type, + config, + ) + return self._row_to_dict(rec) async def delete(self, name: str, preset_type: str) -> bool: - raise stub_not_implemented("Per-partition pipeline presets") + result = await self.pool.execute( + "DELETE FROM pipeline_presets WHERE name = $1 AND preset_type = $2", + name, + preset_type, + ) + return result == "DELETE 1" + + async def count_partitions_using(self, name: str, preset_type: str) -> int: + col = _preset_column(preset_type) + return await self.pool.fetchval( + f"SELECT COUNT(*)::int FROM partitions WHERE {col} = $1", + name, + ) + + +def _preset_column(preset_type: str) -> str: + try: + return _VALID_PRESET_COLUMNS[preset_type] + except KeyError as exc: + raise ValueError(f"Invalid preset_type: {preset_type}") from exc __all__ = ["PgPresetRepository"] diff --git a/openrag/services/persistence/schema.py b/openrag/services/persistence/schema.py index 22ab3ad7b..618946427 100644 --- a/openrag/services/persistence/schema.py +++ b/openrag/services/persistence/schema.py @@ -21,6 +21,7 @@ CheckConstraint, Column, DateTime, + Float, ForeignKey, Index, Integer, @@ -29,17 +30,88 @@ String, Table, UniqueConstraint, + text, ) +from sqlalchemy.dialects.postgresql import JSONB metadata = MetaData() +model_endpoints = Table( + "model_endpoints", + metadata, + Column("name", String, primary_key=True), + Column("model_type", String, primary_key=True), + Column("endpoint", String, nullable=False), + Column("model_name", String, nullable=True), + Column("batch_size", Integer, server_default="32", nullable=False), + Column("timeout", Float, server_default="30.0", nullable=False), + Column("extra", JSONB, server_default=text("'{}'::jsonb"), nullable=False), + Column("is_default", Boolean, server_default="false", nullable=False), + Column( + "created_at", + DateTime(timezone=True), + server_default=text("now()"), + nullable=False, + ), + Column( + "updated_at", + DateTime(timezone=True), + server_default=text("now()"), + nullable=False, + ), + CheckConstraint( + "model_type IN ('embedder','reranker','llm','vlm')", + name="ck_model_endpoint_type", + ), +) + + +pipeline_presets = Table( + "pipeline_presets", + metadata, + Column("name", String, primary_key=True), + Column("preset_type", String, primary_key=True), + Column("config", JSONB, nullable=False), + Column( + "created_at", + DateTime(timezone=True), + server_default=text("now()"), + nullable=False, + ), + Column( + "updated_at", + DateTime(timezone=True), + server_default=text("now()"), + nullable=False, + ), + CheckConstraint( + "preset_type IN ('indexation','retrieval')", + name="ck_pipeline_preset_type", + ), +) + + 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), + Column("description", String, server_default=text("''"), nullable=False), + Column("embedder", String, server_default=text("'default'"), nullable=False), + Column("indexation_preset", String, server_default=text("'default'"), nullable=False), + Column("retrieval_preset", String, server_default=text("'default'"), nullable=False), + Column("dimension", Integer, server_default="1024", nullable=False), + Column("collection_name", String, nullable=True), + Column("chat_history_depth", Integer, server_default="0", nullable=False), + Column("chat_llm", String, nullable=True), + Column( + "updated_at", + DateTime(timezone=True), + server_default=text("now()"), + nullable=False, + ), ) @@ -56,6 +128,7 @@ index=True, ), Column("file_metadata", JSON, nullable=True, default=dict), + Column("indexation_config", JSONB, nullable=True), Column( "created_by", Integer, @@ -195,6 +268,8 @@ __all__ = [ "metadata", + "model_endpoints", + "pipeline_presets", "partitions", "files", "users", diff --git a/openrag/services/storage/milvus_store.py b/openrag/services/storage/milvus_store.py index ab8139845..bd4dca4c0 100644 --- a/openrag/services/storage/milvus_store.py +++ b/openrag/services/storage/milvus_store.py @@ -144,10 +144,10 @@ def __init__(self, config: VectorDBConfig) -> None: self._collection_name = config.collection_name self._hybrid = config.hybrid_search self._uri = f"http://{config.host}:{config.port}" - + self._timeout = 60 try: - self._client = MilvusClient(uri=self._uri) - self._async_client = AsyncMilvusClient(uri=self._uri) + self._client = MilvusClient(uri=self._uri, timeout=self._timeout) + self._async_client = AsyncMilvusClient(uri=self._uri, timeout=self._timeout) except MilvusException as e: raise VDBConnectionError( f"Failed to connect to Milvus: {e!s}", diff --git a/openrag/services/workers/dispatcher.py b/openrag/services/workers/dispatcher.py index b318b9382..80d329b2b 100644 --- a/openrag/services/workers/dispatcher.py +++ b/openrag/services/workers/dispatcher.py @@ -68,6 +68,8 @@ async def dispatch_indexing( user: dict | None, workspace_ids: list[str] | None, replace: bool, + indexation_config: dict | None = None, + embedder_name: str | None = None, ) -> str: task_id = uuid.uuid4().hex @@ -96,6 +98,8 @@ async def dispatch_indexing( user=user, workspace_ids=workspace_ids, replace=replace, + indexation_config=indexation_config, + embedder_name=embedder_name, ) await self._call( diff --git a/openrag/services/workers/indexer_actor.py b/openrag/services/workers/indexer_actor.py index a619b571f..d162e35b3 100644 --- a/openrag/services/workers/indexer_actor.py +++ b/openrag/services/workers/indexer_actor.py @@ -5,6 +5,7 @@ from typing import Any from core.models.document import Document +from services.workers.parsers.doc_serializer_bridge import INDEXATION_CONFIG_METADATA_KEY from services.workers.pipeline_builder import IndexingPipeline @@ -46,6 +47,8 @@ async def process_file( user: dict[str, Any] | None = None, workspace_ids: list[str] | None = None, replace: bool = False, + indexation_config: dict[str, Any] | None = None, + embedder_name: str | None = None, ) -> dict[str, Any]: """Run one file through the indexing pipeline. @@ -55,7 +58,7 @@ async def process_file( """ await self._tsm.set_state.remote(task_id, "SERIALIZING") try: - document = _load_document(path, metadata, partition) + document = _load_document(path, metadata, partition, indexation_config=indexation_config) row: dict[str, Any] = { "document": document, "partition": partition, @@ -64,6 +67,8 @@ async def process_file( "replace": replace, "user": user, "workspace_ids": workspace_ids, + "indexation_config": indexation_config, + "embedder_name": embedder_name, } await self._pipeline.run(row) if self._document_repo is not None: @@ -73,6 +78,7 @@ async def process_file( partition=partition, user=user, replace=replace, + indexation_config=indexation_config, ) await self._tsm.set_state.remote(task_id, "COMPLETED") return {"stored_count": row.get("stored_count", 0), "stage": row.get("stage", "")} @@ -89,9 +95,11 @@ async def _write_catalog_record( partition: str, user: dict[str, Any] | None, replace: bool, + indexation_config: dict[str, Any] | None, ) -> None: file_id = metadata.get("file_id", "") file_metadata = {key: value for key, value in metadata.items() if key != "page"} + config_kwargs = {"indexation_config": indexation_config} if indexation_config is not None else {} if replace: await doc_repo.update_file_in_partition( file_id=file_id, @@ -99,6 +107,7 @@ async def _write_catalog_record( file_metadata=file_metadata, relationship_id=metadata.get("relationship_id"), parent_id=metadata.get("parent_id"), + **config_kwargs, ) return @@ -109,17 +118,27 @@ async def _write_catalog_record( user_id=user.get("id") if user else None, relationship_id=metadata.get("relationship_id"), parent_id=metadata.get("parent_id"), + **config_kwargs, ) -def _load_document(path: str, metadata: dict[str, Any], partition: str) -> Document: +def _load_document( + path: str, + metadata: dict[str, Any], + partition: str, + *, + indexation_config: dict[str, Any] | None = None, +) -> Document: p = Path(path) + document_metadata = dict(metadata) + if indexation_config is not None: + document_metadata[INDEXATION_CONFIG_METADATA_KEY] = dict(indexation_config) 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, + metadata=document_metadata, ) diff --git a/openrag/services/workers/indexer_pool.py b/openrag/services/workers/indexer_pool.py index a458c174e..25087d8a3 100644 --- a/openrag/services/workers/indexer_pool.py +++ b/openrag/services/workers/indexer_pool.py @@ -1,6 +1,8 @@ from __future__ import annotations import asyncio +import threading +from types import SimpleNamespace from typing import Any import ray @@ -12,6 +14,7 @@ class IndexerPool: """Thin Ray actor wrapper around ``IndexerWorker``.""" def __init__(self) -> None: + import services.inference.ollama_client # noqa: F401 import services.inference.vllm_client # noqa: F401 from core.config import load_config from core.embeddings import embedder_registry @@ -24,6 +27,7 @@ def __init__(self) -> None: parser = DocSerializerBridgeParser(config=cfg) chunker = _build_chunker(cfg) + embedder_factory = _build_embedder_factory(cfg) embed_cfg = cfg.embedder embedder = embedder_registry.create( @@ -40,6 +44,8 @@ def __init__(self) -> None: chunker=chunker, embedder=embedder, vector_store=self._vector_store, + chunker_factory=_build_chunker_from_config, + embedder_factory=embedder_factory, ) 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) @@ -65,6 +71,8 @@ async def process_file( user: dict[str, Any] | None = None, workspace_ids: list[str] | None = None, replace: bool = False, + indexation_config: dict[str, Any] | None = None, + embedder_name: str | None = None, ) -> dict[str, Any]: await self._ensure_catalog() result = await self._worker.process_file( @@ -75,6 +83,8 @@ async def process_file( user=user, workspace_ids=workspace_ids, replace=replace, + indexation_config=indexation_config, + embedder_name=embedder_name, ) file_id = metadata.get("file_id", "") if workspace_ids and not replace and file_id: @@ -107,4 +117,42 @@ def _build_chunker(cfg: Any) -> Any: return chunker +def _build_chunker_from_config(chunker_config: Any) -> Any: + return _build_chunker(SimpleNamespace(chunker=chunker_config)) + + +def _build_embedder_factory(cfg: Any) -> Any: + if not getattr(cfg.models, "embedder", None): + return None + + from core.embeddings import embedder_registry + + cache: dict[str, Any] = {} + lock = threading.Lock() + + def factory(name: str = "default") -> Any: + if name in cache: + return cache[name] + with lock: + if name in cache: + return cache[name] + model_cfg = cfg.models.embedder.get(name) + if model_cfg is None: + raise KeyError(f"Unknown embedder '{name}'. Available: {list(cfg.models.embedder)}") + impl_kwargs = {key: value for key, value in model_cfg.extra.items() if key != "implementation"} + impl = model_cfg.extra.get("implementation", "vllm") + instance = embedder_registry.create( + impl, + endpoint=model_cfg.endpoint, + model_name=model_cfg.model_name, + batch_size=model_cfg.batch_size, + timeout=model_cfg.timeout, + **impl_kwargs, + ) + cache[name] = instance + return instance + + return factory + + __all__ = ["IndexerPool", "build_indexer_pool"] diff --git a/openrag/services/workers/parsers/doc_serializer_bridge.py b/openrag/services/workers/parsers/doc_serializer_bridge.py index 87b2e924e..cfe3727e4 100644 --- a/openrag/services/workers/parsers/doc_serializer_bridge.py +++ b/openrag/services/workers/parsers/doc_serializer_bridge.py @@ -7,6 +7,8 @@ from core.indexing.parsers.document_parser import DocumentParser from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock +INDEXATION_CONFIG_METADATA_KEY = "_openrag_indexation_config" + class DocSerializerBridgeParser(DocumentParser): """Transitional parser backed by the legacy loader registry.""" @@ -33,11 +35,12 @@ async def parse(self, document: Document) -> ProcessedDocument: async def _load_via_legacy(self, path: str, document: Document) -> ProcessedDocument: metadata = dict(document.metadata or {}) + indexation_config = metadata.pop(INDEXATION_CONFIG_METADATA_KEY, None) 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) + loader = loader_cls(config=_legacy_loader_config(self._config, indexation_config)) lang_doc = await loader.aload_document( file_path=path, metadata=metadata, @@ -78,4 +81,20 @@ def _suffix_from_document(document: Document) -> str: return f".{document.content_type.value}" if document.content_type else "" +def _legacy_loader_config(config: Any, indexation_config: Any) -> Any: + """Apply per-file indexation overrides to legacy loader config.""" + if not isinstance(indexation_config, dict): + return config + if indexation_config.get("enable_image_captioning", True): + return config + + loader = config.loader.model_copy( + update={ + "image_captioning": False, + "image_captioning_url": False, + } + ) + return config.model_copy(update={"loader": loader}) + + __all__ = ["DocSerializerBridgeParser"] diff --git a/openrag/services/workers/parsers/marker_workers.py b/openrag/services/workers/parsers/marker_workers.py index 255d6fe4f..b08ca8a48 100644 --- a/openrag/services/workers/parsers/marker_workers.py +++ b/openrag/services/workers/parsers/marker_workers.py @@ -25,7 +25,21 @@ def _marker_num_gpus(config) -> float: - return config.loader.marker_num_gpus if torch.cuda.is_available() else 0 + """Return Marker's Ray GPU reservation, falling back to CUDA detection. + + Ray scheduling must see GPU capacity before an actor can request a GPU + fraction. If Ray cannot report cluster resources yet, use local CUDA + availability as the fallback so single-node startup still honors the + configured Marker GPU request. + """ + requested_gpus = config.loader.marker_num_gpus + if requested_gpus <= 0: + return 0 + try: + return requested_gpus if ray.cluster_resources().get("GPU", 0) > 0 else 0 + except Exception as exc: + logger.warning("Failed to query Ray cluster GPU resources; falling back to CUDA check", error=str(exc)) + return requested_gpus if torch.cuda.is_available() else 0 @ray.remote diff --git a/openrag/services/workers/pipeline_builder.py b/openrag/services/workers/pipeline_builder.py index 1bbb41d6a..bcb507fdb 100644 --- a/openrag/services/workers/pipeline_builder.py +++ b/openrag/services/workers/pipeline_builder.py @@ -1,10 +1,11 @@ from __future__ import annotations -from collections.abc import MutableMapping +from collections.abc import Callable, MutableMapping from dataclasses import dataclass from typing import Any from core.chunking.chunking_strategy import ChunkingStrategy +from core.config.indexation_pipeline import IndexationPipelineConfig from core.embeddings.embedder import Embedder from core.indexing.contextualize import ChunkContextualizer from core.indexing.parsers.document_parser import DocumentParser @@ -45,29 +46,42 @@ class IndexingPipeline: vlm: VLM | None = None contextualizer: ChunkContextualizer | None = None timeouts: PipelineTimeouts = PipelineTimeouts() + indexation_config: IndexationPipelineConfig | None = None + parser_factory: Callable[[str], DocumentParser] | None = None + chunker_factory: Callable[[Any], ChunkingStrategy] | None = None + embedder_factory: Callable[[str], Embedder] | None = None + vlm_factory: Callable[[str], VLM] | None = None + contextualizer_factory: Callable[[str], ChunkContextualizer] | None = None 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: + config = self._effective_indexation_config(row) + parser = self._select_parser(config) + chunker = self._select_chunker(config) + embedder = self._select_embedder(row) + vlm = self._select_vlm(config) + contextualizer = self._select_contextualizer(config) + + await parse_stage(row, parser, timeout=self.timeouts.parse) + if vlm is not None: await caption_stage( row, - self.vlm, + 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 chunk_stage(row, chunker, timeout=self.timeouts.chunk) + if contextualizer is not None: await contextualize_stage( row, - self.contextualizer, + contextualizer, timeout=self.timeouts.contextualize, per_chunk_timeout=self.timeouts.contextualize_per_chunk, ) await embed_stage( row, - self.embedder, + embedder, timeout=self.timeouts.embed, per_chunk_timeout=self.timeouts.embed_per_chunk, ) @@ -79,6 +93,48 @@ async def run(self, row: MutableMapping[str, Any]) -> MutableMapping[str, Any]: ) return row + def _effective_indexation_config(self, row: MutableMapping[str, Any]) -> IndexationPipelineConfig | None: + raw_config = row.get("indexation_config", self.indexation_config) + if raw_config is None: + return None + if isinstance(raw_config, IndexationPipelineConfig): + return raw_config + if isinstance(raw_config, dict): + return IndexationPipelineConfig(**raw_config) + raise TypeError("indexation_config must be an IndexationPipelineConfig or dict") + + def _select_parser(self, config: IndexationPipelineConfig | None) -> DocumentParser: + if config is not None and self.parser_factory is not None: + return self.parser_factory(config.parsing_strategy) + return self.parser + + def _select_chunker(self, config: IndexationPipelineConfig | None) -> ChunkingStrategy: + if config is not None and self.chunker_factory is not None: + return self.chunker_factory(config.chunking) + return self.chunker + + def _select_embedder(self, row: MutableMapping[str, Any]) -> Embedder: + embedder_name = row.get("embedder_name") + if embedder_name and self.embedder_factory is not None: + return self.embedder_factory(str(embedder_name)) + return self.embedder + + def _select_vlm(self, config: IndexationPipelineConfig | None) -> VLM | None: + if config is not None: + if not config.enable_image_captioning: + return None + if self.vlm_factory is not None: + return self.vlm_factory(config.vlm or "default") + return self.vlm + + def _select_contextualizer(self, config: IndexationPipelineConfig | None) -> ChunkContextualizer | None: + if config is not None: + if not config.enable_contextualization: + return None + if self.contextualizer_factory is not None: + return self.contextualizer_factory(config.contextualization_llm or "default") + return self.contextualizer + def build_indexing_pipeline( *, @@ -89,6 +145,12 @@ def build_indexing_pipeline( vlm: VLM | None = None, contextualizer: ChunkContextualizer | None = None, timeouts: PipelineTimeouts | None = None, + indexation_config: IndexationPipelineConfig | None = None, + parser_factory: Callable[[str], DocumentParser] | None = None, + chunker_factory: Callable[[Any], ChunkingStrategy] | None = None, + embedder_factory: Callable[[str], Embedder] | None = None, + vlm_factory: Callable[[str], VLM] | None = None, + contextualizer_factory: Callable[[str], ChunkContextualizer] | None = None, ) -> IndexingPipeline: """Build the default sequential indexing pipeline.""" @@ -100,6 +162,12 @@ def build_indexing_pipeline( vlm=vlm, contextualizer=contextualizer, timeouts=timeouts or PipelineTimeouts(), + indexation_config=indexation_config, + parser_factory=parser_factory, + chunker_factory=chunker_factory, + embedder_factory=embedder_factory, + vlm_factory=vlm_factory, + contextualizer_factory=contextualizer_factory, ) diff --git a/scripts/seed_presets.py b/scripts/seed_presets.py new file mode 100644 index 000000000..7be5b3053 --- /dev/null +++ b/scripts/seed_presets.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +"""Seed default model endpoints and pipeline presets from the global Settings. + +One-time migration utility for existing deployments upgrading to the Phase 14 +DB-backed model-endpoint registry and per-partition preset system. It reads the +current YAML / env ``Settings`` and populates the ``model_endpoints``, +``pipeline_presets`` and ``partitions`` tables with their default rows so the +system keeps working without any admin interaction. Subsequent changes go +through the admin API. + +The seeding itself lives in ``ServiceContainer.initialize()`` (the same path the +API runs on boot): it seeds endpoints, then presets, then the default partition, +and is idempotent — re-running skips anything already present. This script is a +thin operator entry point that runs that seed once, standalone (no API / Ray), +and prints a summary, so the migration can be performed without starting the +full app. + +Usage:: + + uv run python scripts/seed_presets.py +""" + +from __future__ import annotations + +import asyncio + +from _bootstrap import ensure_openrag_source_path + +ensure_openrag_source_path() + +from core.config import load_config # noqa: E402 +from di.container import ServiceContainer # noqa: E402 + + +async def main() -> None: + container = ServiceContainer(load_config()) + # initialize() runs the idempotent 3-phase seed: model endpoints → presets + # → default partition (and loads each into the in-memory config). + await container.initialize() + try: + config = container.config + print( + f"Seeded model endpoints: " + f"{len(config.models.embedder)} embedder, " + f"{len(config.models.reranker)} reranker, " + f"{len(config.models.llm)} llm, " + f"{len(config.models.vlm)} vlm" + ) + print( + f"Seeded {len(config.presets.indexation)} indexation presets, " + f"{len(config.presets.retrieval)} retrieval presets" + ) + print(f"Seeded {len(config.partitions)} partition(s)") + finally: + await container.shutdown() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/integration/api/test_model_endpoints.py b/tests/integration/api/test_model_endpoints.py new file mode 100644 index 000000000..f36341ec5 --- /dev/null +++ b/tests/integration/api/test_model_endpoints.py @@ -0,0 +1,81 @@ +"""Model endpoint admin API integration tests.""" + +import uuid + +MOCK_VLLM_ENDPOINT = "http://vllm:8000/v1" +MOCK_CHAT_MODEL = "mock-chat-model" + + +def _assert_success(response, *, context: str) -> None: + assert 200 <= response.status_code < 300, f"{context}: {response.status_code} {response.text}" + + +def _delete_ignore_errors(api_client, path: str) -> None: + try: + api_client.delete(path) + except Exception: + pass + + +def test_model_endpoint_crud_validate_and_default_selection(api_client): + """Create, rename, validate, promote, restore, and delete an LLM endpoint.""" + suffix = uuid.uuid4().hex[:8] + endpoint_name = f"ci-llm-{suffix}" + endpoint_renamed = f"{endpoint_name}-renamed" + original_default_llm = None + + try: + openapi = api_client.get("/openapi.json") + _assert_success(openapi, context="openapi") + assert "/model-endpoints/" in openapi.json()["paths"] + + endpoints = api_client.get("/model-endpoints/", params={"model_type": "llm"}) + _assert_success(endpoints, context="list llm endpoints") + original_default_llm = next(row["name"] for row in endpoints.json() if row["is_default"]) + + create_endpoint = api_client.post( + "/model-endpoints/", + json={ + "name": endpoint_name, + "model_type": "llm", + "endpoint": MOCK_VLLM_ENDPOINT, + "model_name": f"missing-model-{suffix}", + "timeout": 5, + "extra": {"implementation": "vllm", "api_key": "ci-test-key"}, + }, + ) + assert create_endpoint.status_code == 201, create_endpoint.text + assert create_endpoint.json()["extra"]["api_key"] == "ci-test-key" + + validate_missing = api_client.post(f"/model-endpoints/llm/{endpoint_name}/validate") + _assert_success(validate_missing, context="validate missing model") + missing_probe = validate_missing.json() + assert missing_probe["reachable"] is True + assert missing_probe["model_found"] is False + assert MOCK_CHAT_MODEL in missing_probe["models_served"] + + rename_endpoint = api_client.put( + f"/model-endpoints/llm/{endpoint_name}", + json={"name": endpoint_renamed, "model_name": MOCK_CHAT_MODEL}, + ) + _assert_success(rename_endpoint, context="rename endpoint") + assert rename_endpoint.json()["name"] == endpoint_renamed + + validate_real = api_client.post(f"/model-endpoints/llm/{endpoint_renamed}/validate") + _assert_success(validate_real, context="validate mock model") + real_probe = validate_real.json() + assert real_probe["reachable"] is True + assert real_probe["model_found"] is True + + set_default = api_client.post(f"/model-endpoints/llm/{endpoint_renamed}/set-default") + _assert_success(set_default, context="set endpoint default") + assert set_default.json()["is_default"] is True + + restore_default = api_client.post(f"/model-endpoints/llm/{original_default_llm}/set-default") + _assert_success(restore_default, context="restore original default endpoint") + assert restore_default.json()["is_default"] is True + finally: + if original_default_llm is not None: + api_client.post(f"/model-endpoints/llm/{original_default_llm}/set-default") + _delete_ignore_errors(api_client, f"/model-endpoints/llm/{endpoint_renamed}") + _delete_ignore_errors(api_client, f"/model-endpoints/llm/{endpoint_name}") diff --git a/tests/integration/api/test_partition_presets.py b/tests/integration/api/test_partition_presets.py new file mode 100644 index 000000000..ef481a1a3 --- /dev/null +++ b/tests/integration/api/test_partition_presets.py @@ -0,0 +1,136 @@ +"""Partition preset assignment integration tests.""" + +import uuid + +from conftest import wait_for_indexing + + +def _assert_success(response, *, context: str) -> None: + assert 200 <= response.status_code < 300, f"{context}: {response.status_code} {response.text}" + + +def _delete_ignore_errors(api_client, path: str) -> None: + try: + api_client.delete(path) + except Exception: + pass + + +def _indexation_config(chunk_size: int) -> dict: + return { + "chunking": { + "name": "recursive_splitter", + "chunk_size": chunk_size, + "chunk_overlap_rate": 0.1, + }, + "parsing_strategy": "pymupdf", + "enable_image_captioning": False, + "enable_contextualization": False, + "enable_metadata_extraction": False, + "enable_entity_extraction": False, + "enable_topic_tagging": False, + } + + +def _retrieval_config() -> dict: + return { + "type": "single", + "top_k": 5, + "top_n": 3, + "similarity_threshold": 0.0, + "enable_reranker": False, + "include_related": False, + "include_ancestors": False, + } + + +def test_partition_preset_assignment_drives_indexing_and_search(api_client, tmp_path): + """Assign presets to a partition, index a file, and search the resolved partition.""" + suffix = uuid.uuid4().hex[:8] + indexation_preset = f"ci-index-{suffix}" + retrieval_preset = f"ci-retrieval-{suffix}" + partition = f"ci-partition-{suffix}" + file_id = f"ci-file-{suffix}" + marker = f"ci-marker-{suffix}" + + try: + openapi = api_client.get("/openapi.json") + _assert_success(openapi, context="openapi") + assert "/partition/{partition}/config" in openapi.json()["paths"] + + create_indexation = api_client.post( + "/presets/", + json={ + "name": indexation_preset, + "preset_type": "indexation", + "config": _indexation_config(96), + }, + ) + assert create_indexation.status_code == 201, create_indexation.text + + create_retrieval = api_client.post( + "/presets/", + json={ + "name": retrieval_preset, + "preset_type": "retrieval", + "config": _retrieval_config(), + }, + ) + assert create_retrieval.status_code == 201, create_retrieval.text + + create_partition = api_client.post(f"/partition/{partition}") + assert create_partition.status_code == 201, create_partition.text + + update_partition = api_client.patch( + f"/partition/{partition}", + json={ + "indexation_preset": indexation_preset, + "retrieval_preset": retrieval_preset, + }, + ) + _assert_success(update_partition, context="assign partition presets") + partition_config = update_partition.json() + assert partition_config["indexation_preset"] == indexation_preset + assert partition_config["retrieval_preset"] == retrieval_preset + assert partition_config["indexation_pipeline"]["chunking"]["chunk_size"] == 96 + assert partition_config["retrieval_pipeline"]["top_n"] == 3 + + invalid_preset = api_client.patch( + f"/partition/{partition}", + json={"indexation_preset": f"missing-indexation-{suffix}"}, + ) + assert invalid_preset.status_code == 422 + assert "PRESET_NOT_FOUND" in invalid_preset.text + + test_file = tmp_path / "preset-indexing.txt" + test_file.write_text( + f"Preset assignment CI document. Unique marker: {marker}. " + "This file verifies preset-based indexing and search." + ) + + with test_file.open("rb") as handle: + upload = api_client.post( + f"/indexer/partition/{partition}/file/{file_id}", + files={"file": ("preset-indexing.txt", handle, "text/plain")}, + data={"metadata": f'{{"source":"preset-ci","marker":"{marker}"}}'}, + ) + assert upload.status_code in {200, 201, 202}, upload.text + wait_for_indexing(api_client, upload.json()) + + indexed_file = api_client.get(f"/partition/{partition}/file/{file_id}") + _assert_success(indexed_file, context="get indexed file") + assert len(indexed_file.json()["documents"]) > 0 + + search = api_client.get( + f"/search/partition/{partition}", + params={"text": marker, "top_k": 3, "similarity_threshold": 0.0}, + ) + _assert_success(search, context="search indexed file") + documents = search.json()["documents"] + assert documents + assert any(doc["metadata"].get("file_id") == file_id for doc in documents) + finally: + _delete_ignore_errors(api_client, f"/indexer/partition/{partition}/file/{file_id}") + _delete_ignore_errors(api_client, f"/partition/{partition}") + _delete_ignore_errors(api_client, f"/presets/indexation/{indexation_preset}") + _delete_ignore_errors(api_client, f"/presets/retrieval/{retrieval_preset}") diff --git a/tests/integration/api/test_presets.py b/tests/integration/api/test_presets.py new file mode 100644 index 000000000..cdde5b53e --- /dev/null +++ b/tests/integration/api/test_presets.py @@ -0,0 +1,109 @@ +"""Preset admin API integration tests.""" + +import uuid + + +def _assert_success(response, *, context: str) -> None: + assert 200 <= response.status_code < 300, f"{context}: {response.status_code} {response.text}" + + +def _delete_ignore_errors(api_client, path: str) -> None: + try: + api_client.delete(path) + except Exception: + pass + + +def _indexation_config(chunk_size: int) -> dict: + return { + "chunking": { + "name": "recursive_splitter", + "chunk_size": chunk_size, + "chunk_overlap_rate": 0.1, + }, + "parsing_strategy": "pymupdf", + "enable_image_captioning": False, + "enable_contextualization": False, + "enable_metadata_extraction": False, + "enable_entity_extraction": False, + "enable_topic_tagging": False, + } + + +def _retrieval_config(top_n: int = 3) -> dict: + return { + "type": "single", + "top_k": 5, + "top_n": top_n, + "similarity_threshold": 0.0, + "enable_reranker": False, + "include_related": False, + "include_ancestors": False, + } + + +def test_preset_options_crud_and_rename(api_client): + """Create, update, rename, list, and delete indexation/retrieval presets.""" + suffix = uuid.uuid4().hex[:8] + indexation_preset = f"ci-index-{suffix}" + retrieval_preset = f"ci-retrieval-{suffix}" + retrieval_preset_renamed = f"{retrieval_preset}-renamed" + + try: + openapi = api_client.get("/openapi.json") + _assert_success(openapi, context="openapi") + paths = openapi.json()["paths"] + assert "/presets/" in paths + assert "/presets/options" in paths + + options = api_client.get("/presets/options") + _assert_success(options, context="preset options") + option_data = options.json() + assert "recursive_splitter" in option_data["chunking_strategies"] + assert "single" in option_data["retrieval_types"] + + create_indexation = api_client.post( + "/presets/", + json={ + "name": indexation_preset, + "preset_type": "indexation", + "config": _indexation_config(64), + }, + ) + assert create_indexation.status_code == 201, create_indexation.text + + update_indexation = api_client.put( + f"/presets/indexation/{indexation_preset}", + json={"config": _indexation_config(96)}, + ) + _assert_success(update_indexation, context="update indexation preset") + assert update_indexation.json()["config"]["chunking"]["chunk_size"] == 96 + + create_retrieval = api_client.post( + "/presets/", + json={ + "name": retrieval_preset, + "preset_type": "retrieval", + "config": _retrieval_config(), + }, + ) + assert create_retrieval.status_code == 201, create_retrieval.text + + rename_retrieval = api_client.put( + f"/presets/retrieval/{retrieval_preset}", + json={"name": retrieval_preset_renamed}, + ) + _assert_success(rename_retrieval, context="rename retrieval preset") + assert rename_retrieval.json()["name"] == retrieval_preset_renamed + + list_retrieval = api_client.get("/presets/", params={"preset_type": "retrieval"}) + _assert_success(list_retrieval, context="list retrieval presets") + assert retrieval_preset_renamed in {row["name"] for row in list_retrieval.json()} + + get_retrieval = api_client.get(f"/presets/retrieval/{retrieval_preset_renamed}") + _assert_success(get_retrieval, context="get renamed retrieval preset") + assert get_retrieval.json()["config"]["top_n"] == 3 + finally: + _delete_ignore_errors(api_client, f"/presets/indexation/{indexation_preset}") + _delete_ignore_errors(api_client, f"/presets/retrieval/{retrieval_preset_renamed}") + _delete_ignore_errors(api_client, f"/presets/retrieval/{retrieval_preset}") diff --git a/tests/unit/api/mcp/test_server.py b/tests/unit/api/mcp/test_server.py index 07939df4e..7824d3d0d 100644 --- a/tests/unit/api/mcp/test_server.py +++ b/tests/unit/api/mcp/test_server.py @@ -238,3 +238,27 @@ def test_require_container_raises_when_unset(monkeypatch): monkeypatch.setattr(server, "_container", None) with pytest.raises(RuntimeError): server._require_container() + + +@pytest.mark.asyncio +async def test_startup_shuts_down_container_when_initialize_fails(monkeypatch): + calls: list[str] = [] + + class FailingContainer: + async def initialize(self): + calls.append("initialize") + raise RuntimeError("init failed") + + async def shutdown(self): + calls.append("shutdown") + + monkeypatch.setattr(server.ray, "is_initialized", lambda: True) + monkeypatch.setattr(server, "ensure_worker_bootstrap", lambda: calls.append("bootstrap")) + monkeypatch.setattr(server, "ServiceContainer", lambda _config: FailingContainer()) + monkeypatch.setattr(server, "_container", None) + + with pytest.raises(RuntimeError, match="init failed"): + await server._startup() + + assert calls == ["bootstrap", "initialize", "shutdown"] + assert server._container is None diff --git a/tests/unit/api/middleware/test_bypass_config.py b/tests/unit/api/middleware/test_bypass_config.py index 02f5f8285..5af9bbcdf 100644 --- a/tests/unit/api/middleware/test_bypass_config.py +++ b/tests/unit/api/middleware/test_bypass_config.py @@ -9,6 +9,7 @@ from __future__ import annotations +import pytest from api.middleware.auth import ( AuthMiddleware, is_bypass_path, @@ -21,6 +22,7 @@ AuthBypassConfig, ) from fastapi import FastAPI, Request +from starlette.responses import Response # --------------------------------------------------------------------------- # Defaults match the legacy hardcoded sets @@ -163,6 +165,45 @@ def test_auth_middleware_accepts_custom_bypass_config() -> None: assert instance._bypass_config is custom +def _request(headers=None): + raw = [(k.lower().encode(), v.encode()) for k, v in (headers or {}).items()] + scope = {"type": "http", "method": "GET", "path": "/indexer/files", "headers": raw, "query_string": b""} + return Request(scope) + + +async def _unused_call_next(_request): + return Response("ok") + + +@pytest.mark.asyncio +async def test_auth_middleware_returns_503_when_auth_service_is_unavailable(monkeypatch) -> None: + monkeypatch.setenv("AUTH_MODE", "token") + monkeypatch.setenv("AUTH_TOKEN", "secret") + + def unavailable(_request): + raise RuntimeError("container unavailable") + + middleware = AuthMiddleware(lambda scope, receive, send: None, get_auth_service=unavailable) + + response = await middleware.dispatch(_request(headers={"authorization": "Bearer token"}), _unused_call_next) + + assert response.status_code == 503 + + +@pytest.mark.asyncio +async def test_auth_middleware_does_not_swallow_programming_errors(monkeypatch) -> None: + monkeypatch.setenv("AUTH_MODE", "token") + monkeypatch.setenv("AUTH_TOKEN", "secret") + + def broken(_request): + raise ValueError("unexpected bug") + + middleware = AuthMiddleware(lambda scope, receive, send: None, get_auth_service=broken) + + with pytest.raises(ValueError, match="unexpected bug"): + await middleware.dispatch(_request(headers={"authorization": "Bearer token"}), _unused_call_next) + + # --------------------------------------------------------------------------- # Minor: argument is keyword-only so we don't accidentally pass it positionally # --------------------------------------------------------------------------- @@ -171,8 +212,6 @@ def test_auth_middleware_accepts_custom_bypass_config() -> None: def test_is_bypass_path_bypass_config_is_keyword_only() -> None: """Positional misuse should fail loudly rather than silently treating an arbitrary value as the config.""" - import pytest - with pytest.raises(TypeError): is_bypass_path("/docs", AuthBypassConfig()) # type: ignore[misc] diff --git a/tests/unit/api/routers/admin/test_phase14_admin_routers.py b/tests/unit/api/routers/admin/test_phase14_admin_routers.py new file mode 100644 index 000000000..5b5dfcc2c --- /dev/null +++ b/tests/unit/api/routers/admin/test_phase14_admin_routers.py @@ -0,0 +1,361 @@ +from __future__ import annotations + +from typing import Any + +import pytest +from api.dependencies.auth import require_admin +from api.routers.admin import model_endpoints, presets +from di.providers import get_model_endpoint_service, get_preset_service +from fastapi import FastAPI + + +def _model_endpoint_row(**overrides: Any) -> dict[str, Any]: + """Build a model endpoint response row for router tests.""" + row = { + "name": "default", + "model_type": "llm", + "endpoint": "http://llm:8000/v1", + "model_name": "mistral", + "batch_size": 32, + "timeout": 30.0, + "extra": {}, + "is_default": True, + "created_at": "2026-01-01T00:00:00+00:00", + "updated_at": "2026-01-01T00:00:00+00:00", + } + row.update(overrides) + return row + + +def _preset_row(**overrides: Any) -> dict[str, Any]: + """Build a preset response row for router tests.""" + row = { + "name": "default", + "preset_type": "retrieval", + "config": {"type": "single", "top_k": 50}, + "created_at": "2026-01-01T00:00:00+00:00", + "updated_at": "2026-01-01T00:00:00+00:00", + } + row.update(overrides) + return row + + +class FakeModelEndpointService: + """Fake model endpoint service that records router calls.""" + + def __init__(self) -> None: + """Initialize the call log.""" + self.calls: list[tuple[str, dict[str, Any]]] = [] + self.endpoint_extra: dict[str, Any] = {} + + async def create_model_endpoint(self, row: Any) -> dict[str, Any]: + """Record endpoint creation (from a ModelEndpointRow) and echo a row.""" + payload = row.model_dump(exclude={"created_at", "updated_at"}) + self.calls.append(("create", payload)) + return _model_endpoint_row(**payload) + + async def list_model_endpoints(self, model_type: str | None = None) -> list[dict[str, Any]]: + """Record endpoint listing with the optional type filter.""" + self.calls.append(("list", {"model_type": model_type})) + return [_model_endpoint_row(model_type=model_type or "llm")] + + async def get_model_endpoint(self, name: str, model_type: str) -> Any: + """Record a single endpoint lookup.""" + from core.config.model_endpoints import ModelEndpointRow + + self.calls.append(("get", {"name": name, "model_type": model_type})) + return ModelEndpointRow(**_model_endpoint_row(name=name, model_type=model_type, extra=self.endpoint_extra)) + + async def update_model_endpoint(self, name: str, model_type: str, **fields: Any) -> dict[str, Any]: + """Record endpoint updates and echo the merged response row.""" + self.calls.append(("update", {"name": name, "model_type": model_type, **fields})) + return _model_endpoint_row(**{"name": name, "model_type": model_type, **fields}) + + async def delete_model_endpoint(self, name: str, model_type: str) -> None: + """Record endpoint deletion.""" + self.calls.append(("delete", {"name": name, "model_type": model_type})) + + async def set_default(self, model_type: str, name: str) -> None: + """Record default promotion.""" + self.calls.append(("set_default", {"name": name, "model_type": model_type})) + + async def validate_endpoint( + self, + url: str, + model_name: str | None = None, + *, + api_key: str | None = None, + ) -> dict[str, Any]: + """Record endpoint validation.""" + self.calls.append(("validate", {"url": url, "model_name": model_name, "api_key": api_key})) + return {"reachable": True, "model_found": True, "models_served": ["mistral"], "detail": None} + + +class FakePresetService: + """Fake preset service that records router calls.""" + + def __init__(self) -> None: + """Initialize the call log.""" + self.calls: list[tuple[str, dict[str, Any]]] = [] + + async def create_preset(self, name: str, preset_type: str, config: dict[str, Any]) -> dict[str, Any]: + """Record preset creation (from unpacked kwargs) and echo a row.""" + payload = {"name": name, "preset_type": preset_type, "config": config} + self.calls.append(("create", payload)) + return _preset_row(**payload) + + async def list_presets(self, preset_type: str | None = None) -> list[dict[str, Any]]: + """Record preset listing with the optional type filter.""" + self.calls.append(("list", {"preset_type": preset_type})) + return [_preset_row(preset_type=preset_type or "retrieval")] + + async def get_preset(self, name: str, preset_type: str) -> dict[str, Any]: + """Record a single preset lookup.""" + self.calls.append(("get", {"name": name, "preset_type": preset_type})) + return _preset_row(name=name, preset_type=preset_type) + + async def update_preset(self, name: str, preset_type: str, **fields: Any) -> dict[str, Any]: + """Record preset updates and echo the merged response row.""" + self.calls.append(("update", {"name": name, "preset_type": preset_type, **fields})) + return _preset_row(**{"name": name, "preset_type": preset_type, **fields}) + + async def delete_preset(self, name: str, preset_type: str) -> None: + """Record preset deletion.""" + self.calls.append(("delete", {"name": name, "preset_type": preset_type})) + + +def _build_app( + *, + model_service: FakeModelEndpointService | None = None, + preset_service: FakePresetService | None = None, +) -> FastAPI: + """Build a small app with Phase 14 routers and fake dependencies.""" + app = FastAPI() + app.include_router(model_endpoints.router, prefix="/model-endpoints") + app.include_router(presets.router, prefix="/presets") + app.dependency_overrides[require_admin] = lambda: {"id": "admin", "is_admin": True} + if model_service is not None: + app.dependency_overrides[get_model_endpoint_service] = lambda: model_service + if preset_service is not None: + app.dependency_overrides[get_preset_service] = lambda: preset_service + return app + + +@pytest.mark.asyncio +async def test_create_model_endpoint_normalizes_payload(async_client_factory): + """Model endpoint creation should pass normalized schema data to the service.""" + model_service = FakeModelEndpointService() + app = _build_app(model_service=model_service) + + async with async_client_factory(app) as client: + response = await client.post( + "/model-endpoints/", + json={ + "name": " default ", + "model_type": "llm", + "endpoint": " http://llm:8000/v1/ ", + "model_name": "mistral", + }, + ) + + assert response.status_code == 201 + assert model_service.calls == [ + ( + "create", + { + "name": "default", + "model_type": "llm", + "endpoint": "http://llm:8000/v1", + "model_name": "mistral", + "batch_size": 32, + "timeout": 30.0, + "extra": {}, + "is_default": False, + }, + ) + ] + + +@pytest.mark.asyncio +async def test_update_model_endpoint_forwards_only_provided_fields(async_client_factory): + """Model endpoint updates should exclude omitted fields.""" + model_service = FakeModelEndpointService() + app = _build_app(model_service=model_service) + + async with async_client_factory(app) as client: + response = await client.put( + "/model-endpoints/llm/default", + json={"endpoint": "http://llm-next:8000/v1", "timeout": 60}, + ) + + assert response.status_code == 200 + assert model_service.calls == [ + ( + "update", + { + "name": "default", + "model_type": "llm", + "endpoint": "http://llm-next:8000/v1", + "timeout": 60.0, + }, + ) + ] + + +@pytest.mark.asyncio +async def test_update_model_endpoint_maps_name_to_new_name(async_client_factory): + """Model endpoint rename should use the service rename field.""" + model_service = FakeModelEndpointService() + app = _build_app(model_service=model_service) + + async with async_client_factory(app) as client: + response = await client.put( + "/model-endpoints/llm/default", + json={"name": "mistral-small"}, + ) + + assert response.status_code == 200 + assert model_service.calls == [ + ( + "update", + { + "name": "default", + "model_type": "llm", + "new_name": "mistral-small", + }, + ) + ] + + +@pytest.mark.asyncio +async def test_validate_model_endpoint_uses_route_identity(async_client_factory): + """Endpoint validation should resolve route identity before probing.""" + model_service = FakeModelEndpointService() + app = _build_app(model_service=model_service) + + async with async_client_factory(app) as client: + response = await client.post("/model-endpoints/llm/default/validate") + + assert response.status_code == 200 + assert response.json()["reachable"] is True + assert model_service.calls == [ + ("get", {"name": "default", "model_type": "llm"}), + ("validate", {"url": "http://llm:8000/v1", "model_name": "mistral", "api_key": None}), + ] + + +@pytest.mark.asyncio +async def test_validate_model_endpoint_uses_stored_api_key(async_client_factory): + """Endpoint validation should authenticate with the stored endpoint key.""" + model_service = FakeModelEndpointService() + model_service.endpoint_extra = {"api_key": "secret-token"} + app = _build_app(model_service=model_service) + + async with async_client_factory(app) as client: + response = await client.post("/model-endpoints/llm/default/validate") + + assert response.status_code == 200 + assert model_service.calls == [ + ("get", {"name": "default", "model_type": "llm"}), + ("validate", {"url": "http://llm:8000/v1", "model_name": "mistral", "api_key": "secret-token"}), + ] + + +@pytest.mark.asyncio +async def test_set_default_model_endpoint_returns_promoted_endpoint(async_client_factory): + """Default promotion should return the promoted endpoint row.""" + model_service = FakeModelEndpointService() + app = _build_app(model_service=model_service) + + async with async_client_factory(app) as client: + response = await client.post("/model-endpoints/llm/default/set-default") + + assert response.status_code == 200 + assert response.json()["is_default"] is True + assert model_service.calls == [ + ("set_default", {"name": "default", "model_type": "llm"}), + ("get", {"name": "default", "model_type": "llm"}), + ] + + +@pytest.mark.asyncio +async def test_preset_options_return_registered_choices(async_client_factory): + """Preset options should expose available registry choices.""" + app = _build_app() + + async with async_client_factory(app) as client: + response = await client.get("/presets/options") + + assert response.status_code == 200 + body = response.json() + assert body["chunking_strategies"] == ["recursive_splitter"] + assert set(body["retrieval_types"]) == {"single", "multiQuery", "hyde"} + assert body["reranker_providers"] == ["infinity", "openai"] + + +@pytest.mark.asyncio +async def test_create_preset_forwards_schema_payload(async_client_factory): + """Preset creation should pass normalized schema data to the service.""" + preset_service = FakePresetService() + app = _build_app(preset_service=preset_service) + + async with async_client_factory(app) as client: + response = await client.post( + "/presets/", + json={"name": " default ", "preset_type": "retrieval", "config": {"type": "single"}}, + ) + + assert response.status_code == 201 + assert preset_service.calls == [ + ("create", {"name": "default", "preset_type": "retrieval", "config": {"type": "single"}}) + ] + + +@pytest.mark.asyncio +async def test_update_preset_forwards_only_provided_fields(async_client_factory): + """Preset updates should exclude omitted fields.""" + preset_service = FakePresetService() + app = _build_app(preset_service=preset_service) + + async with async_client_factory(app) as client: + response = await client.put( + "/presets/retrieval/default", + json={"config": {"type": "hyde", "top_k": 20}}, + ) + + assert response.status_code == 200 + assert preset_service.calls == [ + ( + "update", + { + "name": "default", + "preset_type": "retrieval", + "config": {"type": "hyde", "top_k": 20}, + }, + ) + ] + + +@pytest.mark.asyncio +async def test_update_preset_maps_name_to_new_name(async_client_factory): + """Preset rename should use the service rename field.""" + preset_service = FakePresetService() + app = _build_app(preset_service=preset_service) + + async with async_client_factory(app) as client: + response = await client.put( + "/presets/retrieval/default", + json={"name": "legal"}, + ) + + assert response.status_code == 200 + assert preset_service.calls == [ + ( + "update", + { + "name": "default", + "preset_type": "retrieval", + "new_name": "legal", + }, + ) + ] diff --git a/tests/unit/api/routers/admin/test_phase14_partition_routes.py b/tests/unit/api/routers/admin/test_phase14_partition_routes.py new file mode 100644 index 000000000..a1845e8ed --- /dev/null +++ b/tests/unit/api/routers/admin/test_phase14_partition_routes.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +from typing import Any + +import pytest +from api.dependencies.auth import require_partition_owner, require_partition_viewer +from api.routers.admin import partitions +from di.providers import get_partition_service +from fastapi import FastAPI + + +def _partition_detail(**overrides: Any) -> dict[str, Any]: + """Build a resolved partition detail response row for router tests.""" + row = { + "name": "legal", + "description": "Legal documents", + "embedder": "default", + "indexation_preset": "default", + "retrieval_preset": "default", + "indexation_pipeline": {"chunking": {"name": "recursive_splitter"}}, + "retrieval_pipeline": {"type": "single", "top_k": 50}, + "dimension": 1024, + "created_at": "2026-01-01T00:00:00+00:00", + } + row.update(overrides) + return row + + +class FakePartitionConfigService: + """Fake partition config service that records router calls.""" + + def __init__(self) -> None: + """Initialize the call log.""" + self.calls: list[tuple[str, dict[str, Any]]] = [] + + async def update_partition_config(self, partition: str, **fields: Any) -> dict[str, Any]: + """Record partition config updates and echo a response row.""" + self.calls.append(("update", {"partition": partition, **fields})) + return _partition_detail(name=partition, **fields) + + async def get_partition_config(self, partition: str) -> dict[str, Any]: + """Record partition config reads.""" + self.calls.append(("get", {"partition": partition})) + return _partition_detail(name=partition) + + +class MissingPartitionConfigService: + """Service double without Phase 14 config methods.""" + + +def _build_app(service) -> FastAPI: + """Build a small app with partition routes and fake dependencies.""" + app = FastAPI() + app.include_router(partitions.router, prefix="/partition") + app.dependency_overrides[require_partition_owner] = lambda: {"id": "admin", "is_admin": True} + app.dependency_overrides[require_partition_viewer] = lambda: {"id": "admin", "is_admin": True} + app.dependency_overrides[get_partition_service] = lambda: service + return app + + +@pytest.mark.asyncio +async def test_update_partition_config_forwards_only_provided_fields(async_client_factory): + """Partition config updates should exclude omitted fields.""" + service = FakePartitionConfigService() + app = _build_app(service) + + async with async_client_factory(app) as client: + response = await client.patch( + "/partition/legal", + json={ + "description": "Updated", + "indexation_preset": "legal-index", + "retrieval_preset": "legal-search", + }, + ) + + assert response.status_code == 200 + assert response.json()["indexation_preset"] == "legal-index" + assert service.calls == [ + ( + "update", + { + "partition": "legal", + "description": "Updated", + "indexation_preset": "legal-index", + "retrieval_preset": "legal-search", + }, + ) + ] + + +@pytest.mark.asyncio +async def test_get_partition_config_returns_resolved_config(async_client_factory): + """Partition config reads should return the resolved service payload.""" + service = FakePartitionConfigService() + app = _build_app(service) + + async with async_client_factory(app) as client: + response = await client.get("/partition/legal/config") + + assert response.status_code == 200 + body = response.json() + assert body["name"] == "legal" + assert body["indexation_pipeline"]["chunking"]["name"] == "recursive_splitter" + assert service.calls == [("get", {"partition": "legal"})] + + +@pytest.mark.asyncio +async def test_partition_config_routes_return_501_until_service_methods_exist(async_client_factory): + """Missing phased service methods should return not-implemented responses.""" + app = _build_app(MissingPartitionConfigService()) + + async with async_client_factory(app) as client: + patch_response = await client.patch("/partition/legal", json={"description": "Updated"}) + get_response = await client.get("/partition/legal/config") + + assert patch_response.status_code == 501 + assert patch_response.json()["detail"] == "update_partition_config is not available." + assert get_response.status_code == 501 + assert get_response.json()["detail"] == "get_partition_config is not available." diff --git a/tests/unit/api/routers/test_router_imports.py b/tests/unit/api/routers/test_router_imports.py index f69b46ab9..bdd9bc627 100644 --- a/tests/unit/api/routers/test_router_imports.py +++ b/tests/unit/api/routers/test_router_imports.py @@ -11,7 +11,9 @@ "api.routers.user.health", "api.routers.user.extract", "api.routers.admin.indexing", + "api.routers.admin.model_endpoints", "api.routers.admin.partitions", + "api.routers.admin.presets", "api.routers.admin.users", "api.routers.admin.workspaces", "api.routers.admin.jobs", @@ -23,6 +25,7 @@ ], ) def test_phase_10_router_module_imports(module_name): + """Router modules should import and expose a FastAPI router.""" module = importlib.import_module(module_name) assert module.router is not None diff --git a/tests/unit/api/schemas/admin/test_phase14_schemas.py b/tests/unit/api/schemas/admin/test_phase14_schemas.py new file mode 100644 index 000000000..105071731 --- /dev/null +++ b/tests/unit/api/schemas/admin/test_phase14_schemas.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +import pytest +from api.schemas.admin.model_endpoint_schemas import CreateModelEndpointRequest, UpdateModelEndpointRequest +from api.schemas.admin.partition_schemas import CreatePartitionRequest, UpdatePartitionRequest +from api.schemas.admin.preset_schemas import CreatePresetRequest, UpdatePresetRequest +from pydantic import ValidationError + + +def test_create_model_endpoint_defaults_and_normalizes_endpoint(): + """Model endpoint creation applies defaults and strips trailing slashes.""" + request = CreateModelEndpointRequest( + name="default", + model_type="embedder", + endpoint="http://embedder:8000/v1/", + model_name="gte", + ) + + assert request.endpoint == "http://embedder:8000/v1" + assert request.batch_size == 32 + assert request.timeout == 30.0 + assert request.extra == {} + assert request.is_default is False + + +@pytest.mark.parametrize("model_type", ["embedding", "chat", "vision", ""]) +def test_create_model_endpoint_rejects_unknown_type(model_type): + """Only supported endpoint types are accepted.""" + with pytest.raises(ValidationError): + CreateModelEndpointRequest(name="default", model_type=model_type, endpoint="http://host") + + +@pytest.mark.parametrize("field,value", [("batch_size", 0), ("timeout", 0)]) +def test_create_model_endpoint_rejects_non_positive_numbers(field, value): + """Batch size and timeout must stay positive.""" + payload = {"name": "default", "model_type": "llm", "endpoint": "http://host", field: value} + + with pytest.raises(ValidationError): + CreateModelEndpointRequest(**payload) + + +@pytest.mark.parametrize("endpoint", ["/", "///", " / "]) +def test_create_model_endpoint_rejects_empty_normalized_endpoint(endpoint): + """Slash-only endpoint values are invalid after normalization.""" + with pytest.raises(ValidationError): + CreateModelEndpointRequest(name="default", model_type="llm", endpoint=endpoint) + + +def test_update_model_endpoint_requires_at_least_one_field(): + """Endpoint updates must contain at least one field.""" + with pytest.raises(ValidationError): + UpdateModelEndpointRequest() + + +def test_create_preset_accepts_indexation_and_retrieval_configs(): + """Preset creation accepts both supported preset families.""" + indexation = CreatePresetRequest(name="fast", preset_type="indexation", config={"chunk_size": 512}) + retrieval = CreatePresetRequest(name="qa", preset_type="retrieval", config={"top_k": 20}) + + assert indexation.config["chunk_size"] == 512 + assert retrieval.config["top_k"] == 20 + + +def test_create_preset_rejects_unknown_type(): + """Preset type must be either indexation or retrieval.""" + with pytest.raises(ValidationError): + CreatePresetRequest(name="bad", preset_type="generation", config={}) + + +def test_update_preset_requires_at_least_one_field(): + """Preset updates must contain at least one field.""" + with pytest.raises(ValidationError): + UpdatePresetRequest() + + +@pytest.mark.parametrize("payload", [{"name": None}, {"config": None}]) +def test_update_preset_rejects_explicit_nulls(payload): + """Preset updates use omission for unchanged fields and reject null.""" + with pytest.raises(ValidationError): + UpdatePresetRequest(**payload) + + +def test_create_partition_defaults_to_default_presets(): + """Partition creation defaults to the default embedder and presets.""" + request = CreatePartitionRequest(name="legal") + + assert request.embedder == "default" + assert request.indexation_preset == "default" + assert request.retrieval_preset == "default" + assert request.chat_history_depth == 0 + + +def test_create_partition_rejects_empty_names(): + """Partition names cannot be blank.""" + with pytest.raises(ValidationError): + CreatePartitionRequest(name=" ", indexation_preset="default", retrieval_preset="default") + + +def test_update_partition_requires_at_least_one_field(): + """Partition updates must contain at least one field.""" + with pytest.raises(ValidationError): + UpdatePartitionRequest() + + +def test_update_partition_rejects_negative_chat_history_depth(): + """Chat history depth cannot be negative.""" + with pytest.raises(ValidationError): + UpdatePartitionRequest(chat_history_depth=-1) + + +@pytest.mark.parametrize( + "field", + [ + "description", + "embedder", + "indexation_preset", + "retrieval_preset", + "chat_history_depth", + ], +) +def test_update_partition_rejects_explicit_nulls(field): + """Partition updates reject null for fields that cannot be cleared.""" + with pytest.raises(ValidationError, match=rf"{field} cannot be null"): + UpdatePartitionRequest(**{field: None}) diff --git a/tests/unit/api/test_main_proxy_headers.py b/tests/unit/api/test_main_proxy_headers.py index 4b1d8725a..dc1ad35fd 100644 --- a/tests/unit/api/test_main_proxy_headers.py +++ b/tests/unit/api/test_main_proxy_headers.py @@ -64,6 +64,37 @@ def test_default_forwarded_allow_ips_env_var_used(): assert "UVICORN_FORWARDED_ALLOW_IPS" in src +def test_phase14_admin_routers_are_mounted(): + """Phase 14 admin routes must be exposed under stable API prefixes.""" + with open(_MAIN_PATH) as f: + tree = ast.parse(f.read()) + + mounted: dict[str, str] = {} + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + func = node.func + if not (isinstance(func, ast.Attribute) and func.attr == "include_router"): + continue + if not node.args or not isinstance(node.args[0], ast.Name): + continue + router_name = node.args[0].id + prefix = None + tag = None + for kw in node.keywords: + if kw.arg == "prefix" and isinstance(kw.value, ast.Constant): + prefix = kw.value.value + if kw.arg == "tags" and isinstance(kw.value, ast.List) and kw.value.elts: + first_tag = kw.value.elts[0] + if isinstance(first_tag, ast.Attribute): + tag = first_tag.attr + if prefix is not None and tag is not None: + mounted[router_name] = f"{prefix}:{tag}" + + assert mounted["model_endpoints_router"] == "/model-endpoints:MODEL_ENDPOINTS" + assert mounted["presets_router"] == "/presets:PRESETS" + + def test_api_package_exports_app_for_legacy_uvicorn_path(monkeypatch): """Older images or overrides may still run ``uvicorn api:app``.""" fake_app = object() diff --git a/tests/unit/core/config/test_model_endpoints.py b/tests/unit/core/config/test_model_endpoints.py new file mode 100644 index 000000000..a02725e33 --- /dev/null +++ b/tests/unit/core/config/test_model_endpoints.py @@ -0,0 +1,64 @@ +"""Tests for ModelsConfig frozen-field + mutable-dict invariant.""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest +from core.config.model_endpoints import ModelEndpointConfig, ModelEndpointRow, ModelsConfig +from pydantic import ValidationError + + +def _row_payload(**overrides): + """Build a valid model endpoint row payload with targeted overrides.""" + now = datetime.now(UTC) + payload = { + "name": "default", + "model_type": "llm", + "endpoint": "http://vllm:8000/v1", + "created_at": now, + "updated_at": now, + } + payload.update(overrides) + return payload + + +def test_models_config_dicts_are_mutable_in_place(): + """Services perform clear()+update() atomic swaps on frozen ModelsConfig fields.""" + cfg = ModelsConfig() + ep = ModelEndpointConfig(endpoint="http://vllm:8000/v1") + + cfg.embedder["default"] = ep + assert "default" in cfg.embedder + + new_ep = ModelEndpointConfig(endpoint="http://new:8000/v1") + cfg.embedder.clear() + cfg.embedder.update({"v2": new_ep}) + assert list(cfg.embedder) == ["v2"] + + +def test_models_config_field_reassignment_raises(): + """Frozen ConfigMixin prevents field reassignment — only in-place mutation works.""" + cfg = ModelsConfig() + with pytest.raises((TypeError, ValidationError)): + cfg.embedder = {} # type: ignore[misc] + + +def test_model_endpoint_row_invalid_model_type(): + """Persisted endpoint rows reject unknown model endpoint types.""" + with pytest.raises(ValidationError): + ModelEndpointRow(**_row_payload(model_type="not-a-model")) + + +@pytest.mark.parametrize("batch_size", [0, -1]) +def test_model_endpoint_row_non_positive_batch_size(batch_size: int): + """Persisted endpoint rows reject non-positive batch sizes.""" + with pytest.raises(ValidationError): + ModelEndpointRow(**_row_payload(batch_size=batch_size)) + + +@pytest.mark.parametrize("timeout", [0, -0.1]) +def test_model_endpoint_row_non_positive_timeout(timeout: float): + """Persisted endpoint rows reject non-positive timeouts.""" + with pytest.raises(ValidationError): + ModelEndpointRow(**_row_payload(timeout=timeout)) diff --git a/tests/unit/core/config/test_pipeline_configs.py b/tests/unit/core/config/test_pipeline_configs.py new file mode 100644 index 000000000..3461274cc --- /dev/null +++ b/tests/unit/core/config/test_pipeline_configs.py @@ -0,0 +1,54 @@ +"""Tests for PresetsConfig frozen-field + mutable-dict invariant.""" + +from __future__ import annotations + +import pytest +from core.config.indexation_pipeline import IndexationPipelineConfig +from core.config.presets import PresetsConfig +from core.config.retrieval_pipeline import RetrievalPipelineConfig +from pydantic import ValidationError + + +def test_presets_config_dicts_are_mutable_in_place(): + """Services perform clear()+update() atomic swaps on frozen PresetsConfig fields.""" + cfg = PresetsConfig() + payload = {"chunking": {"chunk_size": 512}, "parsing_strategy": "marker"} + + cfg.indexation["default"] = payload + assert "default" in cfg.indexation + + cfg.indexation.clear() + cfg.indexation.update({"legal": payload}) + assert list(cfg.indexation) == ["legal"] + + +def test_presets_config_field_reassignment_raises(): + """Frozen ConfigMixin prevents field reassignment — only in-place mutation works.""" + cfg = PresetsConfig() + with pytest.raises((TypeError, ValidationError)): + cfg.indexation = {} # type: ignore[misc] + + +def test_indexation_pipeline_rejects_unknown_parsing_strategy(): + """Indexation presets reject parser names outside the supported set.""" + with pytest.raises(ValidationError): + IndexationPipelineConfig(parsing_strategy="unknown") + + +def test_indexation_pipeline_rejects_unknown_contextualization_mode(): + """Indexation presets reject contextualization modes outside the supported set.""" + with pytest.raises(ValidationError): + IndexationPipelineConfig(contextualization_mode="verbose") + + +def test_retrieval_pipeline_rejects_unknown_type(): + """Retrieval presets reject unsupported retrieval modes.""" + with pytest.raises(ValidationError): + RetrievalPipelineConfig(type="unknown") + + +@pytest.mark.parametrize("threshold", [-0.1, 1.1]) +def test_retrieval_pipeline_rejects_out_of_range_similarity_threshold(threshold: float): + """Retrieval presets keep similarity thresholds in the normalized range.""" + with pytest.raises(ValidationError): + RetrievalPipelineConfig(similarity_threshold=threshold) diff --git a/tests/unit/core/models/test_preset_models.py b/tests/unit/core/models/test_preset_models.py new file mode 100644 index 000000000..8e1d11cd9 --- /dev/null +++ b/tests/unit/core/models/test_preset_models.py @@ -0,0 +1,57 @@ +"""Tests for persisted preset and partition domain row invariants.""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest +from core.config.indexation_pipeline import IndexationPipelineConfig +from core.config.retrieval_pipeline import RetrievalPipelineConfig +from core.models.preset import PartitionConfig, PartitionRow, PresetRow +from pydantic import ValidationError + + +def _now() -> datetime: + """Return a timestamp for persisted row test payloads.""" + return datetime.now(UTC) + + +def test_preset_row_rejects_unknown_preset_type(): + """Persisted preset rows reject unknown preset types.""" + now = _now() + + with pytest.raises(ValidationError): + PresetRow( + name="default", + preset_type="unknown", + config={}, + created_at=now, + updated_at=now, + ) + + +def test_partition_row_rejects_non_positive_dimension(): + """Persisted partition rows reject invalid vector dimensions.""" + now = _now() + + with pytest.raises(ValidationError): + PartitionRow(name="tenant", dimension=0, created_at=now, updated_at=now) + + +def test_partition_row_rejects_negative_chat_history_depth(): + """Persisted partition rows reject negative chat history depth.""" + now = _now() + + with pytest.raises(ValidationError): + PartitionRow(name="tenant", chat_history_depth=-1, created_at=now, updated_at=now) + + +def test_partition_config_rejects_negative_chat_history_depth(): + """Resolved partition configs reject negative chat history depth.""" + with pytest.raises(ValidationError): + PartitionConfig( + name="tenant", + indexation=IndexationPipelineConfig(), + retrieval=RetrievalPipelineConfig(), + chat_history_depth=-1, + ) diff --git a/tests/unit/di/test_container.py b/tests/unit/di/test_container.py index bae9da2c6..4f89b6520 100644 --- a/tests/unit/di/test_container.py +++ b/tests/unit/di/test_container.py @@ -3,10 +3,14 @@ from __future__ import annotations from types import SimpleNamespace +from typing import Any import pytest from core.config.infrastructure import RDBConfig, VectorDBConfig +from core.config.model_endpoints import ModelEndpointConfig, ModelsConfig from core.config.root import Settings +from core.embeddings import embedder_registry +from core.llm import llm_registry from core.ports.audit_log_repo import AuditLogRepository from core.ports.catalog_store import CatalogStore from core.ports.chunk_repo import ChunkRepository @@ -23,6 +27,8 @@ from core.ports.topic_tag_repo import TopicTagRepository from core.ports.user_repo import UserRepository from core.ports.workspace_repo import WorkspaceRepository +from core.rerankers import reranker_registry +from core.vlm import vlm_registry from di.container import ServiceContainer from di.repositories import create_catalog_store from di.vector_stores import create_vector_store @@ -37,6 +43,66 @@ def _settings(database: str | None = None, collection: str = "vdb_test") -> Sett ) +async def _async_call(calls: list, value: object) -> None: + """Record an async lifecycle method call.""" + calls.append(value) + + +class _NamedClient: + """Inference client double built by named component factories.""" + + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + self.closed = False + + async def aclose(self): + """Record shutdown cleanup.""" + self.closed = True + + +def _settings_with_named_models() -> Settings: + """Build settings with one endpoint per model kind.""" + return Settings( + rdb=RDBConfig(password="x"), + vectordb=VectorDBConfig(collection_name="vdb_test"), + models=ModelsConfig( + embedder={ + "embed-a": ModelEndpointConfig( + endpoint="http://embedder:8000/v1", + model_name="embed-model", + batch_size=16, + timeout=12.5, + extra={"implementation": "phase14i-test", "api_key": "embed-key", "dimension": 384}, + ) + }, + reranker={ + "rank-a": ModelEndpointConfig( + endpoint="http://reranker:8000", + model_name="rank-model", + timeout=3.5, + extra={"implementation": "phase14i-test", "api_key": "rank-key"}, + ) + }, + llm={ + "chat-a": ModelEndpointConfig( + endpoint="http://llm:8000/v1", + model_name="chat-model", + timeout=42.0, + extra={"implementation": "phase14i-test", "temperature": 0.2}, + ) + }, + vlm={ + "vision-a": ModelEndpointConfig( + endpoint="http://vlm:8000/v1", + model_name="vision-model", + timeout=7.0, + extra={"implementation": "phase14i-test", "max_tokens": 256}, + ) + }, + ), + ) + + @pytest.fixture(autouse=True) def _stub_milvus_clients(monkeypatch): """Replace the pymilvus gRPC clients with mocks for the whole module. @@ -169,10 +235,31 @@ async def initialize(self): monkeypatch.setenv("AUTH_TOKEN", "admin-token") c = ServiceContainer(_settings()) c._catalog_store = FakeCatalogStore() + c._model_endpoint_service = SimpleNamespace( + seed_defaults=lambda: _async_call(calls, "endpoint.seed"), + load_all=lambda: _async_call(calls, "endpoint.load"), + ) + c._preset_service = SimpleNamespace( + seed_defaults=lambda: _async_call(calls, "preset.seed"), + load_all=lambda: _async_call(calls, "preset.load"), + ) + c._partition_service = SimpleNamespace( + seed_default_partition=lambda: _async_call(calls, "partition.seed"), + load_partitions=lambda: _async_call(calls, "partition.load"), + ) await c.initialize() - assert calls == ["initialize", "admin-token"] + assert calls == [ + "initialize", + "admin-token", + "endpoint.seed", + "endpoint.load", + "preset.seed", + "preset.load", + "partition.seed", + "partition.load", + ] class TestVectorStoreWiring: @@ -240,6 +327,8 @@ def test_does_not_mutate_input_settings(self): ("mcp_service", "get_mcp_service"), ] +_OPTIONAL_PHASE_PROVIDERS = {"get_model_endpoint_service", "get_preset_service"} + class TestPhase8OrchestratorWiring: """8F: all orchestrators wired consistently (container + providers).""" @@ -272,6 +361,7 @@ 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")} + wired -= _OPTIONAL_PHASE_PROVIDERS assert wired == {p for _, p in _ORCHESTRATORS} @@ -328,6 +418,46 @@ def test_uninitialized_container_returns_503_for_service_getter(self): assert exc.value.status_code == 503 + @pytest.mark.parametrize( + ("provider", "attribute_name"), + [ + ("get_model_endpoint_service", "model_endpoint_service"), + ("get_preset_service", "preset_service"), + ], + ) + def test_optional_phase14_provider_resolves_if_present(self, provider, attribute_name): + """Resolve optional Phase 14 services once the container exposes them.""" + from di import providers + + sentinel = object() + fake_container = SimpleNamespace(is_initialized=True, **{attribute_name: sentinel}) + try: + providers.set_container(fake_container) + assert getattr(providers, provider)() is sentinel + finally: + providers.set_container(None) + + @pytest.mark.parametrize( + "provider", + [ + "get_model_endpoint_service", + "get_preset_service", + ], + ) + def test_optional_phase14_provider_missing_returns_503(self, provider): + """Return service-unavailable until optional Phase 14 services are wired.""" + from di import providers + + fake_container = SimpleNamespace(is_initialized=True) + try: + providers.set_container(fake_container) + with pytest.raises(HTTPException) as exc: + getattr(providers, provider)() + finally: + providers.set_container(None) + + assert exc.value.status_code == 503 + class TestPhase11ContainerLifecycle: """Phase 11 introspection + lifecycle the provider bridge relies on.""" @@ -360,6 +490,8 @@ async def test_shutdown_closes_clients_and_resets_state(self): closed = [] class _FakeClient: + """Client double that records successful close calls.""" + async def aclose(self): """Record that the client was closed.""" closed.append(self) @@ -375,17 +507,84 @@ async def aclose(self): assert c._inference_clients == [] assert c.is_initialized is False + +class TestPhase14NamedComponentFactories: + """14I: ServiceContainer exposes named, cached component factories.""" + + @pytest.fixture(autouse=True) + def _register_named_client(self, monkeypatch): + """Register one test implementation in every inference registry.""" + for registry in (embedder_registry, reranker_registry, llm_registry, vlm_registry): + monkeypatch.setitem(registry._registry, "phase14i-test", _NamedClient) + + def test_container_wires_named_factories_from_models_config(self): + """Build each component kind from ``settings.models`` by endpoint name.""" + c = ServiceContainer(_settings_with_named_models()) + + embedder = c.embedder_factory("embed-a") + reranker = c.reranker_factory("rank-a") + llm = c.llm_factory("chat-a") + vlm = c.vlm_factory("vision-a") + + assert embedder.kwargs == { + "endpoint": "http://embedder:8000/v1", + "model_name": "embed-model", + "batch_size": 16, + "timeout": 12.5, + "api_key": "embed-key", + "dimension": 384, + } + assert reranker.kwargs["endpoint"] == "http://reranker:8000" + assert reranker.kwargs["model_name"] == "rank-model" + assert reranker.kwargs["api_key"] == "rank-key" + assert llm.kwargs["endpoint"] == "http://llm:8000/v1" + assert llm.kwargs["temperature"] == 0.2 + assert vlm.kwargs["endpoint"] == "http://vlm:8000/v1" + assert vlm.kwargs["max_tokens"] == 256 + + def test_named_factories_cache_by_endpoint_name(self): + """Repeated factory calls for the same endpoint return one client.""" + c = ServiceContainer(_settings_with_named_models()) + + assert c.embedder_factory("embed-a") is c.embedder_factory("embed-a") + assert c.llm_factory("chat-a") is c.llm_factory("chat-a") + + @pytest.mark.asyncio + async def test_shutdown_closes_named_factory_clients(self): + """Shutdown closes clients created through named factory caches.""" + c = ServiceContainer(_settings_with_named_models()) + embedder = c.embedder_factory("embed-a") + llm = c.llm_factory("chat-a") + c._initialized = True + + await c.shutdown() + + assert embedder.closed is True + assert llm.closed is True + assert c.is_initialized is False + + def test_no_settings_named_factory_fails_with_settings_error(self): + """Legacy no-settings containers reject named factory use clearly.""" + c = ServiceContainer() + + with pytest.raises(RuntimeError, match="without a Settings instance"): + c.embedder_factory("default") + @pytest.mark.asyncio async def test_shutdown_is_best_effort_when_a_client_fails(self): """One client close failure must not skip the rest or the reset.""" closed = [] class _BadClient: + """Client double that fails during close.""" + async def aclose(self): """Fail to close, exercising the best-effort path.""" raise RuntimeError("boom") class _GoodClient: + """Client double that records close calls after a failure.""" + async def aclose(self): """Record a successful close after a prior failure.""" closed.append(self) @@ -400,3 +599,108 @@ async def aclose(self): assert closed == [good] assert c._inference_clients == [] assert c.is_initialized is False + + +class TestPhase14ServiceWiring: + """14K: endpoint/preset services and startup resolution are wired.""" + + def test_model_endpoint_service_is_lazy_cached_and_shares_factory_caches(self): + """Expose ModelEndpointService with the same caches used by named factories.""" + from services.orchestrators.model_endpoint_service import ModelEndpointService + + c = ServiceContainer(_settings()) + + service = c.model_endpoint_service + + assert isinstance(service, ModelEndpointService) + assert c.model_endpoint_service is service + assert service._client_caches["embedder"] is c._embedder_cache + assert service._client_caches["reranker"] is c._reranker_cache + assert service._client_caches["llm"] is c._llm_cache + assert service._client_caches["vlm"] is c._vlm_cache + + def test_preset_service_is_lazy_cached_with_partition_back_reference(self): + """Expose PresetService with the config-aware PartitionService reference.""" + from services.orchestrators.preset_service import PresetService + + settings = _settings() + c = ServiceContainer(settings) + + service = c.preset_service + + assert isinstance(service, PresetService) + assert c.preset_service is service + assert service._partition_service is c.partition_service + assert c.partition_service._config is settings + + @pytest.mark.asyncio + async def test_initialize_loads_phase14_registries_before_partitions(self, monkeypatch): + """Load endpoints, then presets, then partitions after admin seeding.""" + calls = [] + + async def ensure_admin_user(token): + """Record admin token seeding calls.""" + calls.append(("admin", token)) + + class FakeCatalogStore: + """Small catalog-store stand-in for initialization sequencing.""" + + user_repo = SimpleNamespace(ensure_admin_user=ensure_admin_user) + + async def initialize(self): + """Record catalog initialization calls.""" + calls.append("catalog") + + class FakeEndpointService: + """Endpoint registry lifecycle recorder.""" + + async def seed_defaults(self): + """Record endpoint seeding.""" + calls.append("endpoint.seed") + + async def load_all(self): + """Record endpoint loading.""" + calls.append("endpoint.load") + + class FakePresetService: + """Preset registry lifecycle recorder.""" + + async def seed_defaults(self): + """Record preset seeding.""" + calls.append("preset.seed") + + async def load_all(self): + """Record preset loading.""" + calls.append("preset.load") + + class FakePartitionService: + """Partition cache lifecycle recorder.""" + + async def seed_default_partition(self): + """Record default partition seeding.""" + calls.append("partition.seed") + + async def load_partitions(self): + """Record partition config loading.""" + calls.append("partition.load") + + monkeypatch.setenv("AUTH_TOKEN", "admin-token") + c = ServiceContainer(_settings()) + c._catalog_store = FakeCatalogStore() + c._model_endpoint_service = FakeEndpointService() + c._preset_service = FakePresetService() + c._partition_service = FakePartitionService() + + await c.initialize() + + assert calls == [ + "catalog", + ("admin", "admin-token"), + "endpoint.seed", + "endpoint.load", + "preset.seed", + "preset.load", + "partition.seed", + "partition.load", + ] + assert c.is_initialized is True diff --git a/tests/unit/services/orchestrators/test_indexing_service.py b/tests/unit/services/orchestrators/test_indexing_service.py index f719c6342..4566095c5 100644 --- a/tests/unit/services/orchestrators/test_indexing_service.py +++ b/tests/unit/services/orchestrators/test_indexing_service.py @@ -3,6 +3,10 @@ from __future__ import annotations import pytest +from core.config.indexation_pipeline import IndexationPipelineConfig +from core.config.retrieval_pipeline import RetrievalPipelineConfig +from core.models.preset import PartitionConfig +from core.utils.exceptions import PartitionNotFoundError from services.orchestrators.indexing_service import IndexingService @@ -34,7 +38,18 @@ def __init__(self): self.cancelled: list[str] = [] self.cancel_result = True - async def dispatch_indexing(self, *, path, metadata, partition, user, workspace_ids, replace): + async def dispatch_indexing( + self, + *, + path, + metadata, + partition, + user, + workspace_ids, + replace, + indexation_config=None, + embedder_name=None, + ): self.dispatched.append( { "path": path, @@ -43,6 +58,8 @@ async def dispatch_indexing(self, *, path, metadata, partition, user, workspace_ "user": user, "workspace_ids": workspace_ids, "replace": replace, + "indexation_config": indexation_config, + "embedder_name": embedder_name, } ) return "task-abc" @@ -67,11 +84,68 @@ async def cancel_task(self, task_id): return self.cancel_result -def _service(*, doc=None, ws=None, disp=None): +def _config_with_partition(partition: str = "tenant-a"): + return type( + "Config", + (), + { + "partitions": { + partition: PartitionConfig( + name=partition, + embedder="embed-fast", + indexation=IndexationPipelineConfig( + parsing_strategy="pymupdf", + enable_image_captioning=False, + enable_contextualization=True, + contextualization_llm="llm-context", + ), + retrieval=RetrievalPipelineConfig(), + ) + } + }, + )() + + +class FakePartitionService: + """Minimal partition service: records creates and mutates the config cache.""" + + def __init__(self, config, *, db_partitions: set[str] | None = None): + self._config = config + self._db = db_partitions or set() + self.created: list[tuple[str, int]] = [] + self.loaded = 0 + + def _cfg(self, partition: str) -> PartitionConfig: + return PartitionConfig( + name=partition, + embedder="default", + indexation=IndexationPipelineConfig(), + retrieval=RetrievalPipelineConfig(), + ) + + async def partition_exists(self, partition: str) -> bool: + return partition in self._db + + async def create_partition(self, partition: str, *, user_id: int, **_) -> None: + self.created.append((partition, user_id)) + self._db.add(partition) + # Mimic create_partition + load_partitions populating the cache. + self._config.partitions[partition] = self._cfg(partition) + + async def load_partitions(self) -> None: + self.loaded += 1 + # Mimic the cache being rebuilt from all DB rows. + for name in self._db: + self._config.partitions.setdefault(name, self._cfg(name)) + + +def _service(*, doc=None, ws=None, disp=None, config=None, partition_service=None): return IndexingService( document_repo=doc or FakeDocumentRepo(), workspace_repo=ws or FakeWorkspaceRepo(), dispatcher=disp or FakeDispatcher(), + config=config, + partition_service=partition_service, ) @@ -147,6 +221,120 @@ async def test_replace_sets_replace_flag(tmp_path): assert disp.dispatched[0]["replace"] is True +@pytest.mark.asyncio +async def test_add_file_dispatches_partition_indexation_config_and_embedder(tmp_path): + f = tmp_path / "doc.txt" + f.write_text("x") + disp = FakeDispatcher() + svc = _service(disp=disp, config=_config_with_partition("tenant-a")) + + await svc.add_file( + file_path=str(f), + file_id="f1", + partition="tenant-a", + metadata={}, + sanitized_filename="doc.txt", + original_filename="doc.txt", + user=None, + ) + + sent = disp.dispatched[0] + assert sent["embedder_name"] == "embed-fast" + assert sent["indexation_config"]["parsing_strategy"] == "pymupdf" + assert sent["indexation_config"]["enable_image_captioning"] is False + assert sent["indexation_config"]["enable_contextualization"] is True + assert sent["indexation_config"]["contextualization_llm"] == "llm-context" + + +@pytest.mark.asyncio +async def test_add_file_rejects_unknown_partition_when_partition_configs_exist(tmp_path): + f = tmp_path / "doc.txt" + f.write_text("x") + svc = _service(config=_config_with_partition("tenant-a")) + + with pytest.raises(PartitionNotFoundError, match="tenant-b"): + await svc.add_file( + file_path=str(f), + file_id="f1", + partition="tenant-b", + metadata={}, + sanitized_filename="doc.txt", + original_filename="doc.txt", + user=None, + ) + + +@pytest.mark.asyncio +async def test_add_file_auto_creates_unknown_partition_when_service_wired(tmp_path): + f = tmp_path / "doc.txt" + f.write_text("x") + disp = FakeDispatcher() + config = _config_with_partition("tenant-a") + psvc = FakePartitionService(config) + svc = _service(disp=disp, config=config, partition_service=psvc) + + await svc.add_file( + file_path=str(f), + file_id="f1", + partition="tenant-new", + metadata={}, + sanitized_filename="doc.txt", + original_filename="doc.txt", + user={"id": 7}, + ) + + # Partition was created with the uploader as owner, then the file dispatched. + assert psvc.created == [("tenant-new", 7)] + assert disp.dispatched[0]["partition"] == "tenant-new" + + +@pytest.mark.asyncio +async def test_add_file_auto_create_defaults_to_admin_when_user_missing(tmp_path): + f = tmp_path / "doc.txt" + f.write_text("x") + config = _config_with_partition("tenant-a") + psvc = FakePartitionService(config) + svc = _service(config=config, partition_service=psvc) + + await svc.add_file( + file_path=str(f), + file_id="f1", + partition="tenant-new", + metadata={}, + sanitized_filename="doc.txt", + original_filename="doc.txt", + user=None, + ) + + assert psvc.created == [("tenant-new", 1)] + + +@pytest.mark.asyncio +async def test_add_file_refreshes_cache_when_partition_row_exists(tmp_path): + f = tmp_path / "doc.txt" + f.write_text("x") + disp = FakeDispatcher() + config = _config_with_partition("tenant-a") + # Row exists in the DB but is missing from the in-memory cache. + psvc = FakePartitionService(config, db_partitions={"tenant-new"}) + svc = _service(disp=disp, config=config, partition_service=psvc) + + await svc.add_file( + file_path=str(f), + file_id="f1", + partition="tenant-new", + metadata={}, + sanitized_filename="doc.txt", + original_filename="doc.txt", + user=None, + ) + + # No spurious create; the stale cache was refreshed and the file dispatched. + assert psvc.created == [] + assert psvc.loaded == 1 + assert disp.dispatched[0]["partition"] == "tenant-new" + + @pytest.mark.asyncio async def test_delete_file_delegates(): disp = FakeDispatcher() diff --git a/tests/unit/services/orchestrators/test_model_endpoint_service.py b/tests/unit/services/orchestrators/test_model_endpoint_service.py new file mode 100644 index 000000000..548aafb11 --- /dev/null +++ b/tests/unit/services/orchestrators/test_model_endpoint_service.py @@ -0,0 +1,478 @@ +"""Unit tests for ModelEndpointService (Phase 14E).""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest + +_NOW = datetime(2026, 1, 1, tzinfo=UTC) + + +def _row_payload(**kwargs): + """Build a model endpoint row payload with targeted overrides.""" + payload = { + "name": "default", + "model_type": "embedder", + "endpoint": "http://vllm:8000/v1", + "model_name": "jina-v3", + "batch_size": 32, + "timeout": 30.0, + "extra": {}, + "is_default": True, + "created_at": _NOW, + "updated_at": _NOW, + } + payload.update(kwargs) + return payload + + +def _make_row(**kwargs): + from core.config.model_endpoints import ModelEndpointRow + + return ModelEndpointRow(**_row_payload(**kwargs)) + + +def _make_unvalidated_row(**kwargs): + from core.config.model_endpoints import ModelEndpointRow + + return ModelEndpointRow.model_construct(**_row_payload(**kwargs)) + + +class _FakeEndpointRepo: + def __init__(self, rows: list | None = None): + from core.config.model_endpoints import ModelEndpointRow + + self._store: dict[tuple[str, str], ModelEndpointRow] = {} + self.calls: list[tuple[str, tuple]] = [] + for r in rows or []: + self._store[(r.name, r.model_type)] = r + + async def create(self, row): + self._store[(row.name, row.model_type)] = row + self.calls.append(("create", (row.name, row.model_type))) + return row + + async def get(self, name: str, model_type: str): + self.calls.append(("get", (name, model_type))) + return self._store.get((name, model_type)) + + async def list_all(self, model_type: str | None = None): + self.calls.append(("list_all", (model_type,))) + rows = list(self._store.values()) + if model_type is not None: + rows = [r for r in rows if r.model_type == model_type] + return rows + + async def update(self, name: str, model_type: str, **fields): + row = self._store.get((name, model_type)) + if row is None: + return None + updated = row.model_copy(update=fields) + self._store[(name, model_type)] = updated + return updated + + async def rename(self, name: str, model_type: str, new_name: str) -> None: + row = self._store.pop((name, model_type), None) + if row: + self._store[(new_name, model_type)] = row.model_copy(update={"name": new_name}) + + async def delete(self, name: str, model_type: str) -> bool: + return self._store.pop((name, model_type), None) is not None + + async def set_default(self, model_type: str, name: str) -> None: + for key, row in list(self._store.items()): + self._store[key] = row.model_copy(update={"is_default": key[0] == name and key[1] == model_type}) + + +def _make_service(repo=None, rows=None, settings=None): + from core.config.root import Settings + from services.orchestrators.model_endpoint_service import ModelEndpointService + + return ModelEndpointService( + model_endpoint_repo=repo or _FakeEndpointRepo(rows), + config=settings or Settings(), + ) + + +# ------------------------------------------------------------------ +# seed_defaults +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_seed_defaults_inserts_embedder_when_empty(): + repo = _FakeEndpointRepo() + svc = _make_service(repo) + await svc.seed_defaults() + + creates = [c for c in repo.calls if c[0] == "create"] + assert any(name_type[1] == "embedder" for _, name_type in creates) + + +@pytest.mark.asyncio +async def test_seed_defaults_skips_type_when_rows_exist(): + existing = _make_row(name="custom", model_type="embedder", is_default=True) + repo = _FakeEndpointRepo(rows=[existing]) + svc = _make_service(repo) + await svc.seed_defaults() + + creates = [c for c in repo.calls if c[0] == "create"] + assert not any(name_type[1] == "embedder" for _, name_type in creates) + + +@pytest.mark.asyncio +async def test_seed_defaults_skips_empty_endpoint(monkeypatch): + # Settings().llm.base_url defaults to "" and Settings().llm.model to "". + # As long as no LLM_ENDPOINT env var is set, seed_defaults should skip llm. + monkeypatch.delenv("LLM_ENDPOINT", raising=False) + monkeypatch.delenv("LLM_MODEL", raising=False) + + repo = _FakeEndpointRepo() + svc = _make_service(repo) + await svc.seed_defaults() + + creates = [c for c in repo.calls if c[0] == "create"] + assert not any(name_type[1] == "llm" for _, name_type in creates) + + +@pytest.mark.asyncio +async def test_seed_defaults_seeds_disabled_reranker_when_configured(monkeypatch): + from core.config.root import Settings + + monkeypatch.delenv("RERANKER_ENDPOINT", raising=False) + # A disabled reranker is still catalogued as long as it is configured + # (base_url set): registration is about availability, while activation is + # the retrieval preset's enable_reranker kill-switch. + settings = Settings(reranker={"provider": "infinity", "enabled": False}) + repo = _FakeEndpointRepo() + svc = _make_service(repo, settings=settings) + await svc.seed_defaults() + + creates = [c for c in repo.calls if c[0] == "create"] + assert any(name_type[1] == "reranker" for _, name_type in creates) + + +@pytest.mark.asyncio +async def test_seed_defaults_preserves_endpoint_api_keys(monkeypatch): + from core.config.root import Settings + + monkeypatch.delenv("LLM_ENDPOINT", raising=False) + monkeypatch.delenv("LLM_MODEL", raising=False) + + settings = Settings( + embedder={"api_key": "embed-key"}, + llm={"base_url": "http://llm:8000/v1", "model": "mistral", "api_key": "llm-key"}, + vlm={"base_url": "http://vlm:8000/v1", "model": "pixtral", "api_key": "vlm-key"}, + reranker={"provider": "infinity", "api_key": "rerank-key"}, + ) + repo = _FakeEndpointRepo() + svc = _make_service(repo, settings=settings) + + await svc.seed_defaults() + + rows = repo._store.values() + assert {row.model_type: row.extra.get("api_key") for row in rows} == { + "embedder": "embed-key", + "llm": "llm-key", + "vlm": "vlm-key", + "reranker": "rerank-key", + } + + +@pytest.mark.asyncio +async def test_seed_defaults_skips_reranker_when_unconfigured(monkeypatch): + from core.config.root import Settings + + monkeypatch.delenv("RERANKER_ENDPOINT", raising=False) + # No base_url configured → nothing to advertise, so the reranker is skipped. + settings = Settings(reranker={"provider": "openai", "base_url": ""}) + repo = _FakeEndpointRepo() + svc = _make_service(repo, settings=settings) + await svc.seed_defaults() + + creates = [c for c in repo.calls if c[0] == "create"] + assert not any(name_type[1] == "reranker" for _, name_type in creates) + + +# ------------------------------------------------------------------ +# load_all +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_load_all_populates_config_models(): + from core.config.root import Settings + + rows = [ + _make_row(name="jina", model_type="embedder", is_default=True), + _make_row(name="mistral", model_type="llm", is_default=True), + ] + repo = _FakeEndpointRepo(rows=rows) + settings = Settings() + svc = _make_service(repo, settings=settings) + + await svc.load_all() + + assert "jina" in settings.models.embedder + assert "default" in settings.models.embedder + assert "mistral" in settings.models.llm + assert "default" in settings.models.llm + + +@pytest.mark.asyncio +async def test_load_all_no_default_alias_without_is_default(): + from core.config.root import Settings + + rows = [_make_row(name="jina", model_type="embedder", is_default=False)] + repo = _FakeEndpointRepo(rows=rows) + settings = Settings() + svc = _make_service(repo, settings=settings) + + await svc.load_all() + + assert "jina" in settings.models.embedder + assert "default" not in settings.models.embedder + + +# ------------------------------------------------------------------ +# create_model_endpoint +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_create_model_endpoint_rejects_invalid_type(): + from core.utils.exceptions import ValidationError + + svc = _make_service() + row = _make_unvalidated_row(model_type="unknown_type") + with pytest.raises(ValidationError, match="Invalid model_type"): + await svc.create_model_endpoint(row) + + +@pytest.mark.asyncio +async def test_create_model_endpoint_raises_409_on_duplicate(): + from core.utils.exceptions import ValidationError + + existing = _make_row() + repo = _FakeEndpointRepo(rows=[existing]) + svc = _make_service(repo) + + with pytest.raises(ValidationError) as exc_info: + await svc.create_model_endpoint(_make_row()) + assert exc_info.value.status_code == 409 + + +@pytest.mark.asyncio +async def test_create_model_endpoint_inserts_and_returns_row(): + repo = _FakeEndpointRepo() + svc = _make_service(repo) + row = _make_row(name="new-endpoint") + + result = await svc.create_model_endpoint(row) + + assert result.name == "new-endpoint" + assert any(c[0] == "create" for c in repo.calls) + + +# ------------------------------------------------------------------ +# get_model_endpoint +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_get_model_endpoint_raises_404_when_missing(): + from core.utils.exceptions import NotFoundError + + svc = _make_service() + with pytest.raises(NotFoundError): + await svc.get_model_endpoint("ghost", "embedder") + + +@pytest.mark.asyncio +async def test_get_model_endpoint_returns_row_when_found(): + existing = _make_row(name="jina") + svc = _make_service(rows=[existing]) + + row = await svc.get_model_endpoint("jina", "embedder") + assert row.name == "jina" + + +# ------------------------------------------------------------------ +# update_model_endpoint +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_update_model_endpoint_raises_404_when_missing(): + from core.utils.exceptions import NotFoundError + + svc = _make_service() + with pytest.raises(NotFoundError): + await svc.update_model_endpoint("ghost", "embedder", endpoint="http://new:8000/v1") + + +@pytest.mark.asyncio +async def test_update_model_endpoint_renames_and_evicts_cache(): + existing = _make_row(name="old-name") + repo = _FakeEndpointRepo(rows=[existing]) + cache: dict = {"old-name": object()} + svc = _make_service(repo) + svc._client_caches["embedder"] = cache + + await svc.update_model_endpoint("old-name", "embedder", new_name="new-name") + + assert "old-name" not in cache + assert ("old-name", "embedder") not in repo._store + assert ("new-name", "embedder") in repo._store + + +# ------------------------------------------------------------------ +# delete_model_endpoint +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_delete_model_endpoint_raises_404_when_missing(): + from core.utils.exceptions import NotFoundError + + svc = _make_service() + with pytest.raises(NotFoundError): + await svc.delete_model_endpoint("ghost", "embedder") + + +@pytest.mark.asyncio +async def test_delete_model_endpoint_raises_422_when_last(): + from core.utils.exceptions import ValidationError + + existing = _make_row() + svc = _make_service(rows=[existing]) + with pytest.raises(ValidationError, match="last"): + await svc.delete_model_endpoint("default", "embedder") + + +@pytest.mark.asyncio +async def test_delete_model_endpoint_removes_row_and_evicts_cache(): + rows = [ + _make_row(name="jina", is_default=True), + _make_row(name="e5", is_default=False), + ] + repo = _FakeEndpointRepo(rows=rows) + cache: dict = {"jina": object()} + svc = _make_service(repo) + svc._client_caches["embedder"] = cache + + await svc.delete_model_endpoint("jina", "embedder") + + assert ("jina", "embedder") not in repo._store + assert "jina" not in cache + + +# ------------------------------------------------------------------ +# set_default +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_set_default_raises_404_when_missing(): + from core.utils.exceptions import NotFoundError + + svc = _make_service() + with pytest.raises(NotFoundError): + await svc.set_default("embedder", "ghost") + + +@pytest.mark.asyncio +async def test_set_default_calls_repo_and_reloads(): + rows = [ + _make_row(name="jina", is_default=True), + _make_row(name="e5", is_default=False), + ] + repo = _FakeEndpointRepo(rows=rows) + cache: dict = {"default": object()} + svc = _make_service(repo) + svc._client_caches["embedder"] = cache + + await svc.set_default("embedder", "e5") + + assert "default" not in cache + assert any(c[0] == "list_all" for c in repo.calls) + + +# ------------------------------------------------------------------ +# validate_endpoint +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_validate_endpoint_probes_url_and_model_name(monkeypatch): + import httpx + + svc = _make_service() + calls: list[str] = [] + + class FakeResponse: + status_code = 200 + + def json(self): + return {"data": [{"id": "mistral-small"}]} + + class FakeClient: + def __init__(self, *, timeout, headers): + assert timeout == 5.0 + assert headers == {} + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def get(self, url): + calls.append(url) + return FakeResponse() + + monkeypatch.setattr(httpx, "AsyncClient", FakeClient) + + result = await svc.validate_endpoint("http://llm:8000/v1", "mistral-small") + + assert calls == ["http://llm:8000/v1/models"] + assert result == { + "reachable": True, + "model_found": True, + "models_served": ["mistral-small"], + "detail": None, + } + + +@pytest.mark.asyncio +async def test_validate_endpoint_sends_api_key(monkeypatch): + import httpx + + svc = _make_service() + captured_headers: list[dict[str, str]] = [] + + class FakeResponse: + status_code = 200 + + def json(self): + return {"data": [{"id": "mistral-small"}]} + + class FakeClient: + def __init__(self, *, timeout, headers): + captured_headers.append(headers) + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def get(self, url): + return FakeResponse() + + monkeypatch.setattr(httpx, "AsyncClient", FakeClient) + + await svc.validate_endpoint("http://llm:8000/v1", "mistral-small", api_key="secret-token") + + assert captured_headers == [{"Authorization": "Bearer secret-token"}] diff --git a/tests/unit/services/orchestrators/test_partition_preset_resolution.py b/tests/unit/services/orchestrators/test_partition_preset_resolution.py new file mode 100644 index 000000000..582c37c79 --- /dev/null +++ b/tests/unit/services/orchestrators/test_partition_preset_resolution.py @@ -0,0 +1,309 @@ +"""Unit tests for PartitionService preset resolution (Phase 14G).""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest + +_NOW = datetime(2026, 1, 1, tzinfo=UTC) + +_IDX_CONFIG = { + "chunking": {"name": "recursive_splitter", "chunk_size": 512, "chunk_overlap_rate": 0.2}, + "parsing_strategy": "marker", +} +_RET_CONFIG = {"type": "single", "top_k": 50, "top_n": 10} + + +def _full_row(partition: str, **overrides) -> dict: + base = { + "partition": partition, + "description": "", + "embedder": "default", + "indexation_preset": "default", + "retrieval_preset": "default", + "dimension": 1024, + "collection_name": None, + "chat_history_depth": 0, + "chat_llm": None, + "created_at": _NOW, + "updated_at": _NOW, + } + base.update(overrides) + return base + + +class _FakePartitionRepo: + def __init__(self, rows: list[dict] | None = None) -> None: + self._store: dict[str, dict] = {r["partition"]: r for r in (rows or [])} + self.calls: list[tuple[str, tuple]] = [] + + async def partition_exists(self, name: str) -> bool: + return name in self._store + + async def create_partition(self, name: str, user_id: int | None = None) -> dict: + self.calls.append(("create_partition", (name, user_id))) + self._store.setdefault(name, _full_row(name)) + return self._store[name] + + async def get_partition_row(self, name: str) -> dict | None: + return self._store.get(name) + + async def list_partition_rows(self) -> list[dict]: + return list(self._store.values()) + + async def update_partition(self, name: str, **fields) -> dict | None: + self.calls.append(("update_partition", (name,))) + row = self._store.get(name) + if row is None: + return None + row.update(fields) + return row + + async def delete_partition(self, name: str) -> bool: + return self._store.pop(name, None) is not None + + +def _settings(idx=None, ret=None): + from core.config.root import Settings + + s = Settings() + s.presets.indexation.clear() + s.presets.indexation.update(idx if idx is not None else {"default": _IDX_CONFIG}) + s.presets.retrieval.clear() + s.presets.retrieval.update(ret if ret is not None else {"default": _RET_CONFIG}) + return s + + +def _make_service(repo=None, rows=None, settings=None): + from services.orchestrators.partition_service import PartitionService + + return PartitionService( + partition_repo=repo or _FakePartitionRepo(rows), + membership_repo=object(), + document_repo=object(), + vector_store=object(), + user_repo=object(), + collection="vdb", + config=settings if settings is not None else _settings(), + ) + + +# ------------------------------------------------------------------ +# resolve_partition_row +# ------------------------------------------------------------------ + + +def test_resolve_partition_row_builds_config(): + from core.config.indexation_pipeline import IndexationPipelineConfig + from core.config.retrieval_pipeline import RetrievalPipelineConfig + + svc = _make_service() + cfg = svc.resolve_partition_row(_full_row("p1", description="hello")) + + assert cfg.name == "p1" + assert cfg.description == "hello" + assert cfg.embedder == "default" + assert isinstance(cfg.indexation, IndexationPipelineConfig) + assert isinstance(cfg.retrieval, RetrievalPipelineConfig) + assert cfg.retrieval.top_k == 50 + + +def test_resolve_partition_row_missing_indexation_preset_raises(): + from core.utils.exceptions import ConfigError + + svc = _make_service() + with pytest.raises(ConfigError, match="Indexation preset 'ghost'"): + svc.resolve_partition_row(_full_row("p1", indexation_preset="ghost")) + + +def test_resolve_partition_row_missing_retrieval_preset_raises(): + from core.utils.exceptions import ConfigError + + svc = _make_service() + with pytest.raises(ConfigError, match="Retrieval preset 'ghost'"): + svc.resolve_partition_row(_full_row("p1", retrieval_preset="ghost")) + + +def test_resolve_partition_row_without_config_raises(): + from core.utils.exceptions import ConfigError + from services.orchestrators.partition_service import PartitionService + + svc = PartitionService( + partition_repo=_FakePartitionRepo(), + membership_repo=object(), + document_repo=object(), + vector_store=object(), + user_repo=object(), + collection="vdb", + ) + with pytest.raises(ConfigError, match="without a config"): + svc.resolve_partition_row(_full_row("p1")) + + +# ------------------------------------------------------------------ +# load_partitions +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_load_partitions_populates_cache(): + settings = _settings() + repo = _FakePartitionRepo(rows=[_full_row("a"), _full_row("b")]) + svc = _make_service(repo, settings=settings) + + await svc.load_partitions() + + assert set(settings.partitions) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_load_partitions_clears_stale(): + settings = _settings() + settings.partitions["stale"] = object() # type: ignore[assignment] + repo = _FakePartitionRepo(rows=[_full_row("fresh")]) + svc = _make_service(repo, settings=settings) + + await svc.load_partitions() + + assert "stale" not in settings.partitions + assert "fresh" in settings.partitions + + +# ------------------------------------------------------------------ +# seed_default_partition +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_seed_default_partition_creates_when_missing(): + repo = _FakePartitionRepo() + svc = _make_service(repo) + + await svc.seed_default_partition() + + assert await repo.partition_exists("default") + assert any(c[0] == "create_partition" for c in repo.calls) + + +@pytest.mark.asyncio +async def test_seed_default_partition_skips_when_present(): + repo = _FakePartitionRepo(rows=[_full_row("default")]) + svc = _make_service(repo) + + await svc.seed_default_partition() + + assert not any(c[0] == "create_partition" for c in repo.calls) + + +# ------------------------------------------------------------------ +# create_partition (Phase 14 flow) +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_create_partition_validates_presets_before_write(): + from core.utils.exceptions import ValidationError + + repo = _FakePartitionRepo() + svc = _make_service(repo) + + # User-supplied preset name → 422 ValidationError, not a 500 ConfigError. + with pytest.raises(ValidationError, match="Indexation preset 'nope'") as exc: + await svc.create_partition("p1", user_id=1, indexation_preset="nope") + assert exc.value.status_code == 422 + + assert not await repo.partition_exists("p1") + + +@pytest.mark.asyncio +async def test_create_partition_persists_config_and_reloads(): + settings = _settings() + repo = _FakePartitionRepo() + svc = _make_service(repo, settings=settings) + + await svc.create_partition("p1", user_id=1, description="docs") + + assert any(c[0] == "update_partition" for c in repo.calls) + assert "p1" in settings.partitions + assert settings.partitions["p1"].description == "docs" + + +# ------------------------------------------------------------------ +# update_partition +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_update_partition_validates_and_reloads(): + settings = _settings() + repo = _FakePartitionRepo(rows=[_full_row("p1")]) + svc = _make_service(repo, settings=settings) + + await svc.update_partition("p1", description="new") + + assert repo._store["p1"]["description"] == "new" + assert settings.partitions["p1"].description == "new" + + +@pytest.mark.asyncio +async def test_update_partition_rejects_unknown_preset(): + from core.utils.exceptions import ValidationError + + repo = _FakePartitionRepo(rows=[_full_row("p1")]) + svc = _make_service(repo) + + with pytest.raises(ValidationError, match="Retrieval preset 'ghost'") as exc: + await svc.update_partition("p1", retrieval_preset="ghost") + assert exc.value.status_code == 422 + + +@pytest.mark.asyncio +async def test_update_partition_missing_raises_404(): + from core.utils.exceptions import PartitionNotFoundError + + svc = _make_service(_FakePartitionRepo()) + with pytest.raises(PartitionNotFoundError): + await svc.update_partition("ghost", description="x") + + +# ------------------------------------------------------------------ +# get_partition_config / update_partition_config (PartitionDetailResponse) +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_get_partition_config_returns_resolved_detail(): + repo = _FakePartitionRepo(rows=[_full_row("p1", description="docs")]) + svc = _make_service(repo) + + detail = await svc.get_partition_config("p1") + + assert detail["name"] == "p1" + assert detail["description"] == "docs" + assert detail["embedder"] == "default" + assert detail["indexation_preset"] == "default" + assert detail["retrieval_preset"] == "default" + assert detail["dimension"] == 1024 + assert detail["retrieval_pipeline"]["top_k"] == 50 + assert "chunking" in detail["indexation_pipeline"] + + +@pytest.mark.asyncio +async def test_get_partition_config_missing_raises_404(): + from core.utils.exceptions import PartitionNotFoundError + + svc = _make_service(_FakePartitionRepo()) + with pytest.raises(PartitionNotFoundError): + await svc.get_partition_config("ghost") + + +@pytest.mark.asyncio +async def test_update_partition_config_applies_and_returns_detail(): + repo = _FakePartitionRepo(rows=[_full_row("p1")]) + svc = _make_service(repo) + + detail = await svc.update_partition_config("p1", description="new") + + assert detail["description"] == "new" + assert repo._store["p1"]["description"] == "new" diff --git a/tests/unit/services/orchestrators/test_preset_service.py b/tests/unit/services/orchestrators/test_preset_service.py new file mode 100644 index 000000000..9581d8505 --- /dev/null +++ b/tests/unit/services/orchestrators/test_preset_service.py @@ -0,0 +1,350 @@ +"""Unit tests for PresetService (Phase 14F).""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest + +_NOW = datetime(2026, 1, 1, tzinfo=UTC) + +_VALID_IDX_CONFIG = { + "chunking": {"name": "recursive_splitter", "chunk_size": 512, "chunk_overlap_rate": 0.2}, + "parsing_strategy": "marker", + "enable_image_captioning": True, +} + +_VALID_RET_CONFIG = { + "type": "single", + "top_k": 50, + "top_n": 10, + "enable_reranker": True, + "similarity_threshold": 0.6, +} + + +def _make_row(name: str, preset_type: str, config: dict | None = None) -> dict: + return { + "name": name, + "preset_type": preset_type, + "config": config or _VALID_IDX_CONFIG, + "created_at": _NOW, + "updated_at": _NOW, + } + + +class _FakePresetRepo: + def __init__(self, rows: list[dict] | None = None) -> None: + self._store: dict[tuple[str, str], dict] = {} + self._partition_counts: dict[tuple[str, str], int] = {} + self.calls: list[tuple[str, tuple]] = [] + for r in rows or []: + self._store[(r["name"], r["preset_type"])] = r + + async def get(self, name: str, preset_type: str) -> dict | None: + self.calls.append(("get", (name, preset_type))) + return self._store.get((name, preset_type)) + + async def list_all(self, preset_type: str | None = None) -> list[dict]: + self.calls.append(("list_all", (preset_type,))) + rows = list(self._store.values()) + if preset_type is not None: + rows = [r for r in rows if r["preset_type"] == preset_type] + return rows + + async def upsert(self, name: str, preset_type: str, config: dict) -> dict: + self.calls.append(("upsert", (name, preset_type))) + row = _make_row(name, preset_type, config) + self._store[(name, preset_type)] = row + return row + + async def delete(self, name: str, preset_type: str) -> bool: + self.calls.append(("delete", (name, preset_type))) + return self._store.pop((name, preset_type), None) is not None + + async def rename(self, old_name: str, new_name: str, preset_type: str, config: dict) -> dict: + self.calls.append(("rename", (old_name, new_name, preset_type))) + self._store.pop((old_name, preset_type), None) + row = _make_row(new_name, preset_type, config) + self._store[(new_name, preset_type)] = row + return row + + async def count_partitions_using(self, name: str, preset_type: str) -> int: + self.calls.append(("count_partitions_using", (name, preset_type))) + return self._partition_counts.get((name, preset_type), 0) + + +class _FakePartitionService: + def __init__(self) -> None: + self.load_calls = 0 + + async def load_partitions(self) -> None: + self.load_calls += 1 + + +def _make_service(repo=None, rows=None, settings=None, partition_service=None): + from core.config.root import Settings + from services.orchestrators.preset_service import PresetService + + return PresetService( + preset_repo=repo or _FakePresetRepo(rows), + config=settings or Settings(), + partition_service=partition_service, + ) + + +# ------------------------------------------------------------------ +# seed_defaults +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_seed_defaults_inserts_six_presets_when_empty(): + repo = _FakePresetRepo() + svc = _make_service(repo) + await svc.seed_defaults() + + upserts = [c for c in repo.calls if c[0] == "upsert"] + assert len(upserts) == 6 # 3 indexation + 3 retrieval + + +@pytest.mark.asyncio +async def test_seed_defaults_retrieval_inherits_reranker_enabled(): + from core.config.root import Settings + + settings = Settings(reranker={"provider": "infinity", "enabled": False}) + repo = _FakePresetRepo() + svc = _make_service(repo, settings=settings) + await svc.seed_defaults() + + ret = [r for r in repo._store.values() if r["preset_type"] == "retrieval"] + assert ret # sanity: retrieval presets were seeded + assert all(r["config"]["enable_reranker"] is False for r in ret) + + +@pytest.mark.asyncio +async def test_seed_defaults_retrieval_enables_reranker_when_available(): + from core.config.root import Settings + + settings = Settings(reranker={"provider": "infinity", "enabled": True}) + repo = _FakePresetRepo() + svc = _make_service(repo, settings=settings) + await svc.seed_defaults() + + ret = [r for r in repo._store.values() if r["preset_type"] == "retrieval"] + assert all(r["config"]["enable_reranker"] is True for r in ret) + + +@pytest.mark.asyncio +async def test_seed_defaults_skips_type_when_rows_exist(): + existing = _make_row("default", "indexation") + repo = _FakePresetRepo(rows=[existing]) + svc = _make_service(repo) + await svc.seed_defaults() + + upserts = [c for c in repo.calls if c[0] == "upsert"] + upserted_types = {args[1] for _, args in upserts} + assert "indexation" not in upserted_types + assert "retrieval" in upserted_types + + +# ------------------------------------------------------------------ +# load_all +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_load_all_populates_config_presets(): + from core.config.root import Settings + + rows = [ + _make_row("default", "indexation", _VALID_IDX_CONFIG), + _make_row("default", "retrieval", _VALID_RET_CONFIG), + ] + repo = _FakePresetRepo(rows=rows) + settings = Settings() + svc = _make_service(repo, settings=settings) + + await svc.load_all() + + assert "default" in settings.presets.indexation + assert "default" in settings.presets.retrieval + + +@pytest.mark.asyncio +async def test_load_all_clears_stale_entries(): + from core.config.root import Settings + + rows = [_make_row("old", "indexation", _VALID_IDX_CONFIG)] + repo = _FakePresetRepo(rows=rows) + settings = Settings() + settings.presets.indexation["stale"] = {} + svc = _make_service(repo, settings=settings) + + await svc.load_all() + + assert "stale" not in settings.presets.indexation + assert "old" in settings.presets.indexation + + +# ------------------------------------------------------------------ +# create_preset +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_create_preset_rejects_invalid_type(): + from core.utils.exceptions import ValidationError + + svc = _make_service() + with pytest.raises(ValidationError, match="Invalid preset_type"): + await svc.create_preset("my-preset", "unknown_type", {}) + + +@pytest.mark.asyncio +async def test_create_preset_raises_409_on_duplicate(): + from core.utils.exceptions import ValidationError + + existing = _make_row("default", "indexation") + repo = _FakePresetRepo(rows=[existing]) + svc = _make_service(repo) + + with pytest.raises(ValidationError) as exc_info: + await svc.create_preset("default", "indexation", _VALID_IDX_CONFIG) + assert exc_info.value.status_code == 409 + + +@pytest.mark.asyncio +async def test_create_preset_rejects_invalid_config(): + from core.utils.exceptions import ValidationError + + svc = _make_service() + bad_config = {"top_k": -999} # violates gt=0 constraint + with pytest.raises(ValidationError, match="Invalid retrieval preset config"): + await svc.create_preset("bad", "retrieval", bad_config) + + +@pytest.mark.asyncio +async def test_create_preset_inserts_and_returns_row(): + repo = _FakePresetRepo() + svc = _make_service(repo) + + result = await svc.create_preset("legal", "indexation", _VALID_IDX_CONFIG) + + assert result["name"] == "legal" + assert result["preset_type"] == "indexation" + assert any(c[0] == "upsert" for c in repo.calls) + + +# ------------------------------------------------------------------ +# get_preset +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_get_preset_raises_404_when_missing(): + from core.utils.exceptions import NotFoundError + + svc = _make_service() + with pytest.raises(NotFoundError): + await svc.get_preset("ghost", "indexation") + + +@pytest.mark.asyncio +async def test_get_preset_returns_row_when_found(): + existing = _make_row("default", "retrieval", _VALID_RET_CONFIG) + svc = _make_service(rows=[existing]) + + row = await svc.get_preset("default", "retrieval") + assert row["name"] == "default" + assert row["preset_type"] == "retrieval" + + +# ------------------------------------------------------------------ +# update_preset +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_update_preset_raises_404_when_missing(): + from core.utils.exceptions import NotFoundError + + svc = _make_service() + with pytest.raises(NotFoundError): + await svc.update_preset("ghost", "indexation", config=_VALID_IDX_CONFIG) + + +@pytest.mark.asyncio +async def test_update_preset_renames_preset(): + existing = _make_row("old-name", "retrieval", _VALID_RET_CONFIG) + repo = _FakePresetRepo(rows=[existing]) + svc = _make_service(repo) + + await svc.update_preset("old-name", "retrieval", new_name="new-name") + + assert ("old-name", "retrieval") not in repo._store + assert ("new-name", "retrieval") in repo._store + assert ("rename", ("old-name", "new-name", "retrieval")) in repo.calls + assert not any(call[0] == "delete" for call in repo.calls) + + +@pytest.mark.asyncio +async def test_update_preset_calls_partition_service(): + existing = _make_row("default", "retrieval", _VALID_RET_CONFIG) + repo = _FakePresetRepo(rows=[existing]) + part_svc = _FakePartitionService() + svc = _make_service(repo, partition_service=part_svc) + + await svc.update_preset("default", "retrieval", config=_VALID_RET_CONFIG) + + assert part_svc.load_calls == 1 + + +@pytest.mark.asyncio +async def test_update_preset_rejects_invalid_new_config(): + from core.utils.exceptions import ValidationError + + existing = _make_row("default", "retrieval", _VALID_RET_CONFIG) + repo = _FakePresetRepo(rows=[existing]) + svc = _make_service(repo) + + with pytest.raises(ValidationError, match="Invalid retrieval preset config"): + await svc.update_preset("default", "retrieval", config={"top_k": -1}) + + +# ------------------------------------------------------------------ +# delete_preset +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_delete_preset_raises_404_when_missing(): + from core.utils.exceptions import NotFoundError + + svc = _make_service() + with pytest.raises(NotFoundError): + await svc.delete_preset("ghost", "indexation") + + +@pytest.mark.asyncio +async def test_delete_preset_raises_422_when_in_use(): + from core.utils.exceptions import ValidationError + + existing = _make_row("default", "indexation") + repo = _FakePresetRepo(rows=[existing]) + repo._partition_counts[("default", "indexation")] = 2 + svc = _make_service(repo) + + with pytest.raises(ValidationError, match="2 partition"): + await svc.delete_preset("default", "indexation") + + +@pytest.mark.asyncio +async def test_delete_preset_removes_row(): + existing = _make_row("legal", "indexation") + repo = _FakePresetRepo(rows=[existing]) + svc = _make_service(repo) + + await svc.delete_preset("legal", "indexation") + + assert ("legal", "indexation") not in repo._store diff --git a/tests/unit/services/orchestrators/test_retrieval_service.py b/tests/unit/services/orchestrators/test_retrieval_service.py index ed4b570f1..7558af04d 100644 --- a/tests/unit/services/orchestrators/test_retrieval_service.py +++ b/tests/unit/services/orchestrators/test_retrieval_service.py @@ -12,8 +12,12 @@ from types import SimpleNamespace import pytest +from core.config.indexation_pipeline import IndexationPipelineConfig +from core.config.retrieval_pipeline import RetrievalPipelineConfig from core.models.chunk import Chunk +from core.models.preset import PartitionConfig from core.models.query import Query, SearchQueries +from core.utils.exceptions import PartitionNotFoundError from services.orchestrators.retrieval_service import RetrievalService @@ -58,9 +62,33 @@ def _config(rtype: str = "single", reranker_enabled: bool = False) -> SimpleName combine=False, ), reranker=SimpleNamespace(enabled=reranker_enabled, top_k=5), + partitions={}, ) +def _partition( + *, + name: str = "tenant-a", + embedder: str = "embed-a", + retrieval: RetrievalPipelineConfig | None = None, +) -> PartitionConfig: + return PartitionConfig( + name=name, + embedder=embedder, + indexation=IndexationPipelineConfig(), + retrieval=retrieval or RetrievalPipelineConfig(), + ) + + +class FakeReranker: + def __init__(self): + self.calls: list[dict] = [] + + async def rerank(self, *, query, documents, top_k): + self.calls.append({"query": query, "documents": documents, "top_k": top_k}) + return [(idx, 1.0) for idx in range(len(documents))] + + def _svc(searcher, *, rtype="single", reranker_enabled=False) -> RetrievalService: return RetrievalService( searcher=searcher, @@ -169,3 +197,86 @@ def test_fuse_rrf_merges_and_dedupes(): 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 + + +# --------------------------------------------------------------------------- # +# Phase 14J.1 — per-partition retrieval pipeline config +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_retrieve_uses_partition_retrieval_config_and_named_reranker(): + s = FakeSearcher() + s.search_result = [_chunk("a"), _chunk("b"), _chunk("c")] + reranker = FakeReranker() + searcher_calls: list[str] = [] + reranker_calls: list[str] = [] + cfg = _config() + cfg.partitions = { + "tenant-a": _partition( + retrieval=RetrievalPipelineConfig( + top_k=3, + top_n=2, + similarity_threshold=0.77, + include_related=False, + include_ancestors=False, + enable_reranker=True, + reranker="fast-ranker", + ) + ) + } + + svc = RetrievalService( + searcher=s, + reranker=None, + llm=None, + config=cfg, + searcher_factory=lambda name: searcher_calls.append(name) or s, + reranker_factory=lambda name: reranker_calls.append(name) or reranker, + ) + + out = await svc.retrieve(partitions=["tenant-a"], query=Query(query="hello")) + + assert [c.id for c in out] == ["a", "b"] + assert searcher_calls == ["embed-a"] + assert reranker_calls == ["fast-ranker"] + assert reranker.calls[0]["query"] == "hello" + call = s.search_calls[0] + assert call["partition"] == ["tenant-a"] + assert call["top_k"] == 3 + assert call["similarity_threshold"] == 0.77 + + +@pytest.mark.asyncio +async def test_retrieve_uses_partition_searcher_factory_for_named_embedder(): + default_searcher = FakeSearcher() + tenant_searcher = FakeSearcher() + tenant_searcher.search_result = [_chunk("tenant")] + searcher_names: list[str] = [] + cfg = _config() + cfg.partitions = {"tenant-a": _partition(embedder="embed-a")} + + svc = RetrievalService( + searcher=default_searcher, + reranker=None, + llm=None, + config=cfg, + searcher_factory=lambda name: searcher_names.append(name) or tenant_searcher, + ) + + out = await svc.retrieve(partitions=["tenant-a"], query=Query(query="hello")) + + assert [c.id for c in out] == ["tenant"] + assert searcher_names == ["embed-a"] + assert default_searcher.search_calls == [] + assert tenant_searcher.search_calls[0]["partition"] == ["tenant-a"] + + +@pytest.mark.asyncio +async def test_retrieve_rejects_unknown_partition_when_partition_configs_exist(): + cfg = _config() + cfg.partitions = {"tenant-a": _partition()} + svc = RetrievalService(searcher=FakeSearcher(), reranker=None, llm=None, config=cfg) + + with pytest.raises(PartitionNotFoundError, match="tenant-b"): + await svc.retrieve(partitions=["tenant-b"], query=Query(query="hello")) diff --git a/tests/unit/services/persistence/test_add_file_to_partition_user_id.py b/tests/unit/services/persistence/test_add_file_to_partition_user_id.py index 310b92076..cc86c44b5 100644 --- a/tests/unit/services/persistence/test_add_file_to_partition_user_id.py +++ b/tests/unit/services/persistence/test_add_file_to_partition_user_id.py @@ -41,6 +41,15 @@ def acquire(self): return _AsyncContext(self.conn) +class _DirectExecutePool: + def __init__(self): + self.executed: list[tuple[str, tuple]] = [] + + async def execute(self, query: str, *params): + self.executed.append((query, params)) + return "UPDATE 1" + + @pytest.mark.asyncio async def test_auto_create_refuses_when_user_id_is_none(): from services.persistence.document_repo import PgDocumentRepository @@ -63,3 +72,61 @@ async def test_auto_create_succeeds_with_real_user_id(): assert await repo.add_file_to_partition(file_id="f1", partition="new-part", user_id=42) is True assert any("INSERT INTO files" in query for query, _ in pool.conn.executed) + + +@pytest.mark.asyncio +async def test_add_file_to_partition_persists_indexation_config_column(): + from services.persistence.document_repo import PgDocumentRepository + + pool = _FakePool() + repo = PgDocumentRepository(pool_getter=lambda: pool) + snapshot = {"parsing_strategy": "pymupdf"} + + assert ( + await repo.add_file_to_partition( + file_id="f1", + partition="new-part", + user_id=42, + indexation_config=snapshot, + ) + is True + ) + insert_query, insert_params = next( + (query, params) for query, params in pool.conn.executed if "INSERT INTO files" in query + ) + assert "indexation_config" in insert_query + assert snapshot in insert_params + + +@pytest.mark.asyncio +async def test_update_file_in_partition_persists_indexation_config_column(): + from services.persistence.document_repo import PgDocumentRepository + + pool = _DirectExecutePool() + repo = PgDocumentRepository(pool_getter=lambda: pool) + snapshot = {"parsing_strategy": "marker"} + + assert await repo.update_file_in_partition("f1", "tenant-a", indexation_config=snapshot) is True + + query, params = pool.executed[0] + assert "indexation_config" in query + assert snapshot in params + + +def test_row_to_document_exposes_indexation_config_snapshot(): + from services.persistence.document_repo import PgDocumentRepository + + snapshot = {"parsing_strategy": "marker"} + row = { + "file_id": "f1", + "partition_name": "tenant-a", + "file_metadata": {"filename": "doc.txt"}, + "created_by": 42, + "relationship_id": None, + "parent_id": None, + "indexation_config": snapshot, + } + + doc = PgDocumentRepository._row_to_document(row) + + assert doc.indexation_config == snapshot diff --git a/tests/unit/services/persistence/test_model_endpoint_repo.py b/tests/unit/services/persistence/test_model_endpoint_repo.py new file mode 100644 index 000000000..21eadc390 --- /dev/null +++ b/tests/unit/services/persistence/test_model_endpoint_repo.py @@ -0,0 +1,216 @@ +"""Unit tests for PgModelEndpointRepository.""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest + +_NOW = datetime(2026, 1, 1, tzinfo=UTC) + + +def _make_row(**kwargs): + base = { + "name": "default", + "model_type": "embedder", + "endpoint": "http://vllm:8000/v1", + "model_name": "jina-v3", + "batch_size": 32, + "timeout": 30.0, + "extra": {}, + "is_default": True, + "created_at": _NOW, + "updated_at": _NOW, + } + base.update(kwargs) + return base + + +class _AsyncCtx: + def __init__(self, value): + self._value = value + + async def __aenter__(self): + return self._value + + async def __aexit__(self, *_): + return False + + +class _FakeConn: + def __init__(self): + self.executed: list[tuple[str, tuple]] = [] + self._fetchrow_result = None + + def transaction(self): + return _AsyncCtx(self) + + async def execute(self, query: str, *params): + self.executed.append((query, params)) + return "UPDATE 1" + + async def fetchrow(self, query: str, *params): + self.executed.append((query, params)) + return self._fetchrow_result + + +class _FakePool: + def __init__(self): + self.conn = _FakeConn() + self.executed: list[tuple[str, tuple]] = [] + self._fetchrow_result = None + self._fetch_result: list = [] + self._fetchval_result = None + + def acquire(self): + return _AsyncCtx(self.conn) + + async def fetchrow(self, query: str, *params): + self.executed.append((query, params)) + return self._fetchrow_result + + async def fetch(self, query: str, *params): + self.executed.append((query, params)) + return self._fetch_result + + async def fetchval(self, query: str, *params): + self.executed.append((query, params)) + return self._fetchval_result + + async def execute(self, query: str, *params): + self.executed.append((query, params)) + return "DELETE 1" + + +@pytest.mark.asyncio +async def test_create_inserts_and_returns_model(): + from core.config.model_endpoints import ModelEndpointRow + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool._fetchrow_result = _make_row() + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + row = ModelEndpointRow( + name="default", + model_type="embedder", + endpoint="http://vllm:8000/v1", + created_at=_NOW, + updated_at=_NOW, + ) + result = await repo.create(row) + + assert result.name == "default" + assert result.model_type == "embedder" + assert any("INSERT INTO model_endpoints" in q for q, _ in pool.executed) + + +@pytest.mark.asyncio +async def test_get_returns_none_when_missing(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool._fetchrow_result = None + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + assert await repo.get("missing", "embedder") is None + + +@pytest.mark.asyncio +async def test_list_all_no_filter_orders_by_type_name(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool._fetch_result = [_make_row(), _make_row(name="fast", model_type="llm")] + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + results = await repo.list_all() + assert len(results) == 2 + query, params = pool.executed[0] + assert "ORDER BY model_type, name" in query + assert params == () + + +@pytest.mark.asyncio +async def test_list_all_filters_by_model_type(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool._fetch_result = [_make_row()] + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.list_all(model_type="embedder") + query, params = pool.executed[0] + assert "WHERE model_type" in query + assert params == ("embedder",) + + +@pytest.mark.asyncio +async def test_update_builds_set_clause_for_allowed_fields(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool._fetchrow_result = _make_row(endpoint="http://new:8000/v1") + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + result = await repo.update("default", "embedder", endpoint="http://new:8000/v1") + assert result is not None + query, params = pool.executed[0] + assert "UPDATE model_endpoints SET" in query + assert "endpoint = $3" in query + assert "updated_at = now()" in query + assert params[2] == "http://new:8000/v1" + + +@pytest.mark.asyncio +async def test_update_ignores_unknown_fields(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool._fetchrow_result = _make_row() + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + # "unknown_field" should be silently ignored; falls back to a plain GET + await repo.update("default", "embedder", unknown_field="x") + query, _ = pool.executed[0] + assert "SELECT" in query + + +@pytest.mark.asyncio +async def test_delete_returns_true_on_success(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + assert await repo.delete("default", "embedder") is True + + +@pytest.mark.asyncio +async def test_delete_returns_false_when_row_missing(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + + async def _execute(query, *params): + pool.executed.append((query, params)) + return "DELETE 0" + + pool.execute = _execute + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + assert await repo.delete("ghost", "embedder") is False + + +@pytest.mark.asyncio +async def test_set_default_uses_transaction_with_two_updates(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.set_default("embedder", "default") + + queries = [q for q, _ in pool.conn.executed] + assert any("is_default = false" in q for q in queries) + assert any("is_default = true" in q for q in queries) diff --git a/tests/unit/services/persistence/test_preset_repo.py b/tests/unit/services/persistence/test_preset_repo.py new file mode 100644 index 000000000..60b05b02a --- /dev/null +++ b/tests/unit/services/persistence/test_preset_repo.py @@ -0,0 +1,233 @@ +"""Unit tests for PgPresetRepository.""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest + +_NOW = datetime(2026, 1, 1, tzinfo=UTC) + + +def _make_row(**kwargs): + base = { + "name": "default", + "preset_type": "indexation", + "config": {"chunk_size": 512}, + "created_at": _NOW, + "updated_at": _NOW, + } + base.update(kwargs) + return base + + +class _FakePool: + def __init__(self): + self.executed: list[tuple[str, tuple]] = [] + self._fetchrow_result = None + self._fetch_result: list = [] + self._fetchval_result: int = 0 + + async def fetchrow(self, query: str, *params): + self.executed.append((query, params)) + return self._fetchrow_result + + async def fetch(self, query: str, *params): + self.executed.append((query, params)) + return self._fetch_result + + async def fetchval(self, query: str, *params): + self.executed.append((query, params)) + return self._fetchval_result + + async def execute(self, query: str, *params): + self.executed.append((query, params)) + return "DELETE 1" + + def acquire(self): + return _FakeAcquire(self) + + def transaction(self): + return _FakeTransaction(self) + + +class _FakeAcquire: + def __init__(self, pool: _FakePool) -> None: + self.pool = pool + + async def __aenter__(self): + return self.pool + + async def __aexit__(self, exc_type, exc, tb): + return False + + +class _FakeTransaction: + def __init__(self, pool: _FakePool) -> None: + self.pool = pool + self.entered = False + self.exited = False + + async def __aenter__(self): + self.entered = True + self.pool.executed.append(("BEGIN", ())) + return self + + async def __aexit__(self, exc_type, exc, tb): + self.exited = True + self.pool.executed.append(("COMMIT" if exc_type is None else "ROLLBACK", ())) + return False + + +@pytest.mark.asyncio +async def test_get_returns_none_when_missing(): + from services.persistence.preset_repo import PgPresetRepository + + pool = _FakePool() + repo = PgPresetRepository(pool_getter=lambda: pool) + + assert await repo.get("ghost", "indexation") is None + + +@pytest.mark.asyncio +async def test_get_returns_dict_when_found(): + from services.persistence.preset_repo import PgPresetRepository + + pool = _FakePool() + pool._fetchrow_result = _make_row() + repo = PgPresetRepository(pool_getter=lambda: pool) + + result = await repo.get("default", "indexation") + assert result is not None + assert result["name"] == "default" + assert result["config"] == {"chunk_size": 512} + + +@pytest.mark.asyncio +async def test_list_all_no_filter(): + from services.persistence.preset_repo import PgPresetRepository + + pool = _FakePool() + pool._fetch_result = [_make_row(), _make_row(name="legal")] + repo = PgPresetRepository(pool_getter=lambda: pool) + + results = await repo.list_all() + assert len(results) == 2 + query, params = pool.executed[0] + assert "ORDER BY preset_type, name" in query + assert params == () + + +@pytest.mark.asyncio +async def test_list_all_with_type_filter(): + from services.persistence.preset_repo import PgPresetRepository + + pool = _FakePool() + pool._fetch_result = [_make_row()] + repo = PgPresetRepository(pool_getter=lambda: pool) + + await repo.list_all(preset_type="retrieval") + query, params = pool.executed[0] + assert "WHERE preset_type" in query + assert params == ("retrieval",) + + +@pytest.mark.asyncio +async def test_upsert_uses_on_conflict(): + from services.persistence.preset_repo import PgPresetRepository + + pool = _FakePool() + pool._fetchrow_result = _make_row(config={"chunk_size": 256}) + repo = PgPresetRepository(pool_getter=lambda: pool) + + result = await repo.upsert("default", "indexation", {"chunk_size": 256}) + assert result["config"] == {"chunk_size": 256} + query, _ = pool.executed[0] + assert "ON CONFLICT" in query + assert "DO UPDATE" in query + + +@pytest.mark.asyncio +async def test_delete_returns_true_on_success(): + from services.persistence.preset_repo import PgPresetRepository + + pool = _FakePool() + repo = PgPresetRepository(pool_getter=lambda: pool) + + assert await repo.delete("default", "indexation") is True + + +@pytest.mark.asyncio +async def test_delete_returns_false_when_missing(): + from services.persistence.preset_repo import PgPresetRepository + + pool = _FakePool() + + async def _execute(query, *params): + pool.executed.append((query, params)) + return "DELETE 0" + + pool.execute = _execute + repo = PgPresetRepository(pool_getter=lambda: pool) + + assert await repo.delete("ghost", "indexation") is False + + +@pytest.mark.asyncio +async def test_rename_runs_delete_and_upsert_in_one_transaction(): + from services.persistence.preset_repo import PgPresetRepository + + pool = _FakePool() + pool._fetchrow_result = _make_row(name="new", config={"chunk_size": 256}) + repo = PgPresetRepository(pool_getter=lambda: pool) + + result = await repo.rename("old", "new", "indexation", {"chunk_size": 256}) + + assert result["name"] == "new" + operations = [query for query, _params in pool.executed] + assert operations[0] == "BEGIN" + assert "DELETE FROM pipeline_presets" in operations[1] + assert "INSERT INTO pipeline_presets" in operations[2] + assert operations[-1] == "COMMIT" + + +@pytest.mark.asyncio +async def test_count_partitions_using_indexation_queries_correct_column(): + from services.persistence.preset_repo import PgPresetRepository + + pool = _FakePool() + pool._fetchval_result = 3 + repo = PgPresetRepository(pool_getter=lambda: pool) + + count = await repo.count_partitions_using("default", "indexation") + assert count == 3 + query, params = pool.executed[0] + assert "indexation_preset" in query + assert params == ("default",) + + +@pytest.mark.asyncio +async def test_count_partitions_using_retrieval_queries_correct_column(): + from services.persistence.preset_repo import PgPresetRepository + + pool = _FakePool() + pool._fetchval_result = 5 + repo = PgPresetRepository(pool_getter=lambda: pool) + + count = await repo.count_partitions_using("default", "retrieval") + assert count == 5 + query, params = pool.executed[0] + assert "retrieval_preset" in query + assert params == ("default",) + + +@pytest.mark.asyncio +async def test_count_partitions_using_rejects_invalid_type(): + from services.persistence.preset_repo import PgPresetRepository + + pool = _FakePool() + repo = PgPresetRepository(pool_getter=lambda: pool) + + with pytest.raises(ValueError, match="Invalid preset_type"): + await repo.count_partitions_using("default", "bad-type") + assert pool.executed == [] diff --git a/tests/unit/services/workers/parsers/test_doc_serializer_bridge.py b/tests/unit/services/workers/parsers/test_doc_serializer_bridge.py new file mode 100644 index 000000000..9fbb2ebed --- /dev/null +++ b/tests/unit/services/workers/parsers/test_doc_serializer_bridge.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from typing import Any + +import pytest +from core.config.root import Settings +from core.models.document import Document, DocumentType +from langchain_core.documents.base import Document as LangChainDocument +from services.workers.parsers.doc_serializer_bridge import ( + INDEXATION_CONFIG_METADATA_KEY, + DocSerializerBridgeParser, +) + + +class _FakeLoader: + seen_config: Any = None + seen_metadata: dict[str, Any] | None = None + + def __init__(self, *, config: Any) -> None: + type(self).seen_config = config + + async def aload_document(self, *, file_path: str, metadata: dict | None = None, save_markdown: bool = False): + type(self).seen_metadata = dict(metadata or {}) + return LangChainDocument(page_content="hello", metadata=metadata or {}) + + +@pytest.mark.asyncio +async def test_bridge_disables_legacy_captioning_from_indexation_config(monkeypatch): + """Per-file Phase 14 config disables legacy loader image captioning.""" + + def fake_loader_classes(config): + return {".txt": _FakeLoader} + + monkeypatch.setattr("services.workers.parsers.legacy_loaders.get_loader_classes", fake_loader_classes) + + parser = DocSerializerBridgeParser( + Settings( + loader={ + "image_captioning": True, + "image_captioning_url": True, + } + ) + ) + document = Document( + filename="note.txt", + content_type=DocumentType.TEXT, + raw_bytes=b"hello", + metadata={ + "source": "note.txt", + INDEXATION_CONFIG_METADATA_KEY: {"enable_image_captioning": False}, + }, + ) + + await parser.parse(document) + + assert _FakeLoader.seen_config.loader.image_captioning is False + assert _FakeLoader.seen_config.loader.image_captioning_url is False + assert _FakeLoader.seen_metadata == {"source": "note.txt"} diff --git a/tests/unit/services/workers/parsers/test_marker_workers.py b/tests/unit/services/workers/parsers/test_marker_workers.py new file mode 100644 index 000000000..7ecba70f4 --- /dev/null +++ b/tests/unit/services/workers/parsers/test_marker_workers.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from services.workers.parsers import marker_workers + + +def _config(marker_num_gpus: float = 0.25): + """Build the minimal config shape consumed by Marker GPU selection.""" + return SimpleNamespace(loader=SimpleNamespace(marker_num_gpus=marker_num_gpus)) + + +def test_marker_num_gpus_uses_ray_cluster_resources_when_cuda_is_hidden(monkeypatch): + """Marker should request GPUs from Ray even when local CUDA is hidden.""" + monkeypatch.setattr(marker_workers.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(marker_workers.ray, "cluster_resources", lambda: {"GPU": 1.0}) + + assert marker_workers._marker_num_gpus(_config()) == 0.25 diff --git a/tests/unit/services/workers/test_dispatcher.py b/tests/unit/services/workers/test_dispatcher.py index 7da62c7bb..87bc976bc 100644 --- a/tests/unit/services/workers/test_dispatcher.py +++ b/tests/unit/services/workers/test_dispatcher.py @@ -118,6 +118,8 @@ async def test_dispatch_indexing_queues_worker_pool_task_and_records_ref() -> No user={"id": 42}, workspace_ids=["ws-1"], replace=True, + indexation_config={"parsing_strategy": "pymupdf"}, + embedder_name="embed-fast", ) assert task_id == "task-1" @@ -137,6 +139,8 @@ async def test_dispatch_indexing_queues_worker_pool_task_and_records_ref() -> No user={"id": 42}, workspace_ids=["ws-1"], replace=True, + indexation_config={"parsing_strategy": "pymupdf"}, + embedder_name="embed-fast", ) tsm.set_object_ref.remote.assert_called_once_with("task-1", {"ref": ref}) @@ -254,9 +258,13 @@ async def test_delete_file_cleans_database_before_vector_store() -> None: ) call_order = [] - workspace_repo.remove_file_from_all_workspaces = AsyncMock(side_effect=lambda *a, **k: call_order.append("workspace")) + workspace_repo.remove_file_from_all_workspaces = AsyncMock( + side_effect=lambda *a, **k: call_order.append("workspace") + ) document_repo.remove_file_from_partition = AsyncMock(side_effect=lambda *a, **k: call_order.append("document")) - vector_store.query_ids_by_filter = AsyncMock(return_value=["1", "2"], side_effect=lambda *a, **k: call_order.append("query") or ["1", "2"]) + vector_store.query_ids_by_filter = AsyncMock( + return_value=["1", "2"], side_effect=lambda *a, **k: call_order.append("query") or ["1", "2"] + ) vector_store.delete = AsyncMock(side_effect=lambda *a, **k: call_order.append("delete") or None) await dispatcher.delete_file("file-1", "tenant-a") diff --git a/tests/unit/services/workers/test_indexer_worker.py b/tests/unit/services/workers/test_indexer_worker.py index 2785e6a4f..7d2ae1a3a 100644 --- a/tests/unit/services/workers/test_indexer_worker.py +++ b/tests/unit/services/workers/test_indexer_worker.py @@ -8,6 +8,7 @@ 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.parsers.doc_serializer_bridge import INDEXATION_CONFIG_METADATA_KEY from services.workers.pipeline_builder import build_indexing_pipeline # --------------------------------------------------------------------------- @@ -115,6 +116,16 @@ def test_load_document_falls_back_to_filename_when_no_file_id(tmp_path: Path) -> assert doc.filename == "note.txt" +def test_load_document_attaches_internal_indexation_config(tmp_path: Path) -> None: + p = tmp_path / "note.txt" + p.write_bytes(b"hi") + config = {"enable_image_captioning": False} + + doc = _load_document(str(p), {"source": "note.txt"}, "p", indexation_config=config) + + assert doc.metadata[INDEXATION_CONFIG_METADATA_KEY] == config + + # --------------------------------------------------------------------------- # Tests — IndexerWorker.process_file # --------------------------------------------------------------------------- @@ -263,6 +274,32 @@ async def test_process_file_creates_catalog_record_after_successful_pipeline(tmp assert repo.update_calls == [] +@pytest.mark.asyncio +async def test_process_file_stores_indexation_config_snapshot_on_new_file(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, + ) + indexation_config = {"parsing_strategy": "pymupdf", "enable_image_captioning": False} + + await worker.process_file( + task_id="t-new", + path=str(path), + metadata={"file_id": "f1"}, + partition="p", + user={"id": 42}, + indexation_config=indexation_config, + ) + + assert repo.add_calls[0]["indexation_config"] == indexation_config + + @pytest.mark.asyncio async def test_process_file_updates_catalog_record_on_replace(tmp_path: Path) -> None: path = tmp_path / "doc.txt" @@ -296,6 +333,32 @@ async def test_process_file_updates_catalog_record_on_replace(tmp_path: Path) -> assert repo.add_calls == [] +@pytest.mark.asyncio +async def test_process_file_stores_indexation_config_snapshot_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, + ) + indexation_config = {"parsing_strategy": "marker", "enable_contextualization": True} + + await worker.process_file( + task_id="t-replace", + path=str(path), + metadata={"file_id": "f1"}, + partition="p", + replace=True, + indexation_config=indexation_config, + ) + + assert repo.update_calls[0]["indexation_config"] == indexation_config + + @pytest.mark.asyncio async def test_process_file_catalog_failure_sets_failed_state(tmp_path: Path) -> None: path = tmp_path / "doc.txt" diff --git a/tests/unit/services/workers/test_pipeline_builder.py b/tests/unit/services/workers/test_pipeline_builder.py index c6366b0a9..7c30cd59b 100644 --- a/tests/unit/services/workers/test_pipeline_builder.py +++ b/tests/unit/services/workers/test_pipeline_builder.py @@ -1,6 +1,7 @@ import pytest +from core.config.indexation_pipeline import IndexationPipelineConfig from core.models.chunk import Chunk -from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock +from core.models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock from services.workers.pipeline_builder import build_indexing_pipeline @@ -53,6 +54,24 @@ async def ensure_collection(self, name: str, dimension: int, **kwargs) -> None: self.ensure_calls.append((name, dimension)) +class FakeVLM: + def __init__(self) -> None: + self.calls: list[bytes] = [] + + async def caption_image(self, image_bytes: bytes, prompt: str | None = None) -> str: + self.calls.append(image_bytes) + return "caption" + + +class FakeContextualizer: + def __init__(self) -> None: + 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)) + return [chunk.model_copy(update={"text": f"ctx {chunk.text}", "context": "ctx"}) for chunk in chunks] + + @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") @@ -105,3 +124,96 @@ async def test_pipeline_stops_before_later_stages_when_a_stage_fails(): assert row["error"] == "chunk failed" assert vector_store.calls == [] assert "password" not in row + + +@pytest.mark.asyncio +async def test_pipeline_indexation_config_disables_caption_and_contextualization(): + document = Document(filename="note.txt", text="hello", partition="tenant-a") + processed = ProcessedDocument( + document_id=document.id, + text_blocks=[TextBlock(text="hello")], + images=[ImageBlock(image_bytes=b"png")], + ) + chunks = [Chunk(id="c1", text="hello", partition="tenant-a")] + vlm = FakeVLM() + contextualizer = FakeContextualizer() + pipeline = build_indexing_pipeline( + parser=FakeParser(processed), + chunker=FakeChunker(chunks), + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + vlm=vlm, + contextualizer=contextualizer, + indexation_config=IndexationPipelineConfig( + enable_image_captioning=False, + enable_contextualization=False, + ), + ) + + row = {"document": document, "partition": "tenant-a", "filename": "note.txt"} + await pipeline.run(row) + + assert vlm.calls == [] + assert contextualizer.calls == [] + assert row["stage"] == "stored" + + +@pytest.mark.asyncio +async def test_pipeline_row_indexation_config_selects_components(): + document = Document(filename="note.txt", text="hello", partition="tenant-a") + default_processed = ProcessedDocument(document_id="default", text_blocks=[TextBlock(text="default")]) + selected_processed = ProcessedDocument( + document_id="selected", + text_blocks=[TextBlock(text="selected")], + images=[ImageBlock(image_bytes=b"png")], + ) + selected_chunk = Chunk(id="selected", text="selected", partition="tenant-a") + selected_parser = FakeParser(selected_processed) + selected_chunker = FakeChunker([selected_chunk]) + selected_embedder = FakeEmbedder([[0.5]]) + selected_vlm = FakeVLM() + selected_contextualizer = FakeContextualizer() + parser_calls: list[str] = [] + chunker_calls: list[object] = [] + embedder_calls: list[str] = [] + vlm_calls: list[str] = [] + contextualizer_calls: list[str] = [] + + pipeline = build_indexing_pipeline( + parser=FakeParser(default_processed), + chunker=FakeChunker([Chunk(id="default", text="default")]), + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + parser_factory=lambda name: parser_calls.append(name) or selected_parser, + chunker_factory=lambda config: chunker_calls.append(config) or selected_chunker, + embedder_factory=lambda name: embedder_calls.append(name) or selected_embedder, + vlm_factory=lambda name: vlm_calls.append(name) or selected_vlm, + contextualizer_factory=lambda name: contextualizer_calls.append(name) or selected_contextualizer, + ) + row = { + "document": document, + "partition": "tenant-a", + "filename": "note.txt", + "embedder_name": "embed-fast", + "indexation_config": { + "parsing_strategy": "pymupdf", + "chunking": {"name": "recursive_splitter", "chunk_size": 128, "chunk_overlap_rate": 0.1}, + "enable_image_captioning": True, + "vlm": "vlm-fast", + "enable_contextualization": True, + "contextualization_llm": "llm-context", + }, + } + + await pipeline.run(row) + + assert parser_calls == ["pymupdf"] + assert chunker_calls[0].chunk_size == 128 + assert embedder_calls == ["embed-fast"] + assert vlm_calls == ["vlm-fast"] + assert contextualizer_calls == ["llm-context"] + assert selected_parser.calls == [document] + assert selected_chunker.calls == [(row["processed_document"], "tenant-a")] + assert selected_vlm.calls == [b"png"] + assert selected_contextualizer.calls[0][1:] == ("note.txt", "en") + assert row["chunks"][0].embedding == [0.5]