diff --git a/docs/content/docs/documentation/API.mdx b/docs/content/docs/documentation/API.mdx index 5e2a78f06..e800943bb 100644 --- a/docs/content/docs/documentation/API.mdx +++ b/docs/content/docs/documentation/API.mdx @@ -513,6 +513,73 @@ DELETE /presets/{preset_type}/{name} --- +### πŸ“ Prompt Library + +Every prompt the pipeline sends to a model is a stored, editable row rather than a bundled file. On first boot each type is seeded from its bundled template as that type's **default**; an admin can add named variants and select one per preset or per partition. + +All routes are prefixed with **`/prompts`** and require the **admin** role. `prompt_type` is one of `sys_prompt` | `spoken_style_answer` | `query_contextualizer` | `chunk_contextualizer` | `image_captioning` | `hyde` | `multi_query` | `topic_tagger`. + +**Resolution order** for a given type: the name selected for the request β†’ the type's global default β†’ the bundled template. A selection naming a prompt that no longer exists falls back to the default rather than failing. + +**Where a prompt is selected** β€” each setting lives with the thing it configures: + +| Prompt type | Selected on | +|---|---| +| `sys_prompt`, `spoken_style_answer` | partition β€” `generation_prompt_names` | +| `query_contextualizer`, `hyde`, `multi_query` | retrieval preset β€” `*_prompt_name` | +| `chunk_contextualizer`, `image_captioning`, `topic_tagger` | indexation preset β€” `*_prompt_name` | + +#### Create a prompt +```http +POST /prompts/ +``` +**Body:** `prompt_type`, `name`, `content`, `is_default` (default `false`). Returns `201 Created`, or `409` if that `(prompt_type, name)` already exists. + +Types rendered as templates (`sys_prompt`, `spoken_style_answer`, `query_contextualizer`, `hyde`, `multi_query`) accept only their own **plain** `{placeholders}` β€” no conversion (`!r`), format spec (`:>10`) or attribute access β€” and a violation returns `422` at write time rather than failing later at render. Escape a literal brace as `{{` / `}}`. The remaining types are sent to the model verbatim, so any text is valid. + +```bash frame="none" +curl -X POST http://localhost:8080/prompts/ \ + -H "Authorization: Bearer YOUR_AUTH_TOKEN" \ + -H "Content-Type: application/json" \ + -d '{ + "prompt_type": "sys_prompt", + "name": "legal-assistant", + "content": "Answer strictly from the context.\n{context}\nToday is {current_date}." + }' +``` + +#### List / Get / Update / Delete +```http +GET /prompts/ # ?prompt_type= to filter, ?offset= &limit= to page (limit ≀ 500) +GET /prompts/{prompt_id} +PATCH /prompts/{prompt_id} # any of name, content, is_default +DELETE /prompts/{prompt_id} # 204 No Content; refused for a type's current default +``` +List entries carry `used_by` β€” the number of partitions that resolve to that prompt, counting those that fall back to it as the default. + +#### Promote to default for its type +```http +PUT /prompts/{prompt_id}/default +``` +Clears the previous default for that type and promotes this one, atomically. + +#### Selecting a prompt for a partition +```http +PATCH /partition/{partition} +``` +Send `generation_prompt_names`, a map of `{prompt_type: name}` restricted to `sys_prompt` and `spoken_style_answer`. Each name must exist, or the request returns `422`. Send `{}` to clear the selection and fall back to the defaults. + +```bash frame="none" +curl -X PATCH http://localhost:8080/partition/my-partition \ + -H "Authorization: Bearer YOUR_AUTH_TOKEN" \ + -H "Content-Type: application/json" \ + -d '{"generation_prompt_names": {"sys_prompt": "legal-assistant"}}' +``` + +Preset-scoped prompts are selected the same way, by putting the `*_prompt_name` field in the preset's `config` (see [Pipeline Presets](#-pipeline-presets)). + +--- + ### πŸ”Œ Model Endpoints A registry of named inference endpoints (embedder, reranker, LLM, VLM) that partitions and presets can point at, so operators can manage and switch inference backends at runtime instead of via `.env`. Stored API keys are **redacted** in every response and only returned through the explicit reveal action below. diff --git a/openrag/api/main.py b/openrag/api/main.py index e3255a5d8..55f36f1ad 100644 --- a/openrag/api/main.py +++ b/openrag/api/main.py @@ -46,6 +46,7 @@ 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.prompts import router as prompts_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 @@ -117,6 +118,7 @@ class Tags(Enum): PARTITION = "Partitions & files" MODEL_ENDPOINTS = "Model Endpoints" PRESETS = "Presets" + PROMPTS = "Prompts" QUEUE = "Queue management" ACTORS = "Ray Actors" USERS = "User management" @@ -357,6 +359,7 @@ def get_config(): 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(prompts_router, prefix="/prompts", tags=[Tags.PROMPTS]) 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]) diff --git a/openrag/api/routers/admin/prompts.py b/openrag/api/routers/admin/prompts.py new file mode 100644 index 000000000..f7f2a14e0 --- /dev/null +++ b/openrag/api/routers/admin/prompts.py @@ -0,0 +1,77 @@ +"""Admin routes for the DB prompt library. + +Transport-only: auth, request validation, and response shaping live here; +persistence and resolution are delegated to ``PromptService`` from the DI +container. Per-partition assignment routes live alongside the other partition +sub-resources in ``partitions.py``. +""" + +from api.dependencies.auth import require_admin +from api.schemas.admin.prompt_schemas import ( + CreatePromptRequest, + PromptResponse, + PromptTypeName, + UpdatePromptRequest, +) +from di.providers import get_prompt_service +from fastapi import APIRouter, Depends, Query, Response, status + +router = APIRouter(dependencies=[Depends(require_admin)]) + + +@router.post("/", response_model=PromptResponse, status_code=status.HTTP_201_CREATED) +async def create_prompt( + body: CreatePromptRequest, + service=Depends(get_prompt_service), +): + """Add a prompt to the library.""" + return await service.create_prompt(**body.model_dump()) + + +@router.get("/", response_model=list[PromptResponse]) +async def list_prompts( + prompt_type: PromptTypeName | None = None, + offset: int = Query(0, ge=0), + limit: int = Query(100, ge=1, le=500), + service=Depends(get_prompt_service), +): + """List library prompts (optionally by type), each with an override count.""" + return await service.list_prompts(prompt_type=prompt_type, offset=offset, limit=limit) + + +@router.get("/{prompt_id}", response_model=PromptResponse) +async def get_prompt( + prompt_id: str, + service=Depends(get_prompt_service), +): + """Return one library prompt.""" + return await service.get_prompt(prompt_id) + + +@router.patch("/{prompt_id}", response_model=PromptResponse) +async def update_prompt( + prompt_id: str, + body: UpdatePromptRequest, + service=Depends(get_prompt_service), +): + """Edit a prompt's name/content and/or promote it to default.""" + return await service.update_prompt(prompt_id, **body.model_dump(exclude_unset=True)) + + +@router.put("/{prompt_id}/default", response_model=PromptResponse) +async def set_prompt_default( + prompt_id: str, + service=Depends(get_prompt_service), +): + """Promote a prompt to the default for its type.""" + return await service.set_default(prompt_id) + + +@router.delete("/{prompt_id}", status_code=status.HTTP_204_NO_CONTENT) +async def delete_prompt( + prompt_id: str, + service=Depends(get_prompt_service), +): + """Delete a library prompt (rejected if it is the current default).""" + await service.delete_prompt(prompt_id) + return Response(status_code=status.HTTP_204_NO_CONTENT) diff --git a/openrag/api/schemas/admin/partition_schemas.py b/openrag/api/schemas/admin/partition_schemas.py index 335b47a95..fbb99b401 100644 --- a/openrag/api/schemas/admin/partition_schemas.py +++ b/openrag/api/schemas/admin/partition_schemas.py @@ -62,6 +62,24 @@ class UpdatePartitionRequest(BaseModel): retrieval_preset: str | None = None chat_history_depth: int | None = Field(default=None, ge=1) chat_llm: str | None = None + # {prompt_type: library_prompt_name} for this partition's generation prompts. + # Keys are restricted to the generation types; ``{}`` clears all overrides. + generation_prompt_names: dict[str, str] | None = None + + @field_validator("generation_prompt_names") + @classmethod + def validate_generation_prompt_names(cls, value: dict[str, str] | None, info: ValidationInfo) -> dict[str, str]: + value = _reject_explicit_null(info.field_name, value) + # query_contextualizer is a query-side prompt selected on the retrieval + # preset, not a partition generation prompt (see RetrievalPipelineConfig). + allowed = {"sys_prompt", "spoken_style_answer"} + bad = set(value) - allowed + if bad: + raise ValueError(f"generation_prompt_names keys must be one of {sorted(allowed)}; got {sorted(bad)}") + for k, v in value.items(): + if not isinstance(v, str) or not v.strip(): + raise ValueError(f"generation_prompt_names['{k}'] must be a non-empty prompt name") + return value @field_validator("embedder", "indexation_preset", "retrieval_preset") @classmethod @@ -113,6 +131,7 @@ class PartitionDetailResponse(BaseModel): document_count: int = 0 chat_history_depth: int = 4 chat_llm: str | None = None + generation_prompt_names: dict[str, str] = Field(default_factory=dict) __all__ = [ diff --git a/openrag/api/schemas/admin/prompt_schemas.py b/openrag/api/schemas/admin/prompt_schemas.py new file mode 100644 index 000000000..d3ee99aa0 --- /dev/null +++ b/openrag/api/schemas/admin/prompt_schemas.py @@ -0,0 +1,98 @@ +"""Admin schemas for the DB prompt library.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Literal + +from pydantic import BaseModel, ConfigDict, ValidationInfo, field_validator, model_validator + +# The managed prompt types (mirrors core.models.prompt.PromptType). Declaring +# them as a Literal makes FastAPI reject unknown types at the transport edge +# (422) and renders the enum in the OpenAPI schema. +PromptTypeName = Literal[ + "sys_prompt", + "query_contextualizer", + "chunk_contextualizer", + "image_captioning", + "hyde", + "multi_query", + "spoken_style_answer", + "topic_tagger", +] + + +def _require_non_empty(field_name: str, value: str) -> str: + value = value.strip() + if not value: + raise ValueError(f"{field_name} must be non-empty") + return value + + +class CreatePromptRequest(BaseModel): + """Request body for adding a prompt to the library.""" + + model_config = ConfigDict(extra="forbid") + + prompt_type: PromptTypeName + name: str + content: str + is_default: bool = False + + @field_validator("name", "content") + @classmethod + def validate_non_empty(cls, value: str, info: ValidationInfo) -> str: + return _require_non_empty(info.field_name, value) + + +class UpdatePromptRequest(BaseModel): + """Request body for editing a prompt and/or promoting it to default.""" + + model_config = ConfigDict(extra="forbid") + + name: str | None = None + content: str | None = None + is_default: bool | None = None + + @field_validator("content") + @classmethod + def validate_content(cls, value: str | None, info: ValidationInfo) -> str | None: + if value is None: + raise ValueError(f"{info.field_name} cannot be null") + return _require_non_empty("content", value) + + @field_validator("name") + @classmethod + def validate_name(cls, value: str | None, info: ValidationInfo) -> str | None: + if value is None: + raise ValueError(f"{info.field_name} cannot be null") + return _require_non_empty(info.field_name, value) + + @model_validator(mode="after") + def require_at_least_one_update(self) -> UpdatePromptRequest: + if not any(getattr(self, f) is not None for f in ("name", "content", "is_default")): + raise ValueError("at least one field must be provided") + return self + + +class PromptResponse(BaseModel): + """A stored library prompt.""" + + id: str + prompt_type: str + name: str + content: str + is_default: bool + created_at: datetime + updated_at: datetime + # Number of partitions/presets referencing this prompt by name. Populated by + # the list endpoint; 0 on single-item responses where usage isn't computed. + used_by: int = 0 + + +__all__ = [ + "CreatePromptRequest", + "PromptResponse", + "PromptTypeName", + "UpdatePromptRequest", +] diff --git a/openrag/core/config/indexation_pipeline.py b/openrag/core/config/indexation_pipeline.py index 5b53c3908..0ce7aa0bf 100644 --- a/openrag/core/config/indexation_pipeline.py +++ b/openrag/core/config/indexation_pipeline.py @@ -42,9 +42,12 @@ class IndexationPipelineConfig(BaseModel): 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 + # Prompt selection: name a library prompt for this preset's enrichment + # stages (None = fall back to the type's global default, then the disk seed). + # Resolved per file in the indexer via PromptService.resolve_prompt. contextualization_prompt_name: str | None = None + image_captioning_prompt_name: str | None = None + topic_tagging_prompt_name: str | None = None # Entity extraction enable_entity_extraction: bool = True diff --git a/openrag/core/config/retrieval_pipeline.py b/openrag/core/config/retrieval_pipeline.py index 9030a96aa..26ea6e795 100644 --- a/openrag/core/config/retrieval_pipeline.py +++ b/openrag/core/config/retrieval_pipeline.py @@ -29,5 +29,16 @@ class RetrievalPipelineConfig(BaseModel): include_ancestors: bool = True rrf_k: int = Field(default=60, gt=0, le=1000) # Reciprocal Rank Fusion constant + # Prompt selection: name a library prompt for this preset's query-side + # prompts (None = the type's global default, then the disk seed). hyde / + # multi_query drive the query-expansion strategies (resolved per request in + # RetrievalService); query_contextualizer rewrites the user's query before + # retrieval (resolved in QueryService.generate_query). All three are + # query-side concerns, so they live on the retrieval preset rather than the + # partition's generation prompts. + hyde_prompt_name: str | None = None + multi_query_prompt_name: str | None = None + query_contextualizer_prompt_name: str | None = None + __all__ = ["RetrievalPipelineConfig"] diff --git a/openrag/core/indexing/contextualize.py b/openrag/core/indexing/contextualize.py index 897d528c8..fb0e28de0 100644 --- a/openrag/core/indexing/contextualize.py +++ b/openrag/core/indexing/contextualize.py @@ -59,9 +59,10 @@ async def _generate_context( current_chunk: Chunk, filename: str, lang: str, + system_prompt: str, ) -> str: messages = build_messages( - system_prompt=self._system_prompt, + system_prompt=system_prompt, filename=filename, first_chunks_text=[c.text for c in first_chunks], prev_chunks_text=[c.text for c in prev_chunks], @@ -85,6 +86,7 @@ async def contextualize( *, filename: str = "", lang: str = "en", + system_prompt: str | None = None, ) -> list[Chunk]: """Return new chunks with context prepended to ``text``. @@ -100,6 +102,10 @@ async def contextualize( if not chunks: return [] + # A per-call override (the DB-resolved prompt for this file's partition) + # wins over the instance default baked in at construction. + effective_prompt = system_prompt or self._system_prompt + try: first_chunks = chunks[:2] contexts: list[str] = [] @@ -114,6 +120,7 @@ async def contextualize( current_chunk=chunks[i], filename=filename, lang=lang, + system_prompt=effective_prompt, ) for i in range(start, end) ] diff --git a/openrag/core/indexing/topic_tags.py b/openrag/core/indexing/topic_tags.py index eeb67ff5a..84f082b55 100644 --- a/openrag/core/indexing/topic_tags.py +++ b/openrag/core/indexing/topic_tags.py @@ -38,8 +38,13 @@ async def tag( filename: str = "", max_tags: int = 7, lang: str = "en", + system_prompt: str | None = None, ) -> list[str]: - """Return normalized, unique topic tags for a document.""" + """Return normalized, unique topic tags for a document. + + ``system_prompt`` overrides the instance default (the DB-resolved prompt + for this file's partition) when provided. + """ chunks = list(chunks) if not chunks: return [] @@ -48,7 +53,7 @@ async def tag( try: messages = _build_messages( - system_prompt=self._system_prompt, + system_prompt=system_prompt or self._system_prompt, chunks=chunks, filename=filename, max_tags=max_tags, diff --git a/openrag/core/models/preset.py b/openrag/core/models/preset.py index ba35c7a09..13435cfbf 100644 --- a/openrag/core/models/preset.py +++ b/openrag/core/models/preset.py @@ -35,6 +35,10 @@ class PartitionRow(BaseModel): collection_name: str | None = None chat_history_depth: int = Field(default=4, ge=1) chat_llm: str | None = None + # {prompt_type: library_prompt_name} for generation prompts (sys_prompt, + # spoken_style_answer, query_contextualizer). Like chat_llm, generation + # config lives on the partition rather than a preset. + generation_prompt_names: dict[str, str] = Field(default_factory=dict) created_at: datetime updated_at: datetime @@ -54,6 +58,7 @@ class PartitionConfig(BaseModel): collection_name: str | None = None chat_history_depth: int = Field(default=4, ge=1) chat_llm: str | None = None + generation_prompt_names: dict[str, str] = Field(default_factory=dict) def resolve_partition_chat_llm( diff --git a/openrag/core/ports/prompt_repo.py b/openrag/core/ports/prompt_repo.py index 43336f2f5..57d56f4cb 100644 --- a/openrag/core/ports/prompt_repo.py +++ b/openrag/core/ports/prompt_repo.py @@ -1,32 +1,90 @@ -"""Prompt repository interface.""" +"""Prompt repository interface. + +Backs the DB prompt library: a global set of named prompt templates with at +most one ``is_default`` per type. Selection (which prompt a preset/partition +uses) is by name, resolved in ``PromptService`` (named prompt β†’ global default +β†’ disk seed); the repository only exposes the storage primitives. +""" from __future__ import annotations from abc import ABC, abstractmethod -from openrag.core.models.prompt import Prompt +from core.models.prompt import Prompt class PromptRepository(ABC): - """CRUD operations for prompt templates.""" + """CRUD + default-selection + per-partition override storage for prompts.""" + + # ------------------------------------------------------------------ + # Library CRUD + # ------------------------------------------------------------------ @abstractmethod - async def create_prompt(self, prompt: Prompt) -> Prompt: ... + async def create(self, prompt: Prompt) -> Prompt: ... @abstractmethod - async def get_prompt(self, prompt_id: str) -> Prompt | None: ... + async def get(self, prompt_id: str) -> Prompt | None: ... @abstractmethod - async def get_by_type(self, prompt_type: str) -> list[Prompt]: ... + async def list( + self, + *, + prompt_type: str | None = None, + offset: int = 0, + limit: int = 100, + ) -> list[Prompt]: ... @abstractmethod - async def get_active(self, prompt_type: str) -> Prompt | None: ... + async def count(self, *, prompt_type: str | None = None) -> int: ... @abstractmethod - async def list_prompts(self) -> list[Prompt]: ... + async def update(self, prompt_id: str, **fields: object) -> Prompt | None: + """Update whitelisted columns (``name``, ``content``). ``is_default`` is + deliberately not updatable here β€” flip it through :meth:`set_default`, + which clears the previous default in the same transaction (the partial + unique index forbids two defaults per type).""" + ... @abstractmethod - async def update_prompt(self, prompt_id: str, content: str) -> Prompt | None: ... + async def delete(self, prompt_id: str) -> bool: ... + + # ------------------------------------------------------------------ + # Selection by name + # ------------------------------------------------------------------ + + @abstractmethod + async def get_by_name(self, prompt_type: str, name: str) -> Prompt | None: + """Look up a library prompt by (type, name) β€” the selection primitive. + + Presets and partitions select a prompt by naming it; this resolves that + name to the stored prompt (``None`` if no such name exists for the type).""" + ... + + # ------------------------------------------------------------------ + # Usage counts and the global default (one per type) + # ------------------------------------------------------------------ @abstractmethod - async def delete_prompt(self, prompt_id: str) -> bool: ... + async def reference_counts(self) -> dict[tuple[str, str], int]: + """``{(prompt_type, name): partitions_resolving_to_it}`` in one bulk pass. + + Counts *effective* resolution: every partition resolves each prompt type + to a named prompt (when its partition/preset config names an existing one) + or the type's global default. So a default reflects the partitions that + fall back to it, not just those that name it explicitly. Per type the + counts sum to the partition total. Feeds the admin "used by N partitions" + annotation.""" + ... + + @abstractmethod + async def get_default(self, prompt_type: str) -> Prompt | None: ... + + @abstractmethod + async def set_default(self, prompt_id: str) -> Prompt | None: + """Promote ``prompt_id`` to the default for its type, atomically. + + Clears any existing default of the same type and sets this one inside a + single locked transaction. Returns the promoted row, or ``None`` if + ``prompt_id`` does not exist.""" + ... diff --git a/openrag/core/utils/exceptions.py b/openrag/core/utils/exceptions.py index e2fe80c3d..ff0d5e34d 100644 --- a/openrag/core/utils/exceptions.py +++ b/openrag/core/utils/exceptions.py @@ -86,10 +86,16 @@ def to_dict(self) -> dict: class ConfigError(OpenRAGError): - """Configuration-related errors.""" + """Configuration-related errors. - def __init__(self, message: str, **kwargs): - super().__init__(message, code="CONFIG_ERROR", status_code=500, **kwargs) + Accepts a custom ``code`` (same shape as :class:`ValidationError`) so a + caller can name a specific failure. Hard-coding it made ``code=`` collide + with the forwarded ``**kwargs`` and raise ``TypeError`` from the ``raise`` + statement itself, replacing the intended error with an unrelated one. + """ + + def __init__(self, message: str, *, code: str = "CONFIG_ERROR", **kwargs): + super().__init__(message, code=code, status_code=500, **kwargs) class RegistryError(OpenRAGError): diff --git a/openrag/di/container.py b/openrag/di/container.py index 28a8e8ae7..e18efaf8c 100644 --- a/openrag/di/container.py +++ b/openrag/di/container.py @@ -65,6 +65,7 @@ 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.prompt_service import PromptService from services.orchestrators.query_service import QueryService from services.orchestrators.retrieval_service import RetrievalService from services.orchestrators.user_service import UserService @@ -116,6 +117,7 @@ def __init__(self, settings: Settings | None = None) -> None: self._partition_service: PartitionService | None = None self._model_endpoint_service: ModelEndpointService | None = None self._preset_service: PresetService | None = None + self._prompt_service: PromptService | None = None self._workspace_service: WorkspaceService | None = None self._retrieval_service: RetrievalService | None = None self._query_service: QueryService | None = None @@ -210,6 +212,10 @@ async def initialize(self) -> None: 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) + # Prompts resolve request-time from the DB (no in-memory cache to + # load), so seeding the library from the bundled templates is the + # only startup step. + await self._initialize_step("seeding prompts", self.prompt_service.seed_defaults) 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 @@ -424,6 +430,7 @@ def partition_service(self) -> PartitionService: collection=settings.vectordb.collection_name, config=settings, task_state_manager_factory=get_task_state_manager, + prompt_repo=self.prompt_repo, ) return self._partition_service @@ -460,6 +467,18 @@ def preset_service(self) -> PresetService: ) return self._preset_service + @property + def prompt_service(self) -> PromptService: + """PromptService β€” DB-backed prompt library and per-partition overrides.""" + if self._prompt_service is None: + from services.orchestrators.prompt_service import PromptService + + self._prompt_service = PromptService( + prompt_repo=self.prompt_repo, + config=self._require_settings(), + ) + return self._prompt_service + @property def workspace_service(self) -> WorkspaceService: """WorkspaceService β€” lazily built, cached for the container's lifetime.""" @@ -535,6 +554,7 @@ def searcher_factory(embedder_name: str): searcher_factory=searcher_factory, reranker_factory=self.reranker_factory, llm_factory=self.llm_factory, + prompt_service=self.prompt_service, ) return self._retrieval_service @@ -569,6 +589,7 @@ def query_service(self) -> QueryService: config=settings, web_search_service=WebSearchFactory.create_service(settings), workspace_service=self.workspace_service, + prompt_service=self.prompt_service, llm_factory=self.llm_factory, ) return self._query_service diff --git a/openrag/di/providers.py b/openrag/di/providers.py index 866cccba7..7b63a0e55 100644 --- a/openrag/di/providers.py +++ b/openrag/di/providers.py @@ -149,6 +149,11 @@ def get_preset_service(request: Request = None) -> Any: return _get_optional_service(_require_initialized(request), "preset_service") +def get_prompt_service(request: Request = None) -> Any: + """Resolve the prompt-management orchestrator from the active container.""" + return _get_optional_service(_require_initialized(request), "prompt_service") + + def get_config(request: Request = None): """Resolve application configuration from the active container.""" return _require_initialized(request).config @@ -165,6 +170,7 @@ def get_config(request: Request = None): "get_model_endpoint_service", "get_partition_service", "get_preset_service", + "get_prompt_service", "get_query_service", "get_retrieval_service", "get_user_service", diff --git a/openrag/services/inference/_call_log.py b/openrag/services/inference/_call_log.py new file mode 100644 index 000000000..46bd2d33f --- /dev/null +++ b/openrag/services/inference/_call_log.py @@ -0,0 +1,152 @@ +"""DEBUG-level logging of the prompts that actually reach an LLM. + +``PromptService._log_resolution`` proves *which* library prompt each pipeline +stage resolved; this proves the resolved text is what landed in the outbound +request body β€” the other half of the wiring check. One ``llm.call`` line per +request, emitted only when ``LOG_LEVEL=DEBUG``, with every message previewed +rather than dumped so a context-stuffed chat (or a base64 image) cannot flood +the log. +""" + +from __future__ import annotations + +import json +import re +from typing import Any + +from core.utils.logging import get_logger + +logger = get_logger() + +# Long enough to recognise a prompt by its opening sentence, short enough that a +# retrieval context of a dozen documents stays one readable line. +PREVIEW_CHARS = 240 +# A long chat history or a multi-image request would otherwise join into one +# unbounded line, so the whole record is bounded too β€” not just each fragment. +MAX_MESSAGES = 12 +MAX_PARTS = 6 +MAX_DETAIL_CHARS = 2000 +# Identifiers interpolated into the message. ``model`` is client-controllable +# (metadata.llm_override), so it is flattened and bounded like the rest. +MAX_META_CHARS = 120 + +_EMAIL = re.compile(r"[^\s@]+@[^\s@]+\.[^\s@]+") + + +def _redact(text: str) -> str: + """Pseudonymize email addresses before any prompt text reaches a log sink. + + Prompts carry user questions and retrieved document context, both of which + routinely contain addresses; the diagnostic value of this line is the prompt + *shape*, never the personal data inside it. + """ + return _EMAIL.sub("", text) + + +def _clip(text: str, limit: int) -> str: + """Truncate to at most *limit* characters, ellipsis included. + + The ellipsis counts against the budget rather than being appended past it, + so every cap here is the real ceiling on the emitted length β€” otherwise each + clipped span silently ran one character over its limit. + """ + if limit <= 0: + return "" + return text if len(text) <= limit else f"{text[: limit - 1]}…" + + +def _preview(text: str) -> str: + return _clip(_redact(" ".join(text.split())), PREVIEW_CHARS) + + +def _meta(value: object) -> str: + """Flatten a value that is interpolated into the log message. + + Newlines in a client-supplied model name would otherwise let a caller forge + additional log lines, so every identifier is collapsed to one line and + bounded before it is formatted or bound. + """ + return _clip(" ".join(str(value).split()), MAX_META_CHARS) + + +def _render_content(content: Any) -> str: + """Flatten one message's ``content`` to a previewable string. + + Multimodal content arrives as a list of parts whose image entries carry a + base64 data URI β€” those are reduced to a type marker so image bytes never + reach the log. + """ + if isinstance(content, str): + return _preview(content) + if isinstance(content, list): + parts = [] + for part in content[:MAX_PARTS]: + if not isinstance(part, dict): + parts.append(_preview(str(part))) + elif part.get("type") == "text": + parts.append(_preview(str(part.get("text", "")))) + else: + parts.append(f"<{_meta(part.get('type', 'unknown'))}>") + if len(content) > MAX_PARTS: + parts.append(f"(+{len(content) - MAX_PARTS} more parts)") + return " + ".join(parts) + return _preview(json.dumps(content, ensure_ascii=False, default=str)) + + +def _describe(message: Any) -> str: + if not isinstance(message, dict): + return _preview(str(message)) + role = _meta(message.get("role", "?")) + content = message.get("content") + body = _render_content(content) + size = len(content) if isinstance(content, str) else len(body) + return f"{role}[{size}]: {body}" + + +def log_llm_call( + *, + caller: str, + model: str, + endpoint: str, + messages: list | None = None, + prompt: str | None = None, + stream: bool = False, +) -> None: + """Emit one ``llm.call`` line describing an outbound request. + + The previews are built inside a lazily-evaluated argument, so callers pay + nothing for this when the sink level is above DEBUG. Retries log per + attempt, which is deliberate β€” a retried call is a real second request. + + The whole record is bounded: per-message previews, a cap on how many + messages and multimodal parts are rendered, and a final clamp on the joined + result, so no request can turn one call into an unbounded log line. + """ + safe_caller, safe_model, safe_endpoint = _meta(caller), _meta(model), _meta(endpoint) + + def _detail() -> str: + if messages is not None: + rendered = [_describe(m) for m in messages[:MAX_MESSAGES]] + if len(messages) > MAX_MESSAGES: + rendered.append(f"(+{len(messages) - MAX_MESSAGES} more messages)") + return _clip(" || ".join(rendered), MAX_DETAIL_CHARS) + text = prompt or "" + return _clip(f"prompt[{len(text)}]: {_preview(text)}", MAX_DETAIL_CHARS) + + def _line() -> str: + return f"llm.call {safe_caller} model={safe_model} stream={stream} | {_detail()}" + + # The message is a single literal placeholder and everything else is built + # inside the lazy callable. loguru runs ``message.format(*args)``, so an + # identifier interpolated into the format string itself would have its + # braces parsed as format fields β€” and ``model`` is client-controlled via + # ``metadata.llm_override``, so a request naming a model ``gpt{x}`` raised + # ``KeyError`` out of the call path. The substituted value is never + # rescanned, so a brace anywhere in the rendered line is now inert. + # + # Passed positionally, not as a kwarg: loguru copies **kwargs into + # ``record["extra"]`` and the terminal formatter appends every extra, which + # would print the whole payload a second time on each line. + logger.bind(caller=safe_caller, model=safe_model, endpoint=safe_endpoint, stream=stream).opt(lazy=True).debug( + "{}", _line + ) diff --git a/openrag/services/inference/ollama_client.py b/openrag/services/inference/ollama_client.py index f38240825..b0af6bbdf 100644 --- a/openrag/services/inference/ollama_client.py +++ b/openrag/services/inference/ollama_client.py @@ -26,6 +26,7 @@ ) from core.utils.logging import get_logger +from ._call_log import log_llm_call from ._circuit_breaker import with_circuit_breaker from ._retry import with_retry from .vllm_client import _parse_response @@ -83,6 +84,7 @@ def __init__( async def generate(self, prompt: str, **kwargs) -> dict: payload = {**self._defaults, **kwargs, "model": self._model, "prompt": prompt} payload.pop("metadata", None) + log_llm_call(caller="OllamaClient.generate", model=self._model, endpoint=self._endpoint, prompt=prompt) try: resp = await self._client.post(f"{self._endpoint}/completions", json=payload) resp.raise_for_status() @@ -102,6 +104,7 @@ async def generate(self, prompt: str, **kwargs) -> dict: async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: payload = {**self._defaults, **kwargs, "model": self._model, "messages": messages, "stream": False} payload.pop("metadata", None) + log_llm_call(caller="OllamaClient.chat", model=self._model, endpoint=self._endpoint, messages=messages) try: resp = await self._client.post(f"{self._endpoint}/chat/completions", json=payload) resp.raise_for_status() @@ -119,6 +122,13 @@ async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIterator[str]: payload = {**self._defaults, **kwargs, "model": self._model, "messages": messages, "stream": True} payload.pop("metadata", None) + log_llm_call( + caller="OllamaClient.stream_chat", + model=self._model, + endpoint=self._endpoint, + messages=messages, + stream=True, + ) try: async with self._client.stream("POST", f"{self._endpoint}/chat/completions", json=payload) as resp: if resp.status_code >= 400: diff --git a/openrag/services/inference/vllm_client.py b/openrag/services/inference/vllm_client.py index ef5befaf3..0f8ff9c51 100644 --- a/openrag/services/inference/vllm_client.py +++ b/openrag/services/inference/vllm_client.py @@ -33,6 +33,7 @@ from core.vlm import VLM, vlm_registry from tqdm.asyncio import tqdm +from ._call_log import log_llm_call from ._circuit_breaker import with_circuit_breaker from ._retry import with_retry @@ -200,6 +201,7 @@ async def generate(self, prompt: str, **kwargs) -> dict: base_url, model, headers = self._resolve_overrides(kwargs) kwargs.pop("metadata", None) payload = {**self._defaults, **kwargs, "model": model, "prompt": prompt} + log_llm_call(caller="VLLMClient.generate", model=model, endpoint=base_url, prompt=prompt) try: resp = await self._client.post(f"{base_url}/completions", json=payload, headers=headers) resp.raise_for_status() @@ -220,6 +222,7 @@ async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: base_url, model, headers = self._resolve_overrides(kwargs) kwargs.pop("metadata", None) payload = {**self._chat_payload_kwargs(kwargs), "model": model, "messages": messages, "stream": False} + log_llm_call(caller="VLLMClient.chat", model=model, endpoint=base_url, messages=messages) try: resp = await self._client.post(f"{base_url}/chat/completions", json=payload, headers=headers) resp.raise_for_status() @@ -238,6 +241,7 @@ async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIt base_url, model, headers = self._resolve_overrides(kwargs) kwargs.pop("metadata", None) payload = {**self._chat_payload_kwargs(kwargs), "model": model, "messages": messages, "stream": True} + log_llm_call(caller="VLLMClient.stream_chat", model=model, endpoint=base_url, messages=messages, stream=True) try: async with self._client.stream( "POST", f"{base_url}/chat/completions", json=payload, headers=headers @@ -506,6 +510,7 @@ async def caption_image(self, image_bytes: bytes, prompt: str | None = None) -> "messages": messages, "max_tokens": self._max_tokens, } + log_llm_call(caller="VLLMVision.caption_image", model=self._model, endpoint=self._endpoint, messages=messages) try: resp = await self._client.post( f"{self._endpoint}/chat/completions", diff --git a/openrag/services/orchestrators/partition_service.py b/openrag/services/orchestrators/partition_service.py index bdab0d820..b95978562 100644 --- a/openrag/services/orchestrators/partition_service.py +++ b/openrag/services/orchestrators/partition_service.py @@ -102,8 +102,10 @@ def __init__( task_state_manager: Any = None, task_state_manager_factory: Callable[[], Any] | None = None, task_cancel_timeout: float = 60.0, + prompt_repo: Any = None, ) -> None: self._partition_repo = partition_repo + self._prompt_repo = prompt_repo self._membership_repo = membership_repo self._document_repo = document_repo self._vector_store = vector_store @@ -439,6 +441,8 @@ async def update_partition(self, partition: str, **fields: object) -> dict | Non # QueryService falls back to the default LLM for those at runtime. if updates.get("chat_llm"): self._validate_chat_llm_ref(updates["chat_llm"]) + if updates.get("generation_prompt_names"): + await self._validate_generation_prompt_names(updates["generation_prompt_names"]) result = await self._partition_repo.update_partition(partition, **updates) @@ -493,6 +497,22 @@ def _validate_chat_llm_ref(self, chat_llm: str) -> None: code="MODEL_ENDPOINT_NOT_FOUND", ) + async def _validate_generation_prompt_names(self, mapping: dict[str, str]) -> None: + """Assignment-time check: each named generation prompt must exist. + + Mirrors ``_validate_chat_llm_ref`` β€” guards assignment only; a stored + name can go stale later (the prompt may be deleted), which the resolver + tolerates by falling back to the global default at request time. + """ + if self._prompt_repo is None: + return + for prompt_type, name in mapping.items(): + if await self._prompt_repo.get_by_name(prompt_type, name) is None: + raise ValidationError( + f"No '{prompt_type}' prompt named '{name}' exists.", + code="PROMPT_NOT_FOUND", + ) + def _partition_detail(self, row: dict, cfg: PartitionConfig) -> dict: """Shape a resolved row into the ``PartitionDetailResponse`` payload.""" return { @@ -507,6 +527,7 @@ def _partition_detail(self, row: dict, cfg: PartitionConfig) -> dict: "created_at": row.get("created_at"), "chat_history_depth": row.get("chat_history_depth") or self._legacy_chat_history_depth_fallback(), "chat_llm": row.get("chat_llm"), + "generation_prompt_names": row.get("generation_prompt_names") or {}, } # ------------------------------------------------------------------ @@ -544,6 +565,7 @@ def resolve_partition_row(self, row: dict) -> PartitionConfig: # QueryService._resolve_chat_history_depth actually reads at chat time. chat_history_depth=row.get("chat_history_depth") or self._legacy_chat_history_depth_fallback(), chat_llm=row.get("chat_llm"), + generation_prompt_names=row.get("generation_prompt_names") or {}, ) async def load_partitions(self) -> None: diff --git a/openrag/services/orchestrators/prompt_service.py b/openrag/services/orchestrators/prompt_service.py new file mode 100644 index 000000000..944725764 --- /dev/null +++ b/openrag/services/orchestrators/prompt_service.py @@ -0,0 +1,389 @@ +"""PromptService β€” seeding, resolution, and CRUD for the DB prompt library. + +Orchestrates :class:`PromptRepository` to expose the prompt library to the +admin API and to answer the one question the rest of the system asks: + + resolve_prompt(prompt_type, names=[...]) -> str + +Selection is by name: a preset (indexation/retrieval) or a partition +(generation) names a library prompt per type. The caller passes the +precedence-ordered candidate names for a request; the first that resolves +wins, else the global default, else the on-disk seed template. Passing an +ordered list is the extension point β€” e.g. per-user personalization prepends a +user's prompt name ahead of the partition's without changing this signature. +""" + +from __future__ import annotations + +import string +from typing import TYPE_CHECKING + +from core.models.prompt import Prompt, PromptType +from core.prompts.template_loader import load_template_by_key +from core.utils.exceptions import ConfigError, NotFoundError, ValidationError +from core.utils.logging import get_logger + +if TYPE_CHECKING: + from collections.abc import Sequence + + from core.config.root import Settings + from core.ports.prompt_repo import PromptRepository + +logger = get_logger() + +_VALID_TYPES = frozenset(t.value for t in PromptType) + +# Prompt types whose content is a ``str.format`` template rendered on the hot +# path, mapped to the exact placeholders the pipeline substitutes. Content saved +# for these types MUST use only these ``{placeholders}`` (and escape any literal +# brace as ``{{``/``}}``), or the per-request ``.format(...)`` would raise and +# 500 the chat/retrieval path β€” globally if it's the type's default. Validated at +# write time (create/update) so an invalid template can never be stored. +# +# Types NOT listed here (chunk_contextualizer, image_captioning, topic_tagger) +# are sent to the LLM verbatim as a system message β€” never ``.format``-ed β€” so +# they may contain any literal text, braces included, and need no validation. +_PROMPT_FORMAT_FIELDS: dict[str, frozenset[str]] = { + PromptType.SYS_PROMPT.value: frozenset({"context", "current_date"}), + # Rendered by the same call site as sys_prompt (the answer prompt swapped in + # when a request sets metadata.spoken_style_answer), so it takes the same + # placeholders and must be validated identically. + PromptType.SPOKEN_STYLE_ANSWER.value: frozenset({"context", "current_date"}), + PromptType.QUERY_CONTEXTUALIZER.value: frozenset({"query_language", "current_date"}), + PromptType.HYDE.value: frozenset({"question"}), + PromptType.MULTI_QUERY.value: frozenset({"query", "k_queries"}), +} + + +def _validate_template(prompt_type: str, content: str) -> None: + """Reject a format-templated prompt whose ``{placeholders}`` are malformed or + unknown for its type. No-op for verbatim (non-formatted) prompt types. + + Only a *plain* placeholder is accepted β€” the field must be exactly one of the + type's known names, with no conversion (``!r``), format spec (``:>10``), or + attribute/index access (``ctx.attr``, ``ctx[0]``). Reducing such an + expression to its root name would let templates through that this check + calls valid and ``.format()`` then rejects: ``{context!x}`` raises + ``ValueError`` and ``{context.missing}`` raises ``AttributeError`` at render + time. As a type's global default, either would fail every request that falls + back to it β€” exactly what validating at write time exists to prevent. These + prompts are prose with a few injected values, so nothing legitimate is lost. + + Raises ``ValidationError`` (422) so the admin sees a precise message instead + of a later 500 on the chat path. + """ + allowed = _PROMPT_FORMAT_FIELDS.get(prompt_type) + if allowed is None: + return + try: + # Formatter.parse yields (literal, field_name, format_spec, conversion); + # field_name is None for literal text and for escaped {{/}}. It raises + # ValueError on an unbalanced single brace. + parsed = [(f, spec, conv) for _, f, spec, conv in string.Formatter().parse(content) if f is not None] + except ValueError as exc: + raise ValidationError( + f"Prompt template has malformed braces ({exc}). Escape a literal brace as '{{{{' or '}}}}'.", + status_code=422, + code="PROMPT_TEMPLATE_INVALID", + ) from exc + + for field, spec, conversion in parsed: + if conversion or spec: + raise ValidationError( + f"Prompt template placeholder '{{{field}}}' uses a conversion or format spec, " + "which is not supported. Use a plain placeholder such as " + f"'{{{field.split('!')[0].split(':')[0].split('.')[0].split('[')[0]}}}'.", + status_code=422, + code="PROMPT_TEMPLATE_INVALID", + ) + if "." in field or "[" in field: + raise ValidationError( + f"Prompt template placeholder '{{{field}}}' uses attribute or index access, " + "which is not supported. Use a plain placeholder.", + status_code=422, + code="PROMPT_TEMPLATE_INVALID", + ) + + unknown = {field for field, _, _ in parsed if field not in allowed} + if unknown: + raise ValidationError( + f"Prompt template uses unknown placeholder(s) {sorted(unknown)} for type " + f"'{prompt_type}'. Allowed: {sorted(allowed)} (escape a literal brace as '{{{{'/'}}}}').", + status_code=422, + code="PROMPT_TEMPLATE_INVALID", + ) + + +# Canonical ``prompt_type`` (a ``PromptType`` value, and the DB key) β†’ the +# ``PromptsConfig`` attribute the on-disk template loader looks up. Identity for +# every type except image captioning, whose config attribute is historically +# ``image_describer`` while its prompt type is ``image_captioning``. This map is +# the single reconciliation point between the DB type namespace and the disk +# filename namespace; keep it exhaustive over PromptType. +_TYPE_TO_CONFIG_KEY: dict[str, str] = { + PromptType.SYS_PROMPT.value: "sys_prompt", + PromptType.QUERY_CONTEXTUALIZER.value: "query_contextualizer", + PromptType.CHUNK_CONTEXTUALIZER.value: "chunk_contextualizer", + PromptType.IMAGE_CAPTIONING.value: "image_describer", + PromptType.HYDE.value: "hyde", + PromptType.MULTI_QUERY.value: "multi_query", + PromptType.SPOKEN_STYLE_ANSWER.value: "spoken_style_answer", + PromptType.TOPIC_TAGGER.value: "topic_tagger", +} + + +class PromptService: + """CRUD, resolution, and lifecycle for DB-backed prompts.""" + + def __init__(self, *, prompt_repo: PromptRepository, config: Settings) -> None: + self._repo = prompt_repo + self._config = config + + # ------------------------------------------------------------------ + # Startup lifecycle + # ------------------------------------------------------------------ + + async def seed_defaults(self) -> None: + """Create one default library prompt per type from its disk template. + + Idempotent per type: if a default already exists for a type it is left + untouched, so an admin's edits survive restarts. A type whose disk + template is missing is skipped with a warning rather than aborting boot. + """ + for prompt_type in _TYPE_TO_CONFIG_KEY: + if await self._repo.get_default(prompt_type) is not None: + continue + try: + content = self._disk_seed(prompt_type) + except (FileNotFoundError, ValueError) as exc: + logger.warning(f"No disk template to seed prompt type '{prompt_type}': {exc}") + continue + try: + # Seeding writes straight to the repo, so it would otherwise be + # the one path that stores content the CRUD API would reject. A + # bundled template with a bad placeholder must not become a + # type's global default: every request falling back to it would + # raise inside .format() at the point of use. + _validate_template(prompt_type, content) + except ValidationError as exc: + logger.warning(f"Bundled template for '{prompt_type}' is not a valid template; not seeding: {exc}") + continue + try: + await self._repo.create( + Prompt( + prompt_type=prompt_type, + name=f"default_{prompt_type}", + content=content, + is_default=True, + ) + ) + except ValidationError: + # Another replica seeded this type between the check at the top + # of the loop and this insert, and hit the unique index. Losing + # that race is a no-op, not a failure β€” but the 409 mapping makes + # it a ValidationError, which _initialize_step re-raises and + # ServiceContainer.initialize turns into a failed boot. Left + # unhandled, N replicas starting against an empty database + # crash-loop until one wins. + logger.info(f"Default prompt for '{prompt_type}' was seeded concurrently; skipping.") + continue + logger.info(f"Seeded default prompt for '{prompt_type}'.") + + def _disk_seed(self, prompt_type: str) -> str: + """Read a prompt type's bundled template from disk (honours PROMPTS_DIR).""" + config_key = _TYPE_TO_CONFIG_KEY[prompt_type] + return load_template_by_key(self._config.paths.prompts_dir, self._config.prompts, config_key) + + # ------------------------------------------------------------------ + # Resolution β€” the single seam + # ------------------------------------------------------------------ + + async def resolve_prompt(self, prompt_type: str, names: Sequence[str | None] | None = None) -> str: + """Resolve the effective prompt text for ``prompt_type``. + + Tries each candidate ``name`` in order (a preset- or partition-named + library prompt), then the global default, then the on-disk seed. + ``names`` entries may be ``None``/empty (skipped) so callers can pass + optional config values directly. + + Resolution happens per request, which put a Postgres round-trip on the + chat and search paths that did not exist before β€” prompts used to be read + from disk once at construction. A transient repository failure must + therefore not become a 500: lookups are treated as best-effort here and a + failure degrades to the bundled disk template, logged once. Errors are + swallowed at this single choke point rather than at each of the callers, + so chat, query expansion, retrieval and indexing all get the same + guarantee. + + Returns a string in every reachable case: boot seeds a default per type + and deleting a type's default is refused, so reaching the disk seed is + already an anomaly. If even that is unreadable there is no prompt to + return, and inventing one would silently degrade generation β€” so this + raises a typed :class:`ConfigError` naming the type instead of letting a + bare ``FileNotFoundError`` surface as an opaque 500. Callers that must + never fail (the ingest path) catch it and fall back to their own + disk-loaded prompt. + """ + candidates = [n for n in (names or ()) if n] + try: + for name in candidates: + prompt = await self._repo.get_by_name(prompt_type, name) + if prompt is not None: + self._log_resolution(prompt_type, candidates, "named", name, prompt.content) + return prompt.content + default = await self._repo.get_default(prompt_type) + if default is not None: + self._log_resolution(prompt_type, candidates, "default", default.name, default.content) + return default.content + except Exception as exc: # noqa: BLE001 - a DB blip must not fail the request + logger.warning(f"Prompt lookup failed for '{prompt_type}'; falling back to the bundled template: {exc}") + try: + content = self._disk_seed(prompt_type) + except (FileNotFoundError, ValueError, KeyError) as exc: + raise ConfigError( + f"No prompt available for type '{prompt_type}': no library default and " + f"no readable bundled template ({exc}).", + code="PROMPT_UNAVAILABLE", + ) from exc + self._log_resolution(prompt_type, candidates, "disk-seed", None, content) + return content + + @staticmethod + def _log_resolution(prompt_type: str, candidates: list[str], source: str, name: str | None, content: str) -> None: + """Emit one line per resolution so operators can confirm, in the logs, + exactly which library prompt each pipeline stage (indexation / + retrieval / chat) actually used, and preview its text. + + ``source`` is how it resolved: ``named`` (a partition/preset selection), + ``default`` (the type's global default), or ``disk-seed`` (bundled + fallback). ``candidates`` are the names the caller offered, in order. + + DEBUG, not INFO: this fires on every chat request and every indexing + job, and it carries prompt text. It pairs with the ``llm.call`` line + from the inference clients, which sits at the same level β€” turn on + ``LOG_LEVEL=DEBUG`` to see which prompt a stage picked *and* what + actually went to the model. The preview is built lazily so an INFO + deployment pays nothing for it. + """ + + def _line() -> str: + preview = repr(" ".join(content.split())[:80]) + return f"prompt.resolve {prompt_type} <- {source}{f':{name}' if name else ''} | {preview}" + + # Single literal placeholder, everything built inside the lazy callable: + # loguru runs ``message.format(...)``, so interpolating ``name`` into the + # format string made a brace in it a format field. Prompt names are free + # text, so a partition pointed at a prompt named ``my{tmpl}`` raised + # KeyError on *every* request that resolved it. + logger.bind( + prompt_type=prompt_type, + candidates=candidates, + source=source, + resolved_name=name, + length=len(content), + ).opt(lazy=True).debug("{}", _line) + + # ------------------------------------------------------------------ + # Library CRUD + # ------------------------------------------------------------------ + + async def create_prompt(self, *, prompt_type: str, name: str, content: str, is_default: bool = False) -> Prompt: + self._validate_type(prompt_type) + _validate_template(prompt_type, content) + if await self._repo.get_by_name(prompt_type, name) is not None: + raise ValidationError( + f"A '{prompt_type}' prompt named '{name}' already exists.", + status_code=409, + code="PROMPT_EXISTS", + ) + return await self._repo.create( + Prompt(prompt_type=prompt_type, name=name, content=content, is_default=is_default) + ) + + async def get_prompt(self, prompt_id: str) -> Prompt: + prompt = await self._repo.get(prompt_id) + if prompt is None: + raise NotFoundError(f"Prompt '{prompt_id}' not found.") + return prompt + + async def list_prompts(self, *, prompt_type: str | None = None, offset: int = 0, limit: int = 100) -> list[dict]: + """List prompts, each annotated with ``used_by`` β€” the number of + partitions/presets that reference it by name (one bulk aggregate).""" + if prompt_type is not None: + self._validate_type(prompt_type) + prompts = await self._repo.list(prompt_type=prompt_type, offset=offset, limit=limit) + counts = await self._repo.reference_counts() + return [{**p.model_dump(), "used_by": counts.get((p.prompt_type, p.name), 0)} for p in prompts] + + async def update_prompt(self, prompt_id: str, **fields: object) -> Prompt: + """Update ``name``/``content`` and/or promote to default. + + ``is_default=True`` is routed through the repo's atomic set_default + (clear-then-set) rather than a plain column write; a falsey value is a + no-op (you switch the default by promoting another prompt, never by + leaving the type with none). + """ + existing = await self._repo.get(prompt_id) + if existing is None: + raise NotFoundError(f"Prompt '{prompt_id}' not found.") + + new_content = fields.get("content") + if new_content is not None: + _validate_template(existing.prompt_type, str(new_content)) + + new_name = fields.get("name") + if new_name is not None and new_name != existing.name: + clash = await self._repo.get_by_name(existing.prompt_type, new_name) + if clash is not None and clash.id != prompt_id: + raise ValidationError( + f"A '{existing.prompt_type}' prompt named '{new_name}' already exists.", + status_code=409, + code="PROMPT_EXISTS", + ) + + promote_to_default = bool(fields.pop("is_default", None)) + updated = existing + if fields: + updated = await self._repo.update(prompt_id, **fields) or existing + if promote_to_default: + updated = await self._repo.set_default(prompt_id) or updated + return updated + + async def set_default(self, prompt_id: str) -> Prompt: + result = await self._repo.set_default(prompt_id) + if result is None: + raise NotFoundError(f"Prompt '{prompt_id}' not found.") + return result + + async def delete_prompt(self, prompt_id: str) -> None: + """Delete a library prompt. + + Refuses to delete a type's current default β€” removing it would strand + every prompt that resolves to the default on the disk seed and leave the + library with no default for that type. Promote another prompt first. + """ + existing = await self._repo.get(prompt_id) + if existing is None: + raise NotFoundError(f"Prompt '{prompt_id}' not found.") + if existing.is_default: + raise ValidationError( + f"Cannot delete the default '{existing.prompt_type}' prompt. Set another default first.", + ) + await self._repo.delete(prompt_id) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _validate_type(self, prompt_type: str) -> None: + if prompt_type not in _VALID_TYPES: + raise ValidationError( + f"Invalid prompt_type '{prompt_type}'. Must be one of: {sorted(_VALID_TYPES)}", + ) + + +__all__ = ["PromptService", "PROMPT_TYPE_KEYS"] + +# Public view of the canonical type set, for callers that enumerate managed +# prompt types without reaching into the private map. +PROMPT_TYPE_KEYS = tuple(_TYPE_TO_CONFIG_KEY) diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 152cebf7e..7c9287b23 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -48,7 +48,6 @@ SOURCE_SEPARATOR, format_context, format_web_context, - load_template_by_key, prepend_system_prompt, ) from core.utils.exceptions import WorkspaceNotFoundError @@ -65,6 +64,7 @@ if TYPE_CHECKING: from core.config.root import Settings from core.llm.llm import LLM + from services.orchestrators.prompt_service import PromptService from services.orchestrators.retrieval_service import RetrievalService from services.orchestrators.workspace_service import WorkspaceService @@ -117,6 +117,7 @@ def __init__( config: Settings, web_search_service: Any | None, workspace_service: WorkspaceService, + prompt_service: PromptService, llm_factory: Callable[[str], LLM] | None = None, ) -> None: self._retrieval = retrieval_service @@ -124,6 +125,9 @@ def __init__( self._llm_factory = llm_factory self._web = web_search_service self._workspace = workspace_service + # Prompts resolve request-time (override β†’ default β†’ disk seed) so an + # admin's edit takes effect on the next chat without a restart. + self._prompt_service = prompt_service # Keep a live reference so per-partition config (resolved into # ``config.partitions`` and refreshed on every preset change) can be @@ -146,11 +150,6 @@ def __init__( self._mr_expansion = mr.expansion_batch_size self._mr_max = mr.max_total_documents - prompts_dir, mapping = config.paths.prompts_dir, config.prompts - self._query_contextualizer_prompt = load_template_by_key(prompts_dir, mapping, "query_contextualizer") - self._spoken_style_answer_prompt = load_template_by_key(prompts_dir, mapping, "spoken_style_answer") - self._sys_prompt_tmplt = load_template_by_key(prompts_dir, mapping, "sys_prompt") - def _resolve_chat_history_depth(self, partition: list[str] | None) -> int: """Effective chat-history depth for this request. @@ -300,14 +299,53 @@ def _default_llm_name(self) -> str: # Query generation (was RagPipeline.generate_query β€” no LangChain) # ------------------------------------------------------------------ - async def generate_query(self, messages: list[dict], llm: LLM | None = None) -> SearchQueries: + def _generation_prompt_name(self, prompt_type: str, partition: list[str] | None) -> str | None: + """The library prompt this request's partition names for a generation type. + + Honoured only for a single owning partition (same rule as chat_llm / + chat_history_depth); multi-partition and the ``"all"`` sentinel resolve + the global default. Returned as the sole candidate name for + ``PromptService.resolve_prompt`` β€” a future per-user tier prepends ahead + of it. + """ + if not partition or "all" in partition or len(partition) != 1: + return None + cfg = self._config.partitions.get(partition[0]) + if cfg is None: + return None + return getattr(cfg, "generation_prompt_names", {}).get(prompt_type) + + def _retrieval_prompt_name(self, field: str, partition: list[str] | None) -> str | None: + """The library prompt this request's partition names on its retrieval + preset (query-side prompts: query_contextualizer / hyde / multi_query). + + Same single-owning-partition rule as generation prompts; multi-partition + and ``"all"`` resolve the global default. + """ + if not partition or "all" in partition or len(partition) != 1: + return None + cfg = self._config.partitions.get(partition[0]) + if cfg is None: + return None + return getattr(getattr(cfg, "retrieval", None), field, None) + + async def generate_query( + self, + messages: list[dict], + llm: LLM | None = None, + partition: list[str] | None = None, + ) -> SearchQueries: llm = llm or self._llm last_user = messages[-1]["content"] if RAGMODE(self._rag_mode) is RAGMODE.SIMPLERAG: return SearchQueries(query_list=[Query(query=last_user)]) chat_history = "".join(f"{m['role']}: {m['content']}\n" for m in messages) - prompt = self._query_contextualizer_prompt.format( + contextualizer = await self._prompt_service.resolve_prompt( + "query_contextualizer", + names=[self._retrieval_prompt_name("query_contextualizer_prompt_name", partition)], + ) + prompt = contextualizer.format( query_language=detect_language(last_user), current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S"), ) @@ -394,7 +432,7 @@ async def _batch(chunks: list, summaries: list) -> bool: async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: LLM | None = None): messages = payload["messages"][-self._resolve_chat_history_depth(partition) :] - queries = await self.generate_query(messages, llm=llm) + queries = await self.generate_query(messages, llm=llm, partition=partition) metadata = payload.get("metadata") or {} use_map_reduce = metadata.get("use_map_reduce", False) @@ -419,7 +457,14 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L force_retrieval = use_websearch or use_map_reduce if not queries.query_list: if not queries.requires_retrieval and not force_retrieval: - tmpl = self._spoken_style_answer_prompt if spoken_style else self._sys_prompt_tmplt + # Resolved per request from the library (named -> default -> + # bundled), replacing the __init__-time disk snapshots this + # branch removes. The conversational path therefore honours the + # partition's selected answer prompt too. + prompt_type = "spoken_style_answer" if spoken_style else "sys_prompt" + tmpl = await self._prompt_service.resolve_prompt( + prompt_type, names=[self._generation_prompt_name(prompt_type, partition)] + ) payload["messages"] = prepend_system_prompt( messages, tmpl, @@ -479,7 +524,10 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L web_results = [web_results[number - web_start_index] for number in web_source_numbers] new_messages = copy.deepcopy(messages) - tmpl = self._spoken_style_answer_prompt if spoken_style else self._sys_prompt_tmplt + prompt_type = "spoken_style_answer" if spoken_style else "sys_prompt" + tmpl = await self._prompt_service.resolve_prompt( + prompt_type, names=[self._generation_prompt_name(prompt_type, partition)] + ) new_messages.insert( 0, { @@ -507,7 +555,9 @@ async def _gather_rag_and_web(self, queries, partition, top_k, filter_params): async def _prepare_completions(self, partition: list[str], payload: dict, llm: LLM | None = None): prompt = payload["prompt"] - queries = await self.generate_query([{"role": "user", "content": prompt}], llm=llm) + # partition= is ours: the retrieval preset's query_contextualizer is + # resolved per partition. The skip below is from #807. + queries = await self.generate_query([{"role": "user", "content": prompt}], llm=llm, partition=partition) if not queries.query_list: if not queries.requires_retrieval: docs, context = [], "" @@ -524,8 +574,9 @@ async def _prepare_completions(self, partition: list[str], payload: dict, llm: L docs = [docs[i] for i in included] metadata = payload.get("metadata") or {} - tmpl = ( - self._spoken_style_answer_prompt if metadata.get("spoken_style_answer", False) else self._sys_prompt_tmplt + prompt_type = "spoken_style_answer" if metadata.get("spoken_style_answer", False) else "sys_prompt" + tmpl = await self._prompt_service.resolve_prompt( + prompt_type, names=[self._generation_prompt_name(prompt_type, partition)] ) instructions = tmpl.format( context=context, diff --git a/openrag/services/orchestrators/retrieval_service.py b/openrag/services/orchestrators/retrieval_service.py index eea37049a..06ad5f491 100644 --- a/openrag/services/orchestrators/retrieval_service.py +++ b/openrag/services/orchestrators/retrieval_service.py @@ -68,6 +68,7 @@ def __init__( searcher_factory: Callable[[str], RetrievalSearcher] | None = None, reranker_factory: Callable[[str], Reranker] | None = None, llm_factory: Callable[[str], LLM] | None = None, + prompt_service: Any | None = None, ) -> None: self._searcher = searcher self._config = config @@ -76,6 +77,10 @@ def __init__( self._searcher_factory = searcher_factory self._reranker_factory = reranker_factory self._llm_factory = llm_factory + # Resolves a preset's hyde/multi_query prompt by name (named -> default -> + # disk). Optional: when absent (e.g. unit tests, no DB), we fall back to + # the on-disk seed via load_template_by_key, preserving prior behaviour. + self._prompt_service = prompt_service self._pipeline = self._build_legacy_pipeline(reranker=reranker, llm=llm) logger.debug( @@ -131,27 +136,37 @@ def _build_retriever( llm: LLM | None, k_queries: int, combine: bool, + template: str | None = None, ): + # ``template`` is the already-resolved query-expansion prompt for this + # strategy (resolved from the preset's *_prompt_name in _pipeline_for_partition). if rtype == "multiQuery": return MultiQueryRetriever( llm=llm, - multi_query_template=load_template_by_key( - self._config.paths.prompts_dir, - self._config.prompts, - "multi_query", - ), + multi_query_template=template, 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"), + hyde_template=template, combine=combine, **common, ) return SingleRetriever(**common) + async def _resolve_query_template(self, prompt_type: str, name: str | None, disk_key: str) -> str: + """Resolve a query-side prompt (hyde / multi_query) to its text. + + Prefers the library (named preset prompt -> type default) via + PromptService; falls back to the on-disk seed when no PromptService is + wired (unit tests / DB-less runs), so behaviour matches the pre-DB path. + """ + if self._prompt_service is not None: + return await self._prompt_service.resolve_prompt(prompt_type, names=[name]) + return load_template_by_key(self._config.paths.prompts_dir, self._config.prompts, disk_key) + def _partition_configs(self) -> dict[str, Any]: return getattr(self._config, "partitions", {}) or {} @@ -229,7 +244,7 @@ def _default_reranker_name(self) -> str: return name return "default" - def _pipeline_for_partition(self, partition: str) -> tuple[RetrieverPipeline, int | None]: + async def _pipeline_for_partition(self, partition: str) -> tuple[RetrieverPipeline, int | None]: # Callers only ever pass a concrete partition name β€” the "all" sentinel is # expanded to concrete keys by _pipeline_groups_for_partitions before this # runs. With no per-partition configs at all, fall back to the legacy pipeline. @@ -247,10 +262,21 @@ def _pipeline_for_partition(self, partition: str) -> tuple[RetrieverPipeline, in 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") + # Only the expansion strategies need a prompt; type="single" (the common + # case) resolves nothing, so the DB is never touched on that path. + template = None + if rtype == "multiQuery": + template = await self._resolve_query_template( + "multi_query", pipeline_cfg.multi_query_prompt_name, "multi_query" + ) + elif rtype == "hyde": + template = await self._resolve_query_template("hyde", pipeline_cfg.hyde_prompt_name, "hyde") + reranker = self._resolve_reranker(pipeline_cfg.reranker, partition) if pipeline_cfg.enable_reranker else None retriever = self._build_retriever( rtype=rtype, + template=template, common={ "searcher": searcher, "top_k": pipeline_cfg.top_k, @@ -274,7 +300,7 @@ def _pipeline_for_partition(self, partition: str) -> tuple[RetrieverPipeline, in ) return pipeline, pipeline_cfg.top_n - def _pipeline_groups_for_partitions( + async def _pipeline_groups_for_partitions( self, partitions: list[str] ) -> list[tuple[list[str], RetrieverPipeline, int | None]]: configs = self._partition_configs() @@ -293,11 +319,11 @@ def _pipeline_groups_for_partitions( # Nothing to expand (no partitions exist yet) β€” keep the single # legacy pipeline; there is no per-partition config to honour. 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)] - ] + groups: list[tuple[list[str], RetrieverPipeline, int | None]] = [] + for partition in partitions: + pipeline, default_top_k = await self._pipeline_for_partition(partition) + groups.append(([partition], pipeline, default_top_k)) + return groups # ------------------------------------------------------------------ # Raw semantic search (powers routers/search.py β€” was indexer.asearch) @@ -386,6 +412,7 @@ async def retrieve( filter_params: dict | None = None, ) -> list[Chunk]: """Single ``Query`` through retrieve β†’ expand β†’ rerank.""" + groups = await self._pipeline_groups_for_partitions(partitions) ranked_lists = await self._gather_partition_groups( [ pipeline.retrieve_docs( @@ -394,7 +421,7 @@ async def retrieve( 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) + for partition_group, pipeline, default_top_k in groups ] ) return ranked_lists[0] if len(ranked_lists) == 1 else self.fuse(ranked_lists, top_k=top_k) @@ -408,6 +435,7 @@ async def retrieve_multi( filter_params: dict | None = None, ) -> list[Chunk]: """Every sub-query in parallel, fused with RRF.""" + groups = await self._pipeline_groups_for_partitions(partitions) ranked_lists = await self._gather_partition_groups( [ pipeline.get_relevant_docs( @@ -416,7 +444,7 @@ async def retrieve_multi( 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) + for partition_group, pipeline, default_top_k in groups ] ) return ranked_lists[0] if len(ranked_lists) == 1 else self.fuse(ranked_lists, top_k=top_k) diff --git a/openrag/services/persistence/migrations/alembic/versions/e8f9a0b1c2d3_add_prompts_library.py b/openrag/services/persistence/migrations/alembic/versions/e8f9a0b1c2d3_add_prompts_library.py new file mode 100644 index 000000000..240f18d6b --- /dev/null +++ b/openrag/services/persistence/migrations/alembic/versions/e8f9a0b1c2d3_add_prompts_library.py @@ -0,0 +1,101 @@ +"""add_prompts_library + +Adds DB-backed prompt management: + +* ``prompts`` β€” the global prompt library (the menu of named prompts), with a + partial-unique index enforcing at most one ``is_default`` row per prompt type. +* ``partitions.generation_prompt_names`` β€” a JSONB map ``{prompt_type: name}`` + naming the library prompt a partition uses for each generation prompt + (``sys_prompt``, ``spoken_style_answer``). Indexation/retrieval prompts are + named on their presets (JSONB config, no schema change). + +Idempotent: every op is guarded by an inspector check so re-application against +a database that already contains these objects (an older or partially-migrated +deployment) is a safe no-op β€” matching the guarded style of the other +migrations in this tree. + +Revision ID: e8f9a0b1c2d3 +Revises: e5f6a7b8c9d0 +Create Date: 2026-07-24 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from schema_helpers import column_exists, table_exists +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision: str = "e8f9a0b1c2d3" +# Chained after e5f6a7b8c9d0 rather than its original parent d4e5f6a7b8c9: +# that revision landed on develop while this branch was open and took the same +# parent, and two siblings would leave the tree with multiple heads β€” which +# aborts ``alembic upgrade head`` at boot and takes the whole app down, not just +# this feature. +down_revision: str | Sequence[str] | None = "e5f6a7b8c9d0" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +# Kept in lockstep with schema._PROMPT_TYPE_VALUES / core.models.prompt.PromptType. +_PROMPT_TYPE_IN = ( + "prompt_type IN (" + "'sys_prompt','query_contextualizer','chunk_contextualizer','image_captioning'," + "'hyde','multi_query','spoken_style_answer','topic_tagger')" +) + + +def upgrade() -> None: + if not table_exists("prompts"): + op.create_table( + "prompts", + sa.Column("id", sa.String(), nullable=False), + sa.Column("prompt_type", sa.String(), nullable=False), + sa.Column("name", sa.String(), server_default=sa.text("''"), nullable=False), + sa.Column("content", sa.String(), 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(_PROMPT_TYPE_IN, name="ck_prompt_type"), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index("ix_prompts_type", "prompts", ["prompt_type"]) + # Name is the selection key β€” unique per type. + op.create_index("uix_prompts_type_name", "prompts", ["prompt_type", "name"], unique=True) + # At most one global default per type. + op.create_index( + "uix_prompts_default_per_type", + "prompts", + ["prompt_type"], + unique=True, + postgresql_where=sa.text("is_default = true"), + ) + + if not column_exists("partitions", "generation_prompt_names"): + op.add_column( + "partitions", + sa.Column( + "generation_prompt_names", + postgresql.JSONB(astext_type=sa.Text()), + server_default=sa.text("'{}'::jsonb"), + nullable=False, + ), + ) + + +def downgrade() -> None: + if column_exists("partitions", "generation_prompt_names"): + op.drop_column("partitions", "generation_prompt_names") + if table_exists("prompts"): + op.drop_table("prompts") diff --git a/openrag/services/persistence/partition_repo.py b/openrag/services/persistence/partition_repo.py index 1931f2f42..8fbcfeb13 100644 --- a/openrag/services/persistence/partition_repo.py +++ b/openrag/services/persistence/partition_repo.py @@ -55,6 +55,7 @@ "collection_name", "chat_history_depth", "chat_llm", + "generation_prompt_names", } ) _PARTITION_OPERATION_LOCK_NAMESPACE = 20260720 @@ -352,7 +353,7 @@ async def _update_partition_on_conn( sets: list[str] = [] for col, val in updates.items(): idx = len(params) + 1 - sets.append(f"{col} = ${idx}") + sets.append(f"{col} = ${idx}::jsonb" if col == "generation_prompt_names" else f"{col} = ${idx}") params.append(val) sql = f"UPDATE partitions SET {', '.join(sets)}, updated_at = now() WHERE partition = $1 RETURNING *" @@ -437,6 +438,7 @@ def _row_to_full_dict(row: asyncpg.Record) -> dict: "collection_name": row["collection_name"], "chat_history_depth": row["chat_history_depth"], "chat_llm": row["chat_llm"], + "generation_prompt_names": row["generation_prompt_names"], "created_at": row["created_at"], "updated_at": row["updated_at"], } diff --git a/openrag/services/persistence/prompt_repo.py b/openrag/services/persistence/prompt_repo.py index 6c143546b..ea5a3a1af 100644 --- a/openrag/services/persistence/prompt_repo.py +++ b/openrag/services/persistence/prompt_repo.py @@ -1,43 +1,305 @@ -"""Stub :class:`PromptRepository`. - -Prompts are currently disk-based templates (``components/prompts/``). -The post-refactoring P1 feature is DB-stored, per-partition, -versionable prompts that operators can edit without redeploying. When -that lands the on-disk templates become the seed for the table and -this stub becomes a real asyncpg implementation against a ``prompts`` -table. +"""asyncpg-backed :class:`PromptRepository`. + +Manages the ``prompts`` library table. Replaces the earlier stub. Effective +resolution (named prompt β†’ default β†’ disk seed) is the service's job; this +layer is pure storage plus the one invariant that needs SQL-level atomicity: +at most one ``is_default`` row per type β€” held by a partial unique index and +enforced by clear-then-set inside a locked transaction (:meth:`set_default`, +and the default branch of :meth:`create`). """ from __future__ import annotations +from collections.abc import Callable + +import asyncpg from core.models.prompt import Prompt from core.ports.prompt_repo import PromptRepository -from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented +from core.utils.exceptions import ValidationError + +# Only these columns are editable in place. ``is_default`` is excluded on +# purpose: a bare ``UPDATE ... SET is_default = true`` cannot clear the previous +# default in the same statement, so it would collide with the partial unique +# index. Promotion goes through set_default, which clears-then-sets under a lock. +# ``prompt_type`` is immutable β€” retyping a prompt is nonsensical; create a new +# one instead. +_ALLOWED_UPDATE_FIELDS = frozenset({"name", "content"}) + +_COLS = ("id", "prompt_type", "name", "content", "is_default", "created_at", "updated_at") +_SELECT_COLS = ", ".join(_COLS) + + +def _as_conflict(exc: asyncpg.UniqueViolationError, prompt_type: str) -> ValidationError: + """Translate a unique-index violation into the 409 the service intends. + + The service checks for a name clash before writing, but that check and the + write are not one atomic step: two concurrent admins creating (or renaming + to) the same name both pass it, and the loser hits the index. Without this + the loser gets a 500 from the generic exception handler instead of the same + 409 the sequential path returns. Mirrors PgPartitionRepository.create. + """ + if exc.constraint_name == "uix_prompts_default_per_type": + return ValidationError( + f"Another '{prompt_type}' prompt was made the default concurrently; retry.", + status_code=409, + code="PROMPT_DEFAULT_CONFLICT", + ) + return ValidationError( + f"A '{prompt_type}' prompt with that name already exists.", + status_code=409, + code="PROMPT_EXISTS", + ) + + +# Indexation/retrieval preset config field -> the prompt_type it names. Partition +# generation_prompt_names keys ARE prompt_type values, so they need no mapping. +_PRESET_FIELD_TO_TYPE = { + "contextualization_prompt_name": "chunk_contextualizer", + "image_captioning_prompt_name": "image_captioning", + "topic_tagging_prompt_name": "topic_tagger", + "hyde_prompt_name": "hyde", + "multi_query_prompt_name": "multi_query", + "query_contextualizer_prompt_name": "query_contextualizer", +} + + +class PgPromptRepository(PromptRepository): + """asyncpg-backed implementation of :class:`PromptRepository`.""" + + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter + + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() + + @staticmethod + def _to_model(row: asyncpg.Record) -> Prompt: + return Prompt( + id=row["id"], + prompt_type=row["prompt_type"], + name=row["name"], + content=row["content"], + is_default=row["is_default"], + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + + # ------------------------------------------------------------------ + # Library CRUD + # ------------------------------------------------------------------ + + async def create(self, prompt: Prompt) -> Prompt: + # A bare INSERT with is_default=true would collide with the partial + # unique index if a default already exists for the type, so demote the + # current default in the SAME transaction as the insert β€” the new prompt + # becomes the sole default atomically. + async with self.pool.acquire() as conn: + async with conn.transaction(): + if prompt.is_default: + await conn.execute( + "UPDATE prompts SET is_default = false, updated_at = now() " + "WHERE prompt_type = $1 AND is_default = true", + prompt.prompt_type, + ) + try: + rec = await conn.fetchrow( + f""" + INSERT INTO prompts (id, prompt_type, name, content, is_default) + VALUES ($1, $2, $3, $4, $5) + RETURNING {_SELECT_COLS} + """, + prompt.id, + prompt.prompt_type, + prompt.name, + prompt.content, + prompt.is_default, + ) + except asyncpg.UniqueViolationError as exc: + raise _as_conflict(exc, prompt.prompt_type) from exc + return self._to_model(rec) + + async def get(self, prompt_id: str) -> Prompt | None: + rec = await self.pool.fetchrow( + f"SELECT {_SELECT_COLS} FROM prompts WHERE id = $1", + prompt_id, + ) + return self._to_model(rec) if rec else None + + async def list( + self, + *, + prompt_type: str | None = None, + offset: int = 0, + limit: int = 100, + ) -> list[Prompt]: + if prompt_type is not None: + rows = await self.pool.fetch( + f"SELECT {_SELECT_COLS} FROM prompts WHERE prompt_type = $1 " + "ORDER BY prompt_type, name, created_at OFFSET $2 LIMIT $3", + prompt_type, + offset, + limit, + ) + else: + rows = await self.pool.fetch( + f"SELECT {_SELECT_COLS} FROM prompts ORDER BY prompt_type, name, created_at OFFSET $1 LIMIT $2", + offset, + limit, + ) + return [self._to_model(r) for r in rows] + + async def count(self, *, prompt_type: str | None = None) -> int: + if prompt_type is not None: + return await self.pool.fetchval( + "SELECT count(*) FROM prompts WHERE prompt_type = $1", + prompt_type, + ) + return await self.pool.fetchval("SELECT count(*) FROM prompts") + + async def update(self, prompt_id: str, **fields: object) -> Prompt | None: + updates = {k: v for k, v in fields.items() if k in _ALLOWED_UPDATE_FIELDS} + if not updates: + return await self.get(prompt_id) + + params: list = [prompt_id] + sets: list[str] = [] + for col, val in updates.items(): + params.append(val) + sets.append(f"{col} = ${len(params)}") + + try: + rec = await self.pool.fetchrow( + f"UPDATE prompts SET {', '.join(sets)}, updated_at = now() WHERE id = $1 RETURNING {_SELECT_COLS}", + *params, + ) + except asyncpg.UniqueViolationError as exc: + # A rename racing another rename/create onto the same name. + existing = await self.get(prompt_id) + raise _as_conflict(exc, existing.prompt_type if existing else "") from exc + return self._to_model(rec) if rec else None + + async def delete(self, prompt_id: str) -> bool: + # Presets/partitions reference prompts by *name* (soft refs in JSONB), so + # there is no FK cascade: a deleted prompt's stale references simply + # resolve to the global default. The service guards against deleting a + # default; callers surface usage counts before offering delete. + result = await self.pool.execute("DELETE FROM prompts WHERE id = $1", prompt_id) + return result == "DELETE 1" + # ------------------------------------------------------------------ + # Global default (one per type) + # ------------------------------------------------------------------ -class PgPromptRepository(_StubRepositoryBase, PromptRepository): - """TODO: real impl once the ``prompts`` table is added.""" + async def get_by_name(self, prompt_type: str, name: str) -> Prompt | None: + rec = await self.pool.fetchrow( + f"SELECT {_SELECT_COLS} FROM prompts WHERE prompt_type = $1 AND name = $2", + prompt_type, + name, + ) + return self._to_model(rec) if rec else None - async def create_prompt(self, prompt: Prompt) -> Prompt: - raise stub_not_implemented("DB-stored prompts") + async def get_default(self, prompt_type: str) -> Prompt | None: + rec = await self.pool.fetchrow( + f"SELECT {_SELECT_COLS} FROM prompts WHERE prompt_type = $1 AND is_default = true", + prompt_type, + ) + return self._to_model(rec) if rec else None - async def get_prompt(self, prompt_id: str) -> Prompt | None: - raise stub_not_implemented("DB-stored prompts") + async def set_default(self, prompt_id: str) -> Prompt | None: + """Promote ``prompt_id`` to the default for its type, atomically. - async def get_by_type(self, prompt_type: str) -> list[Prompt]: - raise stub_not_implemented("DB-stored prompts") + Locks the type's rows (FOR UPDATE) and confirms ``prompt_id`` still + exists *inside* the transaction before clearing the old default, so a + concurrent delete of ``prompt_id`` can't make the final UPDATE match 0 + rows after the previous default was already cleared β€” which would leave + the type with no default. Same invariant PgModelEndpointRepository + protects. + """ + async with self.pool.acquire() as conn: + async with conn.transaction(): + target = await conn.fetchrow("SELECT prompt_type FROM prompts WHERE id = $1", prompt_id) + if target is None: + return None + prompt_type = target["prompt_type"] + locked = await conn.fetch( + "SELECT id FROM prompts WHERE prompt_type = $1 FOR UPDATE", + prompt_type, + ) + if prompt_id not in {r["id"] for r in locked}: + return None + await conn.execute( + "UPDATE prompts SET is_default = false, updated_at = now() " + "WHERE prompt_type = $1 AND is_default = true", + prompt_type, + ) + rec = await conn.fetchrow( + f"UPDATE prompts SET is_default = true, updated_at = now() WHERE id = $1 RETURNING {_SELECT_COLS}", + prompt_id, + ) + return self._to_model(rec) - async def get_active(self, prompt_type: str) -> Prompt | None: - raise stub_not_implemented("DB-stored prompts") + async def reference_counts(self) -> dict[tuple[str, str], int]: + # Effective resolution count: every partition resolves each prompt type to + # a named library prompt (when its partition/preset config names an + # existing one) or, failing that, to the type's global default. We count + # that resolution β€” so a default correctly shows the partitions that fall + # back to it, not just the (usually zero) partitions that name it + # explicitly. Per type the counts sum to the partition total. + total_partitions = await self.pool.fetchval("SELECT count(*)::int FROM partitions") - async def list_prompts(self) -> list[Prompt]: - raise stub_not_implemented("DB-stored prompts") + prompt_rows = await self.pool.fetch("SELECT prompt_type, name, is_default FROM prompts") + existing = {(r["prompt_type"], r["name"]) for r in prompt_rows} + default_name: dict[str, str] = {r["prompt_type"]: r["name"] for r in prompt_rows if r["is_default"]} - async def update_prompt(self, prompt_id: str, content: str) -> Prompt | None: - raise stub_not_implemented("DB-stored prompts") + # Explicit overrides: partitions that name a prompt directly (generation + # prompts on the partition JSONB) or transitively (their active + # indexation/retrieval preset's *_prompt_name). + overrides: dict[tuple[str, str], int] = {} + part_rows = await self.pool.fetch( + """ + SELECT j.key AS prompt_type, j.value AS name, count(*)::int AS n + FROM partitions p, jsonb_each_text(p.generation_prompt_names) j + WHERE j.value <> '' + GROUP BY 1, 2 + """ + ) + for r in part_rows: + overrides[(r["prompt_type"], r["name"])] = overrides.get((r["prompt_type"], r["name"]), 0) + r["n"] + # count(DISTINCT partition) so a partition is counted once per prompt even + # if two of its presets happened to name it. + preset_rows = await self.pool.fetch( + """ + SELECT c.key AS field, c.value AS name, count(DISTINCT part.partition)::int AS n + FROM partitions part + JOIN pipeline_presets pre + ON (pre.preset_type = 'indexation' AND pre.name = part.indexation_preset) + OR (pre.preset_type = 'retrieval' AND pre.name = part.retrieval_preset) + CROSS JOIN LATERAL jsonb_each_text(pre.config) c + WHERE c.value <> '' + GROUP BY 1, 2 + """ + ) + for r in preset_rows: + prompt_type = _PRESET_FIELD_TO_TYPE.get(r["field"]) + if prompt_type: + key = (prompt_type, r["name"]) + overrides[key] = overrides.get(key, 0) + r["n"] - async def delete_prompt(self, prompt_id: str) -> bool: - raise stub_not_implemented("DB-stored prompts") + # A valid override (names an existing prompt) credits that prompt; a + # dangling one falls through. Each type's default then absorbs every + # partition that didn't validly override it. + counts: dict[tuple[str, str], int] = {} + valid_overrides: dict[str, int] = {} + for (prompt_type, name), n in overrides.items(): + if (prompt_type, name) in existing: + counts[(prompt_type, name)] = counts.get((prompt_type, name), 0) + n + valid_overrides[prompt_type] = valid_overrides.get(prompt_type, 0) + n + for prompt_type, d_name in default_name.items(): + fallback = max(0, (total_partitions or 0) - valid_overrides.get(prompt_type, 0)) + if fallback: + counts[(prompt_type, d_name)] = counts.get((prompt_type, d_name), 0) + fallback + return counts __all__ = ["PgPromptRepository"] diff --git a/openrag/services/persistence/schema.py b/openrag/services/persistence/schema.py index 78b6cfd66..0f56f1472 100644 --- a/openrag/services/persistence/schema.py +++ b/openrag/services/persistence/schema.py @@ -92,6 +92,62 @@ ) +# The 8 canonical prompt types β€” kept in sync with core.models.prompt.PromptType. +# Used by the CHECK constraint on the prompts table so a junk type can never be +# stored. (Mirrors the ck_model_endpoint_type / ck_pipeline_preset_type pattern.) +_PROMPT_TYPE_VALUES = ( + "sys_prompt", + "query_contextualizer", + "chunk_contextualizer", + "image_captioning", + "hyde", + "multi_query", + "spoken_style_answer", + "topic_tagger", +) +_PROMPT_TYPE_IN = "prompt_type IN (" + ",".join(f"'{v}'" for v in _PROMPT_TYPE_VALUES) + ")" + + +prompts = Table( + "prompts", + metadata, + # String (not native UUID) so the column round-trips 1:1 with the + # ``Prompt.id: str`` domain model without asyncpg UUID<->str coercion. + Column("id", String, primary_key=True), + Column("prompt_type", String, nullable=False), + Column("name", String, server_default=text("''"), nullable=False), + Column("content", String, 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(_PROMPT_TYPE_IN, name="ck_prompt_type"), + Index("ix_prompts_type", "prompt_type"), + # Name is the selection key: presets and partitions reference a prompt by + # (type, name), so it must be unique per type for get_by_name to be + # deterministic. + Index("uix_prompts_type_name", "prompt_type", "name", unique=True), + # At most one global default per type β€” the DB-level guardrail behind + # PromptService.set_default's clear-then-set (same shape as the model + # endpoint default invariant, enforced there in application code). + Index( + "uix_prompts_default_per_type", + "prompt_type", + unique=True, + postgresql_where=text("is_default = true"), + ), +) + + partitions = Table( "partitions", metadata, @@ -106,6 +162,15 @@ Column("collection_name", String, nullable=True), Column("chat_history_depth", Integer, server_default="0", nullable=False), Column("chat_llm", String, nullable=True), + # {prompt_type: library_prompt_name} for generation prompts (sys_prompt, + # spoken_style_answer). Like chat_llm, generation config lives on the + # partition; indexation/retrieval prompts are named on their presets instead. + Column( + "generation_prompt_names", + JSONB, + server_default=text("'{}'::jsonb"), + nullable=False, + ), Column( "updated_at", DateTime(timezone=True), @@ -342,6 +407,7 @@ "metadata", "model_endpoints", "pipeline_presets", + "prompts", "topic_tags", "partitions", "files", diff --git a/openrag/services/workers/indexer_actor.py b/openrag/services/workers/indexer_actor.py index 77851ecf9..a4bda8394 100644 --- a/openrag/services/workers/indexer_actor.py +++ b/openrag/services/workers/indexer_actor.py @@ -66,6 +66,7 @@ async def process_file( indexation_config: dict[str, Any] | None = None, embedder_name: str | None = None, require_existing_partition: bool = False, + resolved_prompts: dict[str, str] | None = None, ) -> dict[str, Any]: """Run one file through the indexing pipeline. @@ -92,6 +93,11 @@ async def process_file( "indexation_config": indexation_config, "embedder_name": embedder_name, } + # DB-resolved enrichment prompts (contextualizer/topic-tagger/caption) + # for this partition; each stage prefers its row value over the + # process-wide disk default. Absent keys leave the disk fallback. + if resolved_prompts: + row.update(resolved_prompts) row = await self._pipeline.run(row) indexed_at = row.get("indexed_at") diff --git a/openrag/services/workers/indexer_pool.py b/openrag/services/workers/indexer_pool.py index c8fedb012..75f1e6a2e 100644 --- a/openrag/services/workers/indexer_pool.py +++ b/openrag/services/workers/indexer_pool.py @@ -135,6 +135,7 @@ def __init__(self) -> None: } self._has_default_fallback = self._has_default_fallbacks["llm"] self._model_endpoint_service: Any = None + self._prompt_service: Any = None self._registry_loaded_at: float | None = None self._last_miss_reload_at: float | None = None self._last_miss_reload_key: tuple[tuple[str, tuple[str, ...]], ...] | None = None @@ -231,6 +232,55 @@ def _reload_decision(self, required_model_names: dict[str, list[str]] | list[str missing=missing, ) + # Enrichment stage β†’ (enable flag, prompt_type, preset name-field, row key). + # The name-field is the indexation-preset config key naming a library prompt + # for that stage. Only enabled stages are resolved, so a file that neither + # contextualizes nor tags nor captions pays no prompt-resolution cost. + _INGEST_PROMPTS = ( + ("enable_contextualization", "chunk_contextualizer", "contextualization_prompt_name", "contextualizer_prompt"), + ("enable_topic_tagging", "topic_tagger", "topic_tagging_prompt_name", "topic_tagger_prompt"), + ("enable_image_captioning", "image_captioning", "image_captioning_prompt_name", "caption_prompt"), + ) + + async def _resolve_ingest_prompts(self, partition: str, indexation_config: dict[str, Any]) -> dict[str, str]: + """Resolve the enabled enrichment prompts for this file's indexation preset. + + Returns ``{row_key: prompt_text}`` for each enabled stage. Each is + resolved by the preset's ``*_prompt_name`` (a named library prompt) β†’ + global default β†’ disk seed, so it always yields a string; any failure is + swallowed and the stage falls back to its own disk-loaded prompt rather + than failing the file. + """ + # Fall back to the model's own default for an absent key, not to False: + # enable_image_captioning defaults to True, so a sparse config (one that + # simply omits the flag) still captions during ingest. A bare .get() read + # that as disabled, skipped resolution, and left captioning silently on + # the disk seed β€” ignoring both the preset's *_prompt_name and the type's + # library default, with nothing surfacing the divergence. + enabled = [ + (pt, name_field, key) + for flag, pt, name_field, key in self._INGEST_PROMPTS + if indexation_config.get(flag, _ingest_flag_default(flag)) + ] + if not enabled: + return {} + if self._prompt_service is None: + from services.orchestrators.prompt_service import PromptService + + self._prompt_service = PromptService( + prompt_repo=self._catalog_store.prompt_repo, + config=self._cfg, + ) + resolved: dict[str, str] = {} + for prompt_type, name_field, row_key in enabled: + try: + resolved[row_key] = await self._prompt_service.resolve_prompt( + prompt_type, names=[indexation_config.get(name_field)] + ) + except Exception as exc: # noqa: BLE001 - resolution must never fail a file + self._logger.warning(f"Prompt resolution failed for '{prompt_type}' (partition={partition}): {exc}") + return resolved + async def process_file( self, *, @@ -250,6 +300,11 @@ async def process_file( try: await self._ensure_catalog() await self._ensure_registry_fresh(_required_model_endpoint_names(indexation_config, embedder_name)) + # Resolve the enrichment-stage prompts once for this file (partition + # override β†’ global default β†’ disk seed). Done here, at the job + # boundary, so per-chunk work reuses one resolved string instead of + # hitting the DB per chunk. + resolved_prompts = await self._resolve_ingest_prompts(partition, indexation_config or {}) result = await self._worker.process_file( task_id=task_id, path=path, @@ -261,6 +316,7 @@ async def process_file( indexation_config=indexation_config, embedder_name=embedder_name, require_existing_partition=require_existing_partition, + resolved_prompts=resolved_prompts, ) file_id = metadata.get("file_id", "") if workspace_ids and not replace and file_id: @@ -966,3 +1022,15 @@ def _global_vlm_endpoint_config(cfg: Any) -> Any | None: __all__ = ["IndexerPool", "IndexerWorkerActor", "build_indexer_pool"] + + +def _ingest_flag_default(flag: str) -> bool: + """The IndexationPipelineConfig default for an enrichment flag. + + Read from the model so the two cannot drift: the defaults differ per stage + (captioning is on, contextualization and topic tagging are off). + """ + from core.config.indexation_pipeline import IndexationPipelineConfig + + field = IndexationPipelineConfig.model_fields.get(flag) + return bool(field.default) if field is not None else False diff --git a/openrag/services/workers/stages/contextualize.py b/openrag/services/workers/stages/contextualize.py index 6df3c7625..1b1e0c38c 100644 --- a/openrag/services/workers/stages/contextualize.py +++ b/openrag/services/workers/stages/contextualize.py @@ -24,9 +24,12 @@ async def contextualize_stage( filename = str(row.get("filename") or "") language = str(row.get("language") or row.get("lang") or "en") + # DB-resolved per-partition prompt for this file, if any; otherwise the + # contextualizer falls back to its own (disk-loaded) default. + system_prompt = row.get("contextualizer_prompt") effective_timeout = stage_timeout(timeout, len(chunks), per_item_timeout=per_chunk_timeout) row["chunks"] = await run_with_optional_timeout( - lambda: contextualizer.contextualize(chunks, filename=filename, lang=language), + lambda: contextualizer.contextualize(chunks, filename=filename, lang=language, system_prompt=system_prompt), effective_timeout, ) row["stage"] = "contextualized" diff --git a/openrag/services/workers/stages/topic_tag.py b/openrag/services/workers/stages/topic_tag.py index 265ea6799..3e87068c7 100644 --- a/openrag/services/workers/stages/topic_tag.py +++ b/openrag/services/workers/stages/topic_tag.py @@ -23,8 +23,13 @@ async def topic_tag_stage( filename = str(row.get("filename") or "") language = str(row.get("language") or row.get("lang") or "en") + # DB-resolved per-partition prompt for this file, if any; otherwise the + # tagger falls back to its own (disk-loaded) default. + system_prompt = row.get("topic_tagger_prompt") row["topic_tags"] = await run_with_optional_timeout( - lambda: topic_tagger.tag(chunks, filename=filename, max_tags=max_tags, lang=language), + lambda: topic_tagger.tag( + chunks, filename=filename, max_tags=max_tags, lang=language, system_prompt=system_prompt + ), timeout, ) row["stage"] = "topic_tagged" diff --git a/tests/integration/repos/conftest.py b/tests/integration/repos/conftest.py index a0e9dc155..c3aeb6aab 100644 --- a/tests/integration/repos/conftest.py +++ b/tests/integration/repos/conftest.py @@ -144,6 +144,7 @@ async def postgres_store(test_rdb_config: RDBConfig) -> PostgresStore: workspace_files, workspaces, partition_memberships, + prompts, files, partitions, users diff --git a/tests/integration/repos/test_partition_repo.py b/tests/integration/repos/test_partition_repo.py index 8594e502e..a1f697161 100644 --- a/tests/integration/repos/test_partition_repo.py +++ b/tests/integration/repos/test_partition_repo.py @@ -96,3 +96,16 @@ async def test_delete_cascades_files_and_decrements_uploader_count( assert await partition_repo.get_partition_file_count("cascade-me") == 0 refreshed = await user_repo.get_user_dict_by_id(uploader_id) assert refreshed["file_count"] == 0 + + +class TestGenerationPromptNames: + async def test_round_trip_and_default_empty(self, postgres_store: PostgresStore): + repo = postgres_store.partition_repo + await repo.create_partition("genp") + # Defaults to an empty JSONB map. + row = await repo.get_partition_row("genp") + assert row["generation_prompt_names"] == {} + # Update persists and reads back as a dict (jsonb codec). + await repo.update_partition("genp", generation_prompt_names={"sys_prompt": "legal"}) + row = await repo.get_partition_row("genp") + assert row["generation_prompt_names"] == {"sys_prompt": "legal"} diff --git a/tests/integration/repos/test_prompt_repo.py b/tests/integration/repos/test_prompt_repo.py new file mode 100644 index 000000000..cf3ffad2f --- /dev/null +++ b/tests/integration/repos/test_prompt_repo.py @@ -0,0 +1,143 @@ +"""PgPromptRepository against a real Postgres. + +Also exercises the migration end-to-end: the ``postgres_store`` fixture runs +``e8f9a0b1c2d3_add_prompts_library`` before any of these can pass. +""" + +from __future__ import annotations + +import pytest +from core.models.prompt import Prompt +from core.utils.exceptions import ValidationError +from services.storage.postgres_store import PostgresStore + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +def _prompt( + prompt_type: str = "sys_prompt", *, name: str = "p", content: str = "body", is_default: bool = False +) -> Prompt: + return Prompt(prompt_type=prompt_type, name=name, content=content, is_default=is_default) + + +class TestCrud: + async def test_create_then_get(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + created = await repo.create(_prompt(name="hello", content="world")) + fetched = await repo.get(created.id) + assert fetched is not None + assert (fetched.name, fetched.content, fetched.prompt_type) == ("hello", "world", "sys_prompt") + # Timestamps come from the DB default. + assert fetched.created_at is not None and fetched.updated_at is not None + + async def test_get_missing_returns_none(self, postgres_store: PostgresStore): + assert await postgres_store.prompt_repo.get("no-such-id") is None + + async def test_list_filter_and_paginate(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + await repo.create(_prompt("sys_prompt", name="a")) + await repo.create(_prompt("sys_prompt", name="b")) + await repo.create(_prompt("hyde", name="c")) + assert {p.name for p in await repo.list(prompt_type="sys_prompt")} == {"a", "b"} + assert len(await repo.list()) == 3 + assert len(await repo.list(offset=1, limit=1)) == 1 + + async def test_count(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + await repo.create(_prompt("sys_prompt")) + await repo.create(_prompt("hyde")) + assert await repo.count() == 2 + assert await repo.count(prompt_type="hyde") == 1 + + async def test_update_name_and_content(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + created = await repo.create(_prompt(name="old", content="old-body")) + updated = await repo.update(created.id, name="new", content="new-body") + assert updated is not None + assert (updated.name, updated.content) == ("new", "new-body") + + async def test_update_ignores_non_whitelisted_fields(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + created = await repo.create(_prompt(prompt_type="sys_prompt", is_default=False)) + # is_default / prompt_type must not be writable through update(). + updated = await repo.update(created.id, is_default=True, prompt_type="hyde", name="ok") + assert updated is not None + assert updated.is_default is False + assert updated.prompt_type == "sys_prompt" + assert updated.name == "ok" + + async def test_delete(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + created = await repo.create(_prompt()) + assert await repo.delete(created.id) is True + assert await repo.delete(created.id) is False + assert await repo.get(created.id) is None + + +class TestGetByName: + async def test_get_by_name(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + created = await repo.create(_prompt("sys_prompt", name="legal")) + found = await repo.get_by_name("sys_prompt", "legal") + assert found is not None and found.id == created.id + # Scoped by type, and None when the name doesn't exist. + assert await repo.get_by_name("hyde", "legal") is None + assert await repo.get_by_name("sys_prompt", "missing") is None + + +class TestDefaultPerType: + async def test_get_default_none_when_absent(self, postgres_store: PostgresStore): + assert await postgres_store.prompt_repo.get_default("sys_prompt") is None + + async def test_create_default_is_returned_by_get_default(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + created = await repo.create(_prompt(is_default=True)) + got = await repo.get_default("sys_prompt") + assert got is not None and got.id == created.id + + async def test_second_default_demotes_first(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + first = await repo.create(_prompt(name="first", is_default=True)) + second = await repo.create(_prompt(name="second", is_default=True)) + # Only the second is default now; the invariant (one default/type) holds. + assert (await repo.get_default("sys_prompt")).id == second.id + assert (await repo.get(first.id)).is_default is False + + async def test_set_default_clears_previous(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + a = await repo.create(_prompt(name="a", is_default=True)) + b = await repo.create(_prompt(name="b")) + promoted = await repo.set_default(b.id) + assert promoted is not None and promoted.is_default is True + assert (await repo.get_default("sys_prompt")).id == b.id + assert (await repo.get(a.id)).is_default is False + + async def test_set_default_missing_returns_none(self, postgres_store: PostgresStore): + assert await postgres_store.prompt_repo.set_default("nope") is None + + async def test_default_is_scoped_per_type(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + sysp = await repo.create(_prompt("sys_prompt", is_default=True)) + hyde = await repo.create(_prompt("hyde", is_default=True)) + # Two defaults coexist because they are different types. + assert (await repo.get_default("sys_prompt")).id == sysp.id + assert (await repo.get_default("hyde")).id == hyde.id + + async def test_duplicate_name_raises_409_not_500(self, postgres_store: PostgresStore): + """The service's pre-check is not atomic: a concurrent create can still + reach the unique index. The repo must translate that into the same 409 + the sequential path returns, not let a UniqueViolationError become a 500. + """ + repo = postgres_store.prompt_repo + await repo.create(_prompt(name="clash")) + with pytest.raises(ValidationError) as err: + await repo.create(_prompt(name="clash")) + assert err.value.status_code == 409 + + async def test_rename_onto_an_existing_name_raises_409(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + await repo.create(_prompt(name="taken")) + other = await repo.create(_prompt(name="free")) + with pytest.raises(ValidationError) as err: + await repo.update(other.id, name="taken") + assert err.value.status_code == 409 diff --git a/tests/integration/repos/test_prompt_service_integration.py b/tests/integration/repos/test_prompt_service_integration.py new file mode 100644 index 000000000..cc717ab4f --- /dev/null +++ b/tests/integration/repos/test_prompt_service_integration.py @@ -0,0 +1,96 @@ +"""PromptService against a real Postgres via PgPromptRepository. + +Covers the boot-critical seam end-to-end: seed_defaults writes real rows, and +resolve_prompt resolves named prompt β†’ default β†’ disk against a live DB. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from core.config.infrastructure import PathsConfig, PromptsConfig +from core.models.prompt import Prompt +from services.orchestrators.prompt_service import PROMPT_TYPE_KEYS, PromptService +from services.storage.postgres_store import PostgresStore + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +def _service(store: PostgresStore) -> PromptService: + config = SimpleNamespace(paths=PathsConfig(), prompts=PromptsConfig()) + return PromptService(prompt_repo=store.prompt_repo, config=config) + + +class TestSeedAndResolve: + async def test_seed_defaults_creates_one_default_per_type(self, postgres_store: PostgresStore): + svc = _service(postgres_store) + await svc.seed_defaults() + assert await postgres_store.prompt_repo.count() == len(PROMPT_TYPE_KEYS) + for prompt_type in PROMPT_TYPE_KEYS: + default = await postgres_store.prompt_repo.get_default(prompt_type) + assert default is not None and default.content.strip() + + async def test_seed_is_idempotent(self, postgres_store: PostgresStore): + svc = _service(postgres_store) + await svc.seed_defaults() + await svc.seed_defaults() + assert await postgres_store.prompt_repo.count() == len(PROMPT_TYPE_KEYS) + + async def test_resolution_precedence_end_to_end(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + svc = _service(postgres_store) + await svc.seed_defaults() + + seeded = (await repo.get_default("sys_prompt")).content + + # No candidate names β†’ the seeded default resolves. + assert await svc.resolve_prompt("sys_prompt") == seeded + assert await svc.resolve_prompt("sys_prompt", names=["missing"]) == seeded + + # A named library prompt wins when named. + await svc.create_prompt(prompt_type="sys_prompt", name="legal", content="LEGAL") + assert await svc.resolve_prompt("sys_prompt", names=["legal"]) == "LEGAL" + # First resolvable candidate wins (the per-user tier extension point). + assert await svc.resolve_prompt("sys_prompt", names=["missing", "legal"]) == "LEGAL" + + async def test_reference_counts_are_effective(self, postgres_store: PostgresStore): + repo = postgres_store.prompt_repo + partition_repo = postgres_store.partition_repo + + # Library: a default + a named alternative for two types. + await repo.create(Prompt(prompt_type="sys_prompt", name="d_sys", content="x", is_default=True)) + await repo.create(Prompt(prompt_type="sys_prompt", name="legal", content="x")) + await repo.create(Prompt(prompt_type="chunk_contextualizer", name="d_ctx", content="x", is_default=True)) + await repo.create(Prompt(prompt_type="chunk_contextualizer", name="ctx1", content="x")) + + # An indexation preset naming ctx1, plus one naming a non-existent prompt. + await postgres_store.preset_repo.upsert("legalpreset", "indexation", {"contextualization_prompt_name": "ctx1"}) + await postgres_store.preset_repo.upsert("orphan", "indexation", {"contextualization_prompt_name": "ghost"}) + + # 3 partitions: rc1 overrides sys_prompt=legal and uses legalpreset; rc2 + # uses legalpreset with no generation override; rc3 names a missing + # sys_prompt (dangling -> default) and keeps the default indexation preset. + await partition_repo.create_partition("rc1") + await partition_repo.update_partition( + "rc1", indexation_preset="legalpreset", generation_prompt_names={"sys_prompt": "legal"} + ) + await partition_repo.create_partition("rc2") + await partition_repo.update_partition("rc2", indexation_preset="legalpreset") + await partition_repo.create_partition("rc3") + await partition_repo.update_partition("rc3", generation_prompt_names={"sys_prompt": "missing"}) + + counts = await repo.reference_counts() + + # sys_prompt: rc1 -> legal; rc2 (no override) + rc3 (dangling) fall back to default. + assert counts.get(("sys_prompt", "legal")) == 1 + assert counts.get(("sys_prompt", "d_sys")) == 2 + # chunk_contextualizer: rc1 + rc2 -> ctx1 (via legalpreset); rc3 -> default. + assert counts.get(("chunk_contextualizer", "ctx1")) == 2 + assert counts.get(("chunk_contextualizer", "d_ctx")) == 1 + # The "orphan" preset names a non-existent prompt and is used by no + # partition β€” it contributes to nothing. + assert counts.get(("chunk_contextualizer", "ghost")) is None + # Per type the effective counts sum to the partition total (3). + assert counts[("sys_prompt", "legal")] + counts[("sys_prompt", "d_sys")] == 3 + assert counts[("chunk_contextualizer", "ctx1")] + counts[("chunk_contextualizer", "d_ctx")] == 3 diff --git a/tests/unit/api/routers/admin/test_prompt_routes.py b/tests/unit/api/routers/admin/test_prompt_routes.py new file mode 100644 index 000000000..7cbebc91f --- /dev/null +++ b/tests/unit/api/routers/admin/test_prompt_routes.py @@ -0,0 +1,150 @@ +"""Transport tests for the prompt library router. + +Service behaviour is covered by the PromptService unit tests; here we assert +requestβ†’service forwarding, response shaping, schema validation (422), and that +service-raised domain errors map to the right status via the shared handlers. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from api.dependencies.auth import require_admin +from api.error_handlers import register_error_handlers +from api.routers.admin import prompts +from core.models.prompt import Prompt +from core.utils.exceptions import NotFoundError, ValidationError +from di.providers import get_prompt_service +from fastapi import FastAPI + + +def _prompt(**overrides: Any) -> Prompt: + data = {"prompt_type": "sys_prompt", "name": "p", "content": "body", "is_default": False} + data.update(overrides) + return Prompt(**data) + + +class FakePromptService: + def __init__(self) -> None: + self.calls: list[tuple[str, dict[str, Any]]] = [] + self.error: Exception | None = None + + async def create_prompt(self, *, prompt_type: str, name: str, content: str, is_default: bool = False) -> Prompt: + self.calls.append( + ("create", {"prompt_type": prompt_type, "name": name, "content": content, "is_default": is_default}) + ) + return _prompt(prompt_type=prompt_type, name=name, content=content, is_default=is_default) + + async def list_prompts(self, *, prompt_type=None, offset=0, limit=100) -> list[Prompt]: + self.calls.append(("list", {"prompt_type": prompt_type, "offset": offset, "limit": limit})) + return [_prompt(prompt_type=prompt_type or "sys_prompt")] + + async def get_prompt(self, prompt_id: str) -> Prompt: + if self.error: + raise self.error + self.calls.append(("get", {"prompt_id": prompt_id})) + return _prompt(id=prompt_id) + + async def update_prompt(self, prompt_id: str, **fields: Any) -> Prompt: + self.calls.append(("update", {"prompt_id": prompt_id, **fields})) + return _prompt(id=prompt_id, **{k: v for k, v in fields.items() if k in ("name", "content", "is_default")}) + + async def set_default(self, prompt_id: str) -> Prompt: + self.calls.append(("set_default", {"prompt_id": prompt_id})) + return _prompt(id=prompt_id, is_default=True) + + async def delete_prompt(self, prompt_id: str) -> None: + if self.error: + raise self.error + self.calls.append(("delete", {"prompt_id": prompt_id})) + + +def _build_app(service: FakePromptService) -> FastAPI: + app = FastAPI() + register_error_handlers(app) + app.include_router(prompts.router, prefix="/prompts") + app.dependency_overrides[require_admin] = lambda: {"id": "admin", "is_admin": True} + app.dependency_overrides[get_prompt_service] = lambda: service + return app + + +pytestmark = pytest.mark.asyncio + + +class TestLibraryRoutes: + async def test_create_forwards_and_returns_201(self, async_client_factory): + svc = FakePromptService() + async with async_client_factory(_build_app(svc)) as client: + resp = await client.post( + "/prompts/", + json={"prompt_type": "sys_prompt", "name": "greet", "content": "hi", "is_default": True}, + ) + assert resp.status_code == 201 + assert resp.json()["prompt_type"] == "sys_prompt" + assert svc.calls == [ + ("create", {"prompt_type": "sys_prompt", "name": "greet", "content": "hi", "is_default": True}) + ] + + async def test_create_rejects_empty_content(self, async_client_factory): + svc = FakePromptService() + async with async_client_factory(_build_app(svc)) as client: + resp = await client.post("/prompts/", json={"prompt_type": "sys_prompt", "content": " "}) + assert resp.status_code == 422 + assert svc.calls == [] + + async def test_create_rejects_unknown_type(self, async_client_factory): + svc = FakePromptService() + async with async_client_factory(_build_app(svc)) as client: + resp = await client.post("/prompts/", json={"prompt_type": "bogus", "content": "x"}) + assert resp.status_code == 422 + + async def test_list_forwards_filters(self, async_client_factory): + svc = FakePromptService() + async with async_client_factory(_build_app(svc)) as client: + resp = await client.get("/prompts/?prompt_type=sys_prompt&offset=2&limit=5") + assert resp.status_code == 200 + assert resp.json()[0]["prompt_type"] == "sys_prompt" + assert svc.calls == [("list", {"prompt_type": "sys_prompt", "offset": 2, "limit": 5})] + + async def test_get_missing_maps_to_404(self, async_client_factory): + svc = FakePromptService() + svc.error = NotFoundError("Prompt 'x' not found.") + async with async_client_factory(_build_app(svc)) as client: + resp = await client.get("/prompts/x") + assert resp.status_code == 404 + + async def test_patch_forwards_only_set_fields(self, async_client_factory): + svc = FakePromptService() + async with async_client_factory(_build_app(svc)) as client: + resp = await client.patch("/prompts/pid", json={"content": "new", "is_default": True}) + assert resp.status_code == 200 + assert svc.calls == [("update", {"prompt_id": "pid", "content": "new", "is_default": True})] + + async def test_patch_empty_body_is_422(self, async_client_factory): + svc = FakePromptService() + async with async_client_factory(_build_app(svc)) as client: + resp = await client.patch("/prompts/pid", json={}) + assert resp.status_code == 422 + + async def test_set_default(self, async_client_factory): + svc = FakePromptService() + async with async_client_factory(_build_app(svc)) as client: + resp = await client.put("/prompts/pid/default") + assert resp.status_code == 200 + assert resp.json()["is_default"] is True + assert svc.calls == [("set_default", {"prompt_id": "pid"})] + + async def test_delete_returns_204(self, async_client_factory): + svc = FakePromptService() + async with async_client_factory(_build_app(svc)) as client: + resp = await client.delete("/prompts/pid") + assert resp.status_code == 204 + assert svc.calls == [("delete", {"prompt_id": "pid"})] + + async def test_delete_default_maps_to_422(self, async_client_factory): + svc = FakePromptService() + svc.error = ValidationError("Cannot delete the default 'sys_prompt' prompt.") + async with async_client_factory(_build_app(svc)) as client: + resp = await client.delete("/prompts/pid") + assert resp.status_code == 422 diff --git a/tests/unit/di/test_container.py b/tests/unit/di/test_container.py index 3ae184bcf..51b2b0021 100644 --- a/tests/unit/di/test_container.py +++ b/tests/unit/di/test_container.py @@ -243,6 +243,9 @@ async def initialize(self): seed_defaults=lambda: _async_call(calls, "preset.seed"), load_all=lambda: _async_call(calls, "preset.load"), ) + c._prompt_service = SimpleNamespace( + seed_defaults=lambda: _async_call(calls, "prompt.seed"), + ) c._partition_service = SimpleNamespace( seed_default_partition=lambda: _async_call(calls, "partition.seed"), load_partitions=lambda: _async_call(calls, "partition.load"), @@ -257,6 +260,7 @@ async def initialize(self): "endpoint.load", "preset.seed", "preset.load", + "prompt.seed", "partition.seed", "partition.load", ] @@ -349,7 +353,7 @@ def test_does_not_mutate_input_settings(self): ("mcp_service", "get_mcp_service"), ] -_OPTIONAL_PHASE_PROVIDERS = {"get_model_endpoint_service", "get_preset_service"} +_OPTIONAL_PHASE_PROVIDERS = {"get_model_endpoint_service", "get_preset_service", "get_prompt_service"} class TestPhase8OrchestratorWiring: @@ -690,6 +694,18 @@ def test_preset_service_is_lazy_cached_with_partition_back_reference(self): assert service._partition_service is c.partition_service assert c.partition_service._config is settings + def test_prompt_service_is_lazy_cached_on_shared_prompt_repo(self): + """Expose PromptService wired to the shared prompt_repo.""" + from services.orchestrators.prompt_service import PromptService + + c = ServiceContainer(_settings()) + + service = c.prompt_service + + assert isinstance(service, PromptService) + assert c.prompt_service is service + assert service._repo is c.prompt_repo + @pytest.mark.asyncio async def test_initialize_loads_phase14_registries_before_partitions(self, monkeypatch): """Load endpoints, then presets, then partitions after admin seeding.""" @@ -730,6 +746,13 @@ async def load_all(self): """Record preset loading.""" calls.append("preset.load") + class FakePromptService: + """Prompt library lifecycle recorder (seed only β€” no cache to load).""" + + async def seed_defaults(self): + """Record prompt seeding.""" + calls.append("prompt.seed") + class FakePartitionService: """Partition cache lifecycle recorder.""" @@ -746,6 +769,7 @@ async def load_partitions(self): c._catalog_store = FakeCatalogStore() c._model_endpoint_service = FakeEndpointService() c._preset_service = FakePresetService() + c._prompt_service = FakePromptService() c._partition_service = FakePartitionService() await c.initialize() @@ -757,6 +781,7 @@ async def load_partitions(self): "endpoint.load", "preset.seed", "preset.load", + "prompt.seed", "partition.seed", "partition.load", ] diff --git a/tests/unit/services/inference/test_call_log.py b/tests/unit/services/inference/test_call_log.py new file mode 100644 index 000000000..3a67610ff --- /dev/null +++ b/tests/unit/services/inference/test_call_log.py @@ -0,0 +1,230 @@ +"""The ``llm.call`` line must show the prompt that goes on the wire β€” and only +that. It is the operator-facing half of the prompt-wiring check, so it has to +survive multimodal payloads and stay bounded on a context-stuffed chat. +""" + +import base64 +from contextlib import contextmanager + +from core.utils.logging import get_logger +from loguru import logger +from services.inference._call_log import ( + MAX_DETAIL_CHARS, + MAX_META_CHARS, + MAX_PARTS, + PREVIEW_CHARS, + _clip, + _describe, + _render_content, + log_llm_call, +) + + +@contextmanager +def _only_sink(records: list, *, level: str): + """Swap loguru's handlers for a single sink at *level*, then restore. + + Laziness is a property of the enabled handlers as a whole β€” loguru evaluates + a lazy argument if *any* handler accepts the level β€” so proving the previews + are skipped means owning the handler set for the duration. + """ + logger.remove() + logger.add(records.append, level=level, format="{message}") + try: + yield + finally: + # ``get_logger`` rebuilds the standard handler set from config (it starts + # with its own ``logger.remove()``), so the swap leaves nothing behind. + get_logger() + + +def test_string_content_is_previewed_with_its_true_length(): + assert _describe({"role": "system", "content": "Answer like a pirate."}) == ("system[21]: Answer like a pirate.") + + +def test_long_content_is_truncated_to_the_preview_budget(): + rendered = _describe({"role": "user", "content": "x" * 5000}) + # The bracketed size still reports the real payload, so a truncated preview + # never hides how much context was actually sent. + assert rendered.startswith("user[5000]: ") + assert rendered.endswith("…") + # The preview itself is capped at exactly PREVIEW_CHARS, ellipsis included. + assert len(rendered.split(": ", 1)[1]) == PREVIEW_CHARS + + +def test_newlines_are_flattened_so_one_call_stays_one_line(): + assert "\n" not in _describe({"role": "system", "content": "line one\nline two\n\nline three"}) + + +def test_image_parts_are_reduced_to_a_marker_and_never_logged(): + image_b64 = base64.b64encode(b"\x89PNG" + b"secret-bytes" * 50).decode() + content = [ + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}}, + {"type": "text", "text": "Describe this image in detail."}, + ] + rendered = _render_content(content) + assert rendered == " + Describe this image in detail." + assert image_b64 not in rendered + + +def test_unknown_content_shapes_do_not_raise(): + assert _describe({"role": "user", "content": {"weird": object()}}) + assert _describe("not-a-dict") + + +class _Exploding(str): + """A payload that fails loudly if its preview is ever built. + + ``split`` is what ``_preview`` touches first, so overriding it catches eager + rendering at the earliest point; ``__len__`` covers the size field. + """ + + def split(self, *args, **kwargs): + raise AssertionError("preview built while DEBUG was disabled") + + def __len__(self): + raise AssertionError("preview built while DEBUG was disabled") + + +def test_debug_sink_sees_the_prompt_that_goes_on_the_wire(): + records: list[str] = [] + sink_id = logger.add(records.append, level="DEBUG", format="{message}") + try: + log_llm_call( + caller="VLLMClient.chat", + model="my-model", + endpoint="http://e", + messages=[ + {"role": "system", "content": "Aye! Answer like a pirate."}, + {"role": "user", "content": "What are the office hours?"}, + ], + ) + finally: + logger.remove(sink_id) + + line = "".join(records) + assert "llm.call VLLMClient.chat model=my-model stream=False" in line + assert "system[26]: Aye! Answer like a pirate." in line + assert "user[26]: What are the office hours?" in line + + +def test_payload_is_not_duplicated_into_the_record_extras(): + """Regression: a ``detail=`` kwarg lands in ``record["extra"]``, and the + terminal formatter appends every extra β€” printing the whole payload twice + on every line. The preview must reach the message only. + """ + seen: list = [] + sink_id = logger.add(lambda m: seen.append(m.record), level="DEBUG", format="{message}") + try: + log_llm_call( + caller="VLLMClient.chat", + model="m", + endpoint="http://e", + messages=[{"role": "system", "content": "UNIQUEMARKER pirate instructions"}], + ) + finally: + logger.remove(sink_id) + + extras = seen[0]["extra"] + assert set(extras) == {"caller", "model", "endpoint", "stream"} + assert not any("UNIQUEMARKER" in str(v) for v in extras.values()) + assert "UNIQUEMARKER" in seen[0]["message"] + + +def test_previews_are_not_built_when_no_sink_is_at_debug(): + """Lazy evaluation: an INFO-only sink must pay nothing for the previews.""" + records: list[str] = [] + with _only_sink(records, level="INFO"): + log_llm_call( + caller="VLLMClient.chat", + model="m", + endpoint="http://e", + messages=[{"role": "system", "content": _Exploding("x")}], + ) + assert records == [] + + +def test_email_addresses_are_pseudonymized(): + """Prompts carry user questions and retrieved context; the diagnostic value + is the prompt shape, never the personal data inside it.""" + rendered = _describe({"role": "user", "content": "forward it to alice.smith@example.com please"}) + assert "alice.smith@example.com" not in rendered + assert "" in rendered + + +def test_a_long_conversation_cannot_produce_an_unbounded_line(): + """200 messages of 500 chars is ~100KB of payload; the record stays bounded + by the detail cap plus the fixed caller/model prefix.""" + messages = [{"role": "user", "content": f"message {i} " + "x" * 500} for i in range(200)] + line = _capture_detail(messages).rstrip("\n") + prefix = line.index(" | ") + len(" | ") + # The detail is capped at exactly MAX_DETAIL_CHARS (ellipsis included), and + # the prefix it sits behind is itself bounded by the identifier caps. + assert len(line) - prefix == MAX_DETAIL_CHARS + assert prefix <= 2 * MAX_META_CHARS + 40 + assert line.endswith("…") + + +def test_many_multimodal_parts_are_capped(): + content = [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}} for _ in range(50)] + rendered = _render_content(content) + assert "more parts" in rendered + assert rendered.count("") <= MAX_PARTS + + +def test_a_newline_in_a_client_supplied_model_cannot_forge_log_lines(): + """`model` comes from metadata.llm_override, so it is caller-controlled.""" + records: list[str] = [] + sink_id = logger.add(records.append, level="DEBUG", format="{message}") + try: + log_llm_call( + caller="VLLMClient.chat", + model="evil\nINFO | forged line: everything is fine", + endpoint="http://e", + messages=[{"role": "user", "content": "hi"}], + ) + finally: + logger.remove(sink_id) + # The whole record must remain a single line. + assert "\n" not in "".join(records).rstrip("\n") + + +def _capture_detail(messages: list) -> str: + records: list[str] = [] + sink_id = logger.add(records.append, level="DEBUG", format="{message}") + try: + log_llm_call(caller="c", model="m", endpoint="e", messages=messages) + finally: + logger.remove(sink_id) + return "".join(records) + + +def test_clip_counts_the_ellipsis_against_the_budget(): + """A cap has to be the real ceiling: appending the ellipsis past the limit + made every clipped span one character longer than its stated bound.""" + assert len(_clip("x" * 100, 10)) == 10 + assert _clip("x" * 100, 10).endswith("…") + assert _clip("short", 10) == "short" + assert _clip("anything", 0) == "" + + +def test_a_brace_in_the_client_supplied_model_does_not_raise(): + """`model` comes from metadata.llm_override, and this is called outside the + client's try/except β€” interpolating it into the format string made a brace + a format field, so `gpt{x}` raised KeyError straight out of the request. + """ + records: list[str] = [] + sink_id = logger.add(records.append, level="DEBUG", format="{message}") + try: + log_llm_call( + caller="VLLMClient.chat", + model="gpt{x}", + endpoint="http://e{y}", + messages=[{"role": "user", "content": "brace {in} content too"}], + ) + finally: + logger.remove(sink_id) + + line = "".join(records) + assert "gpt{x}" in line + assert "brace {in} content too" in line diff --git a/tests/unit/services/orchestrators/test_prompt_service.py b/tests/unit/services/orchestrators/test_prompt_service.py new file mode 100644 index 000000000..bebd29334 --- /dev/null +++ b/tests/unit/services/orchestrators/test_prompt_service.py @@ -0,0 +1,383 @@ +"""Unit tests for PromptService. + +Uses an in-memory fake repository so the resolution precedence, seeding, and +validation logic are tested without a database. Seeding runs against the *real* +bundled templates, which also verifies the prompt_type β†’ config-key map lines +up with the on-disk filenames for all managed types. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from core.config.infrastructure import PathsConfig, PromptsConfig +from core.models.prompt import Prompt, PromptType +from core.utils.exceptions import ConfigError, NotFoundError, ValidationError +from loguru import logger +from services.orchestrators.prompt_service import PROMPT_TYPE_KEYS, PromptService + + +class FakePromptRepo: + """Minimal in-memory PromptRepository honouring the same invariants.""" + + def __init__(self) -> None: + self.prompts: dict[str, Prompt] = {} + + async def create(self, prompt: Prompt) -> Prompt: + if prompt.is_default: + for p in self.prompts.values(): + if p.prompt_type == prompt.prompt_type: + p.is_default = False + self.prompts[prompt.id] = prompt + return prompt + + async def get(self, prompt_id: str) -> Prompt | None: + return self.prompts.get(prompt_id) + + async def list(self, *, prompt_type=None, offset=0, limit=100) -> list[Prompt]: + rows = [p for p in self.prompts.values() if prompt_type is None or p.prompt_type == prompt_type] + rows.sort(key=lambda p: (p.prompt_type, p.name, p.created_at)) + return rows[offset : offset + limit] + + async def count(self, *, prompt_type=None) -> int: + return len([p for p in self.prompts.values() if prompt_type is None or p.prompt_type == prompt_type]) + + async def update(self, prompt_id: str, **fields) -> Prompt | None: + p = self.prompts.get(prompt_id) + if p is None: + return None + for k in ("name", "content"): + if k in fields: + setattr(p, k, fields[k]) + return p + + async def delete(self, prompt_id: str) -> bool: + return self.prompts.pop(prompt_id, None) is not None + + async def get_by_name(self, prompt_type: str, name: str) -> Prompt | None: + return next( + (p for p in self.prompts.values() if p.prompt_type == prompt_type and p.name == name), + None, + ) + + async def reference_counts(self) -> dict[tuple[str, str], int]: + return getattr(self, "_ref_counts", {}) + + async def get_default(self, prompt_type: str) -> Prompt | None: + return next((p for p in self.prompts.values() if p.prompt_type == prompt_type and p.is_default), None) + + async def set_default(self, prompt_id: str) -> Prompt | None: + target = self.prompts.get(prompt_id) + if target is None: + return None + for p in self.prompts.values(): + if p.prompt_type == target.prompt_type: + p.is_default = False + target.is_default = True + return target + + +def _service(repo: FakePromptRepo | None = None) -> PromptService: + config = SimpleNamespace(paths=PathsConfig(), prompts=PromptsConfig()) + return PromptService(prompt_repo=repo or FakePromptRepo(), config=config) + + +class TestSeeding: + async def test_seeds_all_eight_types_from_disk(self): + repo = FakePromptRepo() + await _service(repo).seed_defaults() + seeded_types = {p.prompt_type for p in repo.prompts.values()} + assert seeded_types == set(PROMPT_TYPE_KEYS) + assert len(PROMPT_TYPE_KEYS) == 8 + for p in repo.prompts.values(): + assert p.is_default is True + assert p.content.strip() + + async def test_seeding_is_idempotent(self): + repo = FakePromptRepo() + svc = _service(repo) + await svc.seed_defaults() + sysp = await repo.get_default("sys_prompt") + sysp.content = "OPERATOR EDIT" + await svc.seed_defaults() + assert (await repo.get_default("sys_prompt")).content == "OPERATOR EDIT" + assert len(repo.prompts) == 8 + + async def test_type_set_matches_enum(self): + assert set(PROMPT_TYPE_KEYS) == {t.value for t in PromptType} + + +class TestResolution: + async def test_precedence_named_then_default_then_disk(self): + repo = FakePromptRepo() + svc = _service(repo) + + # Nothing in DB β†’ disk seed fallback (never empty). + disk = await svc.resolve_prompt("sys_prompt") + assert disk.strip() + + # Global default β†’ wins over disk. + await repo.create(Prompt(prompt_type="sys_prompt", name="default_sys", content="DEFAULT", is_default=True)) + assert await svc.resolve_prompt("sys_prompt") == "DEFAULT" + + # A named prompt β†’ wins over default when named. + await repo.create(Prompt(prompt_type="sys_prompt", name="legal", content="LEGAL")) + assert await svc.resolve_prompt("sys_prompt", names=["legal"]) == "LEGAL" + # Unknown / None names are skipped, falling through to the default. + assert await svc.resolve_prompt("sys_prompt", names=["missing"]) == "DEFAULT" + assert await svc.resolve_prompt("sys_prompt", names=[None]) == "DEFAULT" + + async def test_first_resolvable_name_wins(self): + repo = FakePromptRepo() + svc = _service(repo) + await repo.create(Prompt(prompt_type="hyde", name="b", content="B")) + # Ordered candidates: first that resolves wins (extension point for a + # future per-user tier prepended ahead of the partition/preset name). + assert await svc.resolve_prompt("hyde", names=["a", "b"]) == "B" + + +class TestCrud: + async def test_create_validates_type(self): + with pytest.raises(ValidationError): + await _service().create_prompt(prompt_type="not_a_type", name="x", content="y") + + async def test_get_missing_raises(self): + with pytest.raises(NotFoundError): + await _service().get_prompt("nope") + + async def test_create_duplicate_name_per_type_is_rejected(self): + repo = FakePromptRepo() + svc = _service(repo) + await svc.create_prompt(prompt_type="sys_prompt", name="formal", content="a") + with pytest.raises(ValidationError): + await svc.create_prompt(prompt_type="sys_prompt", name="formal", content="b") + # Same name under a different type is fine. + await svc.create_prompt(prompt_type="hyde", name="formal", content="c") + + async def test_rename_collision_is_rejected(self): + repo = FakePromptRepo() + svc = _service(repo) + await svc.create_prompt(prompt_type="sys_prompt", name="a", content="a") + b = await svc.create_prompt(prompt_type="sys_prompt", name="b", content="b") + with pytest.raises(ValidationError): + await svc.update_prompt(b.id, name="a") + + async def test_create_accepts_valid_template_placeholders(self): + svc = _service() + # sys_prompt allows {context} and {current_date}; escaped braces are literal. + p = await svc.create_prompt( + prompt_type="sys_prompt", name="ok", content="Use {context} on {current_date}. Literal {{brace}}." + ) + assert p.id + + async def test_create_rejects_unknown_placeholder(self): + svc = _service() + with pytest.raises(ValidationError) as exc: + await svc.create_prompt(prompt_type="sys_prompt", name="bad", content="Answer about {topic}") + assert exc.value.status_code == 422 + + async def test_create_rejects_malformed_braces(self): + svc = _service() + # A stray single brace (e.g. a JSON/code example) would crash str.format at runtime. + with pytest.raises(ValidationError): + await svc.create_prompt(prompt_type="sys_prompt", name="bad", content='return {"a": 1}') + + @pytest.mark.parametrize( + "content", + [ + "{context!x} on {current_date}", # ValueError: unknown conversion + "{context.missing} on {current_date}", # AttributeError at render time + "{context[0]} on {current_date}", # renders, but not a supported form + "{context:>10} on {current_date}", # format spec + ], + ) + async def test_create_rejects_placeholders_str_format_cannot_render(self, content): + """A field reduced to its root name looked valid while `.format()` still + raised β€” and as a type's global default that fails every request that + falls back to it. Only plain placeholders are accepted. + """ + svc = _service() + with pytest.raises(ValidationError) as exc: + await svc.create_prompt(prompt_type="sys_prompt", name="bad", content=content) + assert exc.value.status_code == 422 + + async def test_bundled_templates_all_pass_validation(self): + """Guards the stricter rule against the seed path: a bundled template that + failed validation would be skipped at boot, leaving the type with no + default at all. + """ + from core.prompts.template_loader import load_template_by_key + from services.orchestrators.prompt_service import _TYPE_TO_CONFIG_KEY, _validate_template + + svc = _service() + for prompt_type, config_key in _TYPE_TO_CONFIG_KEY.items(): + content = load_template_by_key(svc._config.paths.prompts_dir, svc._config.prompts, config_key) + _validate_template(prompt_type, content) + + async def test_verbatim_type_allows_any_braces(self): + svc = _service() + # chunk_contextualizer is sent as-is (never str.format-ed), so literal braces are fine. + p = await svc.create_prompt( + prompt_type="chunk_contextualizer", name="ok", content='Emit JSON like {"topic": "x"} with {anything}' + ) + assert p.id + + async def test_update_rejects_bad_template(self): + repo = FakePromptRepo() + svc = _service(repo) + p = await svc.create_prompt(prompt_type="hyde", name="h", content="Hypothetical doc for {question}") + with pytest.raises(ValidationError): + await svc.update_prompt(p.id, content="now with {unknown_var}") + + async def test_update_promotes_default_via_set_default(self): + repo = FakePromptRepo() + svc = _service(repo) + a = await repo.create(Prompt(prompt_type="sys_prompt", name="a", content="a", is_default=True)) + b = await repo.create(Prompt(prompt_type="sys_prompt", name="b", content="b")) + updated = await svc.update_prompt(b.id, content="b2", is_default=True) + assert updated.is_default is True and updated.content == "b2" + assert (await repo.get(a.id)).is_default is False # single-default invariant + + async def test_delete_default_is_rejected(self): + repo = FakePromptRepo() + svc = _service(repo) + d = await repo.create(Prompt(prompt_type="sys_prompt", content="c", is_default=True)) + with pytest.raises(ValidationError): + await svc.delete_prompt(d.id) + + async def test_delete_non_default_ok(self): + repo = FakePromptRepo() + svc = _service(repo) + p = await repo.create(Prompt(prompt_type="sys_prompt", content="c")) + await svc.delete_prompt(p.id) + assert await repo.get(p.id) is None + + async def test_set_default_missing_raises(self): + with pytest.raises(NotFoundError): + await _service().set_default("nope") + + async def test_list_filters_and_annotates_used_by(self): + repo = FakePromptRepo() + svc = _service(repo) + p = await repo.create(Prompt(prompt_type="hyde", name="b", content="c")) + repo._ref_counts = {("hyde", "b"): 3} + listed = await svc.list_prompts(prompt_type="hyde") + assert [row["prompt_type"] for row in listed] == ["hyde"] + assert listed[0]["used_by"] == 3 + assert p.id == listed[0]["id"] + + +class TestLoggingIsBraceSafe: + def test_a_brace_in_a_prompt_name_does_not_raise(self): + """Prompt names are free text. Interpolating one into the log format + string made a brace a format field, so a partition pointed at a prompt + named `my{tmpl}` raised KeyError on *every* request that resolved it. + """ + from services.orchestrators.prompt_service import PromptService + + records: list[str] = [] + sink_id = logger.add(records.append, level="DEBUG", format="{message}") + try: + PromptService._log_resolution("sys_prompt", ["my{tmpl}"], "named", "my{tmpl}", "body {with} braces") + finally: + logger.remove(sink_id) + assert "my{tmpl}" in "".join(records) + + +class TestResolveSurvivesRepositoryFailure: + async def test_a_repo_error_degrades_to_the_disk_seed(self): + """Resolution moved onto the request path, so chat and search now depend + on Postgres per request where they used to read prompts once at boot. A + transient pool error must degrade to the bundled template, not 500. + """ + + class ExplodingRepo(FakePromptRepo): + async def get_by_name(self, prompt_type, name): + raise RuntimeError("connection pool exhausted") + + async def get_default(self, prompt_type): + raise RuntimeError("connection pool exhausted") + + svc = _service(ExplodingRepo()) + content = await svc.resolve_prompt("sys_prompt", names=["whatever"]) + assert "{context}" in content # the bundled sys_prompt template + + +class TestSeedingSurvivesAConcurrentReplica: + async def test_a_lost_seed_race_does_not_fail_boot(self): + """Losing the race is a no-op, but the unique violation maps to a + ValidationError that _initialize_step re-raises β€” so an unhandled one + turns a concurrent boot into a crash-loop instead of a skipped insert. + """ + + class RacingRepo(FakePromptRepo): + async def create(self, prompt): + raise ValidationError("already exists", status_code=409, code="PROMPT_EXISTS") + + await _service(RacingRepo()).seed_defaults() # must not raise + + +class TestUnavailablePromptRaisesTheTypedError: + async def test_missing_default_and_missing_template_raises_configerror(self, tmp_path): + """Exercises the raise itself. ConfigError hard-coded its own code, so + passing code= collided with the forwarded kwargs and the statement threw + TypeError instead β€” the typed error could never be constructed. + """ + # Point the loader at an empty directory β€” no bundled template, and the + # fake repo has no default, which is the only path reaching the raise. + svc = _service() + svc._config = SimpleNamespace( + paths=PathsConfig(prompts_dir=tmp_path), + prompts=PromptsConfig(), + ) + + with pytest.raises(ConfigError) as exc: + await svc.resolve_prompt("hyde") + + assert exc.value.code == "PROMPT_UNAVAILABLE" + assert exc.value.status_code == 500 + assert "hyde" in str(exc.value) + + +class TestErrorPathsAreExercised: + """Every raise in the service reached at least once. + + These paths were reasoned about rather than run, which is how a raise that + itself threw TypeError survived review β€” coverage over the error branches is + the check that catches that class of defect. + """ + + async def test_malformed_braces_raise_at_write_time(self): + svc = _service() + with pytest.raises(ValidationError) as exc: + await svc.create_prompt(prompt_type="hyde", name="bad", content="unbalanced {question") + assert exc.value.status_code == 422 + assert "brace" in str(exc.value).lower() + + @pytest.mark.parametrize("op", ["get", "update", "set_default", "delete"]) + async def test_unknown_id_raises_not_found(self, op): + svc = _service() + with pytest.raises(NotFoundError): + if op == "get": + await svc.get_prompt("nope") + elif op == "update": + await svc.update_prompt("nope", name="x") + elif op == "set_default": + await svc.set_default("nope") + else: + await svc.delete_prompt("nope") + + async def test_deleting_a_types_default_is_refused(self): + repo = FakePromptRepo() + svc = _service(repo) + p = await svc.create_prompt(prompt_type="hyde", name="d", content="{question}", is_default=True) + with pytest.raises(ValidationError): + await svc.delete_prompt(p.id) + + async def test_seeding_skips_a_type_whose_template_is_missing(self, tmp_path): + repo = FakePromptRepo() + svc = _service(repo) + svc._config = SimpleNamespace(paths=PathsConfig(prompts_dir=tmp_path), prompts=PromptsConfig()) + await svc.seed_defaults() # warns per type, never raises + assert repo.prompts == {} diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index 742e8e5f9..3fce8a7ae 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -26,6 +26,23 @@ _PROMPT_CFG = load_config() +class _EmptyPromptRepo: + """No DB rows β†’ PromptService.resolve_prompt falls back to the disk seed, + preserving the pre-DB behaviour these tests assert.""" + + async def get_by_name(self, prompt_type, name): + return None + + async def get_default(self, prompt_type): + return None + + +def _disk_prompt_service(): + from services.orchestrators.prompt_service import PromptService + + return PromptService(prompt_repo=_EmptyPromptRepo(), config=_PROMPT_CFG) + + @pytest.fixture(autouse=True) def _patch_infra(monkeypatch): @asynccontextmanager @@ -123,6 +140,7 @@ def _svc(*, llm=None, retrieval=None, web=None, mode="SimpleRag", llm_factory=No config=_config(mode), web_search_service=web or FakeWeb(), workspace_service=workspace or FakeWorkspace(), + prompt_service=_disk_prompt_service(), llm_factory=llm_factory, ) @@ -186,6 +204,7 @@ def test_default_chat_history_depth_clamps_invalid_global_config(global_depth): config=config, web_search_service=FakeWeb(), workspace_service=FakeWorkspace(), + prompt_service=_disk_prompt_service(), ) assert svc._default_chat_history_depth == 4 assert svc._resolve_chat_history_depth(None) == 4 @@ -722,6 +741,126 @@ async def test_websearch_with_partition_fuses_docs_via_retrieve_multi(): assert web and web[0].url == "https://ex.com" # websearch branch actually taken +@pytest.mark.asyncio +async def test_answer_system_prompt_comes_from_prompt_service(): + # Revert-proves the query seam: the payload's system message is built from + # prompt_service.resolve_prompt("sys_prompt", ...), resolved request-time β€” + # not a startup snapshot. Reverting query_service to load_template_by_key at + # __init__ makes the marker disappear. + class MarkerPromptService: + def __init__(self): + self.seen: list = [] + + async def resolve_prompt(self, prompt_type, names=None): + self.seen.append(prompt_type) + return "MARKER-SYS::{context}" + + svc = _svc(retrieval=FakeRetrieval()) # SimpleRag β†’ no contextualizer call + marker = MarkerPromptService() + svc._prompt_service = marker + + payload, _docs, _web, _citation_protocol_active = await svc._prepare_chat( + ["p"], {"messages": [{"role": "user", "content": "q"}], "metadata": {}} + ) + + assert payload["messages"][0]["role"] == "system" + assert payload["messages"][0]["content"].startswith("MARKER-SYS::") + assert "sys_prompt" in marker.seen + + +@pytest.mark.asyncio +async def test_generation_prompt_name_from_partition_reaches_resolver(): + # Revert-proves #12: a single owning partition's generation_prompt_names is + # passed to resolve_prompt as the candidate name. Multi-partition / "all" + # pass None (global default). + class RecordingPromptService: + def __init__(self): + self.calls: list = [] + + async def resolve_prompt(self, prompt_type, names=None): + self.calls.append((prompt_type, tuple(names or ()))) + return "SYS::{context}" + + svc = _svc(retrieval=FakeRetrieval()) + rec = RecordingPromptService() + svc._prompt_service = rec + svc._config.partitions = { + "p": SimpleNamespace(generation_prompt_names={"sys_prompt": "legal"}, chat_history_depth=4) + } + + await svc._prepare_chat(["p"], {"messages": [{"role": "user", "content": "q"}], "metadata": {}}) + assert ("sys_prompt", ("legal",)) in rec.calls + + rec.calls.clear() + await svc._prepare_chat(["p", "q"], {"messages": [{"role": "user", "content": "q"}], "metadata": {}}) + assert ("sys_prompt", (None,)) in rec.calls # no single owning partition + + +@pytest.mark.asyncio +async def test_spoken_style_metadata_swaps_the_answer_prompt(): + """`metadata.spoken_style_answer` is a public API flag (and a Chainlit + command) that swaps the answer prompt for a voice-friendly one. Nothing + asserted this, so the whole feature could be β€” and briefly was β€” deleted + with the suite still green. + """ + + class RecordingPromptService: + def __init__(self): + self.calls: list = [] + + async def resolve_prompt(self, prompt_type, names=None): + self.calls.append((prompt_type, tuple(names or ()))) + return "SPOKEN::{context}" + + svc = _svc(retrieval=FakeRetrieval()) + rec = RecordingPromptService() + svc._prompt_service = rec + svc._config.partitions = { + "p": SimpleNamespace(generation_prompt_names={"spoken_style_answer": "voice"}, chat_history_depth=4) + } + + await svc._prepare_chat( + ["p"], + {"messages": [{"role": "user", "content": "q"}], "metadata": {"spoken_style_answer": True}}, + ) + # The spoken-style type is resolved, and the partition may name its own. + assert ("spoken_style_answer", ("voice",)) in rec.calls + assert not any(call[0] == "sys_prompt" for call in rec.calls) + + # Without the flag the ordinary answer prompt is used. + rec.calls.clear() + await svc._prepare_chat(["p"], {"messages": [{"role": "user", "content": "q"}], "metadata": {}}) + assert any(call[0] == "sys_prompt" for call in rec.calls) + assert not any(call[0] == "spoken_style_answer" for call in rec.calls) + + +@pytest.mark.asyncio +async def test_query_contextualizer_name_from_retrieval_preset_reaches_resolver(): + # query_contextualizer is selected on the partition's RETRIEVAL preset (not + # generation prompts). A single owning partition's preset name is passed to + # resolve_prompt; multi-partition passes None (global default). + class RecordingPromptService: + def __init__(self): + self.calls: list = [] + + async def resolve_prompt(self, prompt_type, names=None): + self.calls.append((prompt_type, tuple(names or ()))) + return "CTX" + + payload = json.dumps({"query_list": [{"query": "rewritten", "temporal_filters": None}]}) + svc = _svc(llm=FakeLLM(chat_responses=[payload, payload]), mode="ChatBotRag") + rec = RecordingPromptService() + svc._prompt_service = rec + svc._config.partitions = {"p": SimpleNamespace(retrieval=SimpleNamespace(query_contextualizer_prompt_name="myctx"))} + + await svc.generate_query([{"role": "user", "content": "q"}], partition=["p"]) + assert ("query_contextualizer", ("myctx",)) in rec.calls + + rec.calls.clear() + await svc.generate_query([{"role": "user", "content": "q"}], partition=["p", "q"]) + assert ("query_contextualizer", (None,)) in rec.calls # no single owning partition + + @pytest.mark.asyncio async def test_chat_with_valid_workspace_scopes_search_to_file_ids(): scope = WorkspaceScope(workspace_id="w1", partition="p1", file_ids=["fa", "fb"]) @@ -1057,3 +1196,43 @@ def test_sanitize_system_message_preserved(): {"role": "assistant", "content": "hi"}, ] assert QueryService._sanitize_messages(msgs) == msgs + + +@pytest.mark.asyncio +async def test_conversational_reply_resolves_its_prompt_from_the_library(): + """The no-retrieval path (a greeting / capability question) came from #807 + and read an __init__-time snapshot this branch removes. Git auto-merged that + reference without flagging a conflict, so nothing but this test proves the + conversational reply resolves through PromptService at all. + """ + + class RecordingPromptService: + def __init__(self): + self.calls: list = [] + + async def resolve_prompt(self, prompt_type, names=None): + self.calls.append((prompt_type, tuple(names or ()))) + # Each type is rendered with its own placeholders, so the stub has + # to answer in kind rather than with one shared string. + if prompt_type == "query_contextualizer": + return "CTX {query_language} {current_date}" + return "CONVERSATIONAL {context} {current_date}" + + payload = json.dumps({"requires_retrieval": False, "query_list": []}) + svc = _svc(llm=FakeLLM(chat_responses=[payload]), mode="ChatBotRag") + rec = RecordingPromptService() + svc._prompt_service = rec + svc._config.partitions = { + "p": SimpleNamespace(generation_prompt_names={"sys_prompt": "chatty"}, chat_history_depth=4) + } + + out, docs, web, _ = await svc._prepare_chat( + ["p"], {"messages": [{"role": "user", "content": "hello!"}], "metadata": {}} + ) + + # Resolved from the library, honouring the partition's selection, and no + # retrieval happened. + assert ("sys_prompt", ("chatty",)) in rec.calls + assert docs == [] and web == [] + assert out["messages"][0]["role"] == "system" + assert "CONVERSATIONAL" in out["messages"][0]["content"] diff --git a/tests/unit/services/orchestrators/test_retrieval_service.py b/tests/unit/services/orchestrators/test_retrieval_service.py index 67d361f49..c73a3e502 100644 --- a/tests/unit/services/orchestrators/test_retrieval_service.py +++ b/tests/unit/services/orchestrators/test_retrieval_service.py @@ -510,10 +510,58 @@ async def test_retrieve_small_fanout_stays_fully_parallel(): assert state["max"] == 3, "all 3 partitions should run concurrently under the cap" -def test_pipeline_for_partition_threads_rrf_k(): +@pytest.mark.asyncio +async def test_pipeline_for_partition_threads_rrf_k(): """A partition's rrf_k must reach its RetrieverPipeline (#707).""" cfg = _config() cfg.partitions = {"tenant-a": _partition(retrieval=RetrievalPipelineConfig(rrf_k=42))} svc = RetrievalService(searcher=FakeSearcher(), reranker=None, llm=None, config=cfg) - pipeline, _ = svc._pipeline_for_partition("tenant-a") + pipeline, _ = await svc._pipeline_for_partition("tenant-a") assert pipeline.rrf_k == 42 + + +class _RecordingPromptService: + def __init__(self, resolved: str): + self._resolved = resolved + self.calls: list[tuple[str, tuple]] = [] + + async def resolve_prompt(self, prompt_type, names=None): + self.calls.append((prompt_type, tuple(names or ()))) + return self._resolved + + +@pytest.mark.asyncio +async def test_hyde_template_resolved_from_preset_via_prompt_service(): + """A hyde preset's hyde_prompt_name is resolved through PromptService and + threaded into the HyDeRetriever (the #13 retrieval-prompt seam).""" + cfg = _config() + cfg.partitions = {"tenant-a": _partition(retrieval=RetrievalPipelineConfig(type="hyde", hyde_prompt_name="myhyde"))} + rec = _RecordingPromptService("HYDE {question}") + svc = RetrievalService(searcher=FakeSearcher(), reranker=None, llm=object(), config=cfg, prompt_service=rec) + pipeline, _ = await svc._pipeline_for_partition("tenant-a") + assert pipeline.retriever.hyde_template == "HYDE {question}" + assert ("hyde", ("myhyde",)) in rec.calls + + +@pytest.mark.asyncio +async def test_multi_query_template_resolved_from_preset_via_prompt_service(): + cfg = _config() + cfg.partitions = { + "tenant-a": _partition(retrieval=RetrievalPipelineConfig(type="multiQuery", multi_query_prompt_name="mymq")) + } + rec = _RecordingPromptService("MQ {query} {k_queries}") + svc = RetrievalService(searcher=FakeSearcher(), reranker=None, llm=object(), config=cfg, prompt_service=rec) + pipeline, _ = await svc._pipeline_for_partition("tenant-a") + assert pipeline.retriever.multi_query_template == "MQ {query} {k_queries}" + assert ("multi_query", ("mymq",)) in rec.calls + + +@pytest.mark.asyncio +async def test_single_strategy_resolves_no_prompt(): + """type=single needs no expansion prompt β€” PromptService is never called.""" + cfg = _config() + cfg.partitions = {"tenant-a": _partition(retrieval=RetrievalPipelineConfig(type="single"))} + rec = _RecordingPromptService("unused") + svc = RetrievalService(searcher=FakeSearcher(), reranker=None, llm=None, config=cfg, prompt_service=rec) + await svc._pipeline_for_partition("tenant-a") + assert rec.calls == [] diff --git a/tests/unit/services/persistence/test_partition_repo.py b/tests/unit/services/persistence/test_partition_repo.py index 7ca31e504..f54311e02 100644 --- a/tests/unit/services/persistence/test_partition_repo.py +++ b/tests/unit/services/persistence/test_partition_repo.py @@ -167,6 +167,7 @@ def _full_partition_row(**overrides): "collection_name": None, "chat_history_depth": 0, "chat_llm": None, + "generation_prompt_names": {}, "created_at": datetime(2026, 1, 1, tzinfo=UTC), "updated_at": datetime(2026, 1, 1, tzinfo=UTC), } diff --git a/tests/unit/services/workers/stages/test_pipeline_stages.py b/tests/unit/services/workers/stages/test_pipeline_stages.py index 5a789c97a..27c0d54eb 100644 --- a/tests/unit/services/workers/stages/test_pipeline_stages.py +++ b/tests/unit/services/workers/stages/test_pipeline_stages.py @@ -48,9 +48,13 @@ def __init__(self, chunks: list[Chunk], error: Exception | None = None) -> None: self.chunks = chunks self.error = error self.calls: list[tuple[list[Chunk], str, str]] = [] + self.system_prompts: list[str | None] = [] - async def contextualize(self, chunks, *, filename: str = "", lang: str = "en") -> list[Chunk]: + async def contextualize( + self, chunks, *, filename: str = "", lang: str = "en", system_prompt: str | None = None + ) -> list[Chunk]: self.calls.append((list(chunks), filename, lang)) + self.system_prompts.append(system_prompt) if self.error is not None: raise self.error return self.chunks diff --git a/tests/unit/services/workers/test_indexer_pool.py b/tests/unit/services/workers/test_indexer_pool.py index b728a448a..196b85b41 100644 --- a/tests/unit/services/workers/test_indexer_pool.py +++ b/tests/unit/services/workers/test_indexer_pool.py @@ -1448,6 +1448,15 @@ async def process_file(self, **kwargs) -> dict: return {"stored_count": 1, "stage": "stored"} +def _AsyncReturn(value): + """A stub coroutine function returning *value* for any arguments.""" + + async def _call(*_a, **_k): + return value + + return _call + + def _bare_worker_actor(*, save_uploaded_files: bool, worker: _RecordingWorker): """Bare IndexerWorkerActor with only the attributes process_file touches.""" from services.workers.indexer_pool import IndexerWorkerActor @@ -1467,6 +1476,12 @@ async def _noop(*_a, **_k): ) actor._save_uploaded_files = save_uploaded_files actor._logger = SimpleNamespace(debug=lambda *a, **k: None, warning=lambda *a, **k: None) + # These build the actor with __new__, so __init__ never runs. Captioning is + # enabled by default, so ingest now resolves its prompt even for a config + # that omits the flag β€” stub the service these tests don't exercise. + actor._prompt_service = SimpleNamespace( + resolve_prompt=_AsyncReturn("prompt"), + ) return actor diff --git a/tests/unit/services/workers/test_pipeline_builder.py b/tests/unit/services/workers/test_pipeline_builder.py index 6ce4e218f..dc58273b9 100644 --- a/tests/unit/services/workers/test_pipeline_builder.py +++ b/tests/unit/services/workers/test_pipeline_builder.py @@ -83,7 +83,9 @@ 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]: + async def contextualize( + self, chunks, *, filename: str = "", lang: str = "en", system_prompt: str | None = None + ) -> 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] @@ -100,6 +102,7 @@ async def tag( filename: str = "", max_tags: int = 7, lang: str = "en", + system_prompt: str | None = None, ) -> list[str]: self.calls.append((list(chunks), filename, max_tags, lang)) return self.tags @@ -901,3 +904,16 @@ async def test_reindex_with_zero_new_chunks_keeps_old_chunks(): assert row["stored_count"] == 0 assert vs.deleted == [] assert "delete" not in vs.events + + +def test_ingest_flag_defaults_come_from_the_model_not_from_absence(): + """A sparse indexation config that omits enable_image_captioning still + captions during ingest (the model default is True). Reading the flag with a + bare .get() treated that as disabled and skipped prompt resolution, leaving + captioning silently on the disk seed while the preset named another prompt. + """ + from services.workers.indexer_pool import _ingest_flag_default + + assert _ingest_flag_default("enable_image_captioning") is True + assert _ingest_flag_default("enable_contextualization") is False + assert _ingest_flag_default("enable_topic_tagging") is False