diff --git a/.github/RELEASING.md b/.github/RELEASING.md new file mode 100644 index 000000000..b020a8a0c --- /dev/null +++ b/.github/RELEASING.md @@ -0,0 +1,236 @@ +# Releasing OpenRag — verifying the GA images actually published + +Run this **after** pushing a `vX.Y.Z` tag to `main`. It exists because the +v2.0.1 release produced a **green workflow run that built nothing**: the three +build jobs were guarded on `github.event.base_ref`, which is empty for a tag +pushed to a branch-protected `main`, so every job reported `skipped` and the run +was still green. Nothing in our process caught it. + +The rule this checklist enforces: **a green run is not proof. A digest in the +registry is proof.** + +Set the version once, and define the digest helper every step below uses: + +```bash +export VER=v2.1.0 # the tag you just pushed +dg() { docker buildx imagetools inspect --raw "$1" 2>/dev/null | sha256sum | awk '{print "sha256:"$1}'; } +``` + +> **Why `--raw | sha256sum` and not `--format '{{.Manifest.Digest}}'`:** buildx +> v0.30.1 **silently ignores** that `--format` template and prints its default +> human output instead. Comparing those strings makes every image look like it +> drifted. A manifest digest *is* the sha256 of the raw manifest bytes, so this +> form is both correct and self-verifying. Validated against `v2.0.1`, where it +> reproduces the published digest exactly. + +## What a release publishes + +| Image | ghcr.io | Docker Hub | +|---|---|---| +| API | `ghcr.io/linagora/openrag` | `linagoraai/openrag` | +| Ray | `ghcr.io/linagora/openrag-ray` | — (ghcr only, by design) | +| Admin UI | `ghcr.io/linagora/openrag-admin-ui` | `linagoraai/openrag-admin-ui` | + +Each gets two tags: `$VER` and `latest`. + +--- + +## 1. The workflow ran — and did not skip + +A skipped job is the exact v2.0.1 failure mode, and `gh run list` shows such a +run as `success`. Assert on **per-job conclusions**, never on the run's. + +```bash +gh run list --workflow build.yml --limit 5 \ + --json databaseId,headBranch,event,status,conclusion,createdAt \ + --jq '.[] | "\(.databaseId) \(.event) \(.headBranch) \(.status)/\(.conclusion) \(.createdAt)"' +``` + +Take the run id for the tag push, then: + +```bash +RUN_ID= +gh run view "$RUN_ID" --json jobs \ + --jq '.jobs[] | "\(.conclusion)\t\(.name)"' +``` + +**PASS** requires all four jobs `success`: +`verify-tag`, `build-and-push-image`, `build-and-push-image-ray`, +`build-and-push-image-admin-ui`. + +**FAIL** on any `skipped` — that is the v2.0.1 bug recurring. A hard gate: + +```bash +gh run view "$RUN_ID" --json jobs --jq '[.jobs[] | select(.conclusion != "success")] | length' +# must print 0 +``` + +If `verify-tag` failed loudly, the tag is not an ancestor of `origin/main` — +fix the tag placement, do not rerun. + +## 2. The tag exists in every registry + +`imagetools inspect` reads the registry directly (anonymous, no pull, no +`docker login`). If the tag was never pushed, this errors — which is the point. + +```bash +for img in ghcr.io/linagora/openrag ghcr.io/linagora/openrag-ray \ + ghcr.io/linagora/openrag-admin-ui \ + linagoraai/openrag linagoraai/openrag-admin-ui; do + d=$(dg "$img:$VER"); printf '%-42s %s\n' "$img:$VER" "${d:-MISSING}" +done +``` + +**PASS**: five `sha256:…` digests, zero `MISSING`. + +## 3. `latest` points at the release, not something older + +`latest` is published unconditionally by `build.yml`, so a mismatch here means +`latest` is stale and every `docker pull` without a tag gets the wrong build. + +```bash +for img in ghcr.io/linagora/openrag ghcr.io/linagora/openrag-ray \ + ghcr.io/linagora/openrag-admin-ui \ + linagoraai/openrag linagoraai/openrag-admin-ui; do + v=$(dg "$img:$VER"); l=$(dg "$img:latest") + if [ -n "$v" ] && [ "$v" = "$l" ]; then echo "OK $img" + else echo "DRIFT $img"; echo " $VER = ${v:-none}"; echo " latest= ${l:-none}"; fi +done +``` + +**PASS**: all `OK`. + +## 4. ghcr and Docker Hub are the same build + +Both registries are pushed from one `docker/build-push-action` step, so the +digests must be identical. A difference means one push failed and was +back-filled from a different build. + +```bash +for pair in "ghcr.io/linagora/openrag linagoraai/openrag" \ + "ghcr.io/linagora/openrag-admin-ui linagoraai/openrag-admin-ui"; do + set -- $pair; a=$(dg "$1:$VER"); b=$(dg "$2:$VER") + [ "$a" = "$b" ] && echo "OK $1 == $2" || echo "MISMATCH $1=$a $2=$b" +done +``` + +**PASS**: both `OK`. + +## 5. Pull it, and confirm the digest is the one you verified + +Steps 2–4 read metadata. This proves the bytes are actually fetchable. + +```bash +docker pull "linagoraai/openrag:$VER" +docker image inspect "linagoraai/openrag:$VER" --format '{{index .RepoDigests 0}}' +``` + +**PASS**: the printed digest equals the Docker Hub digest from step 2. + +## 6. The image contains the released code + +The strongest check, and the one that catches a build from the wrong commit: +the version baked into the image must equal the tag. `app.version` comes from +`importlib.metadata`, i.e. from `pyproject.toml` at build time. + +```bash +docker run --rm --entrypoint grep "linagoraai/openrag:$VER" -m1 '^version' /app/pyproject.toml +# expect: version = "2.1.0" (tag minus the leading v) +``` + +And on a running stack (the endpoint is unauthenticated): + +```bash +curl -fsS http://:8080/version +# {"version":"2.1.0"} +``` + +**PASS**: both report the release version. A mismatch means the tag sat on a +commit that predates the version bump — the images are mislabelled and must be +rebuilt from a corrected tag. + +## 7. The tag is where it should be + +```bash +git fetch origin main --tags +git tag --contains "$VER" >/dev/null 2>&1 +git merge-base --is-ancestor "$VER" origin/main && echo "OK: $VER is on main" || echo "FAIL: not on main" +git log -1 --format='%H %s' "$VER" +``` + +**PASS**: `OK`, and the commit is the release-branch merge commit. + +## 8. Chart and compose reference the published tags + +The chart and compose pins are part of the release surface; shipping them +pointing at the previous version is a silent regression for anyone deploying +from the tag. + +```bash +git show "$VER:infra/charts/openrag-stack/Chart.yaml" | grep -E '^(version|appVersion)' +git show "$VER:infra/charts/openrag-stack/values.yaml" | grep -nE 'tag: "v[0-9]' +git show "$VER:infra/compose/docker-compose.yaml" | grep -nE 'image: linagoraai/' +``` + +**PASS**: every OpenRag image pin reads `$VER`, `appVersion` matches, chart +`version` was bumped. + +--- + +## What v2.0.1 actually did, and the rules that follow + +Reconstructed from the run log and the tag, 2026-07-30. Three `build.yml` runs +fired for the same tag name: + +| run | time (UTC) | commit | result | +|---|---|---|---| +| 30034799802 | 18:41 | `c08c5e9f` (release/2.0.1 → main merge) | 3 jobs **skipped**, run green | +| 30035356446 | 18:49 | `c08c5e9f` — tag re-pushed unchanged | 3 jobs **skipped** again | +| 30039985149 | 19:55 | `6a18a534` (CI hotfix #764 merge) | `verify-tag` + 3 builds **success** | + +The shipped `v2.0.1` tag therefore sits on the **CI-hotfix merge commit**, not +on the release-branch merge. The tagged tree still carries the version bump +(`c08c5e9f` is its ancestor), so the images are correct — confirmed above. + +Rules this produces: + +1. **A skipped job is a failure.** The first two runs reported `success` at the + run level. Only per-job conclusions revealed the truth. That is step 1. +2. **Never re-push a tag to "retry".** Run 2 proves it is deterministic: the + workflow that executes is the one *at the tagged commit*, so re-pushing the + same tag re-runs the same broken file. Move the tag to a fixed commit, or + fix nothing and diagnose. +3. **Verify images before announcing.** `build.yml` does not create GitHub + Releases. For v2.0.1 the Release was published at 20:01, six minutes after + the images finally landed at 19:55. Keep that order: tag → images verified + → Release notes. +4. **Dry-run the verification itself against the previous release.** Doing that + for v2.0.1 is what exposed the broken `--format` flag above. A checklist + that silently reports nonsense is worse than none. + +### Open risk for the next release + +`main`'s current `build.yml` is **not** the file that successfully built +v2.0.1. PR #767 hardened it afterwards (exact-tag regex, tag name/SHA passed as +env instead of `${{ }}` interpolation, `persist-credentials: false`). The next +GA tag is the **first time that hardened guard ever runs**. + +Pre-flighted locally on 2026-07-30: + +- Regex `^v[0-9]+\.[0-9]+\.[0-9]+$` — `v2.1.0` accepted; `v2.1`, `2.1.0`, + `v1.2.3-rc1`, `v1.0-hardening` rejected loudly; `v2.1.0-rc.N` filtered out by + the job-level `if` and left to `build_rc.yml`. Correct on all six. +- `persist-credentials: false` + `git fetch --no-tags origin main` — verified an + anonymous fetch of this repo succeeds, so the guard can still reach `main`. + This holds only while the repo is **public**; if it is ever made private, + that fetch breaks and every GA build blocks. + +The residual risk is acceptable because the hardened guard's failure mode is +`exit 1` — loud and blocking — not v2.0.1's silent skip. But treat step 1 as +mandatory, not a formality. + +## Result + +Record the outcome on the GitHub Release or the milestone. If any step fails, +the release is **not** done — publishing images is the deliverable, and the tag +alone delivers nothing. diff --git a/docs/content/docs/documentation/API.mdx b/docs/content/docs/documentation/API.mdx index b2fdd675b..e800943bb 100644 --- a/docs/content/docs/documentation/API.mdx +++ b/docs/content/docs/documentation/API.mdx @@ -383,10 +383,13 @@ Manage who can access a partition and with which role — all **owner** only. | Endpoint | Method | Description | |----------|--------|-------------| | `/partition/{partition}/users` | GET | List members → `{ "members": [...] }` | -| `/partition/{partition}/users` | POST | Add a member — form fields `user_id` (int), `role` (default `viewer`) → `201` | +| `/partition/{partition}/users/candidates` | GET | Search non-members by display-name prefix or exact user ID → paginated identities | +| `/partition/{partition}/users` | POST | Add a new member → `201`; returns `409` if already present | | `/partition/{partition}/users/{user_id}` | PATCH | Update a member's role — form field `role` → `200` | | `/partition/{partition}/users/{user_id}` | DELETE | Remove a member → `204` | +`POST` no longer updates an existing member's role. Use the `PATCH` endpoint when changing a role. + #### Document Relationships * Get Files by Relationship @@ -510,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/docs/content/docs/documentation/kubernetes.md b/docs/content/docs/documentation/kubernetes.md index 758e850e2..fe621c50b 100644 --- a/docs/content/docs/documentation/kubernetes.md +++ b/docs/content/docs/documentation/kubernetes.md @@ -21,7 +21,7 @@ This guide explains how to deploy the **OpenRAG** stack on a Kubernetes cluster - Copy or create a new `values.yaml` at the root of your repo. - You can see the full example file inside the chart: - [values.yaml](https://github.com/linagora/openrag/blob/dev/charts/openrag-stack/values.yaml) + [values.yaml](https://github.com/linagora/openrag/blob/dev/infra/charts/openrag-stack/values.yaml) - Customize the values you need (e.g., image tags, resources, ingress host, storage class, environment variables, secrets). 2. **Set environment and secrets**: @@ -36,16 +36,41 @@ This guide explains how to deploy the **OpenRAG** stack on a Kubernetes cluster helm upgrade\ --install openrag oci://ghcr.io/linagora/openrag-stack\ -f ./values.yaml\ - --version 0.1.0 + --version 0.6.0 ``` - `openrag` is the Helm release name. - `oci://ghcr.io/linagora/openrag-stack` is the remote chart location. - `-f ./values.yaml` specifies your custom configuration. - - `--version 0.1.0` ensures you deploy a specific chart version. + - `--version 0.6.0` ensures you deploy a specific chart version — check `Chart.yaml` for the current version before installing. --- +## Upgrading to chart 0.6.0 + +Chart 0.6.0 renames the PVCs, ConfigMap and Secret from fixed `rag-*` names to +`{{ fullname }}-*`, so they follow the release instead of colliding between two +installs in one namespace. With the default `fullnameOverride: "openrag"`: + +| Before | After | +|---|---| +| `rag-model-weights`, `rag-data`, `rag-logs`, `rag-venv` | `openrag-model-weights`, `openrag-data`, `openrag-logs`, `openrag-venv` | +| `rag-env` | `openrag-env` | +| `rag-env-secrets` | `openrag-env-secrets` | + +The old PVCs carry `helm.sh/resource-policy: keep`, so **the upgrade does not +delete them — but it does not mount them either**. It provisions new, empty ones +under the new names, and the release comes up as if it had no indexed data. Pick +one before upgrading: + +- **Keep the existing volumes.** Set `fullnameOverride: "rag"`, which reproduces + the old names exactly. Also set `postgresql.fullnameOverride`, + `milvus.fullnameOverride` and `vllm.hfTokenSecretName` to match (they are kept + in sync by hand — `values.yaml` explains why, and `NOTES.txt` warns on an + HF_TOKEN secret-name mismatch). +- **Migrate to the new names.** Copy the data across (e.g. a Job mounting both + PVCs), then delete the old ones once the release is healthy. + ## Notes - If using a public IP instead of a hostname, you can leave `ingress.host` empty in your `values.yaml`. diff --git a/infra/charts/openrag-stack/Chart.lock b/infra/charts/openrag-stack/Chart.lock index 107a2a026..2fa4d73f1 100644 --- a/infra/charts/openrag-stack/Chart.lock +++ b/infra/charts/openrag-stack/Chart.lock @@ -4,12 +4,12 @@ dependencies: version: 1.4.0 - name: postgresql repository: https://charts.bitnami.com/bitnami - version: 18.1.13 + version: 18.7.3 - name: milvus repository: https://zilliztech.github.io/milvus-helm/ version: 5.0.0 - name: vllm-stack repository: https://vllm-project.github.io/production-stack - version: 0.1.8 -digest: sha256:838fb1143c5471aaf919c128ece414799898727d083103ac2b669083f50b82c2 -generated: "2025-12-02T11:53:42.972907563Z" + version: 0.1.11 +digest: sha256:1ea11f53796e3196848d5d46cdff5a6bdeeb012427c1951ea646a22b176adec6 +generated: "2026-06-10T14:51:21.005020805+02:00" diff --git a/infra/charts/openrag-stack/Chart.yaml b/infra/charts/openrag-stack/Chart.yaml index e96691aa7..557bbbd0a 100644 --- a/infra/charts/openrag-stack/Chart.yaml +++ b/infra/charts/openrag-stack/Chart.yaml @@ -1,39 +1,36 @@ apiVersion: v2 name: openrag-stack description: A Helm chart for Kubernetes - -# A chart can be either an 'application' or a 'library' chart. -# -# Application charts are a collection of templates that can be packaged into versioned archives -# to be deployed. -# -# Library charts provide useful utilities or functions for the chart developer. They're included as -# a dependency of application charts to inject those utilities and functions into the rendering -# pipeline. Library charts do not define any templates and therefore cannot be deployed. type: application -# This is the chart version. This version number should be incremented each time you make changes -# to the chart and its templates, including the app version. -# Versions are expected to follow Semantic Versioning (https://semver.org/) -version: 0.5.2 +version: 0.6.1 +appVersion: "2.1.0" + +maintainers: + - name: linagora + email: openrag@linagora.com + url: https://github.com/linagora/openrag -# This is the version number of the application being deployed. This version number should be -# incremented each time you make changes to the application. Versions are not expected to -# follow Semantic Versioning. They should reflect the version the application is using. -# It is recommended to use it with quotes. -appVersion: "2.0.1" +home: https://github.com/linagora/openrag +sources: + - https://github.com/linagora/openrag dependencies: - name: kuberay-operator version: "1.4.0" repository: "https://ray-project.github.io/kuberay-helm/" + alias: kuberay + condition: kuberay.enabled - name: postgresql - version: ">=17.6.0" + version: ">=17.6.0 <19.0.0" repository: "https://charts.bitnami.com/bitnami" + condition: postgresql.enabled - name: milvus version: "5.0.0" repository: "https://zilliztech.github.io/milvus-helm/" + condition: milvus.enabled - name: vllm-stack alias: vllm - version: ">=0.1.7" + version: ">=0.1.7 <1.0.0" repository: "https://vllm-project.github.io/production-stack" + condition: vllm.enabled diff --git a/infra/charts/openrag-stack/templates/NOTES.txt b/infra/charts/openrag-stack/templates/NOTES.txt new file mode 100644 index 000000000..6a8f9ff0a --- /dev/null +++ b/infra/charts/openrag-stack/templates/NOTES.txt @@ -0,0 +1,50 @@ +OpenRAG stack deployed as release "{{ .Release.Name }}" in namespace "{{ .Release.Namespace }}". + +{{- $secretName := include "openrag-stack.secretName" . }} +{{- $configuredVllmSecret := .Values.vllm.hfTokenSecretName }} +{{- if ne $secretName $configuredVllmSecret }} + +⚠ vLLM HF_TOKEN secret mismatch — action required +--------------------------------------------------- +This release creates the env secret as: {{ $secretName }} +The vLLM sub-chart values still reference: {{ $configuredVllmSecret }} + +vLLM serving pods will fail to mount HF_TOKEN unless you fix this. + +Careful how you override this: Helm never merges *whole lists* across values +sources, it replaces them — so supplying your own +vllm.servingEngineSpec.modelSpec array (a values file, or a --set that +redefines the entire array) silently drops any of the 4 entries you don't +repeat, breaking those deployments. + +An indexed override like +"--set vllm.servingEngineSpec.modelSpec[0].hf_token.secretName=..." is fine — +it only touches that one field of that one entry and leaves the other 3 +entries (and the other fields of that entry) untouched. The wholesale-replace +risk above only applies when the *entire* modelSpec array is redefined. + +To fix this mismatch, either edit hfTokenSecretName's anchor directly (its +value already flows to all 4 entries), or supply the *complete* vllm.servingEngineSpec.modelSpec +array (copy the 4 entries from values.yaml, keep every field, just change +hf_token.secretName — an anchor keeps all 4 in sync from one line), either: + - directly in values.yaml (vllm.hfTokenSecretName's anchor already does this), or + - in your own values file / ArgoCD Application values block, reusing an + anchor the same way, e.g.: + + vllm: + hfTokenSecretName: &hfTokenSecret "{{ $secretName }}" + servingEngineSpec: + modelSpec: + - name: "embedder" + # ...copy every other field from values.yaml unchanged... + hf_token: + secretName: *hfTokenSecret + secretKey: HF_TOKEN + # ...repeat for "whisper", "llm", "vlm"... + +{{- end }} + +Secrets provider: {{ .Values.env.secretsProvider.type | default "values" }} +{{- if .Values.env.existingSecret }} +Using existing secret: {{ .Values.env.existingSecret }} +{{- end }} diff --git a/infra/charts/openrag-stack/templates/_helpers.tpl b/infra/charts/openrag-stack/templates/_helpers.tpl index c5bd0ede7..a86650e03 100644 --- a/infra/charts/openrag-stack/templates/_helpers.tpl +++ b/infra/charts/openrag-stack/templates/_helpers.tpl @@ -47,16 +47,81 @@ Selector labels */}} {{- define "openrag-stack.selectorLabels" -}} app.kubernetes.io/name: {{ include "openrag-stack.name" . }} -app.kubernetes.io/instance: {{ .Release.Name }} +app.kubernetes.io/instance: {{ include "openrag-stack.fullname" . }} {{- end }} {{/* -Create the name of the service account to use +Component-scoped selector labels. app.kubernetes.io/name identifies the +specific workload ("openrag", "admin-ui", "reranker", ...) instead of the +umbrella chart name, so e.g. `kubectl get pods -l app.kubernetes.io/name=admin-ui` +targets one component — every workload template should use this (and +"componentLabels" below) instead of hand-rolling its own label block. +Usage: {{ include "openrag-stack.componentSelectorLabels" (dict "component" "openrag" "context" $) }} */}} -{{- define "openrag-stack.serviceAccountName" -}} -{{- if .Values.serviceAccount.create }} -{{- default (include "openrag-stack.fullname" .) .Values.serviceAccount.name }} +{{- define "openrag-stack.componentSelectorLabels" -}} +app.kubernetes.io/name: {{ .component }} +app.kubernetes.io/instance: {{ include "openrag-stack.fullname" .context }} +{{- end }} + +{{/* +Component-scoped common labels (componentSelectorLabels plus chart/version/managed-by). +Usage: {{ include "openrag-stack.componentLabels" (dict "component" "openrag" "context" $) }} +*/}} +{{- define "openrag-stack.componentLabels" -}} +helm.sh/chart: {{ include "openrag-stack.chart" .context }} +{{ include "openrag-stack.componentSelectorLabels" . }} +{{- if .context.Chart.AppVersion }} +app.kubernetes.io/version: {{ .context.Chart.AppVersion | quote }} +{{- end }} +app.kubernetes.io/managed-by: {{ .context.Release.Service }} +{{- end }} + +{{/* +Environment secret name used by all consumers. +When env.existingSecret is set, that name is returned directly. +*/}} +{{- define "openrag-stack.secretName" -}} +{{- if .Values.env.existingSecret }} +{{- .Values.env.existingSecret }} {{- else }} -{{- default "default" .Values.serviceAccount.name }} +{{- printf "%s-env-secrets" (include "openrag-stack.fullname" .) | trunc 63 | trimSuffix "-" }} {{- end }} {{- end }} + +{{/* +Merge a component's security context override (e.g. just runAsUser/runAsGroup/ +fsGroup, tuned to that component's own Dockerfile) on top of a shared default +from values.yaml's top-level `security` block — component keys win on +conflicts, everything else is inherited from the default. +Deliberately NOT Sprig's `merge` (mergo): mergo treats a zero value (false, 0, +"") as "unset" and overwrites it with the default, so an explicit +`allowPrivilegeEscalation: false` override would be silently discarded the +day the shared default becomes `true`. This does a presence-based (hasKey) +shallow merge instead, so an explicitly-set false/0/"" always wins. +Usage: {{ include "openrag-stack.mergeSecurityContext" (dict "component" .Values.ray.podSecurityContext "default" .Values.security.podSecurityContext) }} +*/}} +{{- define "openrag-stack.mergeSecurityContext" -}} +{{- $result := deepCopy .component -}} +{{- range $key, $value := .default -}} +{{- if not (hasKey $result $key) -}} +{{- $_ := set $result $key $value -}} +{{- end -}} +{{- end -}} +{{- $result | toYaml -}} +{{- end }} + +{{/* +Component override for a single boolean/scalar security field (e.g. +automountServiceAccountToken): uses the component's own value only if the key +is explicitly present, otherwise falls back to the shared default. A plain +`| default` would treat an explicit `false` override as empty and silently +fall back anyway — this checks presence (hasKey) instead of truthiness. +Usage: {{ include "openrag-stack.securityFieldOverride" (dict "component" .Values.ray "key" "automountServiceAccountToken" "default" .Values.security.automountServiceAccountToken) }} +*/}} +{{- define "openrag-stack.securityFieldOverride" -}} +{{- if hasKey .component .key -}} +{{- get .component .key -}} +{{- else -}} +{{- .default -}} +{{- end -}} +{{- end }} diff --git a/infra/charts/openrag-stack/templates/admin-ui.yaml b/infra/charts/openrag-stack/templates/admin-ui.yaml index 5452e4778..8b9a45fd4 100644 --- a/infra/charts/openrag-stack/templates/admin-ui.yaml +++ b/infra/charts/openrag-stack/templates/admin-ui.yaml @@ -2,65 +2,66 @@ apiVersion: apps/v1 kind: Deployment metadata: - name: {{ .Release.Name }}-admin-ui + name: {{ include "openrag-stack.fullname" . }}-admin-ui labels: - app.kubernetes.io/name: admin-ui - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentLabels" (dict "component" "admin-ui" "context" $) | nindent 4 }} spec: replicas: {{ .Values.adminUi.replicaCount }} selector: matchLabels: - app.kubernetes.io/name: admin-ui - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentSelectorLabels" (dict "component" "admin-ui" "context" $) | nindent 6 }} template: metadata: labels: - app.kubernetes.io/name: admin-ui - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentSelectorLabels" (dict "component" "admin-ui" "context" $) | nindent 8 }} spec: - automountServiceAccountToken: {{ .Values.security.automountServiceAccountToken }} securityContext: - {{- toYaml .Values.security.podSecurityContext | nindent 8 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" .Values.adminUi.podSecurityContext "default" .Values.security.podSecurityContext) | nindent 8 }} + {{- with .Values.adminUi.serviceAccountName }} + serviceAccountName: {{ tpl . $ }} + {{- end }} + automountServiceAccountToken: {{ include "openrag-stack.securityFieldOverride" (dict "component" .Values.adminUi "key" "automountServiceAccountToken" "default" .Values.security.automountServiceAccountToken) }} containers: - name: admin-ui - image: "{{ .Values.adminUi.image.repository }}:{{ .Values.adminUi.image.tag }}" - imagePullPolicy: {{ .Values.adminUi.imagePullPolicy }} - # nginx-unprivileged listens on :8080 as a non-root user, so it runs - # under the same hardened context (runAsNonRoot, drop ALL) as the API. + image: "{{ .Values.adminUi.image.registry | default .Values.global.image.registry }}/{{ .Values.adminUi.image.repository }}:{{ .Values.adminUi.image.tag }}" + imagePullPolicy: {{ .Values.adminUi.image.pullPolicy }} securityContext: - {{- toYaml .Values.security.containerSecurityContext | nindent 12 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" (.Values.adminUi.containerSecurityContext | default dict) "default" .Values.security.containerSecurityContext) | nindent 12 }} ports: - containerPort: {{ .Values.adminUi.service.targetPort }} + name: admin-ui # nginx serves the SPA at /app/; probe it so a crash-looping or stuck # pod is restarted and gated out of rollout traffic. + {{- if .Values.adminUi.probes.startup }} + startupProbe: + {{- toYaml .Values.adminUi.probes.startup | nindent 12 }} + {{- end }} + {{- if .Values.adminUi.probes.readiness }} readinessProbe: - httpGet: - path: /app/ - port: {{ .Values.adminUi.service.targetPort }} - initialDelaySeconds: 5 - periodSeconds: 10 + {{- toYaml .Values.adminUi.probes.readiness | nindent 12 }} + {{- end }} + {{- if .Values.adminUi.probes.liveness }} livenessProbe: - httpGet: - path: /app/ - port: {{ .Values.adminUi.service.targetPort }} - initialDelaySeconds: 10 - periodSeconds: 20 + {{- toYaml .Values.adminUi.probes.liveness | nindent 12 }} + {{- end }} resources: {{- toYaml .Values.adminUi.resources | nindent 12 }} --- apiVersion: v1 kind: Service metadata: - name: {{ .Release.Name }}-admin-ui + name: {{ include "openrag-stack.fullname" . }}-admin-ui labels: - app.kubernetes.io/name: admin-ui - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentLabels" (dict "component" "admin-ui" "context" $) | nindent 4 }} spec: type: {{ .Values.adminUi.service.type }} selector: - app.kubernetes.io/name: admin-ui - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentSelectorLabels" (dict "component" "admin-ui" "context" $) | nindent 4 }} ports: - - port: {{ .Values.adminUi.service.port }} + - name: admin-ui + port: {{ .Values.adminUi.service.port }} targetPort: {{ .Values.adminUi.service.targetPort }} + {{- if eq .Values.adminUi.service.type "NodePort" }} + nodePort: {{ .Values.adminUi.service.nodePort }} + {{- end }} {{- end }} diff --git a/infra/charts/openrag-stack/templates/configmap-env.yaml b/infra/charts/openrag-stack/templates/configmap-env.yaml index fa1a5da9e..e7b2483f4 100644 --- a/infra/charts/openrag-stack/templates/configmap-env.yaml +++ b/infra/charts/openrag-stack/templates/configmap-env.yaml @@ -1,8 +1,25 @@ apiVersion: v1 kind: ConfigMap metadata: - name: rag-env + name: {{ include "openrag-stack.fullname" . }}-env data: {{- range $key, $value := .Values.env.config }} {{ $key }}: "{{ tpl (printf "%v" $value) $ }}" {{- end }} +{{- if .Values.proxy.enabled }} +{{- if or (empty .Values.proxy.url) (eq .Values.proxy.url "http://192.168.100.100:80") }} +{{- fail "proxy.enabled is true but proxy.url is not set or still uses the placeholder default. Set a valid proxy URL in proxy.url." }} +{{- end }} +{{- $proxyKeys := list "HTTP_PROXY" "HTTPS_PROXY" "http_proxy" "https_proxy" "NO_PROXY" "no_proxy" }} +{{- range $proxyKeys }} +{{- if hasKey $.Values.env.config . }} +{{- fail (printf "proxy.enabled is true but env.config already defines %q — remove it from env.config to avoid a duplicate ConfigMap key." .) }} +{{- end }} +{{- end }} + HTTP_PROXY: {{ .Values.proxy.url | quote }} + HTTPS_PROXY: {{ .Values.proxy.url | quote }} + http_proxy: {{ .Values.proxy.url | quote }} + https_proxy: {{ .Values.proxy.url | quote }} + NO_PROXY: {{ .Values.proxy.noProxy | quote }} + no_proxy: {{ .Values.proxy.noProxy | quote }} +{{- end }} diff --git a/infra/charts/openrag-stack/templates/extra-objects.yaml b/infra/charts/openrag-stack/templates/extra-objects.yaml new file mode 100644 index 000000000..b058f886d --- /dev/null +++ b/infra/charts/openrag-stack/templates/extra-objects.yaml @@ -0,0 +1,4 @@ +{{- range .Values.extraObjects }} +--- +{{ tpl (toYaml .) $ }} +{{- end }} diff --git a/infra/charts/openrag-stack/templates/infinity.yaml b/infra/charts/openrag-stack/templates/infinity.yaml index dedc1b11f..258198f19 100644 --- a/infra/charts/openrag-stack/templates/infinity.yaml +++ b/infra/charts/openrag-stack/templates/infinity.yaml @@ -1,32 +1,37 @@ +{{- if .Values.reranker.enabled }} apiVersion: apps/v1 kind: Deployment metadata: - name: {{ .Release.Name }}-reranker + name: {{ include "openrag-stack.fullname" . }}-reranker labels: - app.kubernetes.io/name: reranker - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentLabels" (dict "component" "reranker" "context" $) | nindent 4 }} spec: replicas: {{ .Values.reranker.replicas }} selector: matchLabels: - app.kubernetes.io/name: reranker - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentSelectorLabels" (dict "component" "reranker" "context" $) | nindent 6 }} template: metadata: labels: - app.kubernetes.io/name: reranker - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentSelectorLabels" (dict "component" "reranker" "context" $) | nindent 8 }} spec: - automountServiceAccountToken: {{ .Values.security.automountServiceAccountToken }} securityContext: - {{- toYaml .Values.security.podSecurityContext | nindent 8 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" .Values.reranker.podSecurityContext "default" .Values.security.podSecurityContext) | nindent 8 }} + {{- with .Values.reranker.serviceAccountName }} + serviceAccountName: {{ tpl . $ }} + {{- end }} + automountServiceAccountToken: {{ include "openrag-stack.securityFieldOverride" (dict "component" .Values.reranker "key" "automountServiceAccountToken" "default" .Values.security.automountServiceAccountToken) }} + {{- if .Values.reranker.runtimeClassName }} runtimeClassName: {{ .Values.reranker.runtimeClassName }} - nodeSelector: {{- toYaml .Values.reranker.nodeSelector | nindent 8 }} + {{- end }} + {{- with .Values.reranker.nodeSelector }} + nodeSelector: {{- toYaml . | nindent 8 }} + {{- end }} containers: - name: reranker - image: "{{ .Values.reranker.image.repository }}:{{ .Values.reranker.image.tag }}" + image: "{{ .Values.reranker.image.registry | default .Values.global.image.registry }}/{{ .Values.reranker.image.repository }}:{{ .Values.reranker.image.tag }}" securityContext: - {{- toYaml .Values.security.containerSecurityContext | nindent 12 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" (.Values.reranker.containerSecurityContext | default dict) "default" .Values.security.containerSecurityContext) | nindent 12 }} args: - "v2" - "--model-id" @@ -35,32 +40,47 @@ spec: - "{{ .Values.reranker.service.port }}" ports: - containerPort: {{ .Values.reranker.service.port }} + name: reranker resources: {{- toYaml .Values.reranker.resources | nindent 12 }} + readinessProbe: + httpGet: + path: /health + port: reranker + initialDelaySeconds: 30 + periodSeconds: 10 + failureThreshold: 6 + livenessProbe: + httpGet: + path: /health + port: reranker + initialDelaySeconds: 60 + periodSeconds: 30 + failureThreshold: 3 volumeMounts: - name: hf-cache mountPath: /app/.cache/huggingface envFrom: - configMapRef: - name: rag-env + name: {{ include "openrag-stack.fullname" . }}-env volumes: - name: hf-cache persistentVolumeClaim: - claimName: rag-model-weights + claimName: {{ include "openrag-stack.fullname" . }}-model-weights --- apiVersion: v1 kind: Service metadata: - name: {{ .Release.Name }}-reranker + name: {{ include "openrag-stack.fullname" . }}-reranker labels: - app.kubernetes.io/name: reranker - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentLabels" (dict "component" "reranker" "context" $) | nindent 4 }} spec: selector: - app.kubernetes.io/name: reranker - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentSelectorLabels" (dict "component" "reranker" "context" $) | nindent 4 }} ports: - - protocol: TCP + - name: reranker + protocol: TCP port: {{ .Values.reranker.service.port }} - targetPort: {{ .Values.reranker.service.port }} + targetPort: reranker type: {{ .Values.reranker.service.type }} +{{- end }} diff --git a/infra/charts/openrag-stack/templates/ingress.yaml b/infra/charts/openrag-stack/templates/ingress.yaml deleted file mode 100644 index 27f1f5f85..000000000 --- a/infra/charts/openrag-stack/templates/ingress.yaml +++ /dev/null @@ -1,53 +0,0 @@ -{{- if .Values.ingress.enabled }} -apiVersion: networking.k8s.io/v1 -kind: Ingress -metadata: - name: {{ include "openrag-stack.fullname" . }}-ingress - annotations: - nginx.ingress.kubernetes.io/rewrite-target: /$1 - nginx.ingress.kubernetes.io/use-regex: "true" - # Match the API's own upload cap (MAX_UPLOAD_SIZE_MB, default 1024) so large - # document ingestion isn't rejected at the Ingress with a 413 before it ever - # reaches the API. ingress-nginx otherwise defaults this to 1m. - nginx.ingress.kubernetes.io/proxy-body-size: "1024m" - -spec: - ingressClassName: {{ .Values.ingress.className | quote }} - {{- if .Values.ingress.tls.enabled }} - tls: - - hosts: - - {{ required "ingress.host must be set when ingress.tls.enabled=true" .Values.ingress.host | quote }} - {{- with .Values.ingress.tls.secretName }} - secretName: {{ . | quote }} - {{- end }} - {{- end }} - rules: - - host: {{ required "ingress.host must be set when ingress.enabled=true" .Values.ingress.host | quote }} - http: - paths: - {{- if .Values.adminUi.enabled }} - # Admin UI SPA (nginx serves it under /app/). Listed first and more - # specific than the API catch-all below, so only /app and /app/* reach - # the UI pod; the SPA's own same-origin /v1, /auth, … calls fall through - # to the API. The trailing `$` anchors the match so sibling paths like - # /applications don't get misrouted to the UI (whose baked nginx config - # can only reverse-proxy non-/app paths over Docker DNS, unreachable in - # k8s) — every non-/app request must land on the API catch-all below. - - path: /(app(?:/.*)?)$ - pathType: ImplementationSpecific - backend: - service: - name: {{ .Release.Name }}-admin-ui - port: - number: {{ .Values.adminUi.service.port }} - {{- end }} - {{- if .Values.ingress.paths.openrag.enabled }} - - path: /(.*) - pathType: ImplementationSpecific - backend: - service: - name: raycluster-head-svc - port: - number: {{ .Values.env.config.RAY_SERVE_PORT }} - {{- end }} -{{- end }} diff --git a/infra/charts/openrag-stack/templates/openrag.yaml b/infra/charts/openrag-stack/templates/openrag.yaml index bbe5ffb40..ccdba23e2 100644 --- a/infra/charts/openrag-stack/templates/openrag.yaml +++ b/infra/charts/openrag-stack/templates/openrag.yaml @@ -1,30 +1,37 @@ apiVersion: apps/v1 kind: Deployment metadata: - name: {{ .Release.Name }}-openrag + name: {{ include "openrag-stack.fullname" . }}-openrag labels: - app.kubernetes.io/name: openrag - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentLabels" (dict "component" "openrag" "context" $) | nindent 4 }} spec: replicas: {{ .Values.openrag.replicas }} selector: matchLabels: - app.kubernetes.io/name: openrag - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentSelectorLabels" (dict "component" "openrag" "context" $) | nindent 6 }} template: metadata: + annotations: + checksum/config: {{ include (print $.Template.BasePath "/configmap-env.yaml") . | sha256sum }} labels: - app.kubernetes.io/name: openrag - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentSelectorLabels" (dict "component" "openrag" "context" $) | nindent 8 }} spec: - automountServiceAccountToken: {{ .Values.security.automountServiceAccountToken }} securityContext: - {{- toYaml .Values.security.podSecurityContext | nindent 8 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" .Values.openrag.podSecurityContext "default" .Values.security.podSecurityContext) | nindent 8 }} + {{- with .Values.openrag.serviceAccountName }} + serviceAccountName: {{ tpl . $ }} + {{- end }} + automountServiceAccountToken: {{ include "openrag-stack.securityFieldOverride" (dict "component" .Values.openrag "key" "automountServiceAccountToken" "default" .Values.security.automountServiceAccountToken) }} + {{- if or .Values.openrag.extraInitContainers .Values.ray.enabled }} initContainers: + {{- with .Values.openrag.extraInitContainers }} + {{- toYaml . | nindent 8 }} + {{- end }} + {{- if .Values.ray.enabled }} - name: init-venv - image: "{{ .Values.openrag.image.repository }}:{{ .Values.openrag.image.tag }}" + image: "{{ .Values.openrag.image.registry | default .Values.global.image.registry }}/{{ .Values.openrag.image.repository }}:{{ .Values.openrag.image.tag }}" securityContext: - {{- toYaml .Values.security.containerSecurityContext | nindent 12 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" (.Values.openrag.containerSecurityContext | default dict) "default" .Values.security.containerSecurityContext) | nindent 12 }} command: - sh - -c @@ -37,20 +44,40 @@ spec: volumeMounts: - name: venv mountPath: /app/.venv + {{- end }} + {{- end }} containers: - name: openrag - image: "{{ .Values.openrag.image.repository }}:{{ .Values.openrag.image.tag }}" + image: "{{ .Values.openrag.image.registry | default .Values.global.image.registry }}/{{ .Values.openrag.image.repository }}:{{ .Values.openrag.image.tag }}" + imagePullPolicy: {{ .Values.openrag.image.pullPolicy }} securityContext: - {{- toYaml .Values.security.containerSecurityContext | nindent 12 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" (.Values.openrag.containerSecurityContext | default dict) "default" .Values.security.containerSecurityContext) | nindent 12 }} ports: - - containerPort: {{ .Values.openrag.service.port }} + - containerPort: {{ .Values.openrag.service.targetPort }} + name: openrag + {{- if .Values.openrag.rayDashboard.ingress.enabled }} + - containerPort: 8265 + name: ray-dashboard + {{- end }} resources: {{- toYaml .Values.openrag.resources | nindent 12 }} + {{- if .Values.openrag.probes.startup }} + startupProbe: + {{- toYaml .Values.openrag.probes.startup | nindent 12 }} + {{- end }} + {{- if .Values.openrag.probes.readiness }} + readinessProbe: + {{- toYaml .Values.openrag.probes.readiness | nindent 12 }} + {{- end }} + {{- if .Values.openrag.probes.liveness }} + livenessProbe: + {{- toYaml .Values.openrag.probes.liveness | nindent 12 }} + {{- end }} envFrom: - configMapRef: - name: rag-env + name: {{ include "openrag-stack.fullname" . }}-env - secretRef: - name: rag-env-secrets + name: {{ include "openrag-stack.secretName" . }} volumeMounts: - name: model-weights mountPath: /app/model_weights @@ -60,33 +87,133 @@ spec: mountPath: /app/logs - name: venv mountPath: /app/.venv + {{- if .Values.openrag.shmSize }} + - name: dshm + mountPath: /dev/shm + {{- end }} + {{- with .Values.openrag.extraVolumeMounts }} + {{- toYaml . | nindent 12 }} + {{- end }} volumes: - name: model-weights persistentVolumeClaim: - claimName: rag-model-weights + claimName: {{ include "openrag-stack.fullname" . }}-model-weights - name: data persistentVolumeClaim: - claimName: rag-data + claimName: {{ include "openrag-stack.fullname" . }}-data - name: logs persistentVolumeClaim: - claimName: rag-logs + claimName: {{ include "openrag-stack.fullname" . }}-logs - name: venv persistentVolumeClaim: - claimName: rag-venv + claimName: {{ include "openrag-stack.fullname" . }}-venv + {{- if .Values.openrag.shmSize }} + - name: dshm + emptyDir: + medium: Memory + sizeLimit: {{ .Values.openrag.shmSize }} + {{- end }} + {{- with .Values.openrag.extraVolumes }} + {{- toYaml . | nindent 8 }} + {{- end }} --- apiVersion: v1 kind: Service metadata: - name: {{ .Release.Name }}-openrag + name: {{ include "openrag-stack.fullname" . }}-openrag labels: - app.kubernetes.io/name: openrag - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentLabels" (dict "component" "openrag" "context" $) | nindent 4 }} spec: selector: - app.kubernetes.io/name: openrag - app.kubernetes.io/instance: {{ .Release.Name }} + {{- include "openrag-stack.componentSelectorLabels" (dict "component" "openrag" "context" $) | nindent 4 }} ports: - - protocol: TCP + - name: openrag + protocol: TCP port: {{ .Values.openrag.service.port }} - targetPort: {{ .Values.openrag.service.port }} + targetPort: {{ .Values.openrag.service.targetPort }} + {{- if .Values.openrag.rayDashboard.ingress.enabled }} + - name: ray-dashboard + protocol: TCP + port: 8265 + targetPort: 8265 + {{- end }} type: {{ .Values.openrag.service.type }} +{{- if .Values.openrag.ingress.enabled }} +--- +apiVersion: networking.k8s.io/v1 +kind: Ingress +metadata: + name: {{ include "openrag-stack.fullname" . }}-openrag + {{- with .Values.openrag.ingress.annotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} +spec: + {{- if .Values.openrag.ingress.className }} + ingressClassName: {{ .Values.openrag.ingress.className | quote }} + {{- end }} + rules: + - http: + paths: + {{- if .Values.adminUi.ingress.enabled }} + - path: /app/ + pathType: Prefix + backend: + service: + name: {{ include "openrag-stack.fullname" . }}-admin-ui + port: + name: admin-ui + {{- end }} + - path: / + pathType: Prefix + backend: + service: + {{- /* raycluster-head-svc only exists when ray.enabled renders the + RayCluster (templates/raycluster.yaml) — ENABLE_RAY_SERVE alone + isn't enough, otherwise this Ingress can dangle at a Service + that was never created. */}} + {{- if and .Values.ray.enabled (eq (toString .Values.env.config.ENABLE_RAY_SERVE) "true") }} + name: {{ include "openrag-stack.fullname" . }}-raycluster-head-svc + port: + name: ray-serve + {{- else }} + name: {{ include "openrag-stack.fullname" . }}-openrag + port: + name: openrag + {{- end }} + host: {{ required "openrag.ingress.host must be set when openrag.ingress.enabled=true" .Values.openrag.ingress.host | quote }} + {{- with .Values.openrag.ingress.tls }} + tls: + {{- toYaml . | nindent 4 }} + {{- end }} +{{- end }} +{{- if .Values.openrag.rayDashboard.ingress.enabled }} +--- +apiVersion: networking.k8s.io/v1 +kind: Ingress +metadata: + name: {{ include "openrag-stack.fullname" . }}-openrag-ray-dashboard + {{- with .Values.openrag.rayDashboard.ingress.annotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} +spec: + {{- if .Values.openrag.rayDashboard.ingress.className }} + ingressClassName: {{ .Values.openrag.rayDashboard.ingress.className | quote }} + {{- end }} + rules: + - http: + paths: + - path: {{ .Values.openrag.rayDashboard.ingress.path }} + pathType: {{ .Values.openrag.rayDashboard.ingress.pathType }} + backend: + service: + name: {{ include "openrag-stack.fullname" . }}-openrag + port: + name: ray-dashboard + host: {{ required "openrag.rayDashboard.ingress.host must be set when openrag.rayDashboard.ingress.enabled=true" .Values.openrag.rayDashboard.ingress.host | quote }} + {{- with .Values.openrag.rayDashboard.ingress.tls }} + tls: + {{- toYaml . | nindent 4 }} + {{- end }} +{{- end }} diff --git a/infra/charts/openrag-stack/templates/postgres-migration-job.yaml b/infra/charts/openrag-stack/templates/postgres-migration-job.yaml index afc709d95..17ed24a58 100644 --- a/infra/charts/openrag-stack/templates/postgres-migration-job.yaml +++ b/infra/charts/openrag-stack/templates/postgres-migration-job.yaml @@ -14,6 +14,14 @@ Disabled by default via `postgresProvisioning.migrationJob.enabled`. Enable it for managed Postgres and pair it with `runMigrationsInApp: false` so the app no longer migrates at startup. When disabled, the application runs migrations itself during startup instead. + +It runs the same image, UID and security context as the OpenRAG Deployment, so +it also reuses `openrag.serviceAccountName`. That matters on OpenShift: the pod +requests `runAsUser: 10001`, which only a SecurityContextConstraints bound to +that ServiceAccount permits. Falling back to the namespace's `default` SA would +put the Job under `restricted-v2` (MustRunAsRange) instead, and admission would +reject a UID outside the namespace's assigned range — the Job would fail as a +pre-upgrade hook and abort the whole release. */}} {{- if .Values.postgresProvisioning.migrationJob.enabled }} apiVersion: batch/v1 @@ -36,14 +44,17 @@ spec: app.kubernetes.io/component: postgres-migration spec: restartPolicy: Never - automountServiceAccountToken: {{ .Values.security.automountServiceAccountToken }} + {{- with .Values.openrag.serviceAccountName }} + serviceAccountName: {{ tpl . $ }} + {{- end }} + automountServiceAccountToken: {{ include "openrag-stack.securityFieldOverride" (dict "component" .Values.openrag "key" "automountServiceAccountToken" "default" .Values.security.automountServiceAccountToken) }} securityContext: - {{- toYaml .Values.security.podSecurityContext | nindent 8 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" .Values.openrag.podSecurityContext "default" .Values.security.podSecurityContext) | nindent 8 }} containers: - name: migrate - image: "{{ .Values.openrag.image.repository }}:{{ .Values.openrag.image.tag }}" + image: "{{ .Values.openrag.image.registry | default .Values.global.image.registry }}/{{ .Values.openrag.image.repository }}:{{ .Values.openrag.image.tag }}" securityContext: - {{- toYaml .Values.security.containerSecurityContext | nindent 12 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" (.Values.openrag.containerSecurityContext | default dict) "default" .Values.security.containerSecurityContext) | nindent 12 }} workingDir: /app/openrag command: - sh @@ -53,9 +64,9 @@ spec: uv run python -m services.persistence.migrations.run envFrom: - configMapRef: - name: rag-env + name: {{ include "openrag-stack.fullname" . }}-env - secretRef: - name: rag-env-secrets + name: {{ include "openrag-stack.secretName" . }} env: # The runner (services.persistence.migrations.run) always runs # migrations and never auto-creates the database, so it ignores diff --git a/infra/charts/openrag-stack/templates/pvc.yaml b/infra/charts/openrag-stack/templates/pvc.yaml index 8694b9bd3..324e6122b 100644 --- a/infra/charts/openrag-stack/templates/pvc.yaml +++ b/infra/charts/openrag-stack/templates/pvc.yaml @@ -2,57 +2,72 @@ apiVersion: v1 kind: PersistentVolumeClaim metadata: - name: rag-model-weights + name: {{ include "openrag-stack.fullname" . }}-model-weights + {{- with .Values.persistence.annotations }} annotations: - "helm.sh/resource-policy": keep - + {{- toYaml . | nindent 4 }} + {{- end }} spec: accessModes: - {{ .Values.persistence.accessMode }} resources: requests: storage: {{ .Values.persistence.volumes.modelWeights.size }} + {{- if .Values.persistence.storageClass }} storageClassName: {{ .Values.persistence.storageClass }} + {{- end }} --- apiVersion: v1 kind: PersistentVolumeClaim metadata: - name: rag-data + name: {{ include "openrag-stack.fullname" . }}-data + {{- with .Values.persistence.annotations }} annotations: - "helm.sh/resource-policy": keep + {{- toYaml . | nindent 4 }} + {{- end }} spec: accessModes: - {{ .Values.persistence.accessMode }} resources: requests: storage: {{ .Values.persistence.volumes.data.size }} + {{- if .Values.persistence.storageClass }} storageClassName: {{ .Values.persistence.storageClass }} + {{- end }} --- apiVersion: v1 kind: PersistentVolumeClaim metadata: - name: rag-logs + name: {{ include "openrag-stack.fullname" . }}-logs + {{- with .Values.persistence.annotations }} annotations: - "helm.sh/resource-policy": keep + {{- toYaml . | nindent 4 }} + {{- end }} spec: accessModes: - {{ .Values.persistence.accessMode }} resources: requests: storage: {{ .Values.persistence.volumes.logs.size }} + {{- if .Values.persistence.storageClass }} storageClassName: {{ .Values.persistence.storageClass }} + {{- end }} --- apiVersion: v1 kind: PersistentVolumeClaim metadata: - name: rag-venv + name: {{ include "openrag-stack.fullname" . }}-venv + {{- with .Values.persistence.annotations }} annotations: - "helm.sh/resource-policy": keep + {{- toYaml . | nindent 4 }} + {{- end }} spec: accessModes: - {{ .Values.persistence.accessMode }} resources: requests: storage: {{ .Values.persistence.volumes.venv.size }} + {{- if .Values.persistence.storageClass }} storageClassName: {{ .Values.persistence.storageClass }} + {{- end }} {{- end }} diff --git a/infra/charts/openrag-stack/templates/raycluster.yaml b/infra/charts/openrag-stack/templates/raycluster.yaml index 0d81e9e58..af2ed6a80 100644 --- a/infra/charts/openrag-stack/templates/raycluster.yaml +++ b/infra/charts/openrag-stack/templates/raycluster.yaml @@ -1,26 +1,35 @@ +{{- if .Values.ray.enabled }} apiVersion: ray.io/v1 kind: RayCluster metadata: - name: raycluster + name: {{ include "openrag-stack.fullname" . }}-raycluster + labels: + {{- include "openrag-stack.componentLabels" (dict "component" "raycluster" "context" $) | nindent 4 }} + {{- with .Values.ray.annotations }} annotations: - ray.io/overwrite-container-cmd: "true" + {{- . | toYaml | nindent 4 }} + {{- end }} spec: - rayVersion: "2.47.1" - enableInTreeAutoscaling: false - + rayVersion: {{ .Values.ray.rayVersion }} + enableInTreeAutoscaling: {{ .Values.ray.enableInTreeAutoscaling }} headGroupSpec: serviceType: ClusterIP template: spec: - automountServiceAccountToken: {{ .Values.security.automountServiceAccountToken }} + {{- with .Values.ray.nodeSelector }} + nodeSelector: {{- toYaml . | nindent 12 }} + {{- end }} + {{- with .Values.ray.serviceAccountName }} + serviceAccountName: {{ tpl . $ }} + {{- end }} + automountServiceAccountToken: {{ include "openrag-stack.securityFieldOverride" (dict "component" .Values.ray "key" "automountServiceAccountToken" "default" .Values.security.automountServiceAccountToken) }} securityContext: - {{- toYaml .Values.security.podSecurityContext | nindent 10 }} - nodeSelector: {{- toYaml .Values.ray.nodeSelector | nindent 12 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" .Values.ray.podSecurityContext "default" .Values.security.podSecurityContext) | nindent 10 }} initContainers: - name: init-venv-sync - image: {{ .Values.ray.image.repository }}:{{ .Values.ray.image.tag }} + image: "{{ .Values.ray.image.registry | default .Values.global.image.registry }}/{{ .Values.ray.image.repository }}:{{ .Values.ray.image.tag }}" securityContext: - {{- toYaml .Values.security.containerSecurityContext | nindent 14 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" (.Values.ray.containerSecurityContext | default dict) "default" .Values.security.containerSecurityContext) | nindent 14 }} command: - sh - -c @@ -36,12 +45,11 @@ spec: volumeMounts: - name: venv mountPath: /app/.venv - containers: - name: ray-head - image: {{ .Values.ray.image.repository }}:{{ .Values.ray.image.tag }} + image: "{{ .Values.ray.image.registry | default .Values.global.image.registry }}/{{ .Values.ray.image.repository }}:{{ .Values.ray.image.tag }}" securityContext: - {{- toYaml .Values.security.containerSecurityContext | nindent 14 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" (.Values.ray.containerSecurityContext | default dict) "default" .Values.security.containerSecurityContext) | nindent 14 }} command: ["uv"] args: - "run" @@ -50,19 +58,23 @@ spec: - "--head" - "--dashboard-host={{ .Values.ray.dashboardHost }}" - "--dashboard-agent-listen-port=52365" - - "--metrics-export-port=8080" + # Kept off networkPolicy.externalPorts' 8080 (openrag/admin-ui's + # public port) on purpose — that NetworkPolicy rule matches by + # port number across every pod in the namespace, so reusing 8080 + # here would expose these unauthenticated Ray metrics publicly. + - "--metrics-export-port=8090" - "--block" ports: - containerPort: 80 - name: serve + name: ray-serve - containerPort: 8265 - name: dashboard + name: ray-dashboard - containerPort: 10001 - name: client + name: ray-client - containerPort: 6379 - name: gcs - - containerPort: 8080 - name: metrics + name: ray-gcs + - containerPort: 8090 + name: ray-metrics volumeMounts: - name: model-weights mountPath: /app/model_weights @@ -74,26 +86,26 @@ spec: mountPath: /app/.venv envFrom: - configMapRef: - name: rag-env + name: {{ include "openrag-stack.fullname" . }}-env - secretRef: - name: rag-env-secrets + name: {{ include "openrag-stack.secretName" . }} + {{- with .Values.ray.resources.head }} resources: - limits: - nvidia.com/gpu: 1 + {{- toYaml . | nindent 14 }} + {{- end }} volumes: - name: model-weights persistentVolumeClaim: - claimName: rag-model-weights + claimName: {{ include "openrag-stack.fullname" . }}-model-weights - name: data persistentVolumeClaim: - claimName: rag-data + claimName: {{ include "openrag-stack.fullname" . }}-data - name: logs persistentVolumeClaim: - claimName: rag-logs + claimName: {{ include "openrag-stack.fullname" . }}-logs - name: venv persistentVolumeClaim: - claimName: rag-venv - + claimName: {{ include "openrag-stack.fullname" . }}-venv workerGroupSpecs: {{- range $i, $e := until (.Values.ray.workers.count | int) }} - groupName: worker-group-{{ add $i 1 }} @@ -102,31 +114,20 @@ spec: maxReplicas: {{ $.Values.ray.workers.maxReplicas }} template: spec: - automountServiceAccountToken: {{ $.Values.security.automountServiceAccountToken }} + {{- with $.Values.ray.nodeSelector }} + nodeSelector: {{- toYaml . | nindent 14 }} + {{- end }} + {{- with $.Values.ray.serviceAccountName }} + serviceAccountName: {{ tpl . $ }} + {{- end }} + automountServiceAccountToken: {{ include "openrag-stack.securityFieldOverride" (dict "component" $.Values.ray "key" "automountServiceAccountToken" "default" $.Values.security.automountServiceAccountToken) }} securityContext: - {{- toYaml $.Values.security.podSecurityContext | nindent 12 }} - initContainers: - - name: init-venv - image: {{ $.Values.ray.image.repository }}:{{ $.Values.ray.image.tag }} - securityContext: - {{- toYaml $.Values.security.containerSecurityContext | nindent 16 }} - command: - - sh - - -c - - | - echo "Waiting for Python env in /app/.venv..." - while [ ! -f /app/.venv/.ready ]; do - sleep 3 - done - echo "Venv is ready." - volumeMounts: - - name: venv - mountPath: /app/.venv + {{- include "openrag-stack.mergeSecurityContext" (dict "component" $.Values.ray.podSecurityContext "default" $.Values.security.podSecurityContext) | nindent 12 }} containers: - name: ray-worker - image: {{ $.Values.ray.image.repository }}:{{ $.Values.ray.image.tag }} + image: "{{ $.Values.ray.image.registry | default $.Values.global.image.registry }}/{{ $.Values.ray.image.repository }}:{{ $.Values.ray.image.tag }}" securityContext: - {{- toYaml $.Values.security.containerSecurityContext | nindent 16 }} + {{- include "openrag-stack.mergeSecurityContext" (dict "component" ($.Values.ray.containerSecurityContext | default dict) "default" $.Values.security.containerSecurityContext) | nindent 16 }} command: ["/bin/bash", "-lc", "--"] args: - uv run $KUBERAY_GEN_RAY_START_CMD @@ -141,23 +142,55 @@ spec: mountPath: /app/.venv envFrom: - configMapRef: - name: rag-env + name: {{ include "openrag-stack.fullname" $ }}-env - secretRef: - name: rag-env-secrets + name: {{ include "openrag-stack.secretName" $ }} + {{- with $.Values.ray.resources.worker }} resources: - limits: - nvidia.com/gpu: 1 + {{- toYaml . | nindent 16 }} + {{- end }} volumes: - name: model-weights persistentVolumeClaim: - claimName: rag-model-weights + claimName: {{ include "openrag-stack.fullname" $ }}-model-weights - name: data persistentVolumeClaim: - claimName: rag-data + claimName: {{ include "openrag-stack.fullname" $ }}-data - name: logs persistentVolumeClaim: - claimName: rag-logs + claimName: {{ include "openrag-stack.fullname" $ }}-logs - name: venv persistentVolumeClaim: - claimName: rag-venv + claimName: {{ include "openrag-stack.fullname" $ }}-venv + {{- end }} +{{- if .Values.ray.ingress.enabled }} +--- +apiVersion: networking.k8s.io/v1 +kind: Ingress +metadata: + name: {{ include "openrag-stack.fullname" . }}-raycluster + {{- with .Values.ray.ingress.annotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} +spec: + {{- if .Values.ray.ingress.className }} + ingressClassName: {{ .Values.ray.ingress.className | quote }} + {{- end }} + rules: + - http: + paths: + - path: {{ .Values.ray.ingress.path }} + pathType: {{ .Values.ray.ingress.pathType }} + backend: + service: + name: {{ include "openrag-stack.fullname" . }}-raycluster-head-svc + port: + name: ray-dashboard + host: {{ required "ray.ingress.host must be set when ray.ingress.enabled=true" .Values.ray.ingress.host | quote }} + {{- with .Values.ray.ingress.tls }} + tls: + {{- toYaml . | nindent 4 }} {{- end }} +{{- end }} +{{- end }} diff --git a/infra/charts/openrag-stack/templates/secrets-env.yaml b/infra/charts/openrag-stack/templates/secrets-env.yaml index 08a374fd7..4476e0822 100644 --- a/infra/charts/openrag-stack/templates/secrets-env.yaml +++ b/infra/charts/openrag-stack/templates/secrets-env.yaml @@ -1,7 +1,14 @@ +{{- $secretName := include "openrag-stack.secretName" . }} +{{- $type := .Values.env.secretsProvider.type | default "values" }} +{{- if not .Values.env.existingSecret }} + +{{- if eq $type "values" }} +--- apiVersion: v1 kind: Secret metadata: - name: rag-env-secrets + name: {{ $secretName }} + labels: {{- include "openrag-stack.labels" . | nindent 4 }} type: Opaque stringData: {{- $requiredSecrets := list "AUTH_TOKEN" "POSTGRES_PASSWORD" }} @@ -11,16 +18,89 @@ stringData: {{- if not (hasKey $secrets $requiredKey) }} {{- fail (printf "env.secrets.%s must be set before installing the chart" $requiredKey) }} {{- end }} - {{- $requiredRawValue := get $secrets $requiredKey }} - {{- $requiredValue := tpl (printf "%v" $requiredRawValue) $ }} - {{- if or (empty $requiredRawValue) (empty $requiredValue) }} - {{- fail (printf "env.secrets.%s must be set before installing the chart" $requiredKey) }} - {{- end }} {{- end }} -{{- range $key, $value := $secrets }} - {{- $stringValue := tpl (printf "%v" $value) $ }} +{{- range $key, $rawValue := $secrets }} + {{- $stringValue := "" }} + {{- /* postgresql.auth.password is the actual credential Postgres enforces — + read it live here instead of trusting env.secrets.POSTGRES_PASSWORD's YAML + anchor, which only resolves within values.yaml itself and goes stale the + moment postgresql.auth.password is overridden via --set/-f (anchors don't + survive Helm's values merge). Falls back to the literal for external + Postgres (postgresql.enabled=false), where this key is set directly. */}} + {{- if and (eq $key "POSTGRES_PASSWORD") $.Values.postgresql.enabled }} + {{- $stringValue = printf "%v" $.Values.postgresql.auth.password }} + {{- else }} + {{- $stringValue = tpl (printf "%v" $rawValue) $ }} + {{- end }} + {{- if has $key $requiredSecrets }} + {{- if empty $stringValue }} + {{- fail (printf "env.secrets.%s must be set before installing the chart" $key) }} + {{- end }} + {{- end }} {{- if has $stringValue $placeholderSecrets }} {{- fail (printf "env.secrets.%s still uses an unsafe placeholder value" $key) }} {{- end }} {{ $key }}: {{ $stringValue | quote }} {{- end }} + +{{- else if eq $type "externalSecret" }} +--- +# Requires External Secrets Operator >= 0.14 (first release that serves external-secrets.io/v1). +# On older installs only v1beta1 is available — pin apiVersion accordingly or upgrade the operator. +apiVersion: external-secrets.io/v1 +kind: ExternalSecret +metadata: + name: {{ $secretName }} + labels: {{- include "openrag-stack.labels" . | nindent 4 }} +spec: + refreshInterval: {{ .Values.env.secretsProvider.externalSecret.refreshInterval | quote }} + secretStoreRef: + name: {{ .Values.env.secretsProvider.externalSecret.secretStore.name | quote }} + kind: {{ .Values.env.secretsProvider.externalSecret.secretStore.kind | quote }} + target: + name: {{ $secretName }} + creationPolicy: Owner + {{- if .Values.env.secretsProvider.externalSecret.dataFrom }} + dataFrom: + {{- toYaml .Values.env.secretsProvider.externalSecret.dataFrom | nindent 4 }} + {{- else if .Values.env.secretsProvider.externalSecret.data }} + data: + {{- toYaml .Values.env.secretsProvider.externalSecret.data | nindent 4 }} + {{- else }} + {{- fail "externalSecret requires either dataFrom or data to be set" }} + {{- end }} + +{{- else if eq $type "vaultStaticSecret" }} +--- +apiVersion: secrets.hashicorp.com/v1beta1 +kind: VaultStaticSecret +metadata: + name: {{ $secretName }} + labels: {{- include "openrag-stack.labels" . | nindent 4 }} +spec: + type: {{ .Values.env.secretsProvider.vaultStaticSecret.type | quote }} + mount: {{ .Values.env.secretsProvider.vaultStaticSecret.mount | quote }} + path: {{ .Values.env.secretsProvider.vaultStaticSecret.path | quote }} + {{- if .Values.env.secretsProvider.vaultStaticSecret.vaultAuthRef }} + vaultAuthRef: {{ .Values.env.secretsProvider.vaultStaticSecret.vaultAuthRef | quote }} + {{- end }} + refreshAfter: {{ .Values.env.secretsProvider.vaultStaticSecret.refreshAfter | quote }} + hmacSecretData: {{ .Values.env.secretsProvider.vaultStaticSecret.hmacSecretData }} + destination: + create: true + name: {{ $secretName }} + rolloutRestartTargets: + - kind: Deployment + name: {{ include "openrag-stack.fullname" . }}-openrag + {{- if .Values.adminUi.enabled }} + - kind: Deployment + name: {{ include "openrag-stack.fullname" . }}-admin-ui + {{- end }} + # Note: the Ray head and workers also mount this secret via secretRef but + # VSO rolloutRestartTargets only supports Deployment/StatefulSet/DaemonSet — + # RayCluster pods managed by KubeRay must be cycled manually after rotation. + +{{- else }} +{{- fail (printf "invalid env.secretsProvider.type: %q (accepted values: values, externalSecret, vaultStaticSecret)" $type) }} +{{- end }} +{{- end }} diff --git a/infra/charts/openrag-stack/values-linagora.yaml b/infra/charts/openrag-stack/values-linagora.yaml new file mode 100644 index 000000000..6bf7042a7 --- /dev/null +++ b/infra/charts/openrag-stack/values-linagora.yaml @@ -0,0 +1,165 @@ +persistence: + storageClass: longhorn + accessMode: ReadWriteMany + annotations: + "helm.sh/resource-policy": keep + +# values.yaml wires the postgresql/milvus sub-charts' storageClass through a +# YAML alias to persistence.storageClass above — aliases resolve when +# values.yaml is parsed, before Helm merges this overlay, so the top-level +# override alone never reaches them. Set it explicitly per sub-chart here too. +postgresql: + primary: + persistence: + storageClass: longhorn + +milvus: + minio: + persistence: + storageClass: longhorn + etcd: + persistence: + storageClass: longhorn + +openrag: + ingress: + className: nginx + annotations: + nginx.ingress.kubernetes.io/rewrite-target: /$1 + nginx.ingress.kubernetes.io/use-regex: "true" + nginx.ingress.kubernetes.io/proxy-body-size: "500m" + +ray: + # This overlay runs the separate RayCluster (GPU node-selectors/resources + # below) rather than the embedded ray.init() fallback — both this flag and + # env.config.ENABLE_RAY_SERVE must be enabled together (see values.yaml). + enabled: true + nodeSelector: + gpu-role: serving + resources: + head: + limits: + nvidia.com/gpu: 1 + worker: + limits: + nvidia.com/gpu: 1 + +env: + config: + ENABLE_RAY_SERVE: "true" + +reranker: + nodeSelector: + gpu-role: serving + runtimeClassName: nvidia + +vllm: + # Single edit point for the HF_TOKEN secret name — anchor reused in the 4 + # hf_token.secretName entries below. Must stay in sync with fullnameOverride + # ("openrag") in values.yaml; keep the full modelSpec array intact when + # changing this (Helm replaces lists wholesale, it doesn't merge by index — + # see templates/NOTES.txt). + hfTokenSecretName: &hfTokenSecret "openrag-env-secrets" + servingEngineSpec: + modelSpec: + - name: "embedder" + repository: "vllm/vllm-openai" + tag: "latest" + modelURL: "Qwen/Qwen3-Embedding-0.6B" + replicaCount: 1 + requestCPU: 4 + requestMemory: "16Gi" + requestGPU: 1 + limitCPU: 4 + limitMemory: "16Gi" + hf_token: + secretName: *hfTokenSecret + secretKey: HF_TOKEN + vllmConfig: + gpuMemoryUtilization: 0.3 + extraArgs: ["--trust-remote-code", "--task", "embed"] + nodeSelectorTerms: + - matchExpressions: + - key: gpu-role + operator: In + values: ["serving"] + + - name: "whisper" + repository: "ghcr.io/linagora/vllm-whisper" + tag: "latest" + modelURL: "openai/whisper-large-v3-turbo" + replicaCount: 1 + requestCPU: 4 + requestMemory: "16Gi" + requestGPU: 1 + limitCPU: 4 + limitMemory: "16Gi" + hf_token: + secretName: *hfTokenSecret + secretKey: HF_TOKEN + vllmConfig: + gpuMemoryUtilization: 0.5 + extraArgs: ["--trust-remote-code", "--async-scheduling"] + nodeSelectorTerms: + - matchExpressions: + - key: gpu-role + operator: In + values: ["serving"] + + - name: "llm" + repository: "vllm/vllm-openai" + tag: "latest" + modelURL: "RedHatAI/Mistral-Small-3.1-24B-Instruct-2503-quantized.w8a8" + replicaCount: 0 + requestCPU: 8 + requestMemory: "32Gi" + requestGPU: 2 + limitCPU: 8 + limitMemory: "32Gi" + hf_token: + secretName: *hfTokenSecret + secretKey: HF_TOKEN + vllmConfig: + v1: 1 + tensorParallelSize: 2 + pipelineParallelSize: 1 + maxModelLen: 32768 + extraArgs: ["--trust-remote-code"] + shmSize: "20Gi" + raySpec: + headNode: + requestCPU: 8 + requestMemory: "32Gi" + requestGPU: 2 + nodeSelectorTerms: + - matchExpressions: + - key: gpu-role + operator: In + values: ["llm"] + pvcStorage: "50Gi" + + - name: "vlm" + repository: "vllm/vllm-openai" + tag: "v0.11.2" + modelURL: "Qwen/Qwen2.5-VL-7B-Instruct" + replicaCount: 1 + requestCPU: 8 + requestMemory: "32Gi" + requestGPU: 1 + limitCPU: 8 + limitMemory: "32Gi" + hf_token: + secretName: *hfTokenSecret + secretKey: HF_TOKEN + vllmConfig: + v1: 1 + tensorParallelSize: 1 + pipelineParallelSize: 1 + maxModelLen: 8192 + extraArgs: ["--trust-remote-code"] + shmSize: "20Gi" + nodeSelectorTerms: + - matchExpressions: + - key: gpu-role + operator: In + values: ["vlm"] diff --git a/infra/charts/openrag-stack/values.yaml b/infra/charts/openrag-stack/values.yaml index 6e36a781d..23e52e495 100644 --- a/infra/charts/openrag-stack/values.yaml +++ b/infra/charts/openrag-stack/values.yaml @@ -1,18 +1,40 @@ +# Pinned so every resource name is deterministic regardless of the Helm +# release name — required for ArgoCD, which by default sets .Release.Name to +# the Application's own name (or whatever releaseName it's given), not +# something this chart controls. postgresql.fullnameOverride and +# milvus.fullnameOverride below are kept in sync with this value by hand +# (subchart values.yaml is static, it can't reference this one). +# NOTE: the vllm-stack sub-chart has no such override hook — its Service names +# (.../*-engine-service) are permanently derived from the real Release.Name +# (see BASE_URL/EMBEDDER_BASE_URL/VLM_BASE_URL/TRANSCRIBER_BASE_URL below). If +# deploying via ArgoCD, set spec.source.helm.releaseName to this same value +# ("openrag") so those URLs resolve too. +fullnameOverride: "openrag" + +global: + image: + registry: "docker.io" + +# Chart-level toggle for umbrella charts that include openrag-stack as a +# dependency with `condition: openrag-stack.enabled` — has no effect standalone. +enabled: true + # === Network isolation === -# Default-deny ingress for all pods in the release namespace, allowing only -# intra-namespace traffic plus the public HTTP ports below. This isolates the -# unauthenticated Ray dashboard (8265), GCS (6379), Postgres, Milvus, etc. from -# outside the namespace. Disable only if you manage isolation elsewhere. +# Default-deny ingress; allow only same-namespace traffic and the public HTTP +# ports below. Isolates the Ray dashboard, Postgres, Milvus, etc. networkPolicy: enabled: true # Ports reachable from outside the namespace (e.g. via the Ingress controller). externalPorts: - - 8080 # openrag API/UI + - 8080 # openrag API/UI & admin-ui + # - 8265 # ray dashboard - Security issue, the ray dashboard is not protected by any auth, so it should not be exposed to the public internet. # === Global shared persistence === persistence: enabled: true - storageClass: &storageClass longhorn + storageClass: &storageClass "" # empty = use cluster default StorageClass + annotations: + "helm.sh/resource-policy": keep accessMode: ReadWriteMany volumes: modelWeights: { size: 50Gi } @@ -21,27 +43,105 @@ persistence: logs: { size: 1Gi } venv: { size: 10Gi } +# === Shared security defaults === +# Pod Security Standards "restricted" baseline, merged into each workload's +# own podSecurityContext below via the openrag-stack.mergeSecurityContext +# helper — components only set the runAsUser/runAsGroup/fsGroup specific to +# their own Dockerfile (see e.g. openrag's OpenShift arbitrary-UID comment) +# and inherit runAsNonRoot/seccompProfile from here. automountServiceAccountToken +# and containerSecurityContext are identical for every workload today, so +# templates read them directly from here with no per-component override. +security: + automountServiceAccountToken: false + podSecurityContext: + runAsNonRoot: true + seccompProfile: + type: RuntimeDefault + containerSecurityContext: + allowPrivilegeEscalation: false + capabilities: + drop: ["ALL"] + readOnlyRootFilesystem: false # workloads write to runtime paths (openrag /tmp, ray .venv, reranker model cache, ...) + +# === Ray (KubeRay operator + RayCluster) === +kuberay: + enabled: false + ray: + enabled: false + # 0.0.0.0 would expose the unauthenticated Ray dashboard on every pod + # interface (ShadowRay / CVE-2023-48022); loopback-only is safe because the + # dashboard is still reachable through ray.ingress / openrag.rayDashboard.ingress. dashboardHost: "127.0.0.1" + rayVersion: "2.47.1" + enableInTreeAutoscaling: false + annotations: + ray.io/overwrite-container-cmd: "true" workers: count: 1 replicas: 1 minReplicas: 1 maxReplicas: 1 - nodeSelector: - gpu-role: serving + serviceAccountName: "" + # infra/docker/ray.Dockerfile creates a plain uid/gid 10001 "app" user (no + # arbitrary-GID pattern like openrag's image) and chowns /app to it. Merged + # with the shared security.podSecurityContext defaults above. + podSecurityContext: + runAsUser: 10001 + runAsGroup: 10001 + fsGroup: 10001 + # Uncomment to override the shared security.automountServiceAccountToken / + # security.containerSecurityContext defaults for this workload specifically. + # automountServiceAccountToken: false + # containerSecurityContext: {} + nodeSelector: {} image: - repository: ghcr.io/linagora/openrag-ray + registry: "ghcr.io" + repository: "linagora/openrag-ray" # Pin to a release tag (ideally a digest) for reproducible deploys. - tag: "v2.0.1" + tag: "v2.1.0" + resources: + head: + requests: + cpu: "4" + memory: "8Gi" + nvidia.com/gpu: 1 + limits: + cpu: "4" + memory: "8Gi" + nvidia.com/gpu: 1 + worker: + requests: + cpu: "4" + memory: "8Gi" + nvidia.com/gpu: 1 + limits: + cpu: "4" + memory: "8Gi" + nvidia.com/gpu: 1 + ingress: + enabled: false + className: "" + host: "" + annotations: {} + # nginx.ingress.kubernetes.io/rewrite-target: / + # nginx.ingress.kubernetes.io/proxy-body-size: "500m" + tls: [] + path: / + pathType: Prefix # === PostgreSQL (bitnami) === postgresql: enabled: true + # Kept in sync with the root fullnameOverride ("openrag") by hand — see the + # note above it. Without this, the postgresql sub-chart names its own + # Service from the real Release.Name instead, which won't match POSTGRES_HOST below. + fullnameOverride: "openrag-postgresql" auth: username: &pgUser root - # Must be supplied at install time; templates fail closed if it is empty. - password: &pgPass "" + # Must be supplied at install time, e.g. --set postgresql.auth.password=$(openssl rand -hex 16) + # — templates fail closed (secrets-env.yaml's requiredSecrets check) if left empty. + password: "" primary: persistence: enabled: true @@ -50,6 +150,7 @@ postgresql: service: port: &pgPort 5432 +# === PostgreSQL provisioning & migrations === postgresProvisioning: # Keep the bundled PostgreSQL chart self-contained by default. Managed # Postgres deployments should set this to false and point POSTGRES_* values @@ -66,8 +167,12 @@ postgresProvisioning: # === Milvus === milvus: enabled: true + # Kept in sync with the root fullnameOverride ("openrag") by hand — see the + # note on it above. Without this, the milvus sub-chart names its own proxy + # Service from the real Release.Name instead, which won't match VDB_HOST. + fullnameOverride: "openrag-milvus" image: - all: { tag: v2.6.0 } + all: { tag: "v2.6.0" } cluster: { enabled: true } pulsarv3: { enabled: false } woodpecker: { enabled: true } @@ -91,12 +196,25 @@ milvus: type: ClusterIP port: 19530 -# === vLLM (embedder) === +# === vLLM (embedder + whisper + llm + vlm) === vllm: + enabled: true embedderModelName: &embedderModel "Qwen/Qwen3-Embedding-0.6B" vlmModelName: &vlmModel "Qwen/Qwen2.5-VL-7B-Instruct" whisperModelName: &whisperModel "openai/whisper-large-v3-turbo" llmModelName: &llmModel "RedHatAI/Mistral-Small-3.1-24B-Instruct-2503-quantized.w8a8" + + # Single source for the HF_TOKEN secret name passed to all four model specs. + # YAML anchor: editing this value propagates to all four hf_token.secretName entries below. + # Must match the chart env secret name ({{ include "openrag-stack.fullname" . }}-env-secrets by default). + # values.yaml is not templated, so if you override fullnameOverride ("openrag") + # this must be updated to match too — see templates/NOTES.txt for a warning if + # they ever drift, and NEVER use "--set vllm.servingEngineSpec.modelSpec[N]...": + # Helm replaces lists wholesale rather than merging them by index, so a + # partial override there silently wipes repository/tag/modelURL/resources + # for all four models (see NOTES.txt for the safe way to override this). + hfTokenSecretName: &vllmHfTokenSecretName "openrag-env-secrets" + servingEngineSpec: enableEngine: true securityContext: @@ -127,16 +245,12 @@ vllm: limitCPU: 4 limitMemory: "16Gi" hf_token: - secretName: rag-env-secrets + secretName: *vllmHfTokenSecretName secretKey: HF_TOKEN vllmConfig: gpuMemoryUtilization: 0.3 extraArgs: ["--trust-remote-code", "--task", "embed"] - nodeSelectorTerms: - - matchExpressions: - - key: gpu-role - operator: In - values: ["serving"] + nodeSelectorTerms: [] - name: "whisper" repository: "ghcr.io/linagora/vllm-whisper" @@ -149,16 +263,12 @@ vllm: limitCPU: 4 limitMemory: "16Gi" hf_token: - secretName: rag-env-secrets + secretName: *vllmHfTokenSecretName secretKey: HF_TOKEN vllmConfig: gpuMemoryUtilization: 0.5 extraArgs: ["--trust-remote-code", "--async-scheduling"] - nodeSelectorTerms: - - matchExpressions: - - key: gpu-role - operator: In - values: ["serving"] + nodeSelectorTerms: [] - name: "llm" repository: "vllm/vllm-openai" @@ -171,7 +281,7 @@ vllm: limitCPU: 8 limitMemory: "32Gi" hf_token: - secretName: rag-env-secrets + secretName: *vllmHfTokenSecretName secretKey: HF_TOKEN vllmConfig: v1: 1 @@ -185,12 +295,9 @@ vllm: requestCPU: 8 requestMemory: "32Gi" requestGPU: 2 - nodeSelectorTerms: - - matchExpressions: - - key: gpu-role - operator: In - values: ["llm"] + nodeSelectorTerms: [] pvcStorage: "50Gi" + - name: "vlm" repository: "vllm/vllm-openai" tag: "v0.11.2" @@ -202,7 +309,7 @@ vllm: limitCPU: 8 limitMemory: "32Gi" hf_token: - secretName: rag-env-secrets + secretName: *vllmHfTokenSecretName secretKey: HF_TOKEN vllmConfig: v1: 1 @@ -211,25 +318,27 @@ vllm: maxModelLen: 8192 extraArgs: ["--trust-remote-code"] shmSize: "20Gi" - nodeSelectorTerms: - - matchExpressions: - - key: gpu-role - operator: In - values: ["vlm"] + nodeSelectorTerms: [] routerSpec: enableRouter: false # === Infinity reranker === reranker: + # Set to false to skip deploying the local Infinity reranker. Also set + # externalUrl below to point RERANKER_BASE_URL at an external one instead — + # otherwise RERANKER_ENABLED renders "false" and no reranker is used at all. + enabled: true + # URL of an external reranker (vLLM or any OpenAI-compatible reranker). + # When set, this overrides RERANKER_BASE_URL (see env.config below) — + # remember to also set enabled: false so the local Deployment/Service + # aren't created alongside it. + externalUrl: "" rerankerModelName: &rerankerModel "Alibaba-NLP/gte-multilingual-reranker-base" servicePort: &rerankerPort 7997 image: - repository: michaelf34/infinity - # Operator-supplied serving image: pin to a specific infinity release tag - # (or digest) before production use rather than tracking latest. - tag: latest - nodeSelector: - gpu-role: serving + repository: "michaelf34/infinity" + tag: "0.0.75" + nodeSelector: {} model: id: Alibaba-NLP/gte-multilingual-reranker-base command: @@ -239,7 +348,7 @@ reranker: - --port - *rerankerPort replicas: 1 - runtimeClassName: nvidia + runtimeClassName: "" service: type: ClusterIP port: *rerankerPort @@ -250,103 +359,278 @@ reranker: limits: cpu: "4" memory: "16Gi" + serviceAccountName: "" + # Merged with the shared security.podSecurityContext defaults above. + podSecurityContext: + runAsUser: 1000 + fsGroup: 1000 + # Uncomment to override the shared security.automountServiceAccountToken / + # security.containerSecurityContext defaults for this workload specifically. + # automountServiceAccountToken: false + # containerSecurityContext: {} -# === Admin UI (React SPA + nginx) === -# Static SPA served by the admin-ui container (nginx-unprivileged, :8080). It is -# same-origin: the SPA makes relative /v1, /auth, … calls that the Ingress routes -# to the API (see ingress.paths below), so there is no CORS. Runs under the same -# hardened security context as the rest of the stack. +# === Admin UI (replaces indexer-ui) === +# React SPA served by an nginx bundled inside the image. VITE_BASE_PATH=/app/ +# and VITE_API_BASE_URL="" are compiled into the JS bundle at image build time +# (see infra/docker/ui.Dockerfile) — the SPA assumes it is served under /app/ +# and makes same-origin, root-relative API calls (/v1, /auth, /chainlit, ...). +# Changing the path prefix requires rebuilding the image with different +# build-args; this chart cannot do it via runtime env. adminUi: enabled: true image: - repository: linagoraai/openrag-admin-ui + registry: "" + repository: "linagoraai/openrag-admin-ui" # Pin to a release tag (ideally a digest) for reproducible deploys. Must be a # build from infra/docker/ui.Dockerfile (nginx-unprivileged, listens :8080). - tag: "v2.0.1" - imagePullPolicy: IfNotPresent + tag: "v2.1.0" + pullPolicy: IfNotPresent replicaCount: 1 service: type: ClusterIP port: 8080 targetPort: 8080 + # Routing note: admin-ui and openrag/ray-serve share a single optional + # Ingress resource (templates/openrag.yaml) — there is no separate Ingress + # for admin-ui. Off by default; this flag only adds the more specific + # "/app/" path onto openrag.ingress's Ingress ahead of its "/" catch-all, + # and has no effect unless openrag.ingress.enabled=true too. + # openrag.ingress.host/className/annotations/tls apply to both paths — + # there is no automatic "/" -> "/app/" redirect (unlike the compose + # front-door), so users must browse to /app/ directly. + ingress: + enabled: false resources: {} + serviceAccountName: "" + # Uncomment to override the shared security.automountServiceAccountToken / + # security.containerSecurityContext defaults for this workload specifically. + # automountServiceAccountToken: false + # containerSecurityContext: {} + # Base image is nginxinc/nginx-unprivileged:1.27-alpine, listening on 8080 so + # it never needs root/NET_BIND_SERVICE. infra/docker/ui.Dockerfile chowns + # /var/cache/nginx, /etc/nginx/conf.d and /var/run to 10001:0 and makes them + # group-writable, so runAsGroup must stay 0 (not 10001) — the same + # arbitrary-UID pattern as openrag.podSecurityContext below: on OpenShift the + # restricted-v2 SCC replaces runAsUser with a UID from the namespace range but + # always keeps GID 0, and group 0 is what keeps those paths writable then. + # fsGroup would not help here instead (it only affects mounted volumes, not + # paths baked into the image, and this container has no volumeMounts at all). + # Merged with the shared security.podSecurityContext defaults above. + podSecurityContext: + runAsUser: 10001 + runAsGroup: 0 + probes: + startup: {} + readiness: + httpGet: + path: /app/ + port: admin-ui + initialDelaySeconds: 5 + periodSeconds: 10 + liveness: + httpGet: + path: /app/ + port: admin-ui + initialDelaySeconds: 10 + periodSeconds: 20 # === OpenRAG main app === openrag: image: - repository: linagoraai/openrag + registry: "" + repository: "linagoraai/openrag" # Pin to a release tag (ideally a digest) for reproducible deploys. - tag: "v2.0.1" + tag: "v2.1.0" + pullPolicy: IfNotPresent service: type: ClusterIP port: 8080 + targetPort: 8080 + # Off by default: enabling requires an explicit openrag.ingress.host (see + # templates/openrag.yaml's `required` guard) so the chart never exposes an + # unauthenticated wildcard-host Ingress by accident. + ingress: + enabled: false + className: "" + host: "" + annotations: {} + # For OpenShift HAProxy ingress — increase timeout for large file uploads + # (HAProxy default is 30s; uploads over slow links or large files need more): + # haproxy.router.openshift.io/timeout: 600s + # + # For nginx ingress — body size limit (default 1m, raise for large files): + # nginx.ingress.kubernetes.io/proxy-body-size: "500m" + # nginx.ingress.kubernetes.io/proxy-read-timeout: "600" + tls: [] + # Ray dashboard ingress — only relevant when ray.enabled=false (embedded Ray via ray.init()). + # When ray.enabled=true the dashboard is on the RayCluster head and has its own ingress (ray.ingress). + rayDashboard: + ingress: + enabled: false + className: "" + host: "" + annotations: {} + # nginx.ingress.kubernetes.io/rewrite-target: / + tls: [] + path: / + pathType: Prefix resources: requests: cpu: "4" memory: "16Gi" + nvidia.com/gpu: 1 limits: cpu: "4" memory: "16Gi" + nvidia.com/gpu: 1 replicas: 1 - -ingress: - enabled: false - className: "nginx" - host: "" - tls: - enabled: false - secretName: "" - - paths: - openrag: - enabled: true - -security: - automountServiceAccountToken: false + # Size of the /dev/shm emptyDir volume. Ray's plasma object store and Python + # multiprocessing both use /dev/shm for POSIX semaphores and shared memory. + # Kubernetes defaults /dev/shm to 64Mi which causes OOM/ENOSPC errors when + # running Marker workers. Set to ~30% of the container memory limit or more. + # Set to "" or null to disable (use the 64Mi kernel default). + shmSize: "" + extraVolumes: [] + extraVolumeMounts: [] + # Raw Kubernetes initContainer objects, e.g. to merge multiple mounted CA + # bundles into one file on a shared emptyDir before the main container + # starts (see values-openrag.yaml's kc-ca-bundle example). + extraInitContainers: [] + serviceAccountName: "" + # Uncomment to override the shared security.automountServiceAccountToken / + # security.containerSecurityContext defaults for this workload specifically. + # automountServiceAccountToken: false + # containerSecurityContext: {} + # infra/docker/api.Dockerfile bakes an "openrag" user as uid 10001, primary + # gid 0 (`useradd --uid 10001 --gid 0`), and grants group-write on the paths + # it writes at runtime (venv, egg-info, HOME, data/db/logs, HF cache, uv + # cache) via `chgrp -R 0` + `chmod g=u` — the OpenShift arbitrary-UID + # pattern. runAsGroup must stay 0 (not 10001) to match, or e.g. the editable + # install's `openrag.egg-info` isn't writable and uv fails at startup + # ("Cannot update time stamp of directory 'openrag.egg-info'"). fsGroup can + # stay 10001: it's only applied (and added as a supplementary group) to the + # mounted PVCs (data/logs/model_weights/venv), not to the image's own paths. + # Merged with the shared security.podSecurityContext defaults above. podSecurityContext: - runAsNonRoot: true runAsUser: 10001 - runAsGroup: 10001 + runAsGroup: 0 fsGroup: 10001 - seccompProfile: - type: RuntimeDefault - containerSecurityContext: - allowPrivilegeEscalation: false - capabilities: - drop: - - ALL + # Probe configuration. Override with exec probes when ENABLE_RAY_SERVE=true, + # since the Ray Serve HTTP proxy runs on the Ray head (not on this pod). + probes: + startup: + httpGet: + path: /health_check + port: openrag + failureThreshold: 30 + periodSeconds: 10 + readiness: + httpGet: + path: /health_check + port: openrag + initialDelaySeconds: 10 + periodSeconds: 10 + failureThreshold: 3 + liveness: + httpGet: + path: /health_check + port: openrag + initialDelaySeconds: 60 + periodSeconds: 30 + failureThreshold: 3 + +# === Outbound proxy === +# Set enabled: true when the cluster has no direct internet access. +# Both uppercase and lowercase forms are injected (HTTP_PROXY / http_proxy) for broad tool compatibility. +# no_proxy excludes cluster-internal traffic from being routed through the proxy. +proxy: + enabled: false + url: "http://192.168.100.100:80" + # Add your OpenShift apps domain to bypass the proxy for internal routes (e.g. .apps.mycluster.example.com) + noProxy: "localhost,127.0.0.1,.cluster.local,.svc,10.0.0.0/8,172.16.0.0/12" # === Shared env (config + secrets) === env: + # Name of an existing K8s Secret to use as-is. + # When set, the chart creates no Secret/ExternalSecret/VaultStaticSecret resource. + existingSecret: "" + + # Secrets provider — ignored when existingSecret is set. + # type: values (default) | externalSecret | vaultStaticSecret + secretsProvider: + type: values + + # Option A: External Secrets Operator + externalSecret: + refreshInterval: "1h" + secretStore: + name: "vault-backend" + kind: "ClusterSecretStore" + # Option 1: fetch all keys from a path (recommended) + dataFrom: [] + # - extract: + # key: "openrag/secrets" + # Option 2: map individual keys + data: [] + # - secretKey: API_KEY + # remoteRef: + # key: "openrag/secrets" + # property: "api_key" + + # Option B: Vault Secrets Operator + vaultStaticSecret: + vaultAuthRef: "openrag-vault-auth" + mount: "kv" + type: "kv-v2" + path: "openrag/secrets" + refreshAfter: "30s" + hmacSecretData: true + config: # LLM - BASE_URL: "http://{{ .Release.Name }}-llm-engine-service/v1/" - MODEL: *llmModel + # NOTE: postgresql/milvus are sub-charts pinned via their own + # fullnameOverride (see values above) so they honor "fullname" like the + # rest of this chart. vllm-stack has no such override in its own templates + # — its *-engine-service Services are permanently named from the real + # Release.Name, so BASE_URL/VLM_BASE_URL/EMBEDDER_BASE_URL/ + # TRANSCRIBER_BASE_URL below must stay on .Release.Name. Under ArgoCD, set + # spec.source.helm.releaseName to match fullnameOverride ("openrag") so + # these resolve too. + BASE_URL: 'http://{{ .Release.Name }}-llm-engine-service/v1/' + MODEL: '{{ .Values.vllm.llmModelName }}' LLM_SEMAPHORE: "50" # VLM - VLM_BASE_URL: "http://{{ .Release.Name }}-vlm-engine-service/v1/" - VLM_MODEL: *vlmModel + VLM_BASE_URL: 'http://{{ .Release.Name }}-vlm-engine-service/v1/' + VLM_MODEL: '{{ .Values.vllm.vlmModelName }}' VLM_SEMAPHORE: "50" # App APP_PORT: "8080" - ENABLE_RAY_SERVE: "true" + # Requires ray.enabled=true (a separate RayCluster with its own "serve" port + # 80 exposed on {{ fullname }}-raycluster-head-svc — see templates/openrag.yaml + # Ingress). With the default ray.enabled=false, the app is expected to run + # plain uvicorn on APP_PORT instead, which is what container/service/probes + # in templates/openrag.yaml are wired for. Flip both together (see + # values-linagora.yaml for the ray.enabled=true pairing). + ENABLE_RAY_SERVE: "false" RAY_SERVE_NUM_REPLICAS: "4" RAY_SERVE_PORT: "80" WITH_CHAINLIT_UI: "false" - SAVE_UPLOADED_FILES: "false" - # HTTP scaling is handled by Ray Serve above (ENABLE_RAY_SERVE + - # RAY_SERVE_NUM_REPLICAS), not uvicorn workers — see entrypoint.sh. + SAVE_UPLOADED_FILES: "true" # Vector DB - VDB_HOST: "{{ .Release.Name }}-milvus" + # milvus.fullnameOverride above pins this sub-chart's proxy Service name to + # match "fullname" instead of the real Release.Name (see note near the top). + VDB_HOST: '{{ include "openrag-stack.fullname" . }}-milvus' VDB_PORT: "19530" VDB_CONNECTOR_NAME: "milvus" # PostgreSQL - POSTGRES_HOST: "{{ .Release.Name }}-postgresql" + # postgresql.fullnameOverride above pins this sub-chart's Service name to + # match "fullname" instead of the real Release.Name (see note near the top). + POSTGRES_HOST: '{{ include "openrag-stack.fullname" . }}-postgresql' POSTGRES_PORT: *pgPort POSTGRES_USER: *pgUser # POSTGRES_PASSWORD is a secret — defined under env.secrets, not here in the @@ -355,15 +639,18 @@ env: POSTGRES_RUN_MIGRATIONS: "{{ .Values.postgresProvisioning.runMigrationsInApp }}" # Embedder - EMBEDDER_MODEL_NAME: *embedderModel - EMBEDDER_BASE_URL: "http://{{ .Release.Name }}-embedder-engine-service/v1" + EMBEDDER_MODEL_NAME: '{{ .Values.vllm.embedderModelName }}' + EMBEDDER_BASE_URL: 'http://{{ .Release.Name }}-embedder-engine-service/v1' # Reranker - RERANKER_MODEL: *rerankerModel - RERANKER_ENABLED: "true" + RERANKER_MODEL: '{{ .Values.reranker.rerankerModelName }}' + # True when either the local Infinity reranker is deployed or an external + # one is configured — false (and RERANKER_BASE_URL below moot) when + # neither applies, so this never points at a Service that doesn't exist. + RERANKER_ENABLED: '{{ or .Values.reranker.enabled (ne .Values.reranker.externalUrl "") }}' RERANKER_TOP_K: "4" - RERANKER_BASE_URL: "http://{{ .Release.Name }}-reranker:{{ .Values.reranker.servicePort }}" - RERANKER_MODEL_TYPE: "infinity" + RERANKER_BASE_URL: '{{ if .Values.reranker.externalUrl }}{{ .Values.reranker.externalUrl }}{{ else }}http://{{ include "openrag-stack.fullname" . }}-reranker:{{ .Values.reranker.servicePort }}{{ end }}' + RERANKER_PROVIDER: "infinity" # Model Endpoint Registry: when true, re-syncs each type's env-seeded # endpoint (URL/model/api_key/...) from the vars above on every boot, so a @@ -382,24 +669,43 @@ env: MARKER_MAX_PROCESSES: "15" MARKER_POOL_SIZE: "3" MARKER_NUM_GPUS: "0.6" - TRANSCRIBER_BASE_URL: "http://{{ .Release.Name }}-whisper-engine-service/v1" + TRANSCRIBER_BASE_URL: 'http://{{ .Release.Name }}-whisper-engine-service/v1' # Ray + # RAY_ADDRESS: 'ray://{{ include "openrag-stack.fullname" . }}-raycluster-head-svc:10001' + RAY_NUM_GPUS: "0.1" RAY_POOL_SIZE: "3" RAY_MAX_TASKS_PER_WORKER: "50" RAY_DASHBOARD_PORT: "8265" - RAY_ADDRESS: "ray://raycluster-head-svc:10001" RAY_task_retry_delay_ms: "3000" RAY_ENABLE_UV_RUN_RUNTIME_ENV: "0" UV_LINK_MODE: "copy" UV_CACHE_DIR: "/tmp/uv-cache" + UV_PYTHON_PREFERENCE: "only-system" + UV_PYTHON: "python3" # portable: finds python3 in PATH, bypasses .python-version patch-version mismatch secrets: API_KEY: "" # LLM API KEY VLM_API_KEY: "" # VLM API KEY EMBEDDER_API_KEY: "" TRANSCRIBER_API_KEY: "" + # Read live from postgresql.auth.password (not the *pgPass anchor, which + # only resolves within values.yaml and goes stale once overridden via + # --set/-f) so the secret always matches the actual DB credential. POSTGRES_PASSWORD: "{{ .Values.postgresql.auth.password }}" AUTH_TOKEN: "" # API KEY for OpenRAG HF_TOKEN: "" # HuggingFace token + RERANKER_API_KEY: "EMPTY" # Set when using an external reranker that requires authentication + +# Arbitrary additional Kubernetes objects rendered at install time. +# Each entry is a full resource manifest; Go templates are supported. +# Example: +# extraObjects: +# - apiVersion: v1 +# kind: ConfigMap +# metadata: +# name: "{{ include "openrag-stack.fullname" . }}-extra" +# data: +# key: value +extraObjects: [] diff --git a/infra/compose/docker-compose.yaml b/infra/compose/docker-compose.yaml index 5ca821fb5..38d1bb3a4 100644 --- a/infra/compose/docker-compose.yaml +++ b/infra/compose/docker-compose.yaml @@ -18,7 +18,7 @@ x-openrag-env: &openrag_env FONT_PATH: ${FONT_PATH:-/app/data/fonts/GoNotoCurrent-Regular.ttf} x-openrag: &openrag_template - image: linagoraai/openrag:v2.0.1 + image: linagoraai/openrag:v2.1.0 # Start as root so entrypoint.sh can grant GID-0 write on the bind-mounted # writable dirs (data/, logs/, the HF cache) — which a non-root container # can't write when Docker auto-creates them root-owned — then it immediately @@ -113,7 +113,7 @@ x-vllm: &vllm_template services: # ── Admin UI (React SPA + nginx, same-origin reverse proxy to the API) ── admin-ui: - image: linagoraai/openrag-admin-ui:v2.0.1 + image: linagoraai/openrag-admin-ui:v2.1.0 build: context: ../.. dockerfile: infra/docker/ui.Dockerfile diff --git a/infra/docker/ui.Dockerfile b/infra/docker/ui.Dockerfile index 0700ab6fc..207c214c8 100644 --- a/infra/docker/ui.Dockerfile +++ b/infra/docker/ui.Dockerfile @@ -34,12 +34,45 @@ RUN npm run build # nginx-unprivileged listens on :8080 and runs as a non-root user, so the same # image runs under a hardened container security context (runAsNonRoot, # drop ALL capabilities) — required by the Helm chart and good practice in -# compose too. COPY runs as root, then we drop back to the image's non-root user. +# compose too. COPY runs as root, then we drop back to a fixed non-root UID. FROM nginxinc/nginx-unprivileged:1.27-alpine USER root -COPY --from=build /app/dist /usr/share/nginx/html -COPY infra/compose/nginx/openrag-admin.conf /etc/nginx/conf.d/default.conf -USER nginx + +COPY --chown=10001:0 --from=build /app/dist /usr/share/nginx/html +COPY --chown=10001:0 infra/compose/nginx/openrag-admin.conf /etc/nginx/conf.d/default.conf + +# /var/cache/nginx and /var/run come from the base image (not copied above) — +# own them as 10001:0 and make them group-writable, the same arbitrary-UID +# pattern api.Dockerfile uses (`useradd --gid 0` + `chgrp -R 0` + `chmod g=u`): +# OpenShift's restricted-v2 SCC runs the container as an unpredictable UID from +# the namespace range that is always a member of the root group, so group 0 is +# the only ownership every platform agrees on. The base image already ships +# these paths as 101:0 for exactly that reason — chowning them to a *private* +# group would take that away. +# Re-chowning /etc/nginx/conf.d here as well covers the directory itself (COPY +# --chown above only touched the file), which docker-entrypoint.d/10-listen-on- +# ipv6... needs write access to at startup — without it that script logs +# "can not modify /etc/nginx/conf.d/default.conf (read-only file system?)". +# Don't drop these in favour of the base image's own baked-in permissions: a +# stale/cached base layer silently reverted them once already, and the failures +# that follow surface at request time rather than at startup. + +RUN chown -R 10001:0 /var/cache/nginx /etc/nginx/conf.d /var/run && \ + chmod -R g+w /var/cache/nginx /etc/nginx/conf.d /var/run + +# Numeric UID:GID, NOT the base image's `nginx` user: that is uid 101 with only +# gid 101 in its group list, so it has neither owner nor group access to the +# paths chowned above. Pinning 10001:0 makes plain docker/compose (which does +# not remap the user) run as the owner, while an OpenShift arbitrary UID still +# writes through group 0. Matches adminUi.podSecurityContext in the Helm chart +# (runAsUser: 10001, runAsGroup: 0). +# Note the mismatch this replaced was latent rather than actively breaking: +# nginx-unprivileged redirects every *_temp_path (and its pid) to /tmp, which is +# 1777, and openrag-admin.conf sets `proxy_cache off`, so nothing writes under +# /var/cache/nginx today. It only bites once something does — a proxy_cache_path, +# an envsubst template landing in conf.d, or a base image that moves the temp +# paths back — and then it fails at request time, not at startup. +USER 10001:0 EXPOSE 8080 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/partitions.py b/openrag/api/routers/admin/partitions.py index 1f098b78d..f21bda9f3 100644 --- a/openrag/api/routers/admin/partitions.py +++ b/openrag/api/routers/admin/partitions.py @@ -383,8 +383,10 @@ async def get_partition_config( **Response:** Returns list of partition members with: - `user_id`: User identifier +- `display_name`: Human-readable name, when available +- `email`: Account email, when available - `role`: User's role (owner, editor, or viewer) -- Additional user details +- `added_at`: Membership creation time **Permissions:** - Requires partition owner role @@ -401,10 +403,46 @@ async def list_partition_users( service=Depends(get_partition_service), ): """List all users who are members of the given partition.""" - members = await service.list_members(partition=partition) + members = await service.list_members_with_identities(partition=partition) return JSONResponse(status_code=status.HTTP_200_OK, content={"members": members}) +@router.get( + "/{partition}/users/candidates", + description="""List users who can be added to a partition. + +**Parameters:** +- `partition`: The partition name +- `search`: Display-name prefix (at least 3 characters) or exact user ID +- `cursor`: Last user ID from the previous page +- `limit`: Page size (maximum 100) + +**Response:** +Returns a bounded page of non-member users with their display name and email, +plus continuation metadata. + +**Permissions:** +- Requires partition owner role +""", +) +async def list_partition_user_candidates( + partition: str, + search: str = Query(..., max_length=200), + cursor: int | None = Query(default=None, ge=0, le=2_147_483_647), + limit: int = Query(default=25, ge=1, le=100), + partition_owner=Depends(require_partition_owner), + service=Depends(get_partition_service), +): + """List a searchable page of users who are not partition members.""" + page = await service.list_member_candidates( + partition=partition, + search=search, + cursor=cursor, + limit=limit, + ) + return JSONResponse(status_code=status.HTTP_200_OK, content=page) + + @router.post( "/{partition}/users", description="""Add a user to a partition with a specific role. @@ -424,6 +462,7 @@ async def list_partition_users( **Response:** Returns 201 Created on successful addition. +Returns 409 Conflict if the user is already a member; use the role endpoint to change an existing member. """, ) async def add_partition_user( 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/routers/admin/users.py b/openrag/api/routers/admin/users.py index bc3b2b09e..bc86c0a27 100644 --- a/openrag/api/routers/admin/users.py +++ b/openrag/api/routers/admin/users.py @@ -38,6 +38,7 @@ - `id`: User identifier - `display_name`: User's display name - `external_user_id`: External ID (if set) +- `email`: Account email (if set) - `is_admin`: Admin status - `created_at`: Account creation timestamp @@ -141,6 +142,7 @@ async def create_user( - `id`: User identifier - `display_name`: User's display name - `external_user_id`: External ID (if set) +- `email`: Account email (if set) - `is_admin`: Admin status - `created_at`: Account creation timestamp diff --git a/openrag/api/routers/user/chat.py b/openrag/api/routers/user/chat.py index 9e315d110..81f93eeb7 100644 --- a/openrag/api/routers/user/chat.py +++ b/openrag/api/routers/user/chat.py @@ -14,7 +14,6 @@ import asyncio import json from typing import TYPE_CHECKING -from urllib.parse import urlparse import consts from api.dependencies.auth import ( @@ -34,6 +33,7 @@ from core.utils.exceptions import OpenRAGError from core.utils.logging import get_logger from core.utils.text import get_num_tokens, sanitize_text +from core.utils.web_url import normalize_web_url from di.providers import get_config, get_partition_service, get_query_service from fastapi import APIRouter, Body, Depends, HTTPException, Request, status from fastapi.responses import JSONResponse, StreamingResponse @@ -225,8 +225,8 @@ def chunk_url(extract_id) -> str: doc_metadata = dict(doc.metadata) links.append(build_document_source_link(doc_metadata, static_url, chunk_url)) for result in web_results or []: - url = sanitize_text(result.url or "") - if not url or urlparse(url).scheme not in ("http", "https"): + url = normalize_web_url(result.url) + if url is None: continue links.append( { diff --git a/openrag/api/routers/user/source_links.py b/openrag/api/routers/user/source_links.py index 7ace374d9..cb17289f5 100644 --- a/openrag/api/routers/user/source_links.py +++ b/openrag/api/routers/user/source_links.py @@ -33,9 +33,11 @@ def build_document_source_link( encoded_url = None if filename: encoded_url = quote(static_url_builder(doc_metadata["_id"]), safe=":/") - return { - "source_type": "document", - **({"file_url": encoded_url} if encoded_url else {}), - "chunk_url": chunk_url_builder(doc_metadata["_id"]), - **doc_metadata, - } + link = dict(doc_metadata) + link["source_type"] = "document" + link["chunk_url"] = chunk_url_builder(doc_metadata["_id"]) + if encoded_url: + link["file_url"] = encoded_url + else: + link.pop("file_url", None) + return link diff --git a/openrag/api/schemas/admin/model_endpoint_schemas.py b/openrag/api/schemas/admin/model_endpoint_schemas.py index d98dbd6cf..15789957d 100644 --- a/openrag/api/schemas/admin/model_endpoint_schemas.py +++ b/openrag/api/schemas/admin/model_endpoint_schemas.py @@ -2,6 +2,7 @@ from __future__ import annotations +import re from datetime import datetime from typing import Any, Literal @@ -15,12 +16,43 @@ # ``extra`` — validated here so a typo can't persist a nonsensical value. _LLM_TOKEN_EXTRA_KEYS = (LLM_CONTEXT_SIZE_KEY, LLM_OUTPUT_TOKENS_KEY) +# Allowlist, not a denylist: `name` is a single path segment in every +# single-endpoint route (see `_normalize_name`), and enumerating unsafe values +# one at a time as they're discovered — first `/` (#768), then the RFC 3986 +# dot-segments `.`/`..` — never closes the class. Anchoring both ends on +# alphanumeric rules out `/`, `.`, `..`, and any leading/trailing separator by +# construction, while `.`/`_`/`-` stay available in the middle for realistic +# names like `gpt-4.1` or `jina_v3`. +_NAME_PATTERN = re.compile(r"[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?") +_NAME_MAX_LENGTH = 128 + def _normalize_name(value: str) -> str: - """Trim a user-facing registry name and reject blank values.""" + """Trim a user-facing registry name and reject any value unsafe as a URL path segment. + + ``name`` is embedded as a single path segment in every single-endpoint route + (``GET/PUT/DELETE /model-endpoints/{model_type}/{name}``, ``.../set-default``, + ``.../reveal-api-key``, ``.../validate``). A value outside ``_NAME_PATTERN`` + — a ``/`` (splits across path segments), the exact values ``.``/``..`` + (RFC 3986 dot-segments: browsers and HTTP clients normalize these out of + the URL before the request is even sent, resolving to the collection route + or dropping the ``model_type`` segment entirely), or anything else that + doesn't start/end alphanumeric — would leave the row visible in the list + endpoint but permanently unreachable by get/update/delete/set-default, + surfacing as a spurious "not found". Percent-encoding never helps: ASGI + servers decode ``%2F``/dot-segment escapes before Starlette's router sees + the path. + """ value = value.strip() if not value: raise ValueError("name must be non-empty") + if len(value) > _NAME_MAX_LENGTH: + raise ValueError(f"name must be at most {_NAME_MAX_LENGTH} characters") + if not _NAME_PATTERN.fullmatch(value): + raise ValueError( + "name must start and end with a letter or digit, and contain only " + "letters, digits, '.', '_', or '-' (it is used as a URL path segment)" + ) return value 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/app_front.py b/openrag/app_front.py index fd5be7434..7b8efd788 100644 --- a/openrag/app_front.py +++ b/openrag/app_front.py @@ -1,10 +1,13 @@ import json import os +import re import secrets +import string import time +import unicodedata from functools import lru_cache from pathlib import Path -from urllib.parse import urlparse +from urllib.parse import quote, urlparse import chainlit as cl import httpx @@ -19,6 +22,7 @@ CHAINLIT_TOKEN_COOKIE_PATH, ) from core.utils.logging import get_logger, mask_email +from core.utils.web_url import normalize_web_url from dotenv import load_dotenv from openai import AsyncOpenAI @@ -42,12 +46,94 @@ OPENRAG_CHAT_PROFILES_METADATA_KEY = "openrag_chat_profiles" OPENRAG_SESSION_COOKIE_NAME = "openrag_session" _OPENRAG_TOKEN_STORE: dict[str, tuple[str, float]] = {} +_MARKDOWN_ESCAPE_TABLE = str.maketrans({char: f"\\{char}" for char in string.punctuation}) +_MARKDOWN_URL_SAFE_CHARS = ":/?#[]@!$&'+,;=%" +# Chainlit inserts source names into Markdown link labels. Strip only characters +# that can break or restyle that label so ordinary filenames stay recognizable. +_MARKDOWN_UNSAFE_SOURCE_NAME_CHARS = str.maketrans(dict.fromkeys("[]*`\\<>", " ")) +_MARKDOWN_BLOCK_PREFIX_RE = re.compile(r"^(?:(?:#{1,6}|[-+]|\d{1,9}[.)])\s+)+") +_MARKDOWN_THEMATIC_BREAK_RE = re.compile(r"^(?:(?:-\s*){3,}|(?:_\s*){3,}|(?:\*\s*){3,})$") class MissingOpenRAGCredentialError(RuntimeError): pass +def _escape_markdown_text(value: str) -> str: + """Render untrusted source metadata as literal Markdown text.""" + return value.translate(_MARKDOWN_ESCAPE_TABLE) + + +def _neutralize_source_name_delimiters(value: str) -> str: + """Remove active Markdown delimiters while preserving safe filename punctuation.""" + characters = list(value) + underscore_openers: list[tuple[int, int]] = [] + index = 0 + + while index < len(value): + if value[index] != "_": + index += 1 + continue + + end = index + 1 + while end < len(value) and value[end] == "_": + end += 1 + + previous = value[index - 1] if index else None + following = value[end] if end < len(value) else None + previous_whitespace = previous is None or previous.isspace() + following_whitespace = following is None or following.isspace() + previous_punctuation = previous is not None and unicodedata.category(previous)[0] in {"P", "S"} + following_punctuation = following is not None and unicodedata.category(following)[0] in {"P", "S"} + left_flanking = not following_whitespace and ( + not following_punctuation or previous_whitespace or previous_punctuation + ) + right_flanking = not previous_whitespace and ( + not previous_punctuation or following_whitespace or following_punctuation + ) + can_open = left_flanking and (not right_flanking or previous_punctuation) + can_close = right_flanking and (not left_flanking or following_punctuation) + + if can_close and underscore_openers: + opener_start, opener_end = underscore_openers.pop() + characters[opener_start:opener_end] = " " * (opener_end - opener_start) + characters[index:end] = " " * (end - index) + elif can_open: + underscore_openers.append((index, end)) + index = end + + index = 0 + while index < len(value): + if value[index] != "~": + index += 1 + continue + end = index + 1 + while end < len(value) and value[end] == "~": + end += 1 + if end - index >= 2: + characters[index:end] = " " * (end - index) + index = end + + return "".join(characters) + + +def _safe_source_name(value: str, existing: dict) -> str: + """Build a Markdown-inert name that Chainlit can match to its element.""" + value = _neutralize_source_name_delimiters(value) + value = value.lstrip() + if _MARKDOWN_THEMATIC_BREAK_RE.fullmatch(value): + value = "" + else: + value = _MARKDOWN_BLOCK_PREFIX_RE.sub("", value) + base = " ".join(value.translate(_MARKDOWN_UNSAFE_SOURCE_NAME_CHARS).split()) or "source" + candidate = base + suffix = 2 + while candidate in existing: + candidate = f"{base} {suffix}" + suffix += 1 + return candidate + + def get_user_language() -> str: """Return the active language: env override if set, otherwise browser's Accept-Language.""" if DEFAULT_LANGUAGE: @@ -465,26 +551,48 @@ async def __fetch_page_content(chunk_url, headers=None): async def _format_sources(metadata_sources, only_txt=False, api_key=None): - external_url = get_external_url() # used to override the base URL when the front-end requests a file resource - if not metadata_sources: - return None, None + if not isinstance(metadata_sources, list) or not metadata_sources: + return [], [] d = {} headers = get_headers(api_key) + external_url = get_external_url() # used to override the base URL when the front-end requests a file resource for i, s in enumerate(metadata_sources): + if not isinstance(s, dict): + continue + if s.get("source_type") == "web": - title = s.get("title") or s.get("url", f"Web source {i + 1}") - url = s.get("url", "") + title = s.get("title", "") snippet = s.get("snippet", "") - content = f"**[{title}]({url})**\n\n{snippet}" - source_name = title - if source_name in d: - source_name = f"{title} ({i})" + title = title.strip() if isinstance(title, str) else "" + snippet = snippet.strip() if isinstance(snippet, str) else "" + url = normalize_web_url(s.get("url")) + if url is None: + continue + markdown_url = quote(url, safe=_MARKDOWN_URL_SAFE_CHARS) + + source_label = title or url + source_name = _safe_source_name(source_label, d) + content = f"**[{_escape_markdown_text(source_label)}]({markdown_url})**" + if snippet: + content += f"\n\n{_escape_markdown_text(snippet)}" d[source_name] = cl.Text(content=content, name=source_name, display="side") continue - filename = Path(s["filename"]) - file_url = s["file_url"] + filename_value = s.get("filename") + file_url = s.get("file_url") + page = s.get("page") + if ( + not isinstance(filename_value, str) + or not filename_value.strip() + or not isinstance(file_url, str) + or not file_url.strip() + ): + continue + + filename = Path(filename_value.strip()) + suffix = filename.suffix.lower() + file_url = file_url.strip() file_url = file_url.replace(INTERNAL_BASE_URL, external_url) # put the correct base url # Avoid leaking the credential in the URL (browser history, proxy logs, # Referer headers). In OIDC mode the browser already sends the @@ -495,36 +603,51 @@ async def _format_sources(metadata_sources, only_txt=False, api_key=None): # authenticate the fetch. if api_key and (AUTH_MODE != "oidc" or _current_openrag_auth_provider() == "credentials"): file_url = f"{file_url}?token={api_key}" - page = s["page"] - source_name = f"{filename}" + ( - f" (page: {page})" if filename.suffix in [".pdf", ".pptx", ".docx", ".doc"] else "" + page_label = str(page).strip() if page is not None else "" + source_label = f"{filename}" + ( + f" (page: {page_label})" if suffix in [".pdf", ".pptx", ".docx", ".doc"] and page_label else "" ) + source_name = _safe_source_name(source_label, d) - if only_txt: - chunk_content = await __fetch_page_content(chunk_url=s["chunk_url"], headers=headers) - elem = cl.Text(content=chunk_content, name=source_name, display="side") - else: - match filename.suffix.lower(): - case ".pdf": - elem = cl.Pdf( - name=source_name, - url=file_url, - page=int(s["page"]), - display="side", - ) - case suffix if suffix in [".png", ".jpg", ".jpeg"]: - elem = cl.Image(name=source_name, url=file_url, display="side") - case ".mp4": - elem = cl.Video(name=source_name, url=file_url, display="side") - case ".mp3": - elem = cl.Audio(name=source_name, url=file_url, display="side") - case _: - chunk_content = await __fetch_page_content(chunk_url=s["chunk_url"], headers=headers) - elem = cl.Text(content=chunk_content, name=source_name, display="side") + try: + if only_txt: + chunk_url = s.get("chunk_url") + if not isinstance(chunk_url, str) or not chunk_url.strip(): + continue + chunk_content = await __fetch_page_content(chunk_url=chunk_url, headers=headers) + if not isinstance(chunk_content, str) or not chunk_content.strip(): + continue + elem = cl.Text(content=chunk_content, name=source_name, display="side") + else: + match suffix: + case ".pdf": + elem = cl.Pdf( + name=source_name, + url=file_url, + page=int(page) if page_label else None, + display="side", + ) + case suffix if suffix in [".png", ".jpg", ".jpeg"]: + elem = cl.Image(name=source_name, url=file_url, display="side") + case ".mp4": + elem = cl.Video(name=source_name, url=file_url, display="side") + case ".mp3": + elem = cl.Audio(name=source_name, url=file_url, display="side") + case _: + chunk_url = s.get("chunk_url") + if not isinstance(chunk_url, str) or not chunk_url.strip(): + continue + chunk_content = await __fetch_page_content(chunk_url=chunk_url, headers=headers) + if not isinstance(chunk_content, str) or not chunk_content.strip(): + continue + elem = cl.Text(content=chunk_content, name=source_name, display="side") + except (httpx.HTTPError, httpx.InvalidURL, TypeError, ValueError, AttributeError): + logger.warning("Skipping an unavailable source", source_index=i) + continue d[source_name] = elem - source_names = list(d.keys()) + source_names = list(d) elements = list(d.values()) return elements, source_names @@ -582,7 +705,7 @@ async def on_message(message: cl.Message): # Show sources elements, source_names = await _format_sources(sources, api_key=api_key, only_txt=False) msg.elements = elements if elements else [] - if source_names: + if elements and source_names: s = "\n\n" + "-" * 50 + f"\n\n{t('sources_label')}: \n" + "\n".join(source_names) await msg.stream_token(s) await msg.update() 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/models/query.py b/openrag/core/models/query.py index b5dad680d..86c512de5 100644 --- a/openrag/core/models/query.py +++ b/openrag/core/models/query.py @@ -103,6 +103,10 @@ class SearchQueries(BaseModel): """Collection of sub-queries produced by query decomposition.""" query_list: list[Query] = Field(..., description="Search sub-queries to retrieve relevant documents.") + requires_retrieval: bool = Field( + default=True, + description="Whether the user's request needs document retrieval.", + ) def __str__(self) -> str: return " --- ".join(str(q) for q in self.query_list) diff --git a/openrag/core/ports/partition_membership_repo.py b/openrag/core/ports/partition_membership_repo.py index 786d02e47..77666567a 100644 --- a/openrag/core/ports/partition_membership_repo.py +++ b/openrag/core/ports/partition_membership_repo.py @@ -34,3 +34,14 @@ async def update_partition_role(self, user_id: int, partition: str, role: Partit @abstractmethod async def count_partition_users(self, partition: str) -> int: ... + + @abstractmethod + async def list_partition_member_candidates( + self, + partition: str, + *, + search_prefix: str | None, + search_user_id: int | None, + after_id: int | None, + limit: int, + ) -> list[dict]: ... 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/ports/user_repo.py b/openrag/core/ports/user_repo.py index 747eadf0c..dc0dd2e19 100644 --- a/openrag/core/ports/user_repo.py +++ b/openrag/core/ports/user_repo.py @@ -30,6 +30,9 @@ async def create_user(self, user: User) -> User: ... @abstractmethod async def get_user(self, user_id: int) -> User | None: ... + @abstractmethod + async def get_users_by_ids(self, user_ids: list[int]) -> list[User]: ... + @abstractmethod async def get_user_by_email(self, email: str) -> User | None: ... 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/core/utils/source_filtering.py b/openrag/core/utils/source_filtering.py index e64418bda..ff7b52833 100644 --- a/openrag/core/utils/source_filtering.py +++ b/openrag/core/utils/source_filtering.py @@ -17,6 +17,15 @@ re.IGNORECASE, ) _SOURCES_NUMS_RE = re.compile(r"\n?[ \t]*\[?Sources?\]?\s*:\s*\[?([\d,\s]+)\]?[.\s]*?(?=\n|$)", re.IGNORECASE) +_INLINE_SOURCE_NUMS_RE = re.compile( + r"[ \t]*\[\s*Sources?\s+(\d+(?:\s*,\s*\d+)*)\s*\]", + re.IGNORECASE, +) +_UNCLOSED_SOURCE_NUMS_RE = re.compile( + r"[ \t]*\[\s*Sources?\s+(\d+(?:\s*,\s*\d+)*)\s*(?=\n|$)", + re.IGNORECASE, +) +_DANGLING_SOURCE_RE = re.compile(r"[ \t]*\[\s*Sources?\s*(?=\n|$)", re.IGNORECASE) def _sanitize_log_preview(text: str, max_length: int = 150) -> str: @@ -26,22 +35,37 @@ def _sanitize_log_preview(text: str, max_length: int = 150) -> str: return preview -def _strip_sources_tags(text: str) -> tuple[str, set[int], bool]: - """Strip line-terminal source tags and return citations found.""" +def _strip_sources_tags(text: str, *, include_inline_markers: bool = True) -> tuple[str, set[int], bool]: + """Strip source tags and return citations found.""" cited: set[int] = set() - for match in _SOURCES_NUMS_RE.finditer(text): - cited.update(int(n.strip()) for n in match.group(1).split(",") if n.strip().isdigit()) + patterns = [_SOURCES_NUMS_RE] + if include_inline_markers: + patterns.extend((_INLINE_SOURCE_NUMS_RE, _UNCLOSED_SOURCE_NUMS_RE)) + for pattern in patterns: + for match in pattern.finditer(text): + cited.update(int(n.strip()) for n in match.group(1).split(",") if n.strip().isdigit()) saw_none = bool(_SOURCES_NONE_RE.search(text)) cleaned = _SOURCES_NUMS_RE.sub("", text) cleaned = _SOURCES_NONE_RE.sub("", cleaned) + if include_inline_markers: + cleaned = _INLINE_SOURCE_NUMS_RE.sub("", cleaned) + cleaned = _UNCLOSED_SOURCE_NUMS_RE.sub("", cleaned) + cleaned = _DANGLING_SOURCE_RE.sub("", cleaned) return cleaned, cited, saw_none -def extract_and_strip_sources_block(text: str) -> tuple[str, set[int] | None]: - """Strip line-terminal source tags and return merged citations.""" - cleaned, citations, saw_none = _strip_sources_tags(text) +def extract_and_strip_sources_block( + text: str, + *, + include_inline_markers: bool = True, +) -> tuple[str, set[int] | None]: + """Strip source tags and return merged citations.""" + cleaned, citations, saw_none = _strip_sources_tags(text, include_inline_markers=include_inline_markers) if not citations and not saw_none: + if cleaned != text: + logger.debug("Removed incomplete source marker from LLM response") + return cleaned.rstrip(), None tail = text[-150:] if len(text) > 150 else text logger.debug("No [Sources: ...] tag found in LLM response", tail=repr(_sanitize_log_preview(tail))) return text, None @@ -55,14 +79,18 @@ def extract_and_strip_sources_block(text: str) -> tuple[str, set[int] | None]: return cleaned, set() -def filter_sources_by_citations(sources: list, citations: set[int] | None) -> list: +def filter_sources_by_citations( + sources: list, + citations: set[int] | None, + *, + allow_uncited: bool = False, +) -> list: """Keep only sources whose 1-based index was cited.""" if citations is None: - return sources + return sources if allow_uncited else [] if not citations: return [] - filtered = [source for i, source in enumerate(sources, start=1) if i in citations] - return filtered if filtered else sources + return [source for i, source in enumerate(sources, start=1) if i in citations] def _min_sources_tag_buffer_size(n_sources: int) -> int: @@ -83,8 +111,11 @@ async def stream_with_source_filtering( sources: list, model_name: str, buffer_size: int | None = None, + *, + allow_uncited_sources: bool = False, + citation_protocol_active: bool = True, ): - """Process an LLM SSE stream, stripping line-terminal source tags. + """Process an LLM SSE stream and, when active, strip source tags. The terminal flush (tail content + ``extra.sources``) runs exactly once after the loop on *every* termination path — a clean ``data: [DONE]``, the @@ -99,6 +130,7 @@ async def stream_with_source_filtering( """ if buffer_size is None: buffer_size = max(_MIN_STREAM_LOOKAHEAD, _min_sources_tag_buffer_size(len(sources))) + include_inline_markers = citation_protocol_active and bool(sources) pending = "" emitted_len = 0 chunk_template = None @@ -155,7 +187,13 @@ async def stream_with_source_filtering( if len(pending) <= buffer_size: continue - cleaned, _, _ = _strip_sources_tags(pending) + if citation_protocol_active: + cleaned, _, _ = _strip_sources_tags( + pending, + include_inline_markers=include_inline_markers, + ) + else: + cleaned = pending safe_end = max(0, len(cleaned) - buffer_size) if safe_end > emitted_len: out = { @@ -216,10 +254,16 @@ async def stream_with_source_filtering( logger.warning("Upstream stream raised before any content; surfacing error", error=str(stream_error)) raise stream_error - final_clean, citations = extract_and_strip_sources_block(pending) - final_clean = final_clean.rstrip() + if citation_protocol_active: + final_clean, citations = extract_and_strip_sources_block( + pending, + include_inline_markers=include_inline_markers, + ) + final_clean = final_clean.rstrip() + else: + final_clean, citations = pending, None - filtered = filter_sources_by_citations(sources, citations) + filtered = filter_sources_by_citations(sources, citations, allow_uncited=allow_uncited_sources) extra_payload = {"sources": filtered} if not saw_done: extra_payload["truncated"] = True diff --git a/openrag/core/utils/web_url.py b/openrag/core/utils/web_url.py new file mode 100644 index 000000000..8662e8706 --- /dev/null +++ b/openrag/core/utils/web_url.py @@ -0,0 +1,25 @@ +"""Shared validation and normalization for displayable web URLs.""" + +from urllib.parse import urlparse + +import httpx +from core.utils.text import sanitize_text + + +def normalize_web_url(value: object) -> str | None: + """Return a renderable HTTP(S) URL, or None when the value is invalid.""" + if not isinstance(value, str): + return None + url = sanitize_text(value) + if not url: + return None + try: + parsed = urlparse(url) + if parsed.scheme not in {"http", "https"} or not parsed.netloc: + return None + return str(httpx.URL(url)) + except (ValueError, httpx.InvalidURL): + return None + + +__all__ = ["normalize_web_url"] diff --git a/openrag/di/container.py b/openrag/di/container.py index 2510b5638..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 @@ -437,6 +444,7 @@ def model_endpoint_service(self) -> ModelEndpointService: model_endpoint_repo=self.model_endpoint_repo, config=self._require_settings(), partition_service=self.partition_service, + preset_service=self.preset_service, client_caches={ "embedder": self._embedder_cache, "reranker": self._reranker_cache, @@ -459,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.""" @@ -534,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 @@ -568,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/prompts/templates/query_contextualizer_tmpl.txt b/openrag/prompts/templates/query_contextualizer_tmpl.txt index 67a134d3a..a2458377d 100644 --- a/openrag/prompts/templates/query_contextualizer_tmpl.txt +++ b/openrag/prompts/templates/query_contextualizer_tmpl.txt @@ -1,17 +1,18 @@ Produce a JSON object listing sub-queries derived from the user's last message. Output shape (return this JSON object, nothing else): -{{"query_list": [ {{"query": "", "temporal_filters": }} ]}} +{{"requires_retrieval": , "query_list": [ {{"query": "", "temporal_filters": }} ]}} Current date: {current_date} Language for `query` field: {query_language} Timestamps: UTC (`+00:00`). Week starts Monday. # Rewrite rules (`query` field) +- Set `requires_retrieval: false` and return an empty `query_list` when the entire last user message is only a greeting, thanks, casual conversation, or a question about the assistant's general capabilities. +- Set `requires_retrieval: true` for any request that asks for factual or document-backed information. A greeting combined with a factual question still requires retrieval. - Rewrite the last `user:` line as one standalone descriptive sentence. - Use earlier turns only to resolve pronouns or add context directly relevant to the query; do not inject unrelated history. - For independent questions: minimal changes (grammar, missing keywords). -- Greetings / thanks: copy verbatim, `temporal_filters: null`. - Do not answer. Only reformulate. # Sub-queries — when to split @@ -46,19 +47,25 @@ For exclusions, split into two sub-queries covering each remaining range. Never # Examples (Current date = Wednesday, April 15, 2026) User: "Summary of meeting notes uploaded in the past month" -{{"query_list":[{{"query":"Summary of meeting notes uploaded in the past month","temporal_filters":[{{"field":"created_at","operator":">=","value":"2026-03-15T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-04-16T00:00:00+00:00"}}]}}]}} +{{"requires_retrieval":true,"query_list":[{{"query":"Summary of meeting notes uploaded in the past month","temporal_filters":[{{"field":"created_at","operator":">=","value":"2026-03-15T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-04-16T00:00:00+00:00"}}]}}]}} User: "Sales figures for Product A and Product B" -{{"query_list":[{{"query":"Sales figures for Product A","temporal_filters":null}},{{"query":"Sales figures for Product B","temporal_filters":null}}]}} +{{"requires_retrieval":true,"query_list":[{{"query":"Sales figures for Product A","temporal_filters":null}},{{"query":"Sales figures for Product B","temporal_filters":null}}]}} User: "Documents from last year except March" -{{"query_list":[{{"query":"Documents from January or February 2025","temporal_filters":[{{"field":"created_at","operator":">=","value":"2025-01-01T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2025-03-01T00:00:00+00:00"}}]}},{{"query":"Documents from April to December 2025","temporal_filters":[{{"field":"created_at","operator":">=","value":"2025-04-01T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-01-01T00:00:00+00:00"}}]}}]}} +{{"requires_retrieval":true,"query_list":[{{"query":"Documents from January or February 2025","temporal_filters":[{{"field":"created_at","operator":">=","value":"2025-01-01T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2025-03-01T00:00:00+00:00"}}]}},{{"query":"Documents from April to December 2025","temporal_filters":[{{"field":"created_at","operator":">=","value":"2025-04-01T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-01-01T00:00:00+00:00"}}]}}]}} User: "Q3 2024 reporting template" -{{"query_list":[{{"query":"Q3 2024 reporting template","temporal_filters":null}}]}} +{{"requires_retrieval":true,"query_list":[{{"query":"Q3 2024 reporting template","temporal_filters":null}}]}} User: "Evolution of the Department of Justice budget between 2020 and 2022" -{{"query_list":[{{"query":"Department of Justice budget in 2020","temporal_filters":null}},{{"query":"Department of Justice budget in 2021","temporal_filters":null}},{{"query":"Department of Justice budget in 2022","temporal_filters":null}}]}} +{{"requires_retrieval":true,"query_list":[{{"query":"Department of Justice budget in 2020","temporal_filters":null}},{{"query":"Department of Justice budget in 2021","temporal_filters":null}},{{"query":"Department of Justice budget in 2022","temporal_filters":null}}]}} User: "Latest safety bulletins" -{{"query_list":[{{"query":"Latest safety bulletins","temporal_filters":[{{"field":"created_at","operator":">=","value":"2026-01-15T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-04-16T00:00:00+00:00"}}]}}]}} \ No newline at end of file +{{"requires_retrieval":true,"query_list":[{{"query":"Latest safety bulletins","temporal_filters":[{{"field":"created_at","operator":">=","value":"2026-01-15T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-04-16T00:00:00+00:00"}}]}}]}} + +User: "How can you help me?" +{{"requires_retrieval":false,"query_list":[]}} + +User: "Hello, what was Product A revenue in Q1?" +{{"requires_retrieval":true,"query_list":[{{"query":"Product A revenue in Q1","temporal_filters":null}}]}} diff --git a/openrag/prompts/templates/spoken_style_answer_tmpl.txt b/openrag/prompts/templates/spoken_style_answer_tmpl.txt index a2fe6c311..32534293d 100644 --- a/openrag/prompts/templates/spoken_style_answer_tmpl.txt +++ b/openrag/prompts/templates/spoken_style_answer_tmpl.txt @@ -1,4 +1,4 @@ -You are an AI assistant designed for **spoken, conversational answers**. +You are **OpenRAG**, a retrieval-augmented generation system built by **LINAGORA**, designed for spoken, conversational answers. Your goal is to give short (1-2 sentences), clear, and accurate explanations, based only on the retrieved documents in `Context`. # Context @@ -9,10 +9,15 @@ Your goal is to give short (1-2 sentences), clear, and accurate explanations, ba 1. Use only the provided Context * Answer strictly from the information in `Context`. * Do not guess, infer, or use outside knowledge. + * For greetings, thanks, casual conversation, identity questions, or questions about your capabilities, answer briefly without using the Context and end with `[Sources: none]`. + * For identity or capability questions, explain that you are OpenRAG, built by LINAGORA, and that you retrieve and synthesize information from indexed documents with supporting sources. + * Do not present yourself as a general-purpose assistant or list unrelated abilities. + * If a message also asks a factual question, answer that part from the Context as usual. * If the Context lacks enough information, say so briefly and ask the user for more details. 2. Citations * Never place citations, source numbers, or references **inside** your answer text: citation is only at the end + * Do not copy the Context markers such as `[Source 1]` into the answer body. Use source numbers only in the final `[Sources: ...]` line. * Cite sources **only once**, on a **single line** that is the **very last line** of your response, separated from the body by a blank line. * The final line MUST match **exactly one** of these two formats (no other text on that line): - `[Sources: 1, 3]` — when one or more numbered sources from the Context contributed to your answer (comma-separated, ascending order, no duplicates) diff --git a/openrag/prompts/templates/sys_prompt_tmpl.txt b/openrag/prompts/templates/sys_prompt_tmpl.txt index 5e3598e5e..e82a9b1f9 100644 --- a/openrag/prompts/templates/sys_prompt_tmpl.txt +++ b/openrag/prompts/templates/sys_prompt_tmpl.txt @@ -1,4 +1,4 @@ -You are an AI conversational assistant specialized in **information retrieval and synthesis**. +You are **OpenRAG**, a retrieval-augmented generation system built by **LINAGORA**. Your goal is to provide **precise, reliable, and well-structured answers** using **only the retrieved documents** (`Context`). Prioritize **clarity, accuracy, and completeness** in your responses. @@ -10,11 +10,17 @@ Prioritize **clarity, accuracy, and completeness** in your responses. 1. Use only the provided Context * Base your answer **exclusively** on the information contained in the `Context`. * **Never infer**, assume, or rely on any external knowledge. + * For greetings, thanks, casual conversation, identity questions, or questions about your capabilities, answer briefly without using the Context and end with `[Sources: none]`. + * For identity or capability questions, explain that you are OpenRAG, built by LINAGORA, and that you search, retrieve, and synthesize information from indexed documents while citing the supporting sources. + * Present yourself as a document-grounded RAG system, not as a general-purpose assistant. Do not advertise unrelated abilities such as general knowledge, travel advice, creative writing, or coding unless the user asks about indexed documents covering those topics. + * Keep conversational, identity, and capability answers concise: normally 1-3 sentences. + * If a message combines conversation with a factual question, answer the factual part from the Context as usual. * If the context is **insufficient**, **invite the user** to clarify their query or provide additional keywords. * **Always answer with at least one sentence of text.** Your response body must **never be empty**, even when no source is relevant — in that case briefly explain (in the user's language) that the documents do not cover the question, then end with `[Sources: none]`. 2. Citations * Never place citations, source numbers, or references **inside** your answer text: citation is only at the end + * Do not copy the Context markers such as `[Source 1]` into the answer body. Use source numbers only in the final `[Sources: ...]` line. * Cite sources **only once**, on a **single line** that is the **very last line** of your response, separated from the body by a blank line. * The final line MUST match **exactly one** of these two formats (no other text on that line): - `[Sources: 1, 3, 5]` — when one or more numbered sources from the Context contributed to your answer (comma-separated, ascending order, no duplicates) @@ -27,4 +33,4 @@ Prioritize **clarity, accuracy, and completeness** in your responses. * Use **headings**, **bullet points**, **numbered lists**, or **tables** to organize information clearly. * Ensure responses are **concise yet complete**, avoiding omission of key details. -Here are the retrieved documents: `{context}` \ No newline at end of file +Here are the retrieved documents: `{context}` 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/model_endpoint_service.py b/openrag/services/orchestrators/model_endpoint_service.py index 058abfcc1..825e111db 100644 --- a/openrag/services/orchestrators/model_endpoint_service.py +++ b/openrag/services/orchestrators/model_endpoint_service.py @@ -114,11 +114,13 @@ def __init__( model_endpoint_repo: ModelEndpointRepository, config: Settings, partition_service: Any = None, + preset_service: Any = None, client_caches: dict[str, dict[str, Any]] | None = None, ) -> None: self._repo = model_endpoint_repo self._config = config self._partition_service = partition_service + self._preset_service = preset_service self._client_caches: dict[str, dict[str, Any]] = client_caches or {} # ------------------------------------------------------------------ @@ -396,6 +398,29 @@ async def update_model_endpoint(self, name: str, model_type: str, **fields: obje Pass ``new_name=`` to rename. After any change the in-memory config is reloaded and the stale cached client instance is evicted so the next request builds a fresh client against the updated config. + + A rename also cascades to every stored reference — ``partitions.embedder`` + / ``partitions.chat_llm`` and endpoint-name fields embedded in + ``pipeline_presets.config`` — inside the repo's own rename transaction + (see ``PgModelEndpointRepository.rename``, #770). Those writes are + invisible until the referencing services reload their in-memory caches, + which is why a rename also refreshes presets then partitions here — + the same order ``PresetService.update_preset`` uses, since partition + resolution reads the presets dict. + + Both reload calls ``await``, so a concurrent request can run between + them — and the DB rename has *already* committed by that point. Without + ``_alias_renamed_name``, a request landing in that window could resolve + a partition/preset that the cascade already repointed at ``new_name`` + against a registry that (until the final ``load_all()`` below) still + only knows ``name`` — a bare ``KeyError``. The alias makes both ``name`` + and ``new_name`` resolve immediately, built from the row this call just + wrote — not whatever the in-memory bucket held before it — so a rename + combined with a field change (e.g. a new ``endpoint``) aliases the + *updated* config, not a stale pre-update one. That also covers a reload + call above raising: the registry stays queryable under both names, + correctly, instead of the update's failure leaving it stuck on a stale + config until process restart. """ existing = await self._repo.get(name, model_type) if existing is None: @@ -423,6 +448,21 @@ async def update_model_endpoint(self, name: str, model_type: str, **fields: obje await self._repo.rename(name, model_type, new_name) effective_name = new_name renamed_from = name + self._alias_renamed_name(model_type, name, new_name, updated or existing) + # A cached *client instance* under either name would otherwise survive + # this alias — the factory checks its cache before consulting the + # config registry, so a stale pre-rename/pre-update client would keep + # serving until the eviction at the end of this method, which a + # reload call below raising would skip entirely. The config alias + # above is already fresh, so evicting now is safe: anything rebuilt + # from either name resolves through the up-to-date config, not stale + # cached state. + self._invalidate_client_cache(model_type, name) + self._invalidate_client_cache(model_type, new_name) + if self._preset_service is not None: + await self._preset_service.load_all() + if self._partition_service is not None: + await self._partition_service.load_partitions() if promote_to_default: # Clears any prior default and sets this row in one transaction, then @@ -540,6 +580,39 @@ async def validate_endpoint( # Internals # ------------------------------------------------------------------ + def _alias_renamed_name(self, model_type: str, old_name: str, new_name: str, row: ModelEndpointRow) -> None: + """Make both ``old_name`` and ``new_name`` resolve to *row* before any reload runs. + + Runs synchronously right after the rename ``await`` returns — no + further ``await`` happens before this executes, so no concurrent + request can observe the DB already renamed while the registry still + only answers to ``old_name``. + + Built from ``row`` — the just-written DB state — rather than copying + whatever the in-memory bucket currently holds under ``old_name``: a + rename can land in the same call as a field update (e.g. a new + ``endpoint`` URL), applied to the DB *before* this runs, so the stale + in-memory entry would alias both names to the pre-update config. If a + reload below then raises, that staleness would never get corrected + by the final ``load_all()`` this call never reaches — the registry + would keep serving the old endpoint under the new (DB-authoritative) + name until process restart. The next full ``load_all()`` (below, or + from any later CRUD call) rebuilds the bucket straight from DB and + drops the ``old_name`` entry on its own. + """ + bucket: dict[str, Any] | None = getattr(self._config.models, model_type, None) + if bucket is None: + return + cfg = ModelEndpointConfig( + endpoint=row.endpoint, + model_name=row.model_name, + batch_size=row.batch_size, + timeout=row.timeout, + extra=row.extra, + ) + bucket[old_name] = cfg + bucket[new_name] = cfg + def _invalidate_client_cache(self, model_type: str, name: str) -> None: """Evict ``name`` from the component-factory cache for ``model_type``.""" cache = self._client_caches.get(model_type) diff --git a/openrag/services/orchestrators/partition_service.py b/openrag/services/orchestrators/partition_service.py index b0204a9e4..b95978562 100644 --- a/openrag/services/orchestrators/partition_service.py +++ b/openrag/services/orchestrators/partition_service.py @@ -38,6 +38,7 @@ from core.utils.conts import is_internal_metadata_key from core.utils.exceptions import ( ConfigError, + ConflictError, NotFoundError, PartitionNotFoundError, UserNotFoundError, @@ -61,6 +62,9 @@ # and the admin partition-list route would expand it to *every* partition — see # ``list_existant_partitions``. Matched case-insensitively. _RESERVED_PARTITION_NAMES = frozenset({"all"}) +_MAX_MEMBER_CANDIDATE_PAGE_SIZE = 100 +_MAX_POSTGRES_INTEGER = 2_147_483_647 +_MIN_MEMBER_CANDIDATE_SEARCH_LENGTH = 3 # Columns where an explicit ``None`` in a PATCH is a real value (SQL NULL = # "reset to default"), not the omitted-field sentinel that the None-filter @@ -98,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 @@ -435,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) @@ -489,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 { @@ -503,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 {}, } # ------------------------------------------------------------------ @@ -540,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: @@ -673,13 +699,86 @@ def _meta(row: dict[str, Any]) -> dict[str, Any]: # ------------------------------------------------------------------ async def list_members(self, partition: str) -> list[dict]: + """Return role data without identity lookups for authorization callers.""" await self._ensure_partition(partition) return await self._membership_repo.list_partition_members(partition) + async def list_member_candidates( + self, + partition: str, + *, + search: str | None, + cursor: int | None = None, + limit: int = 25, + ) -> dict: + """Return matching non-members without exposing the full user directory.""" + if cursor is not None and (cursor < 0 or cursor > _MAX_POSTGRES_INTEGER): + raise ValidationError("Cursor is outside the supported user ID range.") + if limit < 1 or limit > _MAX_MEMBER_CANDIDATE_PAGE_SIZE: + raise ValidationError( + f"Limit must be between 1 and {_MAX_MEMBER_CANDIDATE_PAGE_SIZE}.", + ) + + await self._ensure_partition(partition) + normalized_search = search.strip() if search else "" + if not normalized_search: + raise ValidationError("Search is required to find users.") + + search_user_id: int | None = None + search_prefix: str | None = None + if normalized_search.isascii() and normalized_search.isdecimal(): + numeric_search = int(normalized_search) + if numeric_search <= _MAX_POSTGRES_INTEGER: + search_user_id = numeric_search + if len(normalized_search) >= _MIN_MEMBER_CANDIDATE_SEARCH_LENGTH: + search_prefix = normalized_search + elif len(normalized_search) < _MIN_MEMBER_CANDIDATE_SEARCH_LENGTH: + raise ValidationError( + f"Enter at least {_MIN_MEMBER_CANDIDATE_SEARCH_LENGTH} characters or an exact user ID.", + ) + else: + search_prefix = normalized_search + + rows = await self._membership_repo.list_partition_member_candidates( + partition, + search_prefix=search_prefix, + search_user_id=search_user_id, + after_id=cursor, + limit=limit + 1, + ) + has_more = len(rows) > limit + candidates = rows[:limit] + return { + "candidates": candidates, + "limit": limit, + "has_more": has_more, + "next_cursor": candidates[-1]["user_id"] if has_more and candidates else None, + } + + async def list_members_with_identities(self, partition: str) -> list[dict]: + """Enrich the admin-facing member list with one bulk user lookup.""" + members = await self.list_members(partition) + users = { + user.id: user for user in await self._user_repo.get_users_by_ids([member["user_id"] for member in members]) + } + for member in members: + user = users.get(member["user_id"]) + member["display_name"] = user.display_name if user else None + member["email"] = user.email if user else None + return members + async def add_member(self, partition: str, user_id: int, role: str) -> None: await self._ensure_partition(partition) await self._ensure_user_exists(user_id) - await self._membership_repo.add_partition_member(partition, user_id, role) + created = await self._membership_repo.add_partition_member(partition, user_id, role) + if not created: + raise ConflictError( + ( + f"User {user_id} is already a member of partition '{partition}'. " + f"Use PATCH /partition/{partition}/users/{user_id} to change their role." + ), + code="PARTITION_MEMBER_EXISTS", + ) logger.info(f"User_id {user_id} added to partition '{partition}'.") async def remove_member(self, partition: str, user_id: int) -> 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 17054fc52..7c9287b23 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -48,7 +48,7 @@ SOURCE_SEPARATOR, format_context, format_web_context, - load_template_by_key, + prepend_system_prompt, ) from core.utils.exceptions import WorkspaceNotFoundError from core.utils.logging import get_logger @@ -58,11 +58,13 @@ stream_with_source_filtering, ) from core.utils.text import get_num_tokens +from core.utils.web_url import normalize_web_url from services.inference.runtime import detect_language, get_llm_semaphore 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 @@ -88,8 +90,10 @@ {query}""" _QUERY_JSON_HINT = ( - "\n\nRespond ONLY with a JSON object of the form " - '{"query_list": [{"query": "", "temporal_filters": null}]}.' + "\n\nRespond ONLY with one of these JSON forms: " + '{"requires_retrieval": false, "query_list": []} when retrieval is not needed, or ' + '{"requires_retrieval": true, ' + '"query_list": [{"query": "", "temporal_filters": null}]} when it is.' ) @@ -113,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 @@ -120,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 @@ -142,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. @@ -296,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"), ) @@ -390,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) @@ -412,6 +454,26 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L partition = [scope.partition] filter_params = {"file_id": scope.file_ids} + force_retrieval = use_websearch or use_map_reduce + if not queries.query_list: + if not queries.requires_retrieval and not force_retrieval: + # 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, + context="", + current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S"), + ) + return payload, [], [], True + queries = SearchQueries(query_list=[Query(query=messages[-1]["content"])]) + web_results: list = [] if partition is not None and use_websearch: chunks, web_lists = await self._gather_rag_and_web(queries, partition, top_k, filter_params) @@ -425,18 +487,19 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L chunks = [] if not chunks and not web_results and partition is None: - return payload, [], [] + return payload, [], [], False docs = [c.to_langchain() for c in chunks] if use_map_reduce and docs: docs = await self._map_reduce(" ".join(q.query for q in queries.query_list), docs) - web_formatted, web_tokens = "", 0 + web_formatted, web_source_numbers, web_tokens = "", [], 0 + web_start_index = 1 if web_results: - web_formatted, _, web_tokens = format_web_context( + web_formatted, web_source_numbers, web_tokens = format_web_context( web_results, length_function=get_num_tokens(), - start_index=1, + start_index=web_start_index, max_tokens=self._web.max_tokens, ) context, included = format_context( @@ -448,18 +511,23 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L if web_results: if docs: - web_formatted, _, _ = format_web_context( + web_start_index = len(docs) + 1 + web_formatted, web_source_numbers, _ = format_web_context( web_results, length_function=get_num_tokens(), - start_index=len(docs) + 1, + start_index=web_start_index, max_tokens=self._web.max_tokens, ) else: context = "" context = f"{context}{SOURCE_SEPARATOR}{web_formatted}" if context else web_formatted + 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, { @@ -470,7 +538,7 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L }, ) payload["messages"] = new_messages - return payload, docs, web_results + return payload, docs, web_results, True async def _gather_rag_and_web(self, queries, partition, top_k, filter_params): # Fuse the doc branch through retrieve_multi so a partition's rrf_k drives @@ -487,21 +555,34 @@ 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) - chunks = await self._retrieval.retrieve_multi(partitions=partition, search_queries=queries) - docs = [c.to_langchain() for c in chunks] - context, included = format_context( - [doc.page_content for doc in docs], - max_context_tokens=self._max_context_tokens, - length_function=get_num_tokens(), - ) - docs = [docs[i] for i in included] - if docs: - payload["prompt"] = ( - f"Given the content\n{context}\nComplete the following prompt: {prompt}\n" - "At the very end of your response, on a new line, list which source numbers " - "you used: [Sources: 1, 3]" + # 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 = [], "" + else: + queries = SearchQueries(query_list=[Query(query=prompt)]) + if queries.query_list: + chunks = await self._retrieval.retrieve_multi(partitions=partition, search_queries=queries) + docs = [c.to_langchain() for c in chunks] + context, included = format_context( + [doc.page_content for doc in docs], + max_context_tokens=self._max_context_tokens, + length_function=get_num_tokens(), ) + docs = [docs[i] for i in included] + + metadata = payload.get("metadata") or {} + 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, + current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S"), + ) + payload["prompt"] = f"{instructions}\n\n# User request\n{prompt}" return payload, docs # ------------------------------------------------------------------ @@ -568,19 +649,35 @@ async def chat( """Non-streaming chat completion → finalized OpenAI dict.""" metadata = payload.get("metadata") or {} llm = self._resolve_llm(partitions) + citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): docs, web_results = [], [] else: - payload, docs, web_results = await self._prepare_chat(partitions, payload, llm) + payload, docs, web_results, citation_protocol_active = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) + structured_output = _allows_uncited_sources(payload) payload["messages"] = self._sanitize_messages(payload["messages"]) chunk = await llm.chat(payload["messages"], **_sampling(payload)) chunk["model"] = model_name content = chunk.get("choices", [{}])[0].get("message", {}).get("content", "") or "" - clean, citations = extract_and_strip_sources_block(content) + if citation_protocol_active and not structured_output: + clean, citations = extract_and_strip_sources_block( + content, + include_inline_markers=bool(sources), + ) + else: + clean, citations = content, None chunk["choices"][0]["message"]["content"] = clean - chunk["extra"] = json.dumps({"sources": filter_sources_by_citations(sources, citations)}) + chunk["extra"] = json.dumps( + { + "sources": filter_sources_by_citations( + sources, + citations, + allow_uncited=structured_output, + ) + } + ) return chunk async def chat_stream( @@ -594,15 +691,23 @@ async def chat_stream( """Streaming chat completion → SSE strings with filtered sources.""" metadata = payload.get("metadata") or {} llm = self._resolve_llm(partitions) + citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): docs, web_results = [], [] else: - payload, docs, web_results = await self._prepare_chat(partitions, payload, llm) + payload, docs, web_results, citation_protocol_active = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) + structured_output = _allows_uncited_sources(payload) payload["messages"] = self._sanitize_messages(payload["messages"]) llm_stream = llm.stream_chat(payload["messages"], **_sampling(payload)) - async for sse_line in stream_with_source_filtering(llm_stream, sources, model_name): + async for sse_line in stream_with_source_filtering( + llm_stream, + sources, + model_name, + allow_uncited_sources=structured_output, + citation_protocol_active=citation_protocol_active and not structured_output, + ): yield sse_line async def complete( @@ -614,17 +719,33 @@ async def complete( ) -> dict: """Non-streaming text completion → finalized OpenAI dict.""" llm = self._resolve_llm(partitions) + citation_protocol_active = partitions is not None if partitions is None: docs = [] else: payload, docs = await self._prepare_completions(partitions, payload, llm) sources = prepare_sources(docs, []) + structured_output = _allows_uncited_sources(payload) resp = await llm.generate(payload["prompt"], **_sampling(payload, key="prompt")) text = resp.get("choices", [{}])[0].get("text", "") or "" - clean, citations = extract_and_strip_sources_block(text) + if citation_protocol_active and not structured_output: + clean, citations = extract_and_strip_sources_block( + text, + include_inline_markers=bool(sources), + ) + else: + clean, citations = text, None resp["choices"][0]["text"] = clean - resp["extra"] = json.dumps({"sources": filter_sources_by_citations(sources, citations)}) + resp["extra"] = json.dumps( + { + "sources": filter_sources_by_citations( + sources, + citations, + allow_uncited=structured_output, + ) + } + ) return resp @@ -649,8 +770,10 @@ def _dedupe_web(web_lists: list[list]) -> list: seen: set[str] = set() out: list = [] for r in (r for lst in web_lists for r in lst): - if r.url not in seen: - seen.add(r.url) + url = normalize_web_url(r.url) + if url is not None and url not in seen: + r.url = url + seen.add(url) out.append(r) return out @@ -666,4 +789,10 @@ def _sampling(payload: dict, key: str = "messages") -> dict: return {k: v for k, v in payload.items() if k not in drop} +def _allows_uncited_sources(payload: dict) -> bool: + """Structured output cannot carry the plain-text citation marker.""" + response_format = payload.get("response_format") + return isinstance(response_format, dict) and response_format.get("type") in {"json_object", "json_schema"} + + __all__ = ["QueryService", "RAGMODE"] diff --git a/openrag/services/orchestrators/retrieval_service.py b/openrag/services/orchestrators/retrieval_service.py index 0fcbf4862..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 {} @@ -164,7 +179,72 @@ def _require_partition_config(self, partition: str): def _legacy_retriever_value(self, name: str, default: Any) -> Any: return getattr(self._config.retriever, name, default) - def _pipeline_for_partition(self, partition: str) -> tuple[RetrieverPipeline, int | None]: + def _resolve_reranker(self, reranker_name: str | None, partition: str) -> Reranker | None: + """Effective reranker for one partition's retrieval pipeline. + + Resolution order — mirrors ``QueryService._resolve_llm``: + + 1. The partition's configured ``reranker`` preset, resolved fresh via + the model-endpoint catalog factory so a rename/promotion of that + endpoint takes effect immediately. + 2. The **catalog default** endpoint (``is_default=True``) when the + partition sets no preset, or its preset name has gone stale (the + endpoint was renamed/deleted after assignment — unlike + ``chat_llm``, this field has no create/PATCH-time validation, so a + stale name reaching here is expected, not a bug). + 3. The static reranker built at startup from ``settings.reranker``, + only when no factory is wired (unit tests) or the catalog has no + default reranker endpoint yet. + + The resolved endpoint name is always logged (at debug), including for + the default, so "which reranker ran?" is answerable from the logs. + """ + if self._reranker_factory is None: + return self._legacy_reranker + if reranker_name: + try: + reranker = self._reranker_factory(reranker_name) + except KeyError: + logger.bind(reranker=reranker_name, partition=partition).warning( + "Partition reranker preset not found in the model-endpoint catalog — " + "falling back to the default reranker" + ) + else: + logger.bind(reranker=reranker_name, partition=partition).debug( + "Reranking with the partition's reranker preset" + ) + return reranker + try: + reranker = self._reranker_factory("default") + except KeyError: + pass + else: + logger.bind(reranker=self._default_reranker_name(), partition=partition).debug( + "Reranking with the default reranker preset" + ) + return reranker + logger.bind(partition=partition).debug( + "Reranking with the static default reranker (no catalog default endpoint)" + ) + return self._legacy_reranker + + def _default_reranker_name(self) -> str: + """Real endpoint name behind the catalog reranker ``"default"`` alias, for logging. + + Same identity-lookup trick as ``QueryService._default_llm_name``: the + ``"default"`` alias config object is the *same* object as its real-named + entry, so the name is recovered by identity. Returns ``"default"`` when + it can't be resolved (e.g. the alias isn't populated yet). + """ + rerankers = self._config.models.reranker + default_cfg = rerankers.get("default") + if default_cfg is not None: + for name, cfg in rerankers.items(): + if name != "default" and cfg is default_cfg: + return name + return "default" + + 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. @@ -182,15 +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") - reranker = None - if pipeline_cfg.enable_reranker: - if self._reranker_factory is not None: - reranker = self._reranker_factory(pipeline_cfg.reranker or "default") - else: - reranker = self._legacy_reranker + # 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, @@ -214,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() @@ -233,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) @@ -326,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( @@ -334,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) @@ -348,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( @@ -356,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/e5f6a7b8c9d0_add_user_display_name_prefix_index.py b/openrag/services/persistence/migrations/alembic/versions/e5f6a7b8c9d0_add_user_display_name_prefix_index.py new file mode 100644 index 000000000..404a860d8 --- /dev/null +++ b/openrag/services/persistence/migrations/alembic/versions/e5f6a7b8c9d0_add_user_display_name_prefix_index.py @@ -0,0 +1,67 @@ +"""add user display name prefix index + +Revision ID: e5f6a7b8c9d0 +Revises: d4e5f6a7b8c9 +Create Date: 2026-07-27 00:00:00.000000 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from schema_helpers import table_exists + +revision: str = "e5f6a7b8c9d0" +down_revision: str | Sequence[str] | None = "d4e5f6a7b8c9" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +_INDEX_NAME = "ix_users_lower_display_name_pattern" + + +def _index_validity() -> bool | None: + result = op.get_bind().execute( + sa.text( + """ + SELECT i.indisvalid + FROM pg_catalog.pg_class index_class + JOIN pg_catalog.pg_index i ON i.indexrelid = index_class.oid + JOIN pg_catalog.pg_class table_class ON table_class.oid = i.indrelid + JOIN pg_catalog.pg_namespace namespace ON namespace.oid = index_class.relnamespace + WHERE namespace.nspname = current_schema() + AND table_class.relname = 'users' + AND index_class.relname = :index_name + """, + ), + {"index_name": _INDEX_NAME}, + ) + validity = result.scalar() + return bool(validity) if validity is not None else None + + +def _drop_index_concurrently() -> None: + with op.get_context().autocommit_block(): + op.execute(sa.text(f"DROP INDEX CONCURRENTLY {_INDEX_NAME}")) + + +def upgrade() -> None: + if not table_exists("users"): + return + + validity = _index_validity() + if validity is False: + _drop_index_concurrently() + + if validity is not True: + with op.get_context().autocommit_block(): + op.execute( + sa.text( + f"CREATE INDEX CONCURRENTLY {_INDEX_NAME} ON users (LOWER(display_name) text_pattern_ops)", + ), + ) + + +def downgrade() -> None: + if table_exists("users") and _index_validity() is not None: + _drop_index_concurrently() 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/model_endpoint_repo.py b/openrag/services/persistence/model_endpoint_repo.py index 933352b88..57b8164b8 100644 --- a/openrag/services/persistence/model_endpoint_repo.py +++ b/openrag/services/persistence/model_endpoint_repo.py @@ -27,6 +27,20 @@ # transaction; ModelEndpointService.update_model_endpoint routes is_default there. _ALLOWED_UPDATE_FIELDS = frozenset({"endpoint", "model_name", "batch_size", "timeout", "extra"}) +# Endpoint names are referenced by value elsewhere, and nothing updates those +# references when an endpoint is renamed (#770) — so ``rename()`` cascades to +# every known reference in the same transaction as the name change. Direct +# endpoint-name columns on ``partitions``, keyed by the model_type they hold: +_PARTITION_COLUMN_BY_TYPE = {"embedder": "embedder", "llm": "chat_llm"} +# Endpoint-name keys embedded in ``pipeline_presets.config`` (JSONB), by the +# preset_type that carries them — see core/config/retrieval_pipeline.py and +# core/config/indexation_pipeline.py for the field definitions. +_RETRIEVAL_PRESET_KEYS_BY_TYPE = {"llm": ("llm",), "reranker": ("reranker",)} +_INDEXATION_PRESET_KEYS_BY_TYPE = { + "llm": ("contextualization_llm", "metadata_extraction_llm", "topic_tagging_llm"), + "vlm": ("vlm",), +} + class PgModelEndpointRepository(ModelEndpointRepository): """asyncpg-backed implementation of :class:`ModelEndpointRepository`.""" @@ -124,12 +138,80 @@ async def update(self, name: str, model_type: str, **fields: object) -> ModelEnd return self._to_model(rec) if rec else None async def rename(self, name: str, model_type: str, new_name: str) -> None: - await self.pool.execute( - "UPDATE model_endpoints SET name = $3, updated_at = now() WHERE name = $1 AND model_type = $2", - name, - model_type, - new_name, - ) + """Rename an endpoint and cascade the new name to every stored reference. + + Left alone, a rename silently strands every partition or preset that + pointed at the old name — ``partitions.embedder`` / ``partitions.chat_llm``, + and the endpoint-name fields embedded in ``pipeline_presets.config`` + (JSONB) — since nothing else in the schema updates those when the + referenced row's name changes (#770). All writes run in this one + transaction so a partial cascade can never leave the registry and its + referents disagreeing. + + The caller (``ModelEndpointService.update_model_endpoint``) still owns + refreshing the in-memory partition/preset caches afterwards — this + method only makes the DB-side references consistent. + + Raises :class:`NotFoundError` if ``name`` vanished between the + service's existence check and this transaction (a concurrent delete) + — mirroring ``PgPipelinePresetRepository.rename``. Without the + ``RETURNING`` check, a lost race would still run the cascade below, + repointing partitions/presets at a ``new_name`` that was never + actually created. + + Also ``LOCK``s ``partitions`` ``IN SHARE MODE`` before touching + anything — the same lock :meth:`PgPresetRepository.delete` takes, and + the same table :meth:`PgPartitionRepository.update_partition` writes + to *before* its own DB-authoritative ``chat_llm`` re-check. Without + this, a partition PATCH could validate ``name`` against the in-memory + catalog, then block behind this transaction's cascade on the exact + row it's about to write, and resume writing the now-renamed-away + ``name`` straight back once this commits — permanently stranding that + partition the moment the service's temporary alias (see + ``ModelEndpointService._alias_renamed_name``) drops. Locking first, in + the same order both call sites use, makes the two block on each other + instead of interleaving: whichever transaction's ``partitions`` write + commits first is the one the other observes. + """ + async with self.pool.acquire() as conn: + async with conn.transaction(): + await conn.execute("LOCK TABLE partitions IN SHARE MODE") + + renamed = await conn.fetchrow( + "UPDATE model_endpoints SET name = $3, updated_at = now() " + "WHERE name = $1 AND model_type = $2 RETURNING name", + name, + model_type, + new_name, + ) + if renamed is None: + raise NotFoundError(f"Endpoint '{name}' of type '{model_type}' not found.") + + partition_col = _PARTITION_COLUMN_BY_TYPE.get(model_type) + if partition_col: + await conn.execute( + f"UPDATE partitions SET {partition_col} = $2 WHERE {partition_col} = $1", + name, + new_name, + ) + + for preset_type, keys in ( + ("retrieval", _RETRIEVAL_PRESET_KEYS_BY_TYPE.get(model_type, ())), + ("indexation", _INDEXATION_PRESET_KEYS_BY_TYPE.get(model_type, ())), + ): + for key in keys: + await conn.execute( + """ + UPDATE pipeline_presets + SET config = jsonb_set(config, $1::text[], to_jsonb($2::text)), updated_at = now() + WHERE preset_type = $3 AND config->>$4 = $5 + """, + [key], + new_name, + preset_type, + key, + name, + ) async def delete(self, name: str, model_type: str) -> bool: result = await self.pool.execute( diff --git a/openrag/services/persistence/partition_membership_repo.py b/openrag/services/persistence/partition_membership_repo.py index c80339be4..977995553 100644 --- a/openrag/services/persistence/partition_membership_repo.py +++ b/openrag/services/persistence/partition_membership_repo.py @@ -123,6 +123,56 @@ async def count_partition_users(self, partition: str) -> int: partition, ) + async def list_partition_member_candidates( + self, + partition: str, + *, + search_prefix: str | None, + search_user_id: int | None, + after_id: int | None, + limit: int, + ) -> list[dict]: + """Return matching users after the cursor who are not partition members.""" + escaped_prefix = None + if search_prefix is not None: + escaped_prefix = search_prefix.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + + rows = await self.pool.fetch( + """ + SELECT u.id AS user_id, u.display_name, u.email + FROM users u + WHERE NOT EXISTS ( + SELECT 1 + FROM partition_memberships m + WHERE m.partition_name = $1 + AND m.user_id = u.id + ) + AND ( + ($2::integer IS NOT NULL AND u.id = $2) + OR ( + $3::text IS NOT NULL + AND LOWER(u.display_name) LIKE LOWER($3) || '%' ESCAPE '\\' + ) + ) + AND ($4::integer IS NULL OR u.id > $4) + ORDER BY u.id + LIMIT $5 + """, + partition, + search_user_id, + escaped_prefix, + after_id, + limit, + ) + return [ + { + "user_id": row["user_id"], + "display_name": row["display_name"], + "email": row["email"], + } + for row in rows + ] + # ── Legacy method names used by the Phase 7C shim ──────────────── async def list_partition_members(self, partition: str) -> list[dict]: @@ -168,19 +218,20 @@ async def add_partition_member(self, partition: str, user_id: int, role: str) -> """, partition, ) - await conn.execute( + created = await conn.fetchval( """ INSERT INTO partition_memberships (partition_name, user_id, role, added_at) VALUES ($1, $2, $3, NOW()) ON CONFLICT (partition_name, user_id) - DO UPDATE SET role = EXCLUDED.role + DO NOTHING + RETURNING 1 """, partition, user_id, role, ) - return True + return created is not None async def remove_partition_member(self, partition: str, user_id: int) -> bool: """TODO(phase-9): remove.""" diff --git a/openrag/services/persistence/partition_repo.py b/openrag/services/persistence/partition_repo.py index 1d4770c33..8fbcfeb13 100644 --- a/openrag/services/persistence/partition_repo.py +++ b/openrag/services/persistence/partition_repo.py @@ -35,6 +35,16 @@ "indexation_preset": "indexation", "retrieval_preset": "retrieval", } +# Partition columns that reference a model_endpoints row, mapped to the +# model_type they point at. Only `chat_llm` is assignment-validated today +# (PartitionService._validate_chat_llm_ref checks the in-memory catalog); +# `embedder` carries no such check, so it is deliberately not listed here. +# Assigning chat_llm must be guarded against a concurrent rename the same way +# a preset assignment is guarded against a concurrent preset delete — see +# update_partition and PgModelEndpointRepository.rename. +_MODEL_ENDPOINT_COLUMN_TYPES = { + "chat_llm": "llm", +} _PARTITION_UPDATE_COLUMNS = frozenset( { "description", @@ -45,6 +55,7 @@ "collection_name", "chat_history_depth", "chat_llm", + "generation_prompt_names", } ) _PARTITION_OPERATION_LOCK_NAMESPACE = 20260720 @@ -56,6 +67,19 @@ def _partition_updates(fields: dict[str, object]) -> dict[str, object]: return {key: value for key, value in fields.items() if key in _PARTITION_UPDATE_COLUMNS} +def _endpoint_refs(updates: dict[str, object]) -> dict[str, str]: + """Model-endpoint-referencing columns in *updates*, mapped to their model_type. + + A ``None`` value clears the (nullable) column rather than pointing it at a + name, so it needs no existence check. + """ + return { + col: model_type + for col, model_type in _MODEL_ENDPOINT_COLUMN_TYPES.items() + if col in updates and updates[col] is not None + } + + class _PartitionOperationGuard: def __init__(self, repo: PgPartitionRepository, conn: asyncpg.Connection) -> None: self._repo = repo @@ -282,24 +306,35 @@ async def update_partition(self, name: str, **fields: object) -> dict | None: """Update a partition's config columns. When the update assigns a preset column (``indexation_preset`` / - ``retrieval_preset``), the write and a DB-authoritative existence check - run in one transaction that touches ``partitions`` before - ``pipeline_presets`` — the same lock order :meth:`PgPresetRepository.delete` - uses (it ``LOCK``s ``partitions`` ``IN SHARE MODE``, then ``DELETE``s the - preset). That makes an assign and a concurrent delete of the same preset + ``retrieval_preset``) or ``chat_llm``, the write and a DB-authoritative + existence check run in one transaction that touches ``partitions`` + before ``pipeline_presets`` / ``model_endpoints`` — the same lock order + :meth:`PgPresetRepository.delete` and :meth:`PgModelEndpointRepository. + rename` use (both ``LOCK`` ``partitions`` ``IN SHARE MODE`` first). That + makes an assign and a concurrent delete/rename of the same reference serialize without deadlocking, so a partition can never end up pointing - at a preset that no longer exists: - - * if this UPDATE commits first, the delete's ``COUNT`` sees the reference - and refuses with 409; - * if the delete commits first, this UPDATE blocks on its ``SHARE`` lock, - then the follow-up ``SELECT`` sees the vanished preset and the - transaction rolls the write back (raising ``PRESET_NOT_FOUND``). + at a preset or model endpoint that no longer exists under that name: + + * if this UPDATE commits first, the delete/rename's own guard against + ``partitions`` (a ``COUNT`` for presets, the ``SHARE`` lock itself for + renames) sees the reference and blocks or refuses accordingly; + * if the delete/rename commits first, this UPDATE blocks on its + ``SHARE``-conflicting write, then the follow-up ``SELECT`` sees the + vanished name and the transaction rolls the write back (raising + ``PRESET_NOT_FOUND`` / ``MODEL_ENDPOINT_NOT_FOUND``) — instead of + silently writing back a name a concurrent rename already moved on + from, which is what a validate-in-memory-then-blind-UPDATE sequence + could otherwise do. + + ``embedder`` carries no such check — it has no assignment-time + validation at all today (see ``_MODEL_ENDPOINT_COLUMN_TYPES``), so + there is nothing here for a concurrent rename to race against. """ updates = _partition_updates(fields) if updates: preset_refs = {col: _PRESET_COLUMN_TYPES[col] for col in updates if col in _PRESET_COLUMN_TYPES} - if preset_refs: + endpoint_refs = _endpoint_refs(updates) + if preset_refs or endpoint_refs: async with self.pool.acquire() as conn: return await self._update_partition_on_conn(conn, name, **fields) return await self._update_partition_on_conn(self.pool, name, **fields) @@ -318,18 +353,19 @@ 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 *" preset_refs = {col: _PRESET_COLUMN_TYPES[col] for col in updates if col in _PRESET_COLUMN_TYPES} - if not preset_refs: + endpoint_refs = _endpoint_refs(updates) + if not preset_refs and not endpoint_refs: row = await conn.fetchrow(sql, *params) return self._row_to_full_dict(row) if row else None transaction = getattr(conn, "transaction", None) if transaction is None: - raise TypeError("preset-reference updates require a connection transaction") + raise TypeError("preset/model-endpoint reference updates require a connection transaction") async with transaction(): row = await conn.fetchrow(sql, *params) if row is None: @@ -345,6 +381,17 @@ async def _update_partition_on_conn( f"{preset_type.capitalize()} preset '{updates[col]}' does not exist.", code="PRESET_NOT_FOUND", ) + for col, model_type in endpoint_refs.items(): + exists = await conn.fetchval( + "SELECT 1 FROM model_endpoints WHERE name = $1 AND model_type = $2", + updates[col], + model_type, + ) + if not exists: + raise ValidationError( + f"{model_type.upper()} endpoint '{updates[col]}' referenced by {col} not found.", + code="MODEL_ENDPOINT_NOT_FOUND", + ) return self._row_to_full_dict(row) # ── Legacy method names used by the Phase 7C shim ──────────────── @@ -391,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 5bd37eb9b..0f56f1472 100644 --- a/openrag/services/persistence/schema.py +++ b/openrag/services/persistence/schema.py @@ -29,6 +29,7 @@ String, Table, UniqueConstraint, + func, text, ) from sqlalchemy.dialects.postgresql import JSONB @@ -91,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, @@ -105,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), @@ -224,6 +290,12 @@ Column("file_count", Integer, nullable=False, default=0), ) +Index( + "ix_users_lower_display_name_pattern", + func.lower(users.c.display_name).label("display_name_lower"), + postgresql_ops={"display_name_lower": "text_pattern_ops"}, +) + oidc_sessions = Table( "oidc_sessions", @@ -335,6 +407,7 @@ "metadata", "model_endpoints", "pipeline_presets", + "prompts", "topic_tags", "partitions", "files", diff --git a/openrag/services/persistence/user_repo.py b/openrag/services/persistence/user_repo.py index 88206850b..21fc3ea11 100644 --- a/openrag/services/persistence/user_repo.py +++ b/openrag/services/persistence/user_repo.py @@ -102,6 +102,15 @@ async def get_user(self, user_id: int) -> User | None: memberships = await self._fetch_memberships(user_id) return self._row_to_user(row, memberships) + async def get_users_by_ids(self, user_ids: list[int]) -> list[User]: + if not user_ids: + return [] + rows = await self.pool.fetch( + "SELECT * FROM users WHERE id = ANY($1::int[])", + user_ids, + ) + return [self._row_to_user(row) for row in rows] + async def get_user_by_email(self, email: str) -> User | None: row = await self.pool.fetchrow( "SELECT * FROM users WHERE email = $1", @@ -332,6 +341,7 @@ async def list_users_dict(self) -> list[dict]: "id": r["id"], "display_name": r["display_name"], "external_user_id": r["external_user_id"], + "email": r["email"], "is_admin": r["is_admin"], "file_quota": r["file_quota"], "file_count": r["file_count"], 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/pyproject.toml b/pyproject.toml index 21e5f5b89..a312607a9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "openrag" -version = "2.0.1" +version = "2.1.0" description = "Add your description here" readme = "README.md" requires-python = ">=3.12" 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_membership_repo.py b/tests/integration/repos/test_partition_membership_repo.py index 157288bba..e4708b2ec 100644 --- a/tests/integration/repos/test_partition_membership_repo.py +++ b/tests/integration/repos/test_partition_membership_repo.py @@ -59,3 +59,85 @@ async def test_remove_partition(self, postgres_store: PostgresStore): ) assert await postgres_store.membership_repo.remove_partition(user.id, "docs") is True assert await postgres_store.membership_repo.list_user_partitions(user.id) == [] + + async def test_candidate_search_is_paginated_and_excludes_members( + self, + postgres_store: PostgresStore, + ): + member = await postgres_store.user_repo.create_user( + _user(display_name="Candidate Member", email="member@example.com"), + ) + candidate = await postgres_store.user_repo.create_user( + _user(display_name="Candidate Sam", email="candidate@example.com"), + ) + third = await postgres_store.user_repo.create_user( + _user(display_name="Candidate Taylor", email="taylor@example.com"), + ) + fourth = await postgres_store.user_repo.create_user( + _user(display_name="Candidate Jordan", email="jordan@example.com"), + ) + await postgres_store.partition_repo.create_partition("docs") + await postgres_store.membership_repo.assign_partition( + UserPartition(user_id=member.id, partition="docs"), + ) + + matches = await postgres_store.membership_repo.list_partition_member_candidates( + "docs", + search_prefix="candidate s", + search_user_id=None, + after_id=None, + limit=10, + ) + assert matches == [ + { + "user_id": candidate.id, + "display_name": "Candidate Sam", + "email": "candidate@example.com", + } + ] + + first_page = await postgres_store.membership_repo.list_partition_member_candidates( + "docs", + search_prefix="candidate", + search_user_id=None, + after_id=None, + limit=2, + ) + second_page = await postgres_store.membership_repo.list_partition_member_candidates( + "docs", + search_prefix="candidate", + search_user_id=None, + after_id=third.id, + limit=2, + ) + assert [row["user_id"] for row in first_page] == [candidate.id, third.id] + assert [row["user_id"] for row in second_page] == [fourth.id] + + exact_match = await postgres_store.membership_repo.list_partition_member_candidates( + "docs", + search_prefix=None, + search_user_id=fourth.id, + after_id=None, + limit=2, + ) + assert [row["user_id"] for row in exact_match] == [fourth.id] + + async def test_add_partition_member_conflict_preserves_existing_role( + self, + postgres_store: PostgresStore, + ): + user = await postgres_store.user_repo.create_user(_user()) + await postgres_store.partition_repo.create_partition("docs") + await postgres_store.membership_repo.assign_partition( + UserPartition(user_id=user.id, partition="docs", role=PartitionRole.VIEWER), + ) + + created = await postgres_store.membership_repo.add_partition_member( + "docs", + user.id, + "owner", + ) + + memberships = await postgres_store.membership_repo.list_user_partitions(user.id) + assert created is False + assert memberships[0].role == PartitionRole.VIEWER 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_phase14_partition_routes.py b/tests/unit/api/routers/admin/test_phase14_partition_routes.py index 0ed5fa01c..957a84d70 100644 --- a/tests/unit/api/routers/admin/test_phase14_partition_routes.py +++ b/tests/unit/api/routers/admin/test_phase14_partition_routes.py @@ -10,7 +10,7 @@ ) from api.routers.admin import partitions from di.providers import get_partition_service -from fastapi import FastAPI +from fastapi import FastAPI, HTTPException, status def _partition_detail(**overrides: Any) -> dict[str, Any]: @@ -91,6 +91,39 @@ async def create_partition(self, **kwargs: Any) -> None: self.created.append(kwargs) +class FakeMemberCandidateService: + """Service double for the partition member candidate route.""" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + async def list_member_candidates( + self, + partition: str, + *, + search: str, + cursor: int | None, + limit: int, + ) -> dict[str, Any]: + self.calls.append( + { + "partition": partition, + "search": search, + "cursor": cursor, + "limit": limit, + } + ) + return { + "candidates": [ + {"user_id": 2, "display_name": "Sam", "email": "sam.2@example.com"}, + {"user_id": 3, "display_name": "Sam", "email": "sam.3@example.com"}, + ], + "limit": limit, + "has_more": True, + "next_cursor": 30, + } + + def _build_list_app(*, is_admin: bool) -> FastAPI: """App for the list route, with a regular user whose sole membership is ``all``.""" app = FastAPI() @@ -141,6 +174,72 @@ async def test_list_partitions_does_not_expand_all_for_non_admin(async_client_fa assert [p["partition"] for p in response.json()["partitions"]] == ["all"] +@pytest.mark.asyncio +async def test_list_partition_member_candidates_returns_stable_identities(async_client_factory): + service = FakeMemberCandidateService() + app = _build_app(service) + + async with async_client_factory(app) as client: + response = await client.get( + "/partition/legal/users/candidates", + params={"search": "sam", "cursor": 20, "limit": 10}, + ) + + assert response.status_code == 200 + assert response.json() == { + "candidates": [ + {"user_id": 2, "display_name": "Sam", "email": "sam.2@example.com"}, + {"user_id": 3, "display_name": "Sam", "email": "sam.3@example.com"}, + ], + "limit": 10, + "has_more": True, + "next_cursor": 30, + } + assert service.calls == [ + { + "partition": "legal", + "search": "sam", + "cursor": 20, + "limit": 10, + } + ] + + +@pytest.mark.asyncio +async def test_list_partition_member_candidates_rejects_non_owner(async_client_factory): + service = FakeMemberCandidateService() + app = _build_app(service) + + async def reject_non_owner(): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Partition owner role required", + ) + + app.dependency_overrides[require_partition_owner] = reject_non_owner + + async with async_client_factory(app) as client: + response = await client.get( + "/partition/legal/users/candidates", + params={"search": "sam"}, + ) + + assert response.status_code == 403 + assert service.calls == [] + + +@pytest.mark.asyncio +async def test_list_partition_member_candidates_requires_search(async_client_factory): + service = FakeMemberCandidateService() + app = _build_app(service) + + async with async_client_factory(app) as client: + response = await client.get("/partition/legal/users/candidates") + + assert response.status_code == 422 + assert service.calls == [] + + @pytest.mark.asyncio async def test_update_partition_config_forwards_only_provided_fields(async_client_factory): """Partition config updates should exclude omitted fields.""" 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/api/routers/user/test_source_links.py b/tests/unit/api/routers/user/test_source_links.py index ed20af42f..9127837c1 100644 --- a/tests/unit/api/routers/user/test_source_links.py +++ b/tests/unit/api/routers/user/test_source_links.py @@ -61,3 +61,25 @@ def test_metadata_is_passed_through(): link = _build({"_id": "x", "source": "doc.pdf", "author": "alice"}) assert link["author"] == "alice" assert link["chunk_url"] == "https://host/extract/x" + + +def test_metadata_cannot_override_authoritative_source_fields(): + link = _build( + { + "_id": "42", + "source": "diagram.png", + "file_url": "https://attacker.example/file", + "chunk_url": "https://attacker.example/chunk", + "source_type": "web", + } + ) + + assert link["source_type"] == "document" + assert link["file_url"] == "https://host/static/42" + assert link["chunk_url"] == "https://host/extract/42" + + +def test_metadata_file_url_is_removed_when_source_is_missing(): + link = _build({"_id": "42", "file_url": "https://attacker.example/file"}) + + assert "file_url" not in link diff --git a/tests/unit/api/schemas/admin/test_phase14_schemas.py b/tests/unit/api/schemas/admin/test_phase14_schemas.py index eb5697d8b..f7cd50f90 100644 --- a/tests/unit/api/schemas/admin/test_phase14_schemas.py +++ b/tests/unit/api/schemas/admin/test_phase14_schemas.py @@ -46,6 +46,45 @@ def test_create_model_endpoint_rejects_empty_normalized_endpoint(endpoint): CreateModelEndpointRequest(name="default", model_type="llm", endpoint=endpoint) +_UNSAFE_NAMES = [ + "owner/model", # splits across the {model_type}/{name} route segment (#768) + ".", # RFC 3986 dot-segment: normalizes to the collection route + "..", # RFC 3986 dot-segment: normalizes away the model_type segment too + "-leading-dash", + "trailing-dash-", + ".leading-dot", + "trailing-dot.", + "_leading_underscore", + "trailing_underscore_", + "has space", + "has%percent", + "a" * 129, # over _NAME_MAX_LENGTH +] + + +@pytest.mark.parametrize("bad_name", _UNSAFE_NAMES) +def test_create_model_endpoint_rejects_unsafe_name(bad_name): + """Any name outside the URL-path-segment allowlist is rejected, not just '/'.""" + with pytest.raises(ValidationError): + CreateModelEndpointRequest(name=bad_name, model_type="reranker", endpoint="http://host") + + +@pytest.mark.parametrize("bad_name", _UNSAFE_NAMES) +def test_update_model_endpoint_rejects_unsafe_name(bad_name): + """Same allowlist applies to renames via the update schema — kept in sync + with the create matrix above (a separate, optional-name validator) so a + regression in one can't go uncovered by the other.""" + with pytest.raises(ValidationError): + UpdateModelEndpointRequest(name=bad_name) + + +@pytest.mark.parametrize("good_name", ["default", "gpt-4.1", "jina_v3", "LocalReranker.prod", "a", "a" * 128]) +def test_create_model_endpoint_accepts_realistic_names(good_name): + """Interior '.', '_', '-' stay available for realistic names.""" + request = CreateModelEndpointRequest(name=good_name, model_type="reranker", endpoint="http://host") + assert request.name == good_name + + def test_update_model_endpoint_requires_at_least_one_field(): """Endpoint updates must contain at least one field.""" with pytest.raises(ValidationError): diff --git a/tests/unit/core/utils/test_source_filtering.py b/tests/unit/core/utils/test_source_filtering.py index 629c0c3eb..877587542 100644 --- a/tests/unit/core/utils/test_source_filtering.py +++ b/tests/unit/core/utils/test_source_filtering.py @@ -135,6 +135,24 @@ def test_tag_inline_in_prose_preserved(self): assert clean == text assert citations is None + def test_context_source_markers_are_recovered_and_stripped(self): + text = "The footprint fell by 28% [Source 7].\nThe partners include Flexis [Source 8][Source 9]." + clean, citations = extract_and_strip_sources_block(text) + assert clean == "The footprint fell by 28%.\nThe partners include Flexis." + assert citations == {7, 8, 9} + + def test_unclosed_numbered_source_marker_is_recovered(self): + text = "The target is 2040 [Source 2" + clean, citations = extract_and_strip_sources_block(text) + assert clean == "The target is 2040" + assert citations == {2} + + def test_dangling_source_marker_is_removed_without_a_citation(self): + text = "Logistics emissions fell by 30% [Source" + clean, citations = extract_and_strip_sources_block(text) + assert clean == "Logistics emissions fell by 30%" + assert citations is None + class TestFilterSourcesByCitations: def test_basic_filtering(self): @@ -142,9 +160,14 @@ def test_basic_filtering(self): result = filter_sources_by_citations(sources, {1, 3, 5}) assert result == ["a", "c", "e"] - def test_none_citations_returns_all(self): + def test_none_citations_returns_empty(self): sources = ["a", "b", "c"] result = filter_sources_by_citations(sources, None) + assert result == [] + + def test_none_citations_can_be_allowed_for_structured_output(self): + sources = ["a", "b", "c"] + result = filter_sources_by_citations(sources, None, allow_uncited=True) assert result == ["a", "b", "c"] def test_empty_citations_returns_empty(self): @@ -152,10 +175,10 @@ def test_empty_citations_returns_empty(self): result = filter_sources_by_citations(sources, set()) assert result == [] - def test_out_of_range_citations_fallback(self): + def test_out_of_range_citations_returns_empty(self): sources = ["a", "b", "c"] result = filter_sources_by_citations(sources, {99}) - assert result == ["a", "b", "c"] + assert result == [] def test_partial_out_of_range(self): sources = ["a", "b", "c"] @@ -370,8 +393,8 @@ async def test_case2_llm_says_sources_none(self): assert _parse_finish_sources(result) == [] @pytest.mark.asyncio - async def test_case3_llm_no_tag_fallback_all(self): - """Case 3: LLM omits tag entirely → fallback to all sources.""" + async def test_case3_llm_no_tag_returns_no_sources(self): + """Case 3: LLM omits tag entirely → no source is attributed.""" lines = [ _make_chunk("Answer without any sources tag."), _make_finish(), @@ -379,8 +402,64 @@ async def test_case3_llm_no_tag_fallback_all(self): ] result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) assert _collect_content(result) == "Answer without any sources tag." + assert _parse_finish_sources(result) == [] + + @pytest.mark.asyncio + async def test_no_tag_keeps_sources_when_uncited_output_is_allowed(self): + lines = [ + _make_chunk('{"answer": "structured"}'), + _make_finish(), + DONE_LINE, + ] + result = await _collect( + stream_with_source_filtering( + _fake_stream(lines), + self.SOURCES, + "test-model", + allow_uncited_sources=True, + ) + ) assert _parse_finish_sources(result) == self.SOURCES + @pytest.mark.asyncio + async def test_structured_output_preserves_source_like_json_values(self): + structured = '{"answer":"Use [Source 1]","literal_format":"[Sources: 1]"}' + lines = [ + _make_chunk(structured), + _make_finish(), + DONE_LINE, + ] + result = await _collect( + stream_with_source_filtering( + _fake_stream(lines), + self.SOURCES, + "test-model", + allow_uncited_sources=True, + citation_protocol_active=False, + ) + ) + assert _collect_content(result) == structured + assert _parse_finish_sources(result) == self.SOURCES + + @pytest.mark.asyncio + async def test_direct_output_preserves_terminal_source_marker(self): + answer = "The requested literal notation is:\n[Sources: 1]" + lines = [ + _make_chunk(answer), + _make_finish(), + DONE_LINE, + ] + result = await _collect( + stream_with_source_filtering( + _fake_stream(lines), + [], + "test-model", + citation_protocol_active=False, + ) + ) + assert _collect_content(result) == answer + assert _parse_finish_sources(result) == [] + @pytest.mark.asyncio async def test_multiple_inline_tags_stripped_from_stream(self): """Bullet-leak: LLM emits [Sources: X] per bullet. All inline tags must be stripped.""" @@ -399,6 +478,30 @@ async def test_multiple_inline_tags_stripped_from_stream(self): assert "Claim two about APEX." in content assert _parse_finish_sources(result) == [{"file": "a.pdf"}, {"file": "c.pdf"}] + @pytest.mark.asyncio + async def test_context_source_markers_are_stripped_and_rendered_as_sources(self): + lines = [ + _make_chunk("First claim [Sour"), + _make_chunk("ce 1]. Second claim [Source 2][Source 3]."), + _make_finish(), + DONE_LINE, + ] + result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) + assert _collect_content(result) == "First claim. Second claim." + assert _parse_finish_sources(result) == self.SOURCES + + @pytest.mark.asyncio + async def test_literal_source_marker_is_preserved_without_sources(self): + lines = [ + _make_chunk("The literal notation [Sour"), + _make_chunk("ce 1] identifies the first source."), + _make_finish(), + DONE_LINE, + ] + result = await _collect(stream_with_source_filtering(_fake_stream(lines), [], "test-model")) + assert _collect_content(result) == "The literal notation [Source 1] identifies the first source." + assert _parse_finish_sources(result) == [] + @pytest.mark.asyncio async def test_inline_prose_tag_preserved_in_stream(self): """Meta-discussion: a [Sources: 1, 3] inside a sentence must NOT be stripped.""" @@ -411,8 +514,8 @@ async def test_inline_prose_tag_preserved_in_stream(self): result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) content = _collect_content(result) assert content == "Use the format [Sources: 1, 3] at the very end of your response." - # No line-terminal tag → fallback to all sources - assert _parse_finish_sources(result) == self.SOURCES + # No line-terminal tag means no source was actually cited. + assert _parse_finish_sources(result) == [] @pytest.mark.asyncio async def test_mid_response_tag_stripped_plus_trailing_tag(self): diff --git a/tests/unit/core/utils/test_web_url.py b/tests/unit/core/utils/test_web_url.py new file mode 100644 index 000000000..f87357128 --- /dev/null +++ b/tests/unit/core/utils/test_web_url.py @@ -0,0 +1,21 @@ +import pytest +from core.utils.web_url import normalize_web_url + + +@pytest.mark.parametrize( + "value", + [ + None, + 42, + "", + "javascript:alert(1)", + "https://", + "http://[::1", + ], +) +def test_normalize_web_url_rejects_unrenderable_values(value): + assert normalize_web_url(value) is None + + +def test_normalize_web_url_returns_canonical_http_url(): + assert normalize_web_url(" https://example.com/a path ") == "https://example.com/a%20path" 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/infra/test_admin_ui_compose.py b/tests/unit/infra/test_admin_ui_compose.py index 7e8737b2d..e9e3a0c57 100644 --- a/tests/unit/infra/test_admin_ui_compose.py +++ b/tests/unit/infra/test_admin_ui_compose.py @@ -39,3 +39,41 @@ def test_admin_ui_nginx_preserves_websocket_upgrades_for_chainlit(): assert re.search(r"proxy_http_version\s+1\.1;", config) assert re.search(r"proxy_set_header\s+Upgrade\s+\$http_upgrade;", config) assert re.search(r"proxy_set_header\s+Connection\s+\$connection_upgrade;", config) + + +def test_admin_ui_image_runs_as_uid_owning_its_writable_paths(): + """The runtime paths nginx writes must be group 0, and the image must run as + a numeric UID that has access to them — the arbitrary-UID pattern + api.Dockerfile uses. + + `USER nginx` is uid 101 with only gid 101, i.e. neither owner nor group on + paths chowned to 10001. Nothing writes under /var/cache/nginx today (the + base image points every *_temp_path at /tmp and openrag-admin.conf sets + `proxy_cache off`), so this stays invisible until something does — and it + then fails at request time, which no smoke test covers. + """ + dockerfile = Path(__file__).resolve().parents[3] / "infra/docker/ui.Dockerfile" + content = dockerfile.read_text(encoding="utf-8") + + assert re.search(r"^USER 10001:0$", content, re.MULTILINE) + assert not re.search(r"^USER nginx$", content, re.MULTILINE) + assert re.search(r"chown -R 10001:0 /var/cache/nginx /etc/nginx/conf\.d /var/run", content) + assert re.search(r"chmod -R g\+w /var/cache/nginx /etc/nginx/conf\.d /var/run", content) + # A private group would undo the chown above for any UID but 10001. + assert "10001:10001" not in content + + +def test_admin_ui_chart_security_context_matches_image_ownership(): + """adminUi.podSecurityContext must keep runAsGroup 0 to match the image's + chown -R 10001:0 — same reason openrag.podSecurityContext does. + """ + values_path = Path(__file__).resolve().parents[3] / "infra/charts/openrag-stack/values.yaml" + + with values_path.open(encoding="utf-8") as handle: + values = yaml.safe_load(handle) + + admin_ui_ctx = values["adminUi"]["podSecurityContext"] + + assert admin_ui_ctx["runAsUser"] == 10001 + assert admin_ui_ctx["runAsGroup"] == 0 + assert values["openrag"]["podSecurityContext"]["runAsGroup"] == 0 diff --git a/tests/unit/infra/test_helm_security_hardening.py b/tests/unit/infra/test_helm_security_hardening.py index 84c8f1335..73fc71261 100644 --- a/tests/unit/infra/test_helm_security_hardening.py +++ b/tests/unit/infra/test_helm_security_hardening.py @@ -60,6 +60,14 @@ def test_secret_template_fails_on_required_or_placeholder_secrets() -> None: def test_chart_workloads_apply_restricted_security_contexts() -> None: + """Shared `security` values block (Pod Security Standards "restricted" + baseline) merged into each workload's own podSecurityContext/ + containerSecurityContext via the openrag-stack.mergeSecurityContext + helper — a component only sets the runAsUser/runAsGroup/fsGroup specific + to its own Dockerfile (see values.yaml's comments, e.g. openrag's + OpenShift arbitrary-UID pattern) and can override automountServiceAccountToken + / containerSecurityContext too, but none currently need to. + """ values = _values() security = values["security"] raycluster = _template("raycluster.yaml") @@ -71,12 +79,21 @@ def test_chart_workloads_apply_restricted_security_contexts() -> None: assert security["containerSecurityContext"]["capabilities"]["drop"] == ["ALL"] assert values["vllm"]["servingEngineSpec"]["containerSecurityContext"]["runAsNonRoot"] is True + # Simulate the template helper's `merge (deepCopy component) default` — + # component keys (e.g. runAsUser) win, the rest is inherited from `security`. + for component in ("openrag", "adminUi", "reranker", "ray"): + block = values[component] + effective_pod_ctx = {**security["podSecurityContext"], **block.get("podSecurityContext", {})} + assert effective_pod_ctx["runAsNonRoot"] is True, component + assert effective_pod_ctx["seccompProfile"]["type"] == "RuntimeDefault", component + assert "runAsUser" in effective_pod_ctx, component + templates = _all_templates() assert templates.count("automountServiceAccountToken:") >= 5 assert templates.count(".Values.security.podSecurityContext") >= 5 assert templates.count(".Values.security.containerSecurityContext") >= 7 assert "ghcr.io/linagora/openrag:dev-latest" not in raycluster - assert "image: {{ $.Values.ray.image.repository }}:{{ $.Values.ray.image.tag }}" in raycluster + assert "{{ .Values.ray.image.repository }}:{{ .Values.ray.image.tag }}" in raycluster def test_ray_dashboard_defaults_to_loopback_in_helm_cluster() -> None: @@ -89,14 +106,22 @@ def test_ray_dashboard_defaults_to_loopback_in_helm_cluster() -> None: def test_ingress_is_not_exposed_by_default_and_supports_tls() -> None: + """No standalone templates/ingress.yaml here — openrag.yaml and + raycluster.yaml each render their own optional Ingress (adminUi.ingress + just toggles a path onto openrag's), gated by their own `required` host + guard so an enabled Ingress can never render with a blank/wildcard host. + """ values = _values() - template = _template("ingress.yaml") - - assert values["ingress"]["enabled"] is False - assert values["ingress"]["host"] == "" - assert values["ingress"]["tls"]["enabled"] is False - - assert "required" in template - assert "ingress.host" in template - assert "tls:" in template - assert "secretName:" in template + openrag_template = _template("openrag.yaml") + raycluster_template = _template("raycluster.yaml") + + assert values["openrag"]["ingress"]["enabled"] is False + assert values["openrag"]["ingress"]["host"] == "" + assert values["adminUi"]["ingress"]["enabled"] is False + assert values["ray"]["ingress"]["enabled"] is False + assert values["ray"]["ingress"]["host"] == "" + + for template in (openrag_template, raycluster_template): + assert "required" in template + assert "ingress.host must be set" in template + assert "tls:" in template diff --git a/tests/unit/infra/test_postgres_migration_job_template.py b/tests/unit/infra/test_postgres_migration_job_template.py index 2ea11be8c..75e93a897 100644 --- a/tests/unit/infra/test_postgres_migration_job_template.py +++ b/tests/unit/infra/test_postgres_migration_job_template.py @@ -7,13 +7,18 @@ def test_postgres_migration_job_uses_secret_ref_instead_of_literal_secret_values() -> None: + """This chart names its ConfigMap/Secret via the openrag-stack.fullname / + secretName helpers (release-scoped, e.g. for ArgoCD) rather than the fixed + "rag-env"/"rag-env-secrets" names — assert it references those same + helpers instead of inlining literal secret values. + """ template = MIGRATION_JOB_TEMPLATE.read_text(encoding="utf-8") assert "envFrom:" in template assert "configMapRef:" in template - assert "name: rag-env" in template + assert 'name: {{ include "openrag-stack.fullname" . }}-env' in template assert "secretRef:" in template - assert "name: rag-env-secrets" in template + assert 'name: {{ include "openrag-stack.secretName" . }}' in template assert ".Values.env.secrets" not in template assert 'value: "{{ $value }}"' not in template @@ -24,6 +29,22 @@ def test_postgres_migration_job_sets_uv_cache_dir() -> None: assert "name: UV_CACHE_DIR" in template +def test_postgres_migration_job_reuses_the_openrag_service_account() -> None: + """The Job runs the OpenRAG image under the same pinned UID as its + Deployment, so it needs the same ServiceAccount to reach the same SCC. + + Without this the Job silently falls back to the namespace's `default` SA. + On OpenShift that means `restricted-v2` (MustRunAsRange), which rejects the + requested runAsUser as outside the namespace's assigned range — and since + this is a pre-install/pre-upgrade hook, the failure aborts the release. + """ + template = MIGRATION_JOB_TEMPLATE.read_text(encoding="utf-8") + + assert "{{- with .Values.openrag.serviceAccountName }}" in template + # tpl(), so a value like "{{ .Release.Name }}-openrag" resolves. + assert "serviceAccountName: {{ tpl . $ }}" in template + + def test_postgres_migration_job_omits_flags_the_runner_ignores() -> None: """The migration runner never reads these, so they must not be set here. 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_model_endpoint_service.py b/tests/unit/services/orchestrators/test_model_endpoint_service.py index a06590457..55ba765b9 100644 --- a/tests/unit/services/orchestrators/test_model_endpoint_service.py +++ b/tests/unit/services/orchestrators/test_model_endpoint_service.py @@ -103,13 +103,15 @@ async def delete_and_promote_default(self, name: str, model_type: str) -> tuple[ return ("ok", promoted) -def _make_service(repo=None, rows=None, settings=None): +def _make_service(repo=None, rows=None, settings=None, partition_service=None, preset_service=None): from core.config.root import Settings from services.orchestrators.model_endpoint_service import ModelEndpointService return ModelEndpointService( model_endpoint_repo=repo or _FakeEndpointRepo(rows), config=settings or Settings(), + partition_service=partition_service, + preset_service=preset_service, ) @@ -913,6 +915,159 @@ async def test_update_model_endpoint_renames_and_evicts_cache(): assert ("new-name", "embedder") in repo._store +class _FakePresetServiceForReload: + def __init__(self): + self.load_all_calls = 0 + + async def load_all(self): + self.load_all_calls += 1 + + +class _FakePartitionServiceForReload: + def __init__(self): + self.load_partitions_calls = 0 + + async def load_partitions(self): + self.load_partitions_calls += 1 + + +@pytest.mark.asyncio +async def test_update_model_endpoint_rename_reloads_presets_then_partitions(): + """A rename cascades DB-side references (#770) inside the repo's own rename + transaction (see PgModelEndpointRepository.rename), but those writes are + invisible until PresetService / PartitionService reload their in-memory + caches — pin that both get refreshed on a rename.""" + existing = _make_row(name="old-name", model_type="llm") + repo = _FakeEndpointRepo(rows=[existing]) + preset_service = _FakePresetServiceForReload() + partition_service = _FakePartitionServiceForReload() + svc = _make_service(repo, partition_service=partition_service, preset_service=preset_service) + + await svc.update_model_endpoint("old-name", "llm", new_name="new-name") + + assert preset_service.load_all_calls == 1 + assert partition_service.load_partitions_calls == 1 + + +@pytest.mark.asyncio +async def test_update_model_endpoint_without_rename_skips_preset_and_partition_reload(): + """A plain field update (no rename) touches no cross-referenced name, so + it must not pay for a presets/partitions reload it doesn't need.""" + existing = _make_row(name="jina") + repo = _FakeEndpointRepo(rows=[existing]) + preset_service = _FakePresetServiceForReload() + partition_service = _FakePartitionServiceForReload() + svc = _make_service(repo, partition_service=partition_service, preset_service=preset_service) + + await svc.update_model_endpoint("jina", "embedder", endpoint="http://new:8000/v1") + + assert preset_service.load_all_calls == 0 + assert partition_service.load_partitions_calls == 0 + + +@pytest.mark.asyncio +async def test_update_model_endpoint_rename_aliases_new_name_before_reload_awaits(): + """A request racing the rename must resolve `new_name` even before the + presets/partitions reload below completes — the DB cascade (#770) has + already repointed partitions/presets at it by the time the rename + `await` returns, so the registry can't lag behind until `load_all()`.""" + existing = _make_row(name="old-name", model_type="llm", endpoint="http://old:8000/v1") + repo = _FakeEndpointRepo(rows=[existing]) + svc = _make_service(repo) + await svc.load_all() + + seen_during_reload = {} + + class _SnoopingPresetService: + async def load_all(self): + seen_during_reload["new-name"] = svc._config.models.llm.get("new-name") + seen_during_reload["old-name"] = svc._config.models.llm.get("old-name") + + svc._preset_service = _SnoopingPresetService() + + await svc.update_model_endpoint("old-name", "llm", new_name="new-name") + + assert seen_during_reload["new-name"].endpoint == "http://old:8000/v1" + assert seen_during_reload["old-name"].endpoint == "http://old:8000/v1" + + +@pytest.mark.asyncio +async def test_update_model_endpoint_rename_keeps_both_names_resolvable_after_failed_reload(): + """If PresetService.load_all() (or PartitionService.load_partitions()) + raises mid-rename, the DB rename has already committed — the registry + must still resolve both the old and the new name afterward, instead of + being stuck answering only to the pre-rename one until process restart.""" + existing = _make_row(name="old-name", model_type="llm", endpoint="http://old:8000/v1") + repo = _FakeEndpointRepo(rows=[existing]) + svc = _make_service(repo) + await svc.load_all() + + class _FailingPresetService: + async def load_all(self): + raise RuntimeError("db blip") + + svc._preset_service = _FailingPresetService() + + with pytest.raises(RuntimeError): + await svc.update_model_endpoint("old-name", "llm", new_name="new-name") + + assert svc._config.models.llm.get("old-name").endpoint == "http://old:8000/v1" + assert svc._config.models.llm.get("new-name").endpoint == "http://old:8000/v1" + + +@pytest.mark.asyncio +async def test_update_model_endpoint_rename_with_field_change_aliases_the_new_values(): + """A rename combined with a field change (e.g. a new endpoint URL) must + alias `new_name`/`old_name` to the row this call just wrote, not to + whatever the in-memory bucket held before the update ran — otherwise a + reload failure right after would leave the registry silently serving the + stale pre-update config under the DB-authoritative new name forever.""" + existing = _make_row(name="old-name", model_type="llm", endpoint="http://old:8000/v1") + repo = _FakeEndpointRepo(rows=[existing]) + svc = _make_service(repo) + await svc.load_all() + + class _FailingPresetService: + async def load_all(self): + raise RuntimeError("db blip") + + svc._preset_service = _FailingPresetService() + + with pytest.raises(RuntimeError): + await svc.update_model_endpoint("old-name", "llm", new_name="new-name", endpoint="http://new:8000/v1") + + assert svc._config.models.llm.get("new-name").endpoint == "http://new:8000/v1" + assert svc._config.models.llm.get("old-name").endpoint == "http://new:8000/v1" + + +@pytest.mark.asyncio +async def test_update_model_endpoint_rename_evicts_stale_client_cache_before_failed_reload(): + """A cached *client instance* under old_name predates this call and can't + know about a field change baked into the same rename — the factory checks + its cache before the config registry, so it must be evicted eagerly (not + only in the post-reload cleanup a failing reload would skip), or a + request through old_name keeps getting the stale pre-update client.""" + existing = _make_row(name="old-name", model_type="llm", endpoint="http://old:8000/v1") + repo = _FakeEndpointRepo(rows=[existing]) + svc = _make_service(repo) + await svc.load_all() + stale_client = object() + cache: dict = {"old-name": stale_client} + svc._client_caches["llm"] = cache + + class _FailingPresetService: + async def load_all(self): + raise RuntimeError("db blip") + + svc._preset_service = _FailingPresetService() + + with pytest.raises(RuntimeError): + await svc.update_model_endpoint("old-name", "llm", new_name="new-name", endpoint="http://new:8000/v1") + + assert "old-name" not in cache + assert "new-name" not in cache + + @pytest.mark.asyncio async def test_update_default_endpoint_evicts_default_alias_cache(): # Updating the current default must evict the 'default' cache key too, not just diff --git a/tests/unit/services/orchestrators/test_partition_service.py b/tests/unit/services/orchestrators/test_partition_service.py index ca2497aeb..d441659cf 100644 --- a/tests/unit/services/orchestrators/test_partition_service.py +++ b/tests/unit/services/orchestrators/test_partition_service.py @@ -7,7 +7,13 @@ from types import SimpleNamespace import pytest -from core.utils.exceptions import NotFoundError, PartitionNotFoundError, UserNotFoundError, ValidationError +from core.utils.exceptions import ( + ConflictError, + NotFoundError, + PartitionNotFoundError, + UserNotFoundError, + ValidationError, +) from services.orchestrators.partition_service import PartitionService @@ -129,12 +135,17 @@ def __init__( self, members: set[tuple[int, str]] | None = None, owned: dict[int, int] | None = None, + candidate_rows: list[dict] | None = None, + add_result: bool = True, ): self._members = members or set() self._owned = owned or {} # user_id -> number of partitions owned + self._candidate_rows = candidate_rows or [] + self._add_result = add_result self.added: list[tuple[str, int, str]] = [] self.removed: list[tuple[str, int]] = [] self.role_updates: list[tuple[str, int, str]] = [] + self.candidate_calls: list[dict] = [] async def list_user_partitions(self, user_id: int): from core.models.user import PartitionRole, UserPartition @@ -150,9 +161,39 @@ async def user_is_partition_member(self, user_id: int, partition: str) -> bool: async def list_partition_members(self, partition: str) -> list[dict]: return [{"user_id": u, "role": "viewer"} for (u, p) in self._members if p == partition] + async def list_partition_member_candidates( + self, + partition: str, + *, + search_prefix: str | None, + search_user_id: int | None, + after_id: int | None, + limit: int, + ) -> list[dict]: + self.candidate_calls.append( + { + "partition": partition, + "search_prefix": search_prefix, + "search_user_id": search_user_id, + "after_id": after_id, + "limit": limit, + } + ) + rows = [ + row + for row in self._candidate_rows + if (row["user_id"], partition) not in self._members + and ( + row["user_id"] == search_user_id + or (search_prefix is not None and (row["display_name"] or "").lower().startswith(search_prefix.lower())) + ) + and (after_id is None or row["user_id"] > after_id) + ] + return rows[:limit] + async def add_partition_member(self, partition: str, user_id: int, role: str) -> bool: self.added.append((partition, user_id, role)) - return True + return self._add_result async def remove_partition_member(self, partition: str, user_id: int) -> bool: self.removed.append((partition, user_id)) @@ -212,12 +253,41 @@ async def query_chunks_by_filter(self, collection, filters, output_fields=None): class FakeUserRepo: - def __init__(self, existing: set[int] | None = None): + def __init__( + self, + existing: set[int] | None = None, + display_names: dict[int, str] | None = None, + emails: dict[int, str] | None = None, + ): self._existing = existing if existing is not None else set() + self._display_names = display_names or {} + self._emails = emails or {} + self.requested_user_id_batches: list[list[int]] = [] async def user_exists(self, user_id: int) -> bool: return user_id in self._existing + async def get_user(self, user_id: int): + if user_id not in self._existing: + return None + return SimpleNamespace( + id=user_id, + display_name=self._display_names.get(user_id), + email=self._emails.get(user_id), + ) + + async def get_users_by_ids(self, user_ids: list[int]): + self.requested_user_id_batches.append(list(user_ids)) + return [ + SimpleNamespace( + id=user_id, + display_name=self._display_names.get(user_id), + email=self._emails.get(user_id), + ) + for user_id in user_ids + if user_id in self._existing + ] + def _svc( *, @@ -830,6 +900,191 @@ async def test_list_members_missing_partition_404(): await _svc(prepo=FakePartitionRepo(set())).list_members("x") +@pytest.mark.asyncio +async def test_list_member_candidates_excludes_existing_members(): + candidate_rows = [ + {"user_id": 1, "display_name": "Partition owner"}, + {"user_id": 2, "display_name": "Sam"}, + {"user_id": 3, "display_name": "Sam"}, + ] + mrepo = FakeMembershipRepo( + {(1, "p"), (3, "p")}, + candidate_rows=candidate_rows, + ) + svc = _svc( + prepo=FakePartitionRepo({"p"}), + mrepo=mrepo, + ) + + assert await svc.list_member_candidates("p", search=" sam ") == { + "candidates": [{"user_id": 2, "display_name": "Sam"}], + "limit": 25, + "has_more": False, + "next_cursor": None, + } + assert mrepo.candidate_calls == [ + { + "partition": "p", + "search_prefix": "sam", + "search_user_id": None, + "after_id": None, + "limit": 26, + } + ] + + +@pytest.mark.asyncio +async def test_list_member_candidates_returns_bounded_page_with_continuation(): + candidate_rows = [{"user_id": user_id, "display_name": f"User {user_id}"} for user_id in range(1, 102)] + mrepo = FakeMembershipRepo(candidate_rows=candidate_rows) + svc = _svc( + prepo=FakePartitionRepo({"p"}), + mrepo=mrepo, + ) + + page = await svc.list_member_candidates("p", search="User", cursor=10, limit=20) + + assert [candidate["user_id"] for candidate in page["candidates"]] == list(range(11, 31)) + assert page == { + "candidates": candidate_rows[10:30], + "limit": 20, + "has_more": True, + "next_cursor": 30, + } + assert mrepo.candidate_calls == [ + { + "partition": "p", + "search_prefix": "User", + "search_user_id": None, + "after_id": 10, + "limit": 21, + } + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("search", [None, "", " ", "sa"]) +async def test_list_member_candidates_requires_a_targeted_search(search): + svc = _svc(prepo=FakePartitionRepo({"p"})) + + with pytest.raises(ValidationError): + await svc.list_member_candidates("p", search=search) + + +@pytest.mark.asyncio +async def test_list_member_candidates_searches_numeric_id_and_name_prefix(): + candidate = {"user_id": 42, "display_name": "Unrelated name"} + mrepo = FakeMembershipRepo(candidate_rows=[candidate]) + svc = _svc(prepo=FakePartitionRepo({"p"}), mrepo=mrepo) + + assert await svc.list_member_candidates("p", search="0042") == { + "candidates": [candidate], + "limit": 25, + "has_more": False, + "next_cursor": None, + } + assert mrepo.candidate_calls[0]["search_user_id"] == 42 + assert mrepo.candidate_calls[0]["search_prefix"] == "0042" + + +@pytest.mark.asyncio +async def test_list_member_candidates_keeps_short_numeric_search_id_only(): + mrepo = FakeMembershipRepo() + svc = _svc(prepo=FakePartitionRepo({"p"}), mrepo=mrepo) + + await svc.list_member_candidates("p", search="42") + + assert mrepo.candidate_calls[0]["search_user_id"] == 42 + assert mrepo.candidate_calls[0]["search_prefix"] is None + + +@pytest.mark.asyncio +async def test_list_member_candidates_uses_name_prefix_for_numeric_values_outside_id_range(): + mrepo = FakeMembershipRepo() + svc = _svc(prepo=FakePartitionRepo({"p"}), mrepo=mrepo) + + await svc.list_member_candidates("p", search="2147483648") + + assert mrepo.candidate_calls[0]["search_user_id"] is None + assert mrepo.candidate_calls[0]["search_prefix"] == "2147483648" + + +@pytest.mark.asyncio +async def test_list_member_candidates_treats_non_ascii_digits_as_a_name_prefix(): + mrepo = FakeMembershipRepo() + svc = _svc(prepo=FakePartitionRepo({"p"}), mrepo=mrepo) + + await svc.list_member_candidates("p", search="١٢٣") + + assert mrepo.candidate_calls[0]["search_user_id"] is None + assert mrepo.candidate_calls[0]["search_prefix"] == "١٢٣" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("search", "cursor"), + [ + ("User", 2_147_483_648), + ("User", -1), + ], +) +async def test_list_member_candidates_rejects_ids_outside_postgres_range(search, cursor): + svc = _svc(prepo=FakePartitionRepo({"p"})) + + with pytest.raises(ValidationError): + await svc.list_member_candidates("p", search=search, cursor=cursor) + + +@pytest.mark.asyncio +async def test_list_members_does_not_lookup_user_identities(): + mrepo = FakeMembershipRepo(members={(9, "p")}) + urepo = FakeUserRepo({9}) + svc = _svc(prepo=FakePartitionRepo({"p"}), mrepo=mrepo, urepo=urepo) + + members = await svc.list_members("p") + + assert members == [{"user_id": 9, "role": "viewer"}] + assert urepo.requested_user_id_batches == [] + + +@pytest.mark.asyncio +async def test_list_members_with_identities_uses_one_lookup(): + mrepo = FakeMembershipRepo(members={(9, "p"), (10, "p")}) + urepo = FakeUserRepo( + {9, 10}, + display_names={9: "Alice", 10: "Bob"}, + emails={9: "alice@example.com", 10: "bob@example.com"}, + ) + svc = _svc(prepo=FakePartitionRepo({"p"}), mrepo=mrepo, urepo=urepo) + members = await svc.list_members_with_identities("p") + assert {member["user_id"]: member for member in members} == { + 9: { + "user_id": 9, + "role": "viewer", + "display_name": "Alice", + "email": "alice@example.com", + }, + 10: { + "user_id": 10, + "role": "viewer", + "display_name": "Bob", + "email": "bob@example.com", + }, + } + assert len(urepo.requested_user_id_batches) == 1 + assert set(urepo.requested_user_id_batches[0]) == {9, 10} + + +@pytest.mark.asyncio +async def test_list_members_missing_user_display_name_is_none(): + mrepo = FakeMembershipRepo(members={(9, "p")}) + urepo = FakeUserRepo(set()) # user_id 9 no longer exists + svc = _svc(prepo=FakePartitionRepo({"p"}), mrepo=mrepo, urepo=urepo) + members = await svc.list_members_with_identities("p") + assert members[0]["display_name"] is None + assert members[0]["email"] is None + + @pytest.mark.asyncio async def test_add_member_checks_partition_and_user(): mrepo = FakeMembershipRepo() @@ -842,6 +1097,22 @@ async def test_add_member_checks_partition_and_user(): assert mrepo.added == [("p", 9, "editor")] +@pytest.mark.asyncio +async def test_add_member_rejects_existing_membership_without_changing_role(): + mrepo = FakeMembershipRepo(add_result=False) + svc = _svc( + prepo=FakePartitionRepo({"p"}), + mrepo=mrepo, + urepo=FakeUserRepo({9}), + ) + + with pytest.raises(ConflictError) as error: + await svc.add_member("p", 9, "editor") + + assert error.value.code == "PARTITION_MEMBER_EXISTS" + assert "PATCH /partition/p/users/9" in error.value.message + + @pytest.mark.asyncio async def test_add_member_unknown_user_404(): svc = _svc(prepo=FakePartitionRepo({"p"}), urepo=FakeUserRepo(set())) 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 62fa1b27d..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 @@ -42,6 +59,7 @@ def __init__(self, *, chat_responses=None, gen_text="answer", stream_lines=None) self._gen_text = gen_text self._stream_lines = stream_lines or ['data: {"choices":[{"delta":{"content":"hi"}}]}\n\n', "data: [DONE]\n\n"] self.chat_calls: list = [] + self.generate_calls: list = [] async def chat(self, messages, **kwargs): self.chat_calls.append((messages, kwargs)) @@ -52,6 +70,7 @@ async def chat(self, messages, **kwargs): return {"choices": [{"message": {"content": content}}]} async def generate(self, prompt, **kwargs): + self.generate_calls.append((prompt, kwargs)) return {"choices": [{"text": self._gen_text}]} async def stream_chat(self, messages, **kwargs): @@ -83,8 +102,10 @@ class FakeWeb: def __init__(self, results=None): self._results = results or [] + self.calls: list[str] = [] async def search(self, query): + self.calls.append(query) return list(self._results) @@ -119,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, ) @@ -182,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 @@ -346,6 +369,16 @@ async def test_generate_query_chatbotrag_parses_json(): svc = _svc(llm=FakeLLM(chat_responses=[payload]), mode="ChatBotRag") sq = await svc.generate_query([{"role": "user", "content": "hi"}]) assert sq.query_list[0].query == "rewritten" + assert sq.requires_retrieval is True + + +@pytest.mark.asyncio +async def test_generate_query_chatbotrag_can_skip_retrieval(): + payload = json.dumps({"requires_retrieval": False, "query_list": []}) + svc = _svc(llm=FakeLLM(chat_responses=[payload]), mode="ChatBotRag") + sq = await svc.generate_query([{"role": "user", "content": "How can you help me?"}]) + assert sq.requires_retrieval is False + assert sq.query_list == [] @pytest.mark.asyncio @@ -379,8 +412,40 @@ async def _spy(**kwargs): ) assert called["n"] == 0 # no retrieval in direct mode assert out["model"] == "m1" - assert out["choices"][0]["message"]["content"] == "hello" # sources tag stripped - assert json.loads(out["extra"])["sources"] == [] # [Sources: none] → no sources + assert out["choices"][0]["message"]["content"] == "hello [Sources: none]" + assert json.loads(out["extra"])["sources"] == [] + + +@pytest.mark.asyncio +async def test_chat_direct_mode_preserves_literal_source_marker(): + answer = "The literal notation [Source 1] identifies the first source." + svc = _svc(llm=FakeLLM(chat_responses=[answer])) + + out = await svc.chat( + partitions=None, + payload={"messages": [{"role": "user", "content": "Explain [Source 1]"}], "metadata": {}}, + prepare_sources=lambda d, w: [], + model_name="m1", + ) + + assert out["choices"][0]["message"]["content"] == answer + assert json.loads(out["extra"])["sources"] == [] + + +@pytest.mark.asyncio +async def test_chat_direct_mode_preserves_literal_terminal_sources_marker(): + answer = "The requested literal notation is:\n[Sources: 1]" + svc = _svc(llm=FakeLLM(chat_responses=[answer])) + + out = await svc.chat( + partitions=None, + payload={"messages": [{"role": "user", "content": "Repeat [Sources: 1]"}], "metadata": {}}, + prepare_sources=lambda d, w: [], + model_name="m1", + ) + + assert out["choices"][0]["message"]["content"] == answer + assert json.loads(out["extra"])["sources"] == [] @pytest.mark.asyncio @@ -397,6 +462,268 @@ async def test_chat_with_partition_retrieves_and_filters_sources(): assert filtered == [{"source_type": "document", "n": 1}] # only cited source 1 +@pytest.mark.asyncio +async def test_chat_recovers_context_markers_as_citations(): + svc = _svc(llm=FakeLLM(chat_responses=["First claim [Source 2]. Second claim [Source 1][Source 2]."])) + sources = [{"source_type": "document", "n": 1}, {"source_type": "document", "n": 2}] + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "q"}], "metadata": {}}, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + + assert out["choices"][0]["message"]["content"] == "First claim. Second claim." + assert json.loads(out["extra"])["sources"] == sources + + +@pytest.mark.asyncio +async def test_chat_conversational_request_skips_partition_retrieval(): + query_json = json.dumps({"requires_retrieval": False, "query_list": []}) + llm = FakeLLM(chat_responses=[query_json, "I can help you search and summarize documents."]) + retrieval = FakeRetrieval() + svc = _svc(mode="ChatBotRag", llm=llm, retrieval=retrieval) + + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "How can you help me?"}], "metadata": {}}, + prepare_sources=lambda d, w: [{"source_type": "document"}], + model_name="m", + ) + + assert retrieval.retrieve_multi_calls == [] + assert out["choices"][0]["message"]["content"] == "I can help you search and summarize documents." + assert json.loads(out["extra"])["sources"] == [] + answer_messages = llm.chat_calls[1][0] + assert answer_messages[0]["role"] == "system" + assert "OpenRAG" in answer_messages[0]["content"] + assert "LINAGORA" in answer_messages[0]["content"] + assert "document-grounded RAG system" in answer_messages[0]["content"] + + +@pytest.mark.asyncio +async def test_chat_conversational_request_keeps_spoken_style_prompt(): + query_json = json.dumps({"requires_retrieval": False, "query_list": []}) + llm = FakeLLM(chat_responses=[query_json, "I'm OpenRAG, built by LINAGORA."]) + svc = _svc(mode="ChatBotRag", llm=llm) + + await svc.chat( + partitions=["p"], + payload={ + "messages": [{"role": "user", "content": "Who are you?"}], + "metadata": {"spoken_style_answer": True}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + + answer_system_prompt = llm.chat_calls[1][0][0]["content"] + assert "OpenRAG" in answer_system_prompt + assert "LINAGORA" in answer_system_prompt + assert "short (1-2 sentences)" in answer_system_prompt + + +@pytest.mark.asyncio +async def test_chat_mixed_request_still_retrieves_documents(): + query_json = json.dumps( + { + "requires_retrieval": True, + "query_list": [{"query": "Product A revenue in Q1", "temporal_filters": None}], + } + ) + llm = FakeLLM(chat_responses=[query_json, "Revenue was 10 million. [Sources: 1]"]) + retrieval = FakeRetrieval() + svc = _svc(mode="ChatBotRag", llm=llm, retrieval=retrieval) + + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "Hello, what was Product A revenue in Q1?"}]}, + prepare_sources=lambda d, w: [{"source_type": "document", "filename": "report.pdf"}], + model_name="m", + ) + + assert len(retrieval.retrieve_multi_calls) == 1 + assert json.loads(out["extra"])["sources"] == [{"source_type": "document", "filename": "report.pdf"}] + + +@pytest.mark.asyncio +async def test_chat_inconsistent_classifier_result_prefers_supplied_query(): + query_json = json.dumps( + { + "requires_retrieval": False, + "query_list": [{"query": "Product A revenue", "temporal_filters": None}], + } + ) + llm = FakeLLM(chat_responses=[query_json, "Revenue was 10 million. [Sources: 1]"]) + retrieval = FakeRetrieval() + svc = _svc(mode="ChatBotRag", llm=llm, retrieval=retrieval) + + await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "What was Product A revenue?"}]}, + prepare_sources=lambda d, w: [{"source_type": "document"}], + model_name="m", + ) + + assert len(retrieval.retrieve_multi_calls) == 1 + + +@pytest.mark.asyncio +async def test_chat_without_citation_does_not_attribute_retrieved_sources(): + svc = _svc(llm=FakeLLM(chat_responses=["A general answer with no citation marker."])) + sources = [{"source_type": "document", "filename": "unrelated.pdf"}] + + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "How can you help me?"}], "metadata": {}}, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + + assert json.loads(out["extra"])["sources"] == [] + + +@pytest.mark.asyncio +async def test_chat_invalid_citation_does_not_fallback_to_unrelated_sources(): + svc = _svc(llm=FakeLLM(chat_responses=["Answer. [Sources: 99]"])) + sources = [{"source_type": "document", "filename": "unrelated.pdf"}] + + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "Question"}], "metadata": {}}, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + + assert json.loads(out["extra"])["sources"] == [] + + +@pytest.mark.asyncio +async def test_chat_structured_output_keeps_retrieved_sources_without_citation_marker(): + structured_answer = '{"answer": "Use [Source 1]", "literal_format": "[Sources: 1]"}' + svc = _svc(llm=FakeLLM(chat_responses=[structured_answer])) + sources = [{"source_type": "document", "filename": "report.pdf"}] + + out = await svc.chat( + partitions=["p"], + payload={ + "messages": [{"role": "user", "content": "Question"}], + "metadata": {}, + "response_format": {"type": "json_object"}, + }, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + + assert out["choices"][0]["message"]["content"] == structured_answer + assert json.loads(out["extra"])["sources"] == sources + + +@pytest.mark.asyncio +async def test_chat_stream_structured_output_preserves_source_like_json_values(): + structured_answer = '{"answer":"Use [Source 1]","literal_format":"[Sources: 1]"}' + stream_lines = [ + "data: " + + json.dumps( + { + "choices": [ + { + "delta": {"content": structured_answer}, + "finish_reason": None, + } + ] + } + ) + + "\n\n", + 'data: {"choices":[{"delta":{},"finish_reason":"stop"}]}\n\n', + "data: [DONE]\n\n", + ] + svc = _svc(llm=FakeLLM(stream_lines=stream_lines)) + sources = [{"source_type": "document", "filename": "report.pdf"}] + + lines = [ + line + async for line in svc.chat_stream( + partitions=["p"], + payload={ + "messages": [{"role": "user", "content": "Question"}], + "metadata": {}, + "response_format": {"type": "json_object"}, + }, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + ] + chunks = [ + json.loads(line[len("data: ") :]) + for line in lines + if line.startswith("data: ") and line.strip() != "data: [DONE]" + ] + content = "".join( + choice.get("delta", {}).get("content", "") for chunk in chunks for choice in chunk.get("choices", []) + ) + extra = next(json.loads(chunk["extra"]) for chunk in reversed(chunks) if chunk.get("extra") not in (None, "{}")) + + assert content == structured_answer + assert extra["sources"] == sources + + +@pytest.mark.asyncio +async def test_structured_websearch_returns_only_sources_included_in_context(): + first = SimpleNamespace( + url="https://example.test/included", + title="Included", + content="short evidence", + snippet="", + ) + excluded = SimpleNamespace( + url="https://example.test/excluded", + title="Excluded", + content="long evidence that does not fit", + snippet="", + ) + web = FakeWeb(results=[first, excluded]) + web.max_tokens = qs.get_num_tokens()("[Source 1]\nIncluded\nshort evidence") + svc = _svc( + llm=FakeLLM(chat_responses=['{"answer": "structured"}']), + retrieval=FakeRetrieval(chunks=[]), + web=web, + ) + + out = await svc.chat( + partitions=None, + payload={ + "messages": [{"role": "user", "content": "Question"}], + "metadata": {"websearch": True}, + "response_format": {"type": "json_object"}, + }, + prepare_sources=lambda _docs, results: [{"url": result.url} for result in results], + model_name="m", + ) + + assert json.loads(out["extra"])["sources"] == [{"url": "https://example.test/included"}] + + +@pytest.mark.asyncio +async def test_explicit_websearch_forces_retrieval_for_conversational_classifier_result(): + query_json = json.dumps({"requires_retrieval": False, "query_list": []}) + llm = FakeLLM(chat_responses=[query_json]) + retrieval = FakeRetrieval(chunks=[]) + web = FakeWeb() + svc = _svc(mode="ChatBotRag", llm=llm, retrieval=retrieval, web=web) + + await svc._prepare_chat( + ["p"], + { + "messages": [{"role": "user", "content": "What is happening today?"}], + "metadata": {"websearch": True}, + }, + ) + + assert len(retrieval.retrieve_multi_calls) == 1 + assert web.calls == ["What is happening today?"] + + @pytest.mark.asyncio async def test_websearch_with_partition_fuses_docs_via_retrieve_multi(): # #707/#740: with a partition AND websearch enabled, the document branch must @@ -406,7 +733,7 @@ async def test_websearch_with_partition_fuses_docs_via_retrieve_multi(): retrieval = FakeRetrieval() web_result = SimpleNamespace(url="https://ex.com", title="T", content="web body", snippet="") svc = _svc(retrieval=retrieval, web=FakeWeb(results=[web_result])) - _payload, _docs, web = await svc._prepare_chat( + _payload, _docs, web, _citation_protocol_active = await svc._prepare_chat( ["p"], {"messages": [{"role": "user", "content": "q"}], "metadata": {"websearch": True}} ) assert len(retrieval.retrieve_multi_calls) == 1 # doc branch fused via the rrf_k-aware retrieve_multi @@ -414,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"]) @@ -513,17 +960,59 @@ async def test_chat_without_workspace_unaffected(): @pytest.mark.asyncio -async def test_complete_strips_and_filters(): - svc = _svc(llm=FakeLLM(gen_text="text body [Sources: none]")) +async def test_complete_direct_mode_preserves_literal_source_marker(): + answer = "text body\n[Sources: none]" + svc = _svc(llm=FakeLLM(gen_text=answer)) out = await svc.complete( partitions=None, payload={"prompt": "do x"}, prepare_sources=lambda d, w: [{"x": 1}], ) - assert out["choices"][0]["text"] == "text body" + assert out["choices"][0]["text"] == answer assert json.loads(out["extra"])["sources"] == [] +@pytest.mark.asyncio +async def test_complete_conversational_request_uses_openrag_prompt_without_retrieval(): + query_json = json.dumps({"requires_retrieval": False, "query_list": []}) + llm = FakeLLM(chat_responses=[query_json], gen_text="I am OpenRAG.\n[Sources: none]") + retrieval = FakeRetrieval() + svc = _svc(mode="ChatBotRag", llm=llm, retrieval=retrieval) + + out = await svc.complete( + partitions=["p"], + payload={"prompt": "Who are you?"}, + prepare_sources=lambda d, w: [], + ) + + assert retrieval.retrieve_multi_calls == [] + assert out["choices"][0]["text"] == "I am OpenRAG." + answer_prompt = llm.generate_calls[0][0] + assert "OpenRAG" in answer_prompt + assert "LINAGORA" in answer_prompt + assert "document-grounded RAG system" in answer_prompt + assert "Who are you?" in answer_prompt + + +@pytest.mark.asyncio +async def test_complete_partition_request_keeps_context_and_filters_citations(): + llm = FakeLLM(gen_text="The answer is grounded.\n[Sources: 1]") + svc = _svc(llm=llm) + sources = [{"source_type": "document", "filename": "report.pdf"}] + + out = await svc.complete( + partitions=["p"], + payload={"prompt": "What does the report say?"}, + prepare_sources=lambda d, w: sources, + ) + + assert out["choices"][0]["text"] == "The answer is grounded." + assert json.loads(out["extra"])["sources"] == sources + answer_prompt = llm.generate_calls[0][0] + assert "ctx" in answer_prompt + assert "What does the report say?" in answer_prompt + + @pytest.mark.asyncio async def test_chat_stream_yields_sse_and_done(): svc = _svc(llm=FakeLLM()) @@ -567,12 +1056,19 @@ def test_json_slice_extracts_object(): def test_dedupe_web_preserves_first_seen(): - a = SimpleNamespace(url="u1") - b = SimpleNamespace(url="u1") - c = SimpleNamespace(url="u2") + a = SimpleNamespace(url="https://example.test/one") + b = SimpleNamespace(url="https://example.test/one") + c = SimpleNamespace(url="https://example.test/two") assert qs._dedupe_web([[a, b], [c]]) == [a, c] +def test_dedupe_web_drops_invalid_urls_before_source_numbering(): + invalid = SimpleNamespace(url="javascript:alert(1)") + valid = SimpleNamespace(url="https://example.test/evidence") + + assert qs._dedupe_web([[invalid, valid]]) == [valid] + + def test_sampling_strips_transport_keys(): out = qs._sampling({"messages": [], "stream": True, "model": "m", "temperature": 0.5}) assert out == {"temperature": 0.5} @@ -700,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 217f737b4..c73a3e502 100644 --- a/tests/unit/services/orchestrators/test_retrieval_service.py +++ b/tests/unit/services/orchestrators/test_retrieval_service.py @@ -65,6 +65,7 @@ def _config(rtype: str = "single", reranker_enabled: bool = False) -> SimpleName ), reranker=SimpleNamespace(enabled=reranker_enabled, top_k=5), partitions={}, + models=SimpleNamespace(reranker={}), ) @@ -249,6 +250,68 @@ async def test_retrieve_uses_partition_retrieval_config_and_named_reranker(): assert call["similarity_threshold"] == 0.77 +@pytest.mark.asyncio +async def test_retrieve_falls_back_to_default_reranker_when_preset_stale(): + """A partition's ``reranker`` preset can go stale (renamed/deleted after + assignment — this field has no create/PATCH-time validation, unlike + ``chat_llm``). The stale name must fall back to the catalog default + instead of raising.""" + s = FakeSearcher() + s.search_result = [_chunk("a")] + default_reranker = FakeReranker() + reranker_calls: list[str] = [] + + def factory(name: str): + reranker_calls.append(name) + if name == "default": + return default_reranker + raise KeyError(name) + + cfg = _config() + cfg.partitions = { + "tenant-a": _partition(retrieval=RetrievalPipelineConfig(enable_reranker=True, reranker="stale-ranker")) + } + + svc = RetrievalService( + searcher=s, + reranker=None, + llm=None, + config=cfg, + reranker_factory=factory, + ) + + out = await svc.retrieve(partitions=["tenant-a"], query=Query(query="hello")) + + assert [c.id for c in out] == ["a"] + assert reranker_calls == ["stale-ranker", "default"] + assert default_reranker.calls[0]["query"] == "hello" + + +@pytest.mark.asyncio +async def test_retrieve_falls_back_to_legacy_reranker_when_no_catalog_default(): + """No ``is_default`` reranker endpoint registered yet — fall back to the + static reranker built at startup instead of raising.""" + s = FakeSearcher() + s.search_result = [_chunk("a")] + legacy_reranker = FakeReranker() + + cfg = _config() + cfg.partitions = {"tenant-a": _partition(retrieval=RetrievalPipelineConfig(enable_reranker=True))} + + svc = RetrievalService( + searcher=s, + reranker=legacy_reranker, + llm=None, + config=cfg, + reranker_factory=lambda name: (_ for _ in ()).throw(KeyError(name)), + ) + + out = await svc.retrieve(partitions=["tenant-a"], query=Query(query="hello")) + + assert [c.id for c in out] == ["a"] + assert legacy_reranker.calls[0]["query"] == "hello" + + @pytest.mark.asyncio async def test_retrieve_uses_partition_searcher_factory_for_named_embedder(): default_searcher = FakeSearcher() @@ -447,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_add_partition_member.py b/tests/unit/services/persistence/test_add_partition_member.py index 1411a20fb..73f972434 100644 --- a/tests/unit/services/persistence/test_add_partition_member.py +++ b/tests/unit/services/persistence/test_add_partition_member.py @@ -15,8 +15,9 @@ async def __aexit__(self, exc_type, exc, tb): class _FakeConn: - def __init__(self, partition_exists: bool): + def __init__(self, partition_exists: bool, membership_created: bool = True): self.partition_exists = partition_exists + self.membership_created = membership_created self.executed: list[tuple[str, tuple]] = [] def transaction(self): @@ -25,6 +26,9 @@ def transaction(self): async def fetchval(self, query: str, *params): if "SELECT 1 FROM partitions" in query: return 1 if self.partition_exists else None + if "INSERT INTO partition_memberships" in query: + self.executed.append((query, params)) + return 1 if self.membership_created else None return None async def execute(self, query: str, *params): @@ -33,8 +37,8 @@ async def execute(self, query: str, *params): class _FakePool: - def __init__(self, partition_exists: bool): - self.conn = _FakeConn(partition_exists) + def __init__(self, partition_exists: bool, membership_created: bool = True): + self.conn = _FakeConn(partition_exists, membership_created) def acquire(self): return _AsyncContext(self.conn) @@ -74,3 +78,17 @@ async def test_add_member_to_existing_partition_allows_any_role(): assert await repo.add_partition_member("existing", 6, "editor") is True assert any("INSERT INTO partition_memberships" in query for query, _ in pool.conn.executed) + + +@pytest.mark.asyncio +async def test_add_member_conflict_does_not_overwrite_existing_role(): + from services.persistence.partition_membership_repo import PgPartitionMembershipRepository + + pool = _FakePool(partition_exists=True, membership_created=False) + repo = PgPartitionMembershipRepository(pool_getter=lambda: pool) + + assert await repo.add_partition_member("existing", 6, "editor") is False + membership_query = next(query for query, _ in pool.conn.executed if "INSERT INTO partition_memberships" in query) + assert "DO NOTHING" in membership_query + assert "DO UPDATE" not in membership_query + assert "RETURNING 1" in membership_query diff --git a/tests/unit/services/persistence/test_display_name_index_migration.py b/tests/unit/services/persistence/test_display_name_index_migration.py new file mode 100644 index 000000000..4f69c2184 --- /dev/null +++ b/tests/unit/services/persistence/test_display_name_index_migration.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import importlib +from contextlib import contextmanager +from pathlib import Path + +import pytest + + +@pytest.fixture +def migration(monkeypatch): + alembic_dir = ( + Path(__file__).resolve().parents[4] / "openrag" / "services" / "persistence" / "migrations" / "alembic" + ) + monkeypatch.syspath_prepend(str(alembic_dir)) + return importlib.import_module( + "services.persistence.migrations.alembic.versions.e5f6a7b8c9d0_add_user_display_name_prefix_index", + ) + + +class _FakeContext: + def __init__(self, calls: list[str]) -> None: + self.calls = calls + + @contextmanager + def autocommit_block(self): + self.calls.append("autocommit") + yield + + +class _FakeOp: + def __init__(self) -> None: + self.calls: list[str] = [] + self.context = _FakeContext(self.calls) + + def get_context(self): + return self.context + + def execute(self, statement) -> None: + self.calls.append(str(statement)) + + +def test_upgrade_rebuilds_an_invalid_concurrent_index(monkeypatch, migration) -> None: + fake_op = _FakeOp() + monkeypatch.setattr(migration, "op", fake_op) + monkeypatch.setattr(migration, "table_exists", lambda _table: True) + monkeypatch.setattr(migration, "_index_validity", lambda: False) + + migration.upgrade() + + statements = "\n".join(fake_op.calls) + assert "DROP INDEX CONCURRENTLY ix_users_lower_display_name_pattern" in statements + assert "CREATE INDEX CONCURRENTLY ix_users_lower_display_name_pattern" in statements + + +def test_upgrade_keeps_a_valid_index(monkeypatch, migration) -> None: + fake_op = _FakeOp() + monkeypatch.setattr(migration, "op", fake_op) + monkeypatch.setattr(migration, "table_exists", lambda _table: True) + monkeypatch.setattr(migration, "_index_validity", lambda: True) + + migration.upgrade() + + assert fake_op.calls == [] + + +def test_upgrade_creates_a_missing_index(monkeypatch, migration) -> None: + fake_op = _FakeOp() + monkeypatch.setattr(migration, "op", fake_op) + monkeypatch.setattr(migration, "table_exists", lambda _table: True) + monkeypatch.setattr(migration, "_index_validity", lambda: None) + + migration.upgrade() + + statements = "\n".join(fake_op.calls) + assert "DROP INDEX" not in statements + assert "CREATE INDEX CONCURRENTLY ix_users_lower_display_name_pattern" in statements diff --git a/tests/unit/services/persistence/test_model_endpoint_repo.py b/tests/unit/services/persistence/test_model_endpoint_repo.py index 660ca6a5a..8544e16ec 100644 --- a/tests/unit/services/persistence/test_model_endpoint_repo.py +++ b/tests/unit/services/persistence/test_model_endpoint_repo.py @@ -237,6 +237,157 @@ async def _execute(query, *params): assert await repo.delete("ghost", "embedder") is False +@pytest.mark.asyncio +async def test_rename_updates_the_model_endpoints_row(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old", "embedder", "new") + + queries = [q for q, _ in pool.conn.executed] + assert any("UPDATE model_endpoints SET name = $3" in q for q in queries) + params = next(p for q, p in pool.conn.executed if "UPDATE model_endpoints SET name" in q) + assert params == ("old", "embedder", "new") + + +@pytest.mark.asyncio +async def test_rename_locks_partitions_table_before_touching_model_endpoints(): + """rename() must LOCK partitions IN SHARE MODE before its own UPDATE — + the same order PgPartitionRepository.update_partition's chat_llm guard + touches partitions (write) then model_endpoints (check), so the two + transactions can only block on each other, never deadlock. Without this + lock, a partition PATCH could validate 'old' in-memory, block on this + transaction's cascade instead, then resume and write 'old' straight back + after this commits — see PgPartitionRepository.update_partition.""" + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old", "llm", "new") + + queries = [q for q, _ in pool.conn.executed] + lock_i = next(i for i, q in enumerate(queries) if q == "LOCK TABLE partitions IN SHARE MODE") + rename_i = next(i for i, q in enumerate(queries) if "UPDATE model_endpoints SET name" in q) + assert lock_i < rename_i + + +@pytest.mark.asyncio +async def test_rename_raises_not_found_and_skips_cascade_when_row_vanished(): + """A concurrent delete between the service's existence check and this + transaction must abort before the cascade — not repoint partitions/presets + at a `new_name` that was never actually created (mirrors + PgPipelinePresetRepository.rename's RETURNING guard).""" + from core.utils.exceptions import NotFoundError + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() # _fetchrow_result defaults to None: row is gone + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + with pytest.raises(NotFoundError): + await repo.rename("old-llm", "llm", "new-llm") + + queries = [q for q, _ in pool.conn.executed] + assert not any("UPDATE partitions SET" in q for q in queries) + assert not any("pipeline_presets" in q for q in queries) + + +@pytest.mark.asyncio +async def test_rename_embedder_cascades_to_partitions_embedder_only(): + """Renaming an embedder must update `partitions.embedder` and touch no + preset JSONB — the embedder name isn't referenced inside any preset.""" + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new-embedder"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old-embedder", "embedder", "new-embedder") + + queries = pool.conn.executed + partition_updates = [(q, p) for q, p in queries if "UPDATE partitions SET" in q] + assert len(partition_updates) == 1 + q, p = partition_updates[0] + assert "embedder = $2 WHERE embedder = $1" in q + assert p == ("old-embedder", "new-embedder") + assert not any("pipeline_presets" in q for q, _ in queries) + + +@pytest.mark.asyncio +async def test_rename_llm_cascades_to_chat_llm_and_both_preset_types(): + """Renaming an LLM endpoint must update `partitions.chat_llm`, the + retrieval preset's `llm` key, and every indexation-preset LLM field.""" + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new-llm"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old-llm", "llm", "new-llm") + + queries = pool.conn.executed + partition_updates = [(q, p) for q, p in queries if "UPDATE partitions SET" in q] + assert len(partition_updates) == 1 + q, p = partition_updates[0] + assert "chat_llm = $2 WHERE chat_llm = $1" in q + assert p == ("old-llm", "new-llm") + + preset_updates = [(q, p) for q, p in queries if "pipeline_presets" in q] + # retrieval.llm + indexation.{contextualization_llm, metadata_extraction_llm, topic_tagging_llm} + assert len(preset_updates) == 4 + keys_by_preset_type = {(p[2], p[3]) for _, p in preset_updates} + assert keys_by_preset_type == { + ("retrieval", "llm"), + ("indexation", "contextualization_llm"), + ("indexation", "metadata_extraction_llm"), + ("indexation", "topic_tagging_llm"), + } + for _, p in preset_updates: + assert p[0] == [p[3]] # jsonb_set path matches the ->> key checked in WHERE + assert p[1] == "new-llm" + assert p[4] == "old-llm" + + +@pytest.mark.asyncio +async def test_rename_reranker_cascades_to_retrieval_preset_only(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new-ranker"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old-ranker", "reranker", "new-ranker") + + queries = pool.conn.executed + assert not any("UPDATE partitions SET" in q for q, _ in queries) + preset_updates = [(q, p) for q, p in queries if "pipeline_presets" in q] + assert len(preset_updates) == 1 + q, p = preset_updates[0] + assert p == (["reranker"], "new-ranker", "retrieval", "reranker", "old-ranker") + + +@pytest.mark.asyncio +async def test_rename_vlm_cascades_to_indexation_preset_only(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new-vlm"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old-vlm", "vlm", "new-vlm") + + queries = pool.conn.executed + assert not any("UPDATE partitions SET" in q for q, _ in queries) + preset_updates = [(q, p) for q, p in queries if "pipeline_presets" in q] + assert len(preset_updates) == 1 + q, p = preset_updates[0] + assert p == (["vlm"], "new-vlm", "indexation", "vlm", "old-vlm") + + def _row(name, is_default): return {"name": name, "is_default": is_default} diff --git a/tests/unit/services/persistence/test_partition_member_candidates.py b/tests/unit/services/persistence/test_partition_member_candidates.py new file mode 100644 index 000000000..97bd07759 --- /dev/null +++ b/tests/unit/services/persistence/test_partition_member_candidates.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import pytest + + +class _FakePool: + def __init__(self): + self.query = "" + self.params: tuple = () + + async def fetch(self, query: str, *params): + self.query = query + self.params = params + return [] + + +@pytest.mark.asyncio +async def test_candidate_search_escapes_like_wildcards_and_uses_a_cursor(): + from services.persistence.partition_membership_repo import PgPartitionMembershipRepository + + pool = _FakePool() + repo = PgPartitionMembershipRepository(pool_getter=lambda: pool) + + await repo.list_partition_member_candidates( + "legal", + search_prefix=r"Sam_%\\", + search_user_id=None, + after_id=42, + limit=26, + ) + + assert pool.params == ("legal", None, r"Sam\_\%\\\\", 42, 26) + assert "u.id > $4" in pool.query + assert "OFFSET" not in pool.query diff --git a/tests/unit/services/persistence/test_partition_repo.py b/tests/unit/services/persistence/test_partition_repo.py index 0aaf6209c..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), } @@ -180,10 +181,13 @@ class _UpdateFakeConn: ``preset_exists`` models whether the referenced preset row is still present when the guard's follow-up SELECT runs (False simulates a concurrent delete_preset committing while this UPDATE was blocked on its SHARE lock). + ``model_endpoint_exists`` is the same, for a ``chat_llm`` reference racing + a concurrent ``PgModelEndpointRepository.rename``. """ - def __init__(self, *, preset_exists: bool = True): + def __init__(self, *, preset_exists: bool = True, model_endpoint_exists: bool = True): self.preset_exists = preset_exists + self.model_endpoint_exists = model_endpoint_exists self.operations: list[tuple[str, tuple]] = [] self.transactions = 0 @@ -209,6 +213,8 @@ async def fetchval(self, query: str, *params): self.operations.append((query, params)) if "FROM pipeline_presets" in query: return 1 if self.preset_exists else None + if "FROM model_endpoints" in query: + return 1 if self.model_endpoint_exists else None return None # conn interface @@ -267,3 +273,76 @@ async def test_update_partition_without_preset_change_skips_the_guard(): assert result["description"] == "notes" assert conn.transactions == 0 assert not any("FROM pipeline_presets" in q for q, _ in conn.operations) + + +# ── update_partition chat_llm-assignment race guard ────────────────── + + +@pytest.mark.asyncio +async def test_update_partition_rolls_back_when_chat_llm_endpoint_renamed_concurrently(): + """A concurrent PgModelEndpointRepository.rename() moving 'old-name' away + while this UPDATE was blocked on the LOCK it also takes must not let the + partition silently keep pointing at the now-nonexistent name.""" + from services.persistence.partition_repo import PgPartitionRepository + + conn = _UpdateFakeConn(model_endpoint_exists=False) + repo = PgPartitionRepository(pool_getter=lambda: conn) + + with pytest.raises(ValidationError) as exc: + await repo.update_partition("p1", chat_llm="old-name") + + assert exc.value.code == "MODEL_ENDPOINT_NOT_FOUND" + # The write and the existence check share one transaction, and the write + # (partitions) happens before the check (model_endpoints) — the same lock + # order PgModelEndpointRepository.rename uses, keeping the two deadlock-free. + assert conn.transactions == 1 + queries = [q for q, _ in conn.operations] + update_i = next(i for i, q in enumerate(queries) if "UPDATE partitions" in q) + check_i = next(i for i, q in enumerate(queries) if "FROM model_endpoints" in q) + assert update_i < check_i + + +@pytest.mark.asyncio +async def test_update_partition_commits_when_chat_llm_endpoint_exists(): + from services.persistence.partition_repo import PgPartitionRepository + + conn = _UpdateFakeConn(model_endpoint_exists=True) + repo = PgPartitionRepository(pool_getter=lambda: conn) + + result = await repo.update_partition("p1", chat_llm="gpt-4.1") + + assert result["chat_llm"] == "gpt-4.1" + assert conn.transactions == 1 + assert any("FROM model_endpoints" in q for q, _ in conn.operations) + + +@pytest.mark.asyncio +async def test_update_partition_clearing_chat_llm_skips_the_guard(): + """chat_llm=None clears the (nullable) column — it names no endpoint to + validate, so this must take the fast, non-transactional path.""" + from services.persistence.partition_repo import PgPartitionRepository + + conn = _UpdateFakeConn(model_endpoint_exists=False) + repo = PgPartitionRepository(pool_getter=lambda: conn) + + result = await repo.update_partition("p1", chat_llm=None) + + assert result["chat_llm"] is None + assert conn.transactions == 0 + assert not any("FROM model_endpoints" in q for q, _ in conn.operations) + + +@pytest.mark.asyncio +async def test_update_partition_embedder_change_skips_the_guard(): + """embedder carries no assignment-time validation today, so assigning it + alone must not pay for a transaction or a model_endpoints lookup.""" + from services.persistence.partition_repo import PgPartitionRepository + + conn = _UpdateFakeConn(model_endpoint_exists=False) + repo = PgPartitionRepository(pool_getter=lambda: conn) + + result = await repo.update_partition("p1", embedder="some-embedder") + + assert result["embedder"] == "some-embedder" + assert conn.transactions == 0 + assert not any("FROM model_endpoints" in q for q, _ in conn.operations) diff --git a/tests/unit/services/persistence/test_user_repo_external_id.py b/tests/unit/services/persistence/test_user_repo_external_id.py index f582db6ef..201d38242 100644 --- a/tests/unit/services/persistence/test_user_repo_external_id.py +++ b/tests/unit/services/persistence/test_user_repo_external_id.py @@ -30,10 +30,14 @@ def __init__(self): self.last_query: str | None = None self.last_params: tuple = () self._next_row: _FakeRow | None = None + self._rows: list[_FakeRow] = [] def set_next_row(self, **fields): self._next_row = _FakeRow(fields) + def set_rows(self, *rows: _FakeRow): + self._rows = list(rows) + async def fetchrow(self, query: str, *params): self.last_query = query self.last_params = params @@ -49,7 +53,9 @@ async def execute(self, query: str, *params): async def fetch(self, query: str, *params): self.last_query = query self.last_params = params - return [] + if "partition_memberships" in query: + return [] + return self._rows def _make_user_with_ext(ext: str | None): @@ -141,3 +147,65 @@ async def test_create_legacy_user_coerces_empty_external_id_to_none(): ) # Same column position (display_name, external_user_id, ...) assert pool.last_params[1] is None + + +@pytest.mark.asyncio +async def test_list_users_dict_includes_email(): + from services.persistence.user_repo import PgUserRepository + + pool = _FakePool() + pool.set_rows( + _FakeRow( + id=42, + display_name="Alice", + external_user_id="kc-alice", + email="alice@example.com", + is_admin=False, + file_quota=None, + file_count=0, + created_at=__import__("datetime").datetime(2026, 1, 1), + ) + ) + repo = PgUserRepository(pool_getter=lambda: pool) + + users = await repo.list_users_dict() + + assert users[0]["email"] == "alice@example.com" + + +@pytest.mark.asyncio +async def test_get_users_by_ids_fetches_all_users_in_one_query(): + from services.persistence.user_repo import PgUserRepository + + pool = _FakePool() + pool.set_rows( + _FakeRow( + id=42, + display_name="Alice", + external_user_id="kc-alice", + email="alice@example.com", + token=None, + is_admin=False, + file_quota=None, + file_count=0, + created_at=__import__("datetime").datetime(2026, 1, 1), + ), + _FakeRow( + id=84, + display_name="Bob", + external_user_id="kc-bob", + email="bob@example.com", + token=None, + is_admin=False, + file_quota=None, + file_count=0, + created_at=__import__("datetime").datetime(2026, 1, 2), + ), + ) + repo = PgUserRepository(pool_getter=lambda: pool) + + users = await repo.get_users_by_ids([42, 84]) + + assert {user.id for user in users} == {42, 84} + assert pool.last_params == ([42, 84],) + assert "ANY($1::int[])" in (pool.last_query or "") 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 diff --git a/tests/unit/test_app_front_secret.py b/tests/unit/test_app_front_secret.py index b0a9a5e37..49565bd88 100644 --- a/tests/unit/test_app_front_secret.py +++ b/tests/unit/test_app_front_secret.py @@ -29,6 +29,16 @@ def _load_app_front(monkeypatch, *, auth_mode: str, module_name: str): return module +def _stub_chainlit_elements(module): + module.cl = SimpleNamespace( + Pdf=lambda **kwargs: SimpleNamespace(**kwargs), + Text=lambda **kwargs: SimpleNamespace(**kwargs), + Image=lambda **kwargs: SimpleNamespace(**kwargs), + Video=lambda **kwargs: SimpleNamespace(**kwargs), + Audio=lambda **kwargs: SimpleNamespace(**kwargs), + ) + + def test_no_hardcoded_default_secret_assignment_in_source(): """The fall-through to a literal default secret must be gone. @@ -398,6 +408,303 @@ async def fake_load_model_ids(_client, api_key): assert module._OPENRAG_TOKEN_STORE[auth_handle][0] == "handoff-token" +@pytest.mark.parametrize( + "sources", + [ + None, + [], + 42, + "not-a-source-list", + {"filename": "notes.txt"}, + [{}], + [None], + [{"source_type": "web", "title": "", "url": "", "snippet": ""}], + [{"source_type": "web", "title": "Invalid URL", "url": "http://[invalid"}], + [{"filename": "", "file_url": "", "page": ""}], + ], +) +@pytest.mark.asyncio +async def test_chainlit_hides_sources_when_none_are_displayable(monkeypatch, sources): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_empty_sources_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + elements, source_names = await module._format_sources(sources) + + assert elements == [] + assert source_names == [] + + +@pytest.mark.asyncio +async def test_chainlit_keeps_page_less_non_pdf_sources(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_page_less_sources_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + async def available_chunk(*_args, **_kwargs): + return "Page-less text content" + + monkeypatch.setattr(module, "__fetch_page_content", available_chunk) + + elements, source_names = await module._format_sources( + [ + { + "filename": "diagram.png", + "file_url": "https://openrag.example/static/image-id", + }, + { + "filename": "demo.mp4", + "file_url": "https://openrag.example/static/video-id", + }, + { + "filename": "recording.mp3", + "file_url": "https://openrag.example/static/audio-id", + }, + { + "filename": "notes.txt", + "file_url": "https://openrag.example/static/text-id", + "chunk_url": "https://openrag.example/chunks/text-id", + }, + ] + ) + + assert source_names == ["diagram.png", "demo.mp4", "recording.mp3", "notes.txt"] + assert [element.name for element in elements] == ["diagram.png", "demo.mp4", "recording.mp3", "notes.txt"] + assert source_names == [element.name for element in elements] + assert elements[-1].content == "Page-less text content" + + +@pytest.mark.asyncio +async def test_chainlit_keeps_page_less_pdf_sources(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_page_less_pdf_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + elements, source_names = await module._format_sources( + [ + { + "filename": "report_[draft].pdf", + "file_url": "https://openrag.example/static/pdf-id", + } + ] + ) + + assert source_names == ["report_ draft .pdf"] + assert source_names == [element.name for element in elements] + assert elements[0].page is None + + +@pytest.mark.asyncio +async def test_chainlit_keeps_valid_web_sources(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_web_source_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + elements, source_names = await module._format_sources( + [ + { + "source_type": "web", + "title": "Example reference", + "url": "https://example.test/reference", + "snippet": "Supporting evidence", + } + ] + ) + + assert source_names == ["Example reference"] + assert elements[0].name == "Example reference" + assert elements[0].content == ("**[Example reference](https://example.test/reference)**\n\nSupporting evidence") + + +@pytest.mark.asyncio +async def test_chainlit_escapes_untrusted_web_source_markdown(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_web_source_escaping_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + title = "Reference ](https://spoof.test) **trusted**" + snippet = "Evidence [click here](https://spoof.test) or *ignore this*." + url = "https://example.test/reference_(draft)" + elements, source_names = await module._format_sources( + [ + { + "source_type": "web", + "title": title, + "url": url, + "snippet": snippet, + } + ] + ) + + assert source_names == ["Reference (https://spoof.test) trusted"] + assert elements[0].name == source_names[0] + assert elements[0].content == ( + r"**[Reference \]\(https\:\/\/spoof\.test\) \*\*trusted\*\*]" + "(https://example.test/reference_%28draft%29)**\n\n" + r"Evidence \[click here\]\(https\:\/\/spoof\.test\) or \*ignore this\*\." + ) + + +@pytest.mark.asyncio +async def test_chainlit_keeps_source_names_unique_after_sanitizing(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_collision_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + elements, source_names = await module._format_sources( + [ + {"source_type": "web", "title": "Reference [draft]", "url": "https://example.test/one"}, + {"source_type": "web", "title": "Reference *draft*", "url": "https://example.test/two"}, + ] + ) + + assert source_names == ["Reference draft", "Reference draft 2"] + assert source_names == [element.name for element in elements] + + +@pytest.mark.parametrize( + ("source_name", "expected"), + [ + ("rapport_annuel_2026.pdf", "rapport_annuel_2026.pdf"), + ("report.pdf (page: 3)", "report.pdf (page: 3)"), + ("C++_style_guide.pdf", "C++_style_guide.pdf"), + ], +) +def test_chainlit_preserves_safe_source_name_punctuation(monkeypatch, source_name, expected): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_punctuation_test") + + assert module._safe_source_name(source_name, {}) == expected + + +@pytest.mark.parametrize( + ("source_name", "expected"), + [ + ("_x_", "x"), + ("_trusted_", "trusted"), + ("__trusted__", "trusted"), + ("foo _trusted_ bar", "foo trusted bar"), + ("~~irrelevant~~", "irrelevant"), + ("# Trusted source", "Trusted source"), + ("1. Official result", "Official result"), + ("- Search result", "Search result"), + ("+ Search result", "Search result"), + ("# 1. Official result", "Official result"), + ("___", "source"), + ], +) +def test_chainlit_neutralizes_active_source_name_markdown(monkeypatch, source_name, expected): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_markdown_test") + + assert module._safe_source_name(source_name, {}) == expected + + +@pytest.mark.parametrize("source_name", ["#report.pdf", "1.report.pdf"]) +def test_chainlit_preserves_non_block_source_name_prefixes(monkeypatch, source_name): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_prefix_test") + + assert module._safe_source_name(source_name, {}) == source_name + + +def test_chainlit_handles_long_source_name_delimiter_runs(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_long_run_test") + source_name = "_" * 10_000 + + assert module._safe_source_name(source_name, {}) == "source" + + +@pytest.mark.parametrize( + "hostile_name", + [ + "x](https://evil.test)", + "nested[a](b)c", + "trailing\\", + "**bold**", + "back`tick`", + ">quote", + ], +) +def test_chainlit_source_name_cannot_break_markdown_link(monkeypatch, hostile_name): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_safety_test") + + safe_name = module._safe_source_name(hostile_name, {}) + + assert not set(safe_name) & set("[]*`\\>") + assert safe_name.strip() == safe_name + + +@pytest.mark.parametrize( + "source_error", + [ + pytest.param(httpx.ConnectError("source unavailable"), id="connection-error"), + pytest.param(httpx.InvalidURL("invalid source URL"), id="invalid-url"), + ], +) +@pytest.mark.asyncio +async def test_chainlit_skips_unavailable_text_sources(monkeypatch, source_error): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_unavailable_source_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + async def unavailable_chunk(*_args, **_kwargs): + raise source_error + + monkeypatch.setattr(module, "__fetch_page_content", unavailable_chunk) + + elements, source_names = await module._format_sources( + [ + { + "filename": "notes.txt", + "file_url": "http://internal:8080/static/source-id", + "page": "1", + "chunk_url": "http://internal:8080/chunks/source-id", + } + ] + ) + + assert elements == [] + assert source_names == [] + + +@pytest.mark.asyncio +async def test_chainlit_skips_text_source_with_non_object_json(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_malformed_source_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + class FakeResponse: + def raise_for_status(self): + return None + + def json(self): + return [{"page_content": "unexpected list response"}] + + class FakeAsyncClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + async def get(self, *_args, **_kwargs): + return FakeResponse() + + monkeypatch.setattr(module.httpx, "AsyncClient", FakeAsyncClient) + + elements, source_names = await module._format_sources( + [ + { + "filename": "notes.txt", + "file_url": "http://internal:8080/static/source-id", + "page": "1", + "chunk_url": "http://internal:8080/chunks/source-id", + } + ] + ) + + assert elements == [] + assert source_names == [] + + @pytest.mark.asyncio async def test_oidc_token_handoff_keeps_bearer_on_static_source_urls(monkeypatch): module = _load_app_front(monkeypatch, auth_mode="oidc", module_name="app_front_source_token_test") @@ -410,14 +717,8 @@ def get(self, key): return SimpleNamespace(metadata={"provider": "credentials"}) return None - module.cl = SimpleNamespace( - user_session=UserSession(), - Pdf=lambda **kwargs: SimpleNamespace(**kwargs), - Text=lambda **kwargs: SimpleNamespace(**kwargs), - Image=lambda **kwargs: SimpleNamespace(**kwargs), - Video=lambda **kwargs: SimpleNamespace(**kwargs), - Audio=lambda **kwargs: SimpleNamespace(**kwargs), - ) + _stub_chainlit_elements(module) + module.cl.user_session = UserSession() elements, _ = await module._format_sources( [ @@ -445,14 +746,8 @@ def get(self, key): return SimpleNamespace(metadata={"provider": "oidc"}) return None - module.cl = SimpleNamespace( - user_session=UserSession(), - Pdf=lambda **kwargs: SimpleNamespace(**kwargs), - Text=lambda **kwargs: SimpleNamespace(**kwargs), - Image=lambda **kwargs: SimpleNamespace(**kwargs), - Video=lambda **kwargs: SimpleNamespace(**kwargs), - Audio=lambda **kwargs: SimpleNamespace(**kwargs), - ) + _stub_chainlit_elements(module) + module.cl.user_session = UserSession() elements, _ = await module._format_sources( [ diff --git a/ui/src/components/layout/sidebar.tsx b/ui/src/components/layout/sidebar.tsx index 7c003df4e..53af3efc5 100644 --- a/ui/src/components/layout/sidebar.tsx +++ b/ui/src/components/layout/sidebar.tsx @@ -6,6 +6,7 @@ import { Clock, Cpu, Settings, + MessageSquareText, Users, Activity, UserCog, @@ -40,6 +41,7 @@ const navItems: NavItem[] = [ { title: "Jobs", href: "/jobs", icon: Clock }, { title: "Models", href: "/models", icon: Cpu, requires: (p) => p.canManageModels }, { title: "Presets", href: "/presets", icon: Settings, requires: (p) => p.canManagePresets }, + { title: "Prompts", href: "/prompts", icon: MessageSquareText, requires: (p) => p.canManagePrompts }, { title: "Users", href: "/users", icon: Users, requires: (p) => p.canManageUsers }, { title: "System", href: "/system", icon: Activity, requires: (p) => p.canViewSystem }, ]; diff --git a/ui/src/components/shared/data-table.tsx b/ui/src/components/shared/data-table.tsx index 71fb6668f..2a19d0a10 100644 --- a/ui/src/components/shared/data-table.tsx +++ b/ui/src/components/shared/data-table.tsx @@ -28,7 +28,10 @@ interface BaseDataTableProps { columns: ColumnDef[]; data: TData[]; pageSize?: number; + emptyMessage?: string; initialSorting?: SortingState; + /** Reset pagination to the first page whenever this value changes. */ + pageResetKey?: unknown; /** Render a leading checkbox column. */ enableSelection?: boolean; /** Optional row-level selection guard for pages with state-dependent bulk actions. */ @@ -54,7 +57,9 @@ export function DataTable({ columns, data, pageSize = 10, + emptyMessage = "No results.", initialSorting = [], + pageResetKey, enableSelection = false, canSelectRow, getRowId, @@ -68,6 +73,12 @@ export function DataTable({ const rowSelection = controlledRowSelection ?? internalRowSelection; const setRowSelection = onRowSelectionChange ?? setInternalRowSelection; + useEffect(() => { + setPagination((previous) => + previous.pageIndex === 0 ? previous : { ...previous, pageIndex: 0 }, + ); + }, [pageResetKey]); + // Prepend a checkbox column when selection is enabled. const tableColumns = useMemo[]>(() => { if (!enableSelection) return columns; @@ -177,7 +188,7 @@ export function DataTable({ colSpan={columnCount} className="h-24 text-center text-muted-foreground" > - No results. + {emptyMessage} )} diff --git a/ui/src/components/ui/tabs.tsx b/ui/src/components/ui/tabs.tsx index 7bf18aa7e..e41440631 100644 --- a/ui/src/components/ui/tabs.tsx +++ b/ui/src/components/ui/tabs.tsx @@ -24,7 +24,7 @@ function Tabs({ } const tabsListVariants = cva( - "rounded-lg p-[3px] group-data-[orientation=horizontal]/tabs:h-9 data-[variant=line]:rounded-none group/tabs-list text-muted-foreground inline-flex w-fit items-center justify-center group-data-[orientation=vertical]/tabs:h-fit group-data-[orientation=vertical]/tabs:flex-col", + "rounded-lg p-[3px] group-data-[orientation=horizontal]/tabs:h-8 data-[variant=line]:rounded-none group/tabs-list text-muted-foreground inline-flex w-fit items-center justify-center group-data-[orientation=vertical]/tabs:h-fit group-data-[orientation=vertical]/tabs:flex-col", { variants: { variant: { @@ -62,7 +62,7 @@ function TabsTrigger({ { expect(fetchMock.mock.calls.map(methodOf)).not.toContain("DELETE"); }); }); + +describe("listPartitionMemberCandidates", () => { + it("requests the owner-protected candidate endpoint", async () => { + fetchMock.mockResolvedValue( + fakeResponse({ + body: JSON.stringify({ + candidates: [{ user_id: 2, display_name: "Sam", email: "sam@example.com" }], + limit: 10, + has_more: true, + next_cursor: 30, + }), + }), + ); + + await expect( + listPartitionMemberCandidates("legal docs", { + search: "Sam Lee", + cursor: 20, + limit: 10, + }), + ).resolves.toEqual({ + candidates: [{ user_id: 2, display_name: "Sam", email: "sam@example.com" }], + limit: 10, + has_more: true, + next_cursor: 30, + }); + const requestedUrl = String(fetchMock.mock.calls[0][0]); + expect(requestedUrl).toContain("/partition/legal%20docs/users/candidates?"); + expect(Array.from(new URLSearchParams(requestedUrl.split("?")[1]).entries())).toEqual( + expect.arrayContaining([ + ["search", "Sam Lee"], + ["cursor", "20"], + ["limit", "10"], + ]), + ); + }); +}); diff --git a/ui/src/lib/api/partitions.ts b/ui/src/lib/api/partitions.ts index 13743a141..d28c6134c 100644 --- a/ui/src/lib/api/partitions.ts +++ b/ui/src/lib/api/partitions.ts @@ -8,7 +8,7 @@ import { request } from "./client"; // POST /partition/{p} create (name in path, NO body; caller becomes owner) → 201 // PATCH /partition/{p} update config → PartitionDetailResponse // DELETE /partition/{p} delete → 204 -// GET /partition/{p}/users members → { members: [{ user_id, role, added_at }] } +// GET /partition/{p}/users members → { members: [{ user_id, display_name, email, role, added_at }] } // POST /partition/{p}/users add member (multipart: user_id, role) // PATCH /partition/{p}/users/{user_id} change role (multipart: role) // DELETE /partition/{p}/users/{user_id} remove member @@ -66,6 +66,9 @@ export interface PartitionConfig { document_count: number; chat_history_depth: number; chat_llm: string | null; + // The final-answer prompt selected for this partition, keyed by prompt type + // (sys_prompt). Absent = the type's global default. + generation_prompt_names: Record; } export interface UpdatePartitionRequest { @@ -75,6 +78,7 @@ export interface UpdatePartitionRequest { retrieval_preset?: string; chat_history_depth?: number; chat_llm?: string | null; + generation_prompt_names?: Record; /** Accepted for compat but never sent (server has no such column). */ collection_name?: string; } @@ -90,6 +94,7 @@ const _PATCH_FIELDS = [ "retrieval_preset", "chat_history_depth", "chat_llm", + "generation_prompt_names", ] as const; function _toRow(r: Record): PartitionResponse { @@ -194,14 +199,48 @@ export function listPartitionFiles(name: string, limit?: number): Promise<{ file export interface PartitionMember { user_id: number; + display_name: string | null; + email: string | null; role: PartitionRole; added_at: string | null; } +export interface PartitionMemberCandidate { + user_id: number; + display_name: string | null; + email: string | null; +} + +export interface PartitionMemberCandidatePage { + candidates: PartitionMemberCandidate[]; + limit: number; + has_more: boolean; + next_cursor: number | null; +} + +interface ListPartitionMemberCandidatesOptions { + search: string; + cursor?: number; + limit?: number; +} + export function listPartitionMembers(name: string): Promise<{ members: PartitionMember[] }> { return request<{ members: PartitionMember[] }>(`${P}/${enc(name)}/users`); } +export function listPartitionMemberCandidates( + name: string, + options: ListPartitionMemberCandidatesOptions, +): Promise { + const query = new URLSearchParams(); + query.set("search", options.search.trim()); + if (options.cursor !== undefined) query.set("cursor", String(options.cursor)); + if (options.limit !== undefined) query.set("limit", String(options.limit)); + return request( + `${P}/${enc(name)}/users/candidates?${query.toString()}`, + ); +} + export function addPartitionMember(name: string, userId: number, role: PartitionRole): Promise { const form = new FormData(); form.append("user_id", String(userId)); diff --git a/ui/src/lib/api/prompts.test.ts b/ui/src/lib/api/prompts.test.ts new file mode 100644 index 000000000..34d951236 --- /dev/null +++ b/ui/src/lib/api/prompts.test.ts @@ -0,0 +1,138 @@ +import { afterEach, beforeEach, describe, it, expect, vi } from "vitest"; +import { + listAllPrompts, + listPrompts, + getPrompt, + createPrompt, + updatePrompt, + setPromptDefault, + deletePrompt, +} from "./prompts"; + +function fakeResponse({ + status = 200, + body = "{}", +}: { status?: number; body?: string } = {}): Response { + return { + status, + ok: status >= 200 && status < 300, + headers: { get: (key: string) => (key.toLowerCase() === "content-type" ? "application/json" : null) }, + json: async () => JSON.parse(body), + text: async () => body, + } as unknown as Response; +} + +const fetchMock = vi.fn(); + +beforeEach(() => { + vi.stubGlobal("fetch", fetchMock); + fetchMock.mockReset(); +}); + +afterEach(() => { + vi.unstubAllGlobals(); +}); + +function lastCall() { + const [url, init] = fetchMock.mock.calls.at(-1)!; + return { url: url as string, init: (init ?? {}) as RequestInit }; +} + +describe("listPrompts", () => { + it("hits /prompts/ and returns the bare array", async () => { + fetchMock.mockResolvedValue(fakeResponse({ body: '[{"id":"1","name":"a","used_by":2}]' })); + const result = await listPrompts(); + expect(lastCall().url).toBe("/prompts/"); + expect(result).toEqual([{ id: "1", name: "a", used_by: 2 }]); + }); + + it("serializes type/offset/limit into the query string", async () => { + fetchMock.mockResolvedValue(fakeResponse({ body: "[]" })); + await listPrompts({ prompt_type: "sys_prompt", offset: 10, limit: 50 }); + expect(lastCall().url).toBe("/prompts/?prompt_type=sys_prompt&offset=10&limit=50"); + }); +}); + +describe("CRUD verbs and paths", () => { + it("getPrompt → GET /prompts/{id}", async () => { + fetchMock.mockResolvedValue(fakeResponse({ body: '{"id":"abc"}' })); + await getPrompt("abc"); + const { url, init } = lastCall(); + expect(url).toBe("/prompts/abc"); + expect(init.method ?? "GET").toBe("GET"); + }); + + it("createPrompt → POST /prompts/ with a JSON body", async () => { + fetchMock.mockResolvedValue(fakeResponse({ status: 201, body: '{"id":"new"}' })); + await createPrompt({ prompt_type: "hyde", name: "x", content: "c" }); + const { url, init } = lastCall(); + expect(url).toBe("/prompts/"); + expect(init.method).toBe("POST"); + expect(JSON.parse(init.body as string)).toEqual({ prompt_type: "hyde", name: "x", content: "c" }); + }); + + it("updatePrompt → PATCH /prompts/{id}", async () => { + fetchMock.mockResolvedValue(fakeResponse({ body: '{"id":"abc"}' })); + await updatePrompt("abc", { content: "new" }); + const { url, init } = lastCall(); + expect(url).toBe("/prompts/abc"); + expect(init.method).toBe("PATCH"); + expect(JSON.parse(init.body as string)).toEqual({ content: "new" }); + }); + + it("setPromptDefault → PUT /prompts/{id}/default", async () => { + fetchMock.mockResolvedValue(fakeResponse({ body: '{"id":"abc","is_default":true}' })); + await setPromptDefault("abc"); + const { url, init } = lastCall(); + expect(url).toBe("/prompts/abc/default"); + expect(init.method).toBe("PUT"); + }); + + it("deletePrompt → DELETE /prompts/{id} (204)", async () => { + fetchMock.mockResolvedValue(fakeResponse({ status: 204, body: "" })); + await deletePrompt("abc"); + const { url, init } = lastCall(); + expect(url).toBe("/prompts/abc"); + expect(init.method).toBe("DELETE"); + }); + + it("url-encodes the id", async () => { + fetchMock.mockResolvedValue(fakeResponse({ body: "{}" })); + await getPrompt("a b/c"); + expect(lastCall().url).toBe("/prompts/a%20b%2Fc"); + }); +}); + +// A single capped request hid prompts past the cap and reported a partial count +// as the total, leaving them unmanageable in the library and unselectable in +// every picker. +describe("listAllPrompts", () => { + const page = (n: number, from: number) => + JSON.stringify(Array.from({ length: n }, (_, i) => ({ id: `p${from + i}`, name: `p${from + i}` }))); + + it("follows pagination until a short page", async () => { + fetchMock + .mockResolvedValueOnce(fakeResponse({ body: page(200, 0) })) + .mockResolvedValueOnce(fakeResponse({ body: page(200, 200) })) + .mockResolvedValueOnce(fakeResponse({ body: page(37, 400) })); + + const all = await listAllPrompts(); + + expect(all).toHaveLength(437); + expect(fetchMock).toHaveBeenCalledTimes(3); + expect(fetchMock.mock.calls[1][0]).toContain("offset=200"); + expect(fetchMock.mock.calls[2][0]).toContain("offset=400"); + }); + + it("stops after one request when the first page is short", async () => { + fetchMock.mockResolvedValueOnce(fakeResponse({ body: page(7, 0) })); + expect(await listAllPrompts()).toHaveLength(7); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it("returns an empty library without looping", async () => { + fetchMock.mockResolvedValueOnce(fakeResponse({ body: "[]" })); + expect(await listAllPrompts()).toEqual([]); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/src/lib/api/prompts.ts b/ui/src/lib/api/prompts.ts index cee5882d6..96a586ef7 100644 --- a/ui/src/lib/api/prompts.ts +++ b/ui/src/lib/api/prompts.ts @@ -1,34 +1,61 @@ import { request } from "./client"; -export interface PromptUsedBy { - count: number; - partitions: string[]; -} +// OpenRag prompt library (admin-only). Mounted at `/prompts`. +// Shapes verified against openrag/api/schemas/admin/prompt_schemas.py + routers/admin/prompts.py: +// GET /prompts/ list (bare array; optional ?prompt_type=&offset=&limit=) +// POST /prompts/ create → 201 +// GET /prompts/{id} get one +// PATCH /prompts/{id} update (name/content/is_default) +// PUT /prompts/{id}/default promote to default for its type +// DELETE /prompts/{id} delete → 204 +// +// Prompts are *named*; a preset or partition selects one by naming it. There is +// no partition-assignment endpoint — selection lives in the preset/partition +// editors (see presets.tsx / partitions). + +// The 8 managed prompt types (mirrors PromptTypeName on the backend). +export type PromptType = + | "sys_prompt" + | "query_contextualizer" + | "chunk_contextualizer" + | "image_captioning" + | "hyde" + | "multi_query" + | "topic_tagger"; export interface PromptResponse { id: string; - prompt_type: string; + prompt_type: PromptType; name: string; content: string; is_default: boolean; created_at: string; updated_at: string; - used_by: PromptUsedBy | null; + // Number of partitions that reference this prompt by name (directly for + // generation prompts, transitively via preset for indexation/retrieval). + // Only populated by listPrompts(); single-item responses default it to 0. + used_by: number; +} + +export interface CreatePromptRequest { + prompt_type: PromptType; + name: string; + content: string; + is_default?: boolean; } -export interface PromptListResponse { - prompts: PromptResponse[]; - offset: number; - limit: number; +export interface UpdatePromptRequest { + name?: string; + content?: string; + is_default?: boolean; } -const BASE = "/api/v1/admin/prompts"; +const BASE = "/prompts"; +const enc = encodeURIComponent; -// Prompts is a dropped feature; presets only reads the prompt list (to pick a -// chat prompt). The CRUD/partition-assignment surface was removed — re-add from -// git history if a prompt-management page is reinstated. +/** List library prompts (bare array). Optionally filter by type. */ export function listPrompts(params?: { - prompt_type?: string; + prompt_type?: PromptType; offset?: number; limit?: number; }) { @@ -37,5 +64,56 @@ export function listPrompts(params?: { if (params?.offset !== undefined) search.set("offset", String(params.offset)); if (params?.limit !== undefined) search.set("limit", String(params.limit)); const qs = search.toString(); - return request(`${BASE}${qs ? `?${qs}` : ""}`); + return request(`${BASE}/${qs ? `?${qs}` : ""}`); +} + +/** Page size used when walking the whole library. The API caps `limit` at 500. */ +const PAGE_SIZE = 200; + +/** Fetch every library prompt, following offset pagination to the end. + * + * The management page and every prompt picker need the complete library: a + * single capped request would silently hide prompts past the cap and report a + * partial count as if it were the total, leaving those prompts unmanageable + * and unselectable. + */ +export async function listAllPrompts(params?: { prompt_type?: PromptType }): Promise { + const all: PromptResponse[] = []; + for (let offset = 0; ; offset += PAGE_SIZE) { + const page = await listPrompts({ ...params, offset, limit: PAGE_SIZE }); + all.push(...page); + // A short page means the end; a full page means there may be more. + if (page.length < PAGE_SIZE) return all; + } +} + +export function getPrompt(id: string) { + return request(`${BASE}/${enc(id)}`); +} + +export function createPrompt(data: CreatePromptRequest) { + return request(`${BASE}/`, { + method: "POST", + body: JSON.stringify(data), + }); +} + +export function updatePrompt(id: string, data: UpdatePromptRequest) { + return request(`${BASE}/${enc(id)}`, { + method: "PATCH", + body: JSON.stringify(data), + }); +} + +/** Promote a prompt to the default for its type. */ +export function setPromptDefault(id: string) { + return request(`${BASE}/${enc(id)}/default`, { + method: "PUT", + }); +} + +export function deletePrompt(id: string) { + return request(`${BASE}/${enc(id)}`, { + method: "DELETE", + }); } diff --git a/ui/src/lib/permissions.ts b/ui/src/lib/permissions.ts index 3c14df7b5..e3e0deb17 100644 --- a/ui/src/lib/permissions.ts +++ b/ui/src/lib/permissions.ts @@ -35,6 +35,7 @@ export interface Permissions { canManageUsers: boolean; canManageModels: boolean; canManagePresets: boolean; + canManagePrompts: boolean; canManagePartitions: boolean; // Any authenticated user may create (and thereby own) a partition — the backend's // POST /partition/{name} has no admin guard. Distinct from canManagePartitions, @@ -74,6 +75,7 @@ export function usePermissions(): Permissions { canManageUsers: isAdmin, canManageModels: isAdmin, canManagePresets: isAdmin, + canManagePrompts: isAdmin, canManagePartitions: isAdmin, canCreatePartition: true, canRead: (role) => superAdmin || !!role, diff --git a/ui/src/lib/prompt-meta.test.ts b/ui/src/lib/prompt-meta.test.ts new file mode 100644 index 000000000..7aefd95a3 --- /dev/null +++ b/ui/src/lib/prompt-meta.test.ts @@ -0,0 +1,148 @@ +import { describe, it, expect } from "vitest"; +import { + PROMPT_DEFAULT_OPTION, + promptOptionToName, + promptOptionValue, + promptSelectValue, + renderPreview, + scanTemplate, + validatePlaceholders, +} from "./prompt-meta"; + +// The editor must reach the same verdict as the API's `string.Formatter` +// validation, or a template looks fine here and comes back 422 on save. +describe("scanTemplate", () => { + it("collects simple fields", () => { + expect(scanTemplate("Answer with {context} on {current_date}").fields).toEqual([ + "context", + "current_date", + ]); + }); + + it("treats doubled braces as literal text, not placeholders", () => { + const scan = scanTemplate('emit {{"a": 1}} verbatim'); + expect(scan.malformed).toBe(false); + expect(scan.fields).toEqual([]); + }); + + it("flags a lone opening brace the way Python does", () => { + expect(scanTemplate("what { is this").malformed).toBe(true); + }); + + it("flags a lone closing brace", () => { + expect(scanTemplate("what } is this").malformed).toBe(true); + }); + + it("reduces conversions, format specs and attribute/index access to the root", () => { + expect(scanTemplate("{context!r} {context:>10} {context.attr} {context[0]}").fields).toEqual([ + "context", + ]); + }); + + it("treats auto-numbering as an unnamed field, which the API rejects", () => { + expect(validatePlaceholders("hello {}", "sys_prompt").unknown).toEqual([""]); + }); +}); + +describe("validatePlaceholders", () => { + it("accepts a template using only its type's variables", () => { + const v = validatePlaceholders("Use {context} on {current_date}", "sys_prompt"); + expect(v.unknown).toEqual([]); + expect(v.malformed).toBe(false); + }); + + it("reports an unknown variable", () => { + expect(validatePlaceholders("About {topic}", "sys_prompt").unknown).toEqual(["topic"]); + }); + + it("accepts any literal text for verbatim prompt types", () => { + // chunk_contextualizer is sent to the model as-is and never `.format`-ed, + // so braces in it are just characters. + const v = validatePlaceholders('Return {"json": true} exactly', "chunk_contextualizer"); + expect(v.unknown).toEqual([]); + expect(v.malformed).toBe(false); + }); +}); + +describe("prompt picker option values", () => { + it("round-trips a name", () => { + expect(promptOptionToName(promptOptionValue("legal-assistant"))).toBe("legal-assistant"); + }); + + it("maps an unset name to the default sentinel", () => { + expect(promptSelectValue(undefined)).toBe(PROMPT_DEFAULT_OPTION); + expect(promptSelectValue("")).toBe(PROMPT_DEFAULT_OPTION); + expect(promptOptionToName(PROMPT_DEFAULT_OPTION)).toBe(""); + }); + + it("keeps a prompt whose name looks like the sentinel selectable", () => { + // Names are free text, so the sentinel must not be able to collide with one. + for (const name of ["__default__", "__use_default__"]) { + const value = promptOptionValue(name); + expect(value).not.toBe(PROMPT_DEFAULT_OPTION); + expect(promptOptionToName(value)).toBe(name); + } + }); +}); + +// The preview must agree with what the pipeline actually renders — it shares +// the tokenizer with validation so the two can't drift apart. +describe("renderPreview", () => { + const flat = (content: string, type = "sys_prompt") => + renderPreview(content, type) + .map((s) => s.value) + .join(""); + + it("substitutes a known variable with its sample", () => { + const segments = renderPreview("Today is {current_date}.", "sys_prompt"); + const substituted = segments.find((s) => s.isVariable); + expect(substituted?.varName).toBe("current_date"); + expect(flat("Today is {current_date}.")).toContain("2026-07-27"); + }); + + it("unescapes doubled braces the way str.format does", () => { + expect(flat('emit {{"a": 1}} verbatim')).toBe('emit {"a": 1} verbatim'); + }); + + it("still previews a modified field, though such a template cannot be saved", () => { + // Modifiers are rejected at validation (see below), so this only governs + // what the editor shows while the author is mid-edit: the sample lands + // where the value would, with no attempt to emulate padding or `!r`. + expect(flat("Today is {current_date:>12}.")).toBe("Today is 2026-07-27."); + expect(flat("Today is {current_date!r}.")).toBe("Today is 2026-07-27."); + }); + + it("leaves an unknown field visible rather than dropping it", () => { + expect(flat("About {topic}")).toBe("About {topic}"); + }); + + it("previews a verbatim prompt type unchanged", () => { + const raw = 'Return {"json": true} exactly'; + expect(flat(raw, "chunk_contextualizer")).toBe(raw); + }); +}); + + +// Reducing an expression to its root name made unrenderable templates look +// valid: `{context!x}` raises ValueError and `{context.missing}` raises +// AttributeError when the pipeline formats the prompt. As a type's default that +// breaks every request, so only plain placeholders are accepted — mirroring +// `_validate_template` on the API. +describe("validatePlaceholders rejects non-plain placeholders", () => { + it.each([ + "{context!x} on {current_date}", + "{context.missing} on {current_date}", + "{context[0]} on {current_date}", + "{context:>10} on {current_date}", + ])("rejects %s", (content) => { + const v = validatePlaceholders(content, "sys_prompt"); + expect(v.malformed).toBe(true); + expect(v.error).toMatch(/not supported/); + }); + + it("still accepts the plain form", () => { + const v = validatePlaceholders("Use {context} on {current_date}", "sys_prompt"); + expect(v.malformed).toBe(false); + expect(v.unknown).toEqual([]); + }); +}); diff --git a/ui/src/lib/prompt-meta.ts b/ui/src/lib/prompt-meta.ts new file mode 100644 index 000000000..6959628b7 --- /dev/null +++ b/ui/src/lib/prompt-meta.ts @@ -0,0 +1,282 @@ +import type { PromptType } from "@/lib/api/prompts"; + +// Shared metadata for the managed prompt types: their human labels, how they +// group by concern, and the `{variables}` each template understands. Consumed by +// the Prompt Library page and by the preset/partition editors that select a +// prompt by name. + +export interface PromptTypeEntry { + value: PromptType; + label: string; +} + +export interface PromptGroup { + // Concern this group of prompts belongs to. + name: string; + // Where a prompt of this concern is *selected* (presets vs partition). + description: string; + types: PromptTypeEntry[]; +} + +export const PROMPT_GROUPS: PromptGroup[] = [ + { + name: "Final Answer", + description: "The final-answer prompt — selected per partition", + types: [{ value: "sys_prompt", label: "Final answer prompt" }], + }, + { + name: "Indexation", + description: "Document-enrichment prompts — selected on the indexation preset", + types: [ + { value: "chunk_contextualizer", label: "Contextualization" }, + { value: "image_captioning", label: "Image captioning" }, + { value: "topic_tagger", label: "Topic tagging" }, + ], + }, + { + name: "Retrieval", + description: "Query-transformation prompts — selected on the retrieval preset", + types: [ + { value: "query_contextualizer", label: "Query contextualizer" }, + { value: "hyde", label: "HyDE" }, + { value: "multi_query", label: "Multi-query" }, + ], + }, +]; + +export const PROMPT_TYPES: PromptTypeEntry[] = PROMPT_GROUPS.flatMap((g) => g.types); + +const PROMPT_TYPE_LABELS: Record = Object.fromEntries( + PROMPT_TYPES.map((t) => [t.value, t.label]), +); + +export function promptTypeLabel(type: string): string { + return PROMPT_TYPE_LABELS[type] ?? type; +} + +export interface TemplateVariable { + name: string; + description: string; + sample: string; +} + +// Placeholders each template understands, keyed by prompt type. Empty arrays are +// intentional: those prompts are system messages sent with the chunk/image +// attached as a separate message — they take no inline placeholders. +export const PROMPT_TYPE_VARIABLES: Record = { + sys_prompt: [ + { name: "context", description: "Retrieved document chunks injected by the pipeline", sample: "[Source 1] Employees are entitled to 20 days of paid vacation per year, accrued monthly." }, + { name: "current_date", description: "Today's date, injected at request time", sample: "2026-07-27" }, + ], + query_contextualizer: [ + { name: "current_date", description: "Today's date, injected at request time", sample: "2026-07-27" }, + { name: "query_language", description: "Detected language of the user's query", sample: "English" }, + ], + chunk_contextualizer: [], + image_captioning: [], + topic_tagger: [], + hyde: [ + { name: "question", description: "The user's search query", sample: "How do I configure SSL certificates?" }, + ], + multi_query: [ + { name: "k_queries", description: "Number of alternative queries to generate", sample: "3" }, + { name: "query", description: "The user's original search query", sample: "What are the termination clauses in the contract?" }, + ], +}; + +export interface TemplateScan { + /** Root field names, reduced the way the backend reduces them. */ + fields: string[]; + /** Raw expressions carrying a conversion/format spec/attribute access, which + * the API rejects as not-plain placeholders. */ + modified: string[]; + /** True when Python's parser would raise — i.e. the API would return 422. */ + malformed: boolean; + error?: string; +} + +type TemplateToken = + | { kind: "literal"; text: string } + | { kind: "field"; root: string; raw: string; hasModifier: boolean }; + +interface TokenizeResult { + tokens: TemplateToken[]; + malformed: boolean; + error?: string; +} + +/** Tokenize a template the way Python's ``string.Formatter().parse`` does. + * + * Single source of grammar for both validation and the preview, so the two can + * never disagree about what a template means: `{{`/`}}` are literal braces, a + * lone brace is an error, and a field may carry a conversion (`!r`), a format + * spec (`:>10`) or attribute/index access (`a.b`, `a[0]`), all of which reduce + * to the root name the backend validates against. + */ +function tokenizeTemplate(content: string): TokenizeResult { + const tokens: TemplateToken[] = []; + let literal = ""; + let i = 0; + const flush = () => { + if (literal) { + tokens.push({ kind: "literal", text: literal }); + literal = ""; + } + }; + while (i < content.length) { + const ch = content[i]; + if (ch === "{") { + if (content[i + 1] === "{") { + literal += "{"; + i += 2; + continue; + } + const end = content.indexOf("}", i + 1); + if (end === -1) { + flush(); + return { tokens, malformed: true, error: "Single '{' encountered in format string" }; + } + const raw = content.slice(i + 1, end); + const root = raw.split("!")[0].split(":")[0].split(".")[0].split("[")[0].trim(); + // A conversion, format spec or attribute/index access is not a plain + // placeholder. The API rejects those because reducing them to a root name + // hides templates `.format()` cannot render (`{context!x}` raises + // ValueError, `{context.missing}` raises AttributeError). + const hasModifier = raw !== root; + flush(); + tokens.push({ kind: "field", root, raw, hasModifier }); + i = end + 1; + continue; + } + if (ch === "}") { + if (content[i + 1] === "}") { + literal += "}"; + i += 2; + continue; + } + flush(); + return { tokens, malformed: true, error: "Single '}' encountered in format string" }; + } + literal += ch; + i += 1; + } + flush(); + return { tokens, malformed: false }; +} + +export function scanTemplate(content: string): TemplateScan { + const { tokens, malformed, error } = tokenizeTemplate(content); + const fields: string[] = []; + const modified: string[] = []; + for (const token of tokens) { + if (token.kind !== "field") continue; + if (!fields.includes(token.root)) fields.push(token.root); + if (token.hasModifier && !modified.includes(token.raw)) modified.push(token.raw); + } + return { fields, modified, malformed, error }; +} + +export function extractPlaceholders(content: string): string[] { + return scanTemplate(content).fields; +} + +/** Types the backend renders with `str.format` and therefore validates. + * + * Mirrors `_PROMPT_FORMAT_FIELDS` in `prompt_service.py`. It cannot be derived + * from `PROMPT_TYPE_VARIABLES`, because a type there may legitimately have an + * empty variable list; the others are sent to the model verbatim, so braces in + * them are ordinary characters and must not be validated at all. + */ +const FORMATTED_PROMPT_TYPES = new Set([ + "sys_prompt", + "query_contextualizer", + "hyde", + "multi_query", +]); + +export function validatePlaceholders(content: string, promptType: string) { + const known = new Set((PROMPT_TYPE_VARIABLES[promptType] ?? []).map((v) => v.name)); + if (!FORMATTED_PROMPT_TYPES.has(promptType)) { + return { used: [], unknown: [], missing: [], malformed: false, error: undefined }; + } + const scan = scanTemplate(content); + const used = scan.fields; + const unknown = used.filter((v) => !known.has(v)); + const missing = [...known].filter((v) => !used.includes(v)); + if (!scan.malformed && scan.modified.length > 0) { + return { + used, + unknown, + missing, + malformed: true, + error: + `Placeholder {${scan.modified[0]}} uses a conversion, format spec or attribute access, ` + + `which is not supported. Use a plain placeholder.`, + }; + } + return { used, unknown, missing, malformed: scan.malformed, error: scan.error }; +} + +export interface PreviewSegment { + value: string; + isVariable: boolean; + varName?: string; +} + +/** Render the template the way the pipeline will: substitute each known + * `{var}` with its sample value, unescape `{{`/`}}`, and track which spans were + * substituted so the preview can highlight them. + * + * Shares `tokenizeTemplate` with validation so what an author sees here matches + * what the model receives — a separate regex used to leave escaped braces + * doubled and skip formatted fields entirely. + * + * A field's conversion (`!r`) and format spec (`:>12`) are recognised but + * deliberately NOT emulated: the substituted value is the illustrative sample, + * not the runtime one, so reimplementing Python's format mini-language in + * TypeScript would add a second thing to keep in sync for no gain. The preview + * shows where a value lands, not its exact padding or quoting. + */ +export function renderPreview(content: string, promptType: string): PreviewSegment[] { + // Verbatim types are never `.format`-ed at runtime, so they preview as-is. + if (!FORMATTED_PROMPT_TYPES.has(promptType)) { + return content ? [{ value: content, isVariable: false }] : []; + } + const vars = PROMPT_TYPE_VARIABLES[promptType] ?? []; + const varMap = Object.fromEntries(vars.map((v) => [v.name, v.sample])); + const { tokens } = tokenizeTemplate(content); + + return tokens.map((token) => { + if (token.kind === "literal") return { value: token.text, isVariable: false }; + const sample = varMap[token.root]; + return sample !== undefined + ? { value: sample, isVariable: true, varName: token.root } + : { value: `{${token.raw}}`, isVariable: false }; + }); +} + +/* ---------- Prompt picker option values ---------- */ + +/** Sentinel for the "use the type's global default" choice. + * + * Prompt names are free text, so any bare sentinel is a name a prompt could + * legitimately have — a prompt actually called `__default__` would then be + * impossible to select. Real options are therefore namespaced with a prefix + * and the sentinel is the only unprefixed value, which makes a collision + * structurally impossible rather than merely unlikely. + */ +export const PROMPT_DEFAULT_OPTION = "__use_default__"; +const PROMPT_OPTION_PREFIX = "name:"; + +export function promptOptionValue(name: string): string { + return `${PROMPT_OPTION_PREFIX}${name}`; +} + +export function promptOptionToName(value: string): string { + return value.startsWith(PROMPT_OPTION_PREFIX) ? value.slice(PROMPT_OPTION_PREFIX.length) : ""; +} + +/** The Select value for a stored prompt name ("" / unset → the default). */ +export function promptSelectValue(name: string | undefined | null): string { + return name ? promptOptionValue(name) : PROMPT_DEFAULT_OPTION; +} diff --git a/ui/src/pages/admin/jobs/list.test.tsx b/ui/src/pages/admin/jobs/list.test.tsx index b8d3f42bd..2659ebbcd 100644 --- a/ui/src/pages/admin/jobs/list.test.tsx +++ b/ui/src/pages/admin/jobs/list.test.tsx @@ -128,7 +128,7 @@ describe("JobListPage filters", () => { await userEvent.type(search, "completed"); expect(search.value).toBe("completed"); - await userEvent.click(screen.getByRole("tab", { name: "FAILED" })); + await userEvent.click(screen.getByRole("tab", { name: "Failed" })); await waitFor(() => expect(search.value).toBe("")); expect(await screen.findByText("failed.pdf")).not.toBeNull(); @@ -245,7 +245,7 @@ describe("JobListPage filters", () => { [expect.objectContaining({ task_id: "all-named-task" })], ); - await userEvent.click(screen.getByRole("tab", { name: "FAILED" })); + await userEvent.click(screen.getByRole("tab", { name: "Failed" })); await waitFor(() => expect(screen.getByText("docs.pdf")).not.toBeNull()); }); diff --git a/ui/src/pages/admin/jobs/list.tsx b/ui/src/pages/admin/jobs/list.tsx index 32643580f..f2dd84043 100644 --- a/ui/src/pages/admin/jobs/list.tsx +++ b/ui/src/pages/admin/jobs/list.tsx @@ -351,7 +351,7 @@ export default function JobListPage() { {STATUS_TABS.map((tab) => ( - {tab} + {tab.charAt(0) + tab.slice(1).toLowerCase()} ))} diff --git a/ui/src/pages/admin/models.tsx b/ui/src/pages/admin/models.tsx index 8811682db..aa1016cd1 100644 --- a/ui/src/pages/admin/models.tsx +++ b/ui/src/pages/admin/models.tsx @@ -72,6 +72,14 @@ type RevealedApiKey = { const normalizeEndpointUrl = (value: string) => value.trim().replace(/\/+$/, ""); +// Mirrors the backend's `_NAME_PATTERN` allowlist (api/schemas/admin/model_endpoint_schemas.py): +// `name` is a single path segment in every single-endpoint route. An allowlist, not a +// denylist — `/` splits across path segments, and `.`/`..` are RFC 3986 dot-segments that +// browsers normalize out of the URL before the request is even sent — so anchoring both +// ends on alphanumeric rules out all of that (and any leading/trailing separator) at once. +const NAME_PATTERN = /^[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?$/; +const NAME_MAX_LENGTH = 128; + // Placeholder shown when a per-endpoint budget is left blank. The real // fallback (the MAX_LLM_CONTEXT_SIZE / MAX_OUTPUT_TOKENS env vars — see // core/config/endpoints.py:LLMContextConfig) is environment-configurable, so @@ -210,7 +218,7 @@ export default function ModelsPage() {
{isDefault && ( - Default + Default )}
@@ -316,6 +324,20 @@ function EndpointDialog({ const [vendor, setVendor] = useState(""); const [extraJson, setExtraJson] = useState("{}"); + // The backend trims `name` before validating it (and before persisting it), + // so validate — and submit — the same trimmed value here, not the raw + // input; otherwise e.g. " gpt-4.1 " would be accepted by the backend but + // blocked by this form. Checked client-side too because the resulting 422 + // has no readable message — FastAPI's validation `detail` is a list, and + // ApiError only unwraps string detail. + const trimmedName = name.trim(); + const nameError = + trimmedName !== "" && trimmedName.length > NAME_MAX_LENGTH + ? `Name must be at most ${NAME_MAX_LENGTH} characters.` + : trimmedName !== "" && !NAME_PATTERN.test(trimmedName) + ? "Name must start/end with a letter or digit, and contain only letters, digits, '.', '_', or '-' — it's used in the endpoint's URL path." + : null; + const modelType = (editing ? editing.model_type : activeTab) as ModelType; // LLM token-budget fields (max context / max output) apply to LLM endpoints only. const isLlm = modelType === "llm"; @@ -572,6 +594,10 @@ function EndpointDialog({ const handleSubmit = (e: React.FormEvent) => { e.preventDefault(); + if (nameError) { + toast.error(nameError); + return; + } let extra: Record = {}; try { extra = mergeModelEndpointApiKeyExtra(JSON.parse(extraJson), apiKey, { @@ -594,13 +620,13 @@ function EndpointDialog({ timeout: numOr(timeout, 30), extra, }; - if (name !== editing.name) { - updateData.name = name; + if (trimmedName !== editing.name) { + updateData.name = trimmedName; } onUpdate(editing.model_type, editing.name, updateData); } else { onCreate({ - name, + name: trimmedName, model_type: activeTab as ModelType, endpoint, model_name: modelName || undefined, @@ -633,6 +659,7 @@ function EndpointDialog({ onChange={(e) => setName(e.target.value)} required /> + {nameError &&

{nameError}

}
@@ -800,7 +827,7 @@ function EndpointDialog({ diff --git a/ui/src/pages/admin/partitions/detail.tsx b/ui/src/pages/admin/partitions/detail.tsx index 704573d3f..371eb5219 100644 --- a/ui/src/pages/admin/partitions/detail.tsx +++ b/ui/src/pages/admin/partitions/detail.tsx @@ -1,10 +1,11 @@ -import { useState, useEffect, useCallback } from "react"; +import { useState, useEffect, useCallback, useMemo } from "react"; import { useParams, Link } from "react-router-dom"; -import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query"; +import { useInfiniteQuery, useQuery, useMutation, useQueryClient } from "@tanstack/react-query"; import { ArrowLeft, Save, UserPlus, Trash2, CheckCircle, XCircle, Loader2, Info } from "lucide-react"; import { toast } from "sonner"; import { Button } from "@/components/ui/button"; +import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { Textarea } from "@/components/ui/textarea"; @@ -42,15 +43,37 @@ import { updatePartition, listPartitions, listPartitionMembers, - addPartitionMember, + listPartitionMemberCandidates, removePartitionMember, updatePartitionMemberRole, } from "@/lib/api/partitions"; -import type { PartitionConfig, PartitionRole } from "@/lib/api/partitions"; +import type { PartitionConfig, PartitionMemberCandidate, PartitionRole } from "@/lib/api/partitions"; import { listPresets } from "@/lib/api/presets"; import { listModelEndpoints, validateStoredModelEndpoint, resolveEmbedderName } from "@/lib/api/models"; +import { listAllPrompts } from "@/lib/api/prompts"; +import type { PromptResponse } from "@/lib/api/prompts"; +import { + PROMPT_DEFAULT_OPTION, + PROMPT_GROUPS, + promptOptionToName, + promptOptionValue, + promptSelectValue, +} from "@/lib/prompt-meta"; import { usePermissions } from "@/lib/permissions"; import { formatDate, intOr } from "@/lib/utils"; +import { MemberPicker } from "./member-picker"; +import { candidateLabel, candidateSecondaryLabel } from "./member-candidate"; +import { addPartitionMembers } from "./member-batch"; +import type { MemberAddFailure } from "./member-batch"; +import { + PartitionMemberEmail, + PartitionMemberIdentity, +} from "./partition-member-identity"; +import { describePartitionMember } from "./partition-member"; + +// The answer prompt is selected on the partition (keyed by prompt type). Mirrors +// the "Final Answer" concern in the prompt library. +const GENERATION_PROMPT_TYPES = PROMPT_GROUPS.find((g) => g.name === "Final Answer")!.types; // --- General Tab --- @@ -71,6 +94,16 @@ function GeneralTab({ partition }: { partition: PartitionConfig }) { const [indexationPreset, setIndexationPreset] = useState(partition.indexation_preset); const [retrievalPreset, setRetrievalPreset] = useState(partition.retrieval_preset); const [chatLlm, setChatLlm] = useState(partition.chat_llm ?? "__default__"); + // Only carry the generation types this editor manages — a stale key (e.g. a + // pre-move query_contextualizer) would be rejected by the partition PATCH. + const initialGenerationPrompts = useMemo(() => { + const allowed = new Set(GENERATION_PROMPT_TYPES.map((t) => t.value)); + return Object.fromEntries( + Object.entries(partition.generation_prompt_names ?? {}).filter(([k]) => allowed.has(k)), + ); + }, [partition.generation_prompt_names]); + const [generationPrompts, setGenerationPrompts] = + useState>(initialGenerationPrompts); const [llmValidated, setLlmValidated] = useState( partition.chat_llm ? null : true, ); @@ -96,6 +129,28 @@ function GeneralTab({ partition }: { partition: PartitionConfig }) { enabled: isAdmin, }); + const { data: promptsData } = useQuery({ + queryKey: ["prompts-library"], + queryFn: () => listAllPrompts(), + enabled: isAdmin, + }); + const promptsByType = (type: string): PromptResponse[] => + (promptsData ?? []).filter((p) => p.prompt_type === type); + + // Compared against what the partition was loaded with, so an untouched + // mapping is omitted from the PATCH entirely. + const generationPromptsChanged = + JSON.stringify(generationPrompts) !== JSON.stringify(initialGenerationPrompts); + + const setGenerationPrompt = (type: string, name: string) => { + setGenerationPrompts((prev) => { + const next = { ...prev }; + if (name) next[type] = name; + else delete next[type]; + return next; + }); + }; + const validateLlm = useCallback( async (name: string) => { const ep = llmEndpoints?.find((e) => e.name === name); @@ -145,6 +200,12 @@ function GeneralTab({ partition }: { partition: PartitionConfig }) { indexation_preset: indexationPreset, retrieval_preset: retrievalPreset, chat_llm: chatLlm === "__default__" ? null : chatLlm, + // Only sent when this editor actually changed it. The backend validates + // every name in the mapping, so resubmitting an untouched-but-stale + // reference (its prompt since renamed or deleted) would 422 the whole + // PATCH and block unrelated edits like the description — and a non-admin + // owner, for whom this picker is disabled, could never clear it. + ...(generationPromptsChanged ? { generation_prompt_names: generationPrompts } : {}), }), onSuccess: () => { toast.success("Partition updated"); @@ -176,7 +237,7 @@ function GeneralTab({ partition }: { partition: PartitionConfig }) { General Settings -
+ {!canEdit && (

You have read-only access to this partition's settings. Only an owner can change them. @@ -197,7 +258,7 @@ function GeneralTab({ partition }: { partition: PartitionConfig }) { disabled={!canEdit} />

-
+

{partition.dimension}

@@ -213,11 +274,11 @@ function GeneralTab({ partition }: { partition: PartitionConfig }) {

{partition.document_count}

-
+
- + @@ -256,7 +317,7 @@ function GeneralTab({ partition }: { partition: PartitionConfig }) { )}
+ {GENERATION_PROMPT_TYPES.map((t) => { + const options = promptsByType(t.value); + return ( +
+ + +
+ ); + })}
) : usersQuery.data && usersQuery.data.members.length > 0 ? ( -
+
- User ID + User + Email Role Added {canManage && Actions} @@ -476,8 +643,11 @@ function UsersTab({ partitionName }: { partitionName: string }) { {usersQuery.data.members.map((user) => ( - - {user.user_id} + + + + + {canManage ? ( @@ -508,13 +678,14 @@ function UsersTab({ partitionName }: { partitionName: string }) { removeMutation.mutate(user.user_id)} > @@ -532,26 +703,70 @@ function UsersTab({ partitionName }: { partitionName: string }) {

)} - + - Add User to Partition + Add Users to Partition - Assign a user to this partition with a specific role. + Select users and assign the same partition role to each of them.
-
- - setUserId(e.target.value)} - /> -
+ { + setCandidateSearch(search); + setAddFailures([]); + }} + onRetry={() => { + void candidatesQuery.refetch(); + }} + hasMore={Boolean(candidatesQuery.hasNextPage)} + isLoadingMore={candidatesQuery.isFetchingNextPage} + onLoadMore={() => { + void candidatesQuery.fetchNextPage(); + }} + selectedCandidates={selectedCandidates} + onSelectionChange={(selection) => { + setSelectedCandidates(selection); + setAddFailures([]); + }} + /> + {addFailures.length > 0 && ( + + Some users were not added + +
    + {addFailures.map((failure) => ( +
  • + {candidateLabel(failure.candidate)} ({candidateSecondaryLabel(failure.candidate)}): + {" "} + {failure.message} +
    • + ))} +
+ + + )}
- { + setRole(value); + setAddFailures([]); + }} + > @@ -567,12 +782,20 @@ function UsersTab({ partitionName }: { partitionName: string }) { - diff --git a/ui/src/pages/admin/partitions/member-batch.test.ts b/ui/src/pages/admin/partitions/member-batch.test.ts new file mode 100644 index 000000000..8d12be939 --- /dev/null +++ b/ui/src/pages/admin/partitions/member-batch.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, it, vi } from "vitest"; + +import { ApiError } from "@/lib/api/client"; +import { addPartitionMembers } from "./member-batch"; + +const candidates = [ + { user_id: 2, display_name: "Sam Lee", email: "sam@example.com" }, + { user_id: 3, display_name: "Alex Morgan", email: "alex@example.com" }, + { user_id: 4, display_name: null, email: null }, +]; + +describe("addPartitionMembers", () => { + it("preserves the server reason for each failed user", async () => { + const addMember = vi.fn() + .mockResolvedValueOnce(undefined) + .mockRejectedValueOnce(new Error("User is already a member")) + .mockRejectedValueOnce(new Error("Account is inactive")); + + const result = await addPartitionMembers({ + partitionName: "legal", + candidates, + role: "viewer", + addMember, + }); + + expect(result.addedCandidates).toEqual([candidates[0]]); + expect(result.failures).toEqual([ + { + candidate: candidates[1], + message: "User is already a member", + }, + { + candidate: candidates[2], + message: "Account is inactive", + }, + ]); + }); + + it("stops sending requests after authorization is lost", async () => { + const addMember = vi.fn().mockRejectedValueOnce( + new ApiError(403, { detail: "Partition owner role required" }), + ); + + const result = await addPartitionMembers({ + partitionName: "legal", + candidates, + role: "editor", + addMember, + }); + + expect(addMember).toHaveBeenCalledOnce(); + expect(result.failures).toHaveLength(3); + expect(result.failures[0]).toMatchObject({ + candidate: candidates[0], + message: "Partition owner role required", + }); + expect(result.failures[1]).toMatchObject({ + candidate: candidates[1], + message: "Not attempted because permission to manage members was lost.", + }); + }); +}); diff --git a/ui/src/pages/admin/partitions/member-batch.ts b/ui/src/pages/admin/partitions/member-batch.ts new file mode 100644 index 000000000..5bc5c5f95 --- /dev/null +++ b/ui/src/pages/admin/partitions/member-batch.ts @@ -0,0 +1,50 @@ +import { addPartitionMember } from "@/lib/api/partitions"; +import type { PartitionMemberCandidate, PartitionRole } from "@/lib/api/partitions"; +import { ApiError } from "@/lib/api/client"; + +export interface MemberAddFailure { + candidate: PartitionMemberCandidate; + message: string; +} + +interface AddPartitionMembersOptions { + partitionName: string; + candidates: PartitionMemberCandidate[]; + role: PartitionRole; + addMember?: typeof addPartitionMember; +} + +export async function addPartitionMembers({ + partitionName, + candidates, + role, + addMember = addPartitionMember, +}: AddPartitionMembersOptions): Promise<{ + addedCandidates: PartitionMemberCandidate[]; + failures: MemberAddFailure[]; +}> { + const addedCandidates: PartitionMemberCandidate[] = []; + const failures: MemberAddFailure[] = []; + + for (const [index, candidate] of candidates.entries()) { + try { + await addMember(partitionName, candidate.user_id, role); + addedCandidates.push(candidate); + } catch (error) { + const message = error instanceof Error ? error.message : "Unknown error"; + failures.push({ candidate, message }); + + if (error instanceof ApiError && (error.status === 401 || error.status === 403)) { + for (const remainingCandidate of candidates.slice(index + 1)) { + failures.push({ + candidate: remainingCandidate, + message: "Not attempted because permission to manage members was lost.", + }); + } + break; + } + } + } + + return { addedCandidates, failures }; +} diff --git a/ui/src/pages/admin/partitions/member-candidate.ts b/ui/src/pages/admin/partitions/member-candidate.ts new file mode 100644 index 000000000..eeb541e64 --- /dev/null +++ b/ui/src/pages/admin/partitions/member-candidate.ts @@ -0,0 +1,9 @@ +import type { PartitionMemberCandidate } from "@/lib/api/partitions"; + +export function candidateLabel(candidate: PartitionMemberCandidate): string { + return candidate.display_name?.trim() || "Unnamed user"; +} + +export function candidateSecondaryLabel(candidate: PartitionMemberCandidate): string { + return candidate.email?.trim() || `User ID ${candidate.user_id}`; +} diff --git a/ui/src/pages/admin/partitions/member-picker.test.tsx b/ui/src/pages/admin/partitions/member-picker.test.tsx new file mode 100644 index 000000000..4a6f76ed6 --- /dev/null +++ b/ui/src/pages/admin/partitions/member-picker.test.tsx @@ -0,0 +1,124 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import { MemberPicker } from "./member-picker"; +import type { PartitionMemberCandidate } from "@/lib/api/partitions"; + +const candidates = [ + { user_id: 2, display_name: "Sam Lee", email: "sam.lee@example.com" }, + { user_id: 3, display_name: "Sam Lee", email: "sam.lee@linagora.com" }, + { user_id: 4, display_name: null, email: null }, +]; + +function renderPicker( + overrides: Partial> = {}, +) { + const props: React.ComponentProps = { + candidates, + isLoading: false, + isInitialError: false, + isRefreshError: false, + isLoadMoreError: false, + search: "Sam", + searchReady: true, + onSearchChange: vi.fn(), + onRetry: vi.fn(), + hasMore: false, + isLoadingMore: false, + onLoadMore: vi.fn(), + selectedCandidates: [], + onSelectionChange: vi.fn(), + ...overrides, + }; + return { ...render(), props }; +} + +describe("MemberPicker", () => { + it("shows email addresses so duplicate names remain distinguishable", () => { + renderPicker(); + + expect(screen.getAllByText("Sam Lee")).toHaveLength(2); + expect(screen.getByText("sam.lee@example.com")).not.toBeNull(); + expect(screen.getByText("sam.lee@linagora.com")).not.toBeNull(); + expect(screen.getByText("Unnamed user")).not.toBeNull(); + expect(screen.getByText("User ID 4")).not.toBeNull(); + }); + + it("forwards search changes to the server-backed query", async () => { + const onSearchChange = vi.fn(); + const user = userEvent.setup(); + renderPicker({ search: "", searchReady: false, onSearchChange }); + + const search = screen.getByRole("textbox", { name: "Users" }); + await user.type(search, "3"); + + expect(onSearchChange).toHaveBeenCalledWith("3"); + }); + + it("returns selected identities when a candidate is checked", async () => { + const onSelectionChange = vi.fn(); + const user = userEvent.setup(); + renderPicker({ selectedCandidates: [candidates[0]], onSelectionChange }); + + await user.click(screen.getByRole("checkbox", { name: /select sam lee, sam.lee@linagora.com/i })); + + expect(onSelectionChange).toHaveBeenCalledWith([candidates[0], candidates[1]]); + }); + + it("loads the next server page on request", async () => { + const onLoadMore = vi.fn(); + const user = userEvent.setup(); + renderPicker({ hasMore: true, onLoadMore }); + + await user.click(screen.getByRole("button", { name: "Load more users" })); + + expect(onLoadMore).toHaveBeenCalledOnce(); + }); + + it("keeps selected identities visible when they are absent from current results", async () => { + const selected: PartitionMemberCandidate = { + user_id: 9, + display_name: "Alex Morgan", + email: "alex@example.com", + }; + const onSelectionChange = vi.fn(); + const user = userEvent.setup(); + renderPicker({ + candidates: [candidates[0]], + selectedCandidates: [selected], + onSelectionChange, + }); + + expect(screen.getByRole("region", { name: "Selected users" })).not.toBeNull(); + expect(screen.getByText("Alex Morgan")).not.toBeNull(); + await user.click(screen.getByRole("button", { name: /remove alex morgan, alex@example.com/i })); + + expect(onSelectionChange).toHaveBeenCalledWith([]); + }); + + it("asks for a targeted search before showing candidates", () => { + renderPicker({ search: "Sa", searchReady: false }); + + expect(screen.getByText("Enter at least 3 characters or an exact user ID.")).not.toBeNull(); + expect(screen.queryByText("Sam Lee")).toBeNull(); + }); + + it("keeps cached candidates visible after a refresh error", () => { + renderPicker({ isRefreshError: true }); + + expect(screen.getByText(/previous results remain available/i)).not.toBeNull(); + expect(screen.getAllByText("Sam Lee")).toHaveLength(2); + }); + + it("offers a page-specific retry without discarding loaded candidates", async () => { + const onLoadMore = vi.fn(); + const user = userEvent.setup(); + renderPicker({ isLoadMoreError: true, onLoadMore }); + + expect(screen.getByText(/retry without losing this page/i)).not.toBeNull(); + await user.click(screen.getByRole("button", { name: "Retry loading more" })); + + expect(onLoadMore).toHaveBeenCalledOnce(); + }); +}); diff --git a/ui/src/pages/admin/partitions/member-picker.tsx b/ui/src/pages/admin/partitions/member-picker.tsx new file mode 100644 index 000000000..e36c9f56e --- /dev/null +++ b/ui/src/pages/admin/partitions/member-picker.tsx @@ -0,0 +1,203 @@ +import { useMemo } from "react"; +import { Search, X } from "lucide-react"; + +import { Alert, AlertDescription } from "@/components/ui/alert"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Skeleton } from "@/components/ui/skeleton"; +import type { PartitionMemberCandidate } from "@/lib/api/partitions"; +import { candidateLabel, candidateSecondaryLabel } from "./member-candidate"; + +interface MemberPickerProps { + candidates: PartitionMemberCandidate[]; + isLoading: boolean; + isInitialError: boolean; + isRefreshError: boolean; + isLoadMoreError: boolean; + search: string; + searchReady: boolean; + onSearchChange: (search: string) => void; + onRetry: () => void; + hasMore: boolean; + isLoadingMore: boolean; + onLoadMore: () => void; + selectedCandidates: PartitionMemberCandidate[]; + onSelectionChange: (candidates: PartitionMemberCandidate[]) => void; +} + +export function MemberPicker({ + candidates, + isLoading, + isInitialError, + isRefreshError, + isLoadMoreError, + search, + searchReady, + onSearchChange, + onRetry, + hasMore, + isLoadingMore, + onLoadMore, + selectedCandidates, + onSelectionChange, +}: MemberPickerProps) { + const selected = useMemo( + () => new Set(selectedCandidates.map((candidate) => candidate.user_id)), + [selectedCandidates], + ); + + const toggleCandidate = (candidate: PartitionMemberCandidate, checked: boolean) => { + if (checked) { + onSelectionChange( + selected.has(candidate.user_id) ? selectedCandidates : [...selectedCandidates, candidate], + ); + } else { + onSelectionChange( + selectedCandidates.filter((selectedCandidate) => selectedCandidate.user_id !== candidate.user_id), + ); + } + }; + + return ( +
+
+ + + {selectedCandidates.length} selected + +
+
+
+ + {selectedCandidates.length > 0 && ( +
+

Selected users

+ {selectedCandidates.map((candidate) => { + const label = candidateLabel(candidate); + const secondaryLabel = candidateSecondaryLabel(candidate); + return ( +
+ + {label} ({secondaryLabel}) + + +
+ ); + })} +
+ )} + + {isRefreshError && ( + + + Results could not be refreshed. The previous results remain available. + + + + )} + +
+ {!searchReady ? ( +

+ Enter at least 3 characters or an exact user ID. +

+ ) : isLoading ? ( +
+ {Array.from({ length: 3 }).map((_, index) => ( + + ))} +
+ ) : isInitialError ? ( +
+

Users could not be loaded.

+ +
+ ) : candidates.length === 0 ? ( +

+ No available users match this search. +

+ ) : ( + <> +
+ {candidates.map((candidate) => { + const label = candidateLabel(candidate); + const secondaryLabel = candidateSecondaryLabel(candidate); + const checkboxId = `partition-member-${candidate.user_id}`; + return ( + + ); + })} +
+ {(hasMore || isLoadMoreError) && ( +
+ {isLoadMoreError && ( +

+ More users could not be loaded. You can retry without losing this page. +

+ )} + +
+ )} + + )} +
+
+ ); +} diff --git a/ui/src/pages/admin/partitions/partition-member-identity.test.tsx b/ui/src/pages/admin/partitions/partition-member-identity.test.tsx new file mode 100644 index 000000000..2abe9928c --- /dev/null +++ b/ui/src/pages/admin/partitions/partition-member-identity.test.tsx @@ -0,0 +1,54 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; +import type { PartitionMember } from "@/lib/api/partitions"; +import { + PartitionMemberEmail, + PartitionMemberIdentity, +} from "./partition-member-identity"; +import { describePartitionMember } from "./partition-member"; + +function member(overrides: Partial = {}): PartitionMember { + return { + user_id: 9, + display_name: "Alice", + email: "alice@example.com", + role: "viewer", + added_at: null, + ...overrides, + }; +} + +describe("PartitionMemberIdentity", () => { + it("shows the display name and stable user ID", () => { + render(); + + expect(screen.getByText("Alice")).not.toBeNull(); + expect(screen.getByText("User ID 9")).not.toBeNull(); + }); + + it("retains a useful identity when the display name is missing", () => { + render(); + + expect(screen.getByText("User 9")).not.toBeNull(); + expect(screen.getByText("User ID 9")).not.toBeNull(); + }); + + it("shows email explicitly with a clear missing-value state", () => { + const { rerender } = render(); + + expect(screen.getByText("alice@example.com")).not.toBeNull(); + + rerender(); + + expect(screen.getByText("Not available")).not.toBeNull(); + }); + + it("describes a member unambiguously in destructive actions", () => { + expect(describePartitionMember(member())).toBe( + "Alice, alice@example.com (user ID 9)", + ); + expect(describePartitionMember(member({ display_name: null, email: null }))).toBe( + "user ID 9", + ); + }); +}); diff --git a/ui/src/pages/admin/partitions/partition-member-identity.tsx b/ui/src/pages/admin/partitions/partition-member-identity.tsx new file mode 100644 index 000000000..bd5a8d714 --- /dev/null +++ b/ui/src/pages/admin/partitions/partition-member-identity.tsx @@ -0,0 +1,26 @@ +import type { PartitionMember } from "@/lib/api/partitions"; + +export function PartitionMemberIdentity({ member }: { member: PartitionMember }) { + const primaryIdentity = member.display_name || `User ${member.user_id}`; + + return ( +
+
+ {primaryIdentity} +
+
User ID {member.user_id}
+
+ ); +} + +export function PartitionMemberEmail({ member }: { member: PartitionMember }) { + if (!member.email) { + return Not available; + } + + return ( + + {member.email} + + ); +} diff --git a/ui/src/pages/admin/partitions/partition-member.ts b/ui/src/pages/admin/partitions/partition-member.ts new file mode 100644 index 000000000..05223240e --- /dev/null +++ b/ui/src/pages/admin/partitions/partition-member.ts @@ -0,0 +1,8 @@ +import type { PartitionMember } from "@/lib/api/partitions"; + +export function describePartitionMember(member: PartitionMember): string { + const knownIdentity = [member.display_name, member.email].filter(Boolean).join(", "); + return knownIdentity + ? `${knownIdentity} (user ID ${member.user_id})` + : `user ID ${member.user_id}`; +} diff --git a/ui/src/pages/admin/presets.test.tsx b/ui/src/pages/admin/presets.test.tsx index 75f1a7260..f8bc5063b 100644 --- a/ui/src/pages/admin/presets.test.tsx +++ b/ui/src/pages/admin/presets.test.tsx @@ -26,7 +26,7 @@ vi.mock("@/lib/api/presets", async () => { }); vi.mock("@/lib/api/prompts", () => ({ - listPrompts: vi.fn().mockResolvedValue({ prompts: [] }), + listPrompts: vi.fn().mockResolvedValue([]), })); vi.mock("@/lib/api/models", () => ({ diff --git a/ui/src/pages/admin/presets.tsx b/ui/src/pages/admin/presets.tsx index f6c5de3cd..d92b8b375 100644 --- a/ui/src/pages/admin/presets.tsx +++ b/ui/src/pages/admin/presets.tsx @@ -10,7 +10,7 @@ import { getPresetOptions, } from "@/lib/api/presets"; import type { PresetResponse, PresetType } from "@/lib/api/presets"; -import { listPrompts } from "@/lib/api/prompts"; +import { listAllPrompts } from "@/lib/api/prompts"; import type { PromptResponse } from "@/lib/api/prompts"; import { listModelEndpoints, pickDefaultEndpoint } from "@/lib/api/models"; import { PageHeader } from "@/components/shared/page-header"; @@ -39,6 +39,12 @@ import { Separator } from "@/components/ui/separator"; import { Badge } from "@/components/ui/badge"; import { Skeleton } from "@/components/ui/skeleton"; import { formatDate, intOr, numOr } from "@/lib/utils"; +import { + PROMPT_DEFAULT_OPTION, + promptOptionToName, + promptOptionValue, + promptSelectValue, +} from "@/lib/prompt-meta"; import { type Config, configGet, @@ -387,9 +393,9 @@ function IndexationPresetForm({ onModelChange={(v) => set("vlm", v)} models={vlms} promptLabel="Caption prompt" - promptValue={configGet(config, "vlm_caption_prompt_name", "")} - onPromptChange={(v) => set("vlm_caption_prompt_name", v || null)} - prompts={promptsByType("vlm_caption")} + promptValue={configGet(config, "image_captioning_prompt_name", "")} + onPromptChange={(v) => set("image_captioning_prompt_name", v || null)} + prompts={promptsByType("image_captioning")} /> set("contextualization_prompt_name", v || null)} - prompts={promptsByType("contextualization")} + prompts={promptsByType("chunk_contextualizer")} /> set("topic_tagging_llm", v)} models={llms} + promptLabel="Prompt" + promptValue={configGet(config, "topic_tagging_prompt_name", "")} + onPromptChange={(v) => set("topic_tagging_prompt_name", v || null)} + prompts={promptsByType("topic_tagger")} numberLabel="Max tags" numberValue={configGet(config, "max_topic_tags", 7)} onNumberChange={(v) => set("max_topic_tags", v)} @@ -512,24 +522,21 @@ function FeatureToggle({ }) { const handleToggle = (on: boolean) => { onToggle(on); - if (on) { - // Auto-select when only one option available - if (models.length === 1 && !modelValue) onModelChange(models[0]); - if (prompts?.length === 1 && onPromptChange && !promptValue) onPromptChange(prompts[0].name); - } + // Auto-select the sole model, but never the sole prompt: an empty prompt + // value is a real choice ("use the type's global default"), and after + // seeding a type usually has exactly one prompt — the default itself. + // Auto-selecting it would pin that *name* into the preset just by opening + // and saving, so promoting a different global default later would silently + // stop affecting this preset. Models have no such fallback, so they keep it. + if (on && models.length === 1 && !modelValue) onModelChange(models[0]); }; - // Auto-select if toggled on and model/prompt list resolves to single item later + // Same for a list that resolves to a single item after the query settles. useEffect(() => { if (!enabled) return; if (models.length === 1 && !modelValue) onModelChange(models[0]); }, [enabled, models, modelValue, onModelChange]); - useEffect(() => { - if (!enabled || !prompts || !onPromptChange) return; - if (prompts.length === 1 && !promptValue) onPromptChange(prompts[0].name); - }, [enabled, prompts, promptValue, onPromptChange]); - return (
@@ -575,17 +582,17 @@ function FeatureToggle({
onChange(promptOptionToName(v))} + > + + + + + Use default + {prompts.map((p) => ( + + {p.name}{p.is_default ? " (default)" : ""} + + ))} + + +
+ ); +} + /* ---------- Retrieval form ---------- */ function RetrievalPresetForm({ @@ -606,16 +652,19 @@ function RetrievalPresetForm({ retrievalPipelines, rerankers, llms, + prompts, }: { config: Config; onChange: (c: Config) => void; retrievalPipelines: string[]; rerankers: string[]; llms: string[]; + prompts: PromptResponse[]; }) { const set = (key: string, value: unknown) => onChange(configSet(config, key, value)); const pipelineType: string = configGet(config, "type", "single"); const [advancedOpen, setAdvancedOpen] = useState(false); + const promptsByType = (type: string) => prompts.filter((p) => p.prompt_type === type); return (
@@ -633,6 +682,12 @@ function RetrievalPresetForm({
+ set("query_contextualizer_prompt_name", v || null)} + /> {pipelineType !== "single" && (
@@ -654,6 +709,22 @@ function RetrievalPresetForm({
)} + {pipelineType === "hyde" && ( + set("hyde_prompt_name", v || null)} + /> + )} + {pipelineType === "multiQuery" && ( + set("multi_query_prompt_name", v || null)} + /> + )}
listPrompts({ limit: 200 }), - enabled: open && presetType === "indexation", + queryFn: () => listAllPrompts(), + enabled: open, }); - const allPrompts = promptData?.prompts ?? []; + const allPrompts = promptData ?? []; const handleSubmit = (e: React.FormEvent) => { e.preventDefault(); @@ -901,6 +972,7 @@ function PresetDialog({ retrievalPipelines={options?.retrieval_types ?? []} rerankers={rerankers} llms={llms} + prompts={allPrompts} /> )} diff --git a/ui/src/pages/admin/prompts.tsx b/ui/src/pages/admin/prompts.tsx new file mode 100644 index 000000000..e0219f24e --- /dev/null +++ b/ui/src/pages/admin/prompts.tsx @@ -0,0 +1,636 @@ +import { useState, useEffect, useRef, useMemo } from "react"; +import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query"; +import { toast } from "sonner"; +import { + Plus, + Trash2, + Pencil, + Star, + Code2, + Eye, + AlertTriangle, + Circle, + Users, +} from "lucide-react"; +import { + listAllPrompts, + createPrompt, + updatePrompt, + deletePrompt, + setPromptDefault, +} from "@/lib/api/prompts"; +import type { PromptResponse, PromptType } from "@/lib/api/prompts"; +import { + PROMPT_GROUPS, + PROMPT_TYPES, + promptTypeLabel, + PROMPT_TYPE_VARIABLES, + validatePlaceholders, + renderPreview, +} from "@/lib/prompt-meta"; +import { PageHeader } from "@/components/shared/page-header"; +import { ConfirmDialog } from "@/components/shared/confirm-dialog"; +import { Button } from "@/components/ui/button"; +import { Badge } from "@/components/ui/badge"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Textarea } from "@/components/ui/textarea"; +import { Skeleton } from "@/components/ui/skeleton"; +import { Card, CardContent, CardHeader } from "@/components/ui/card"; +import { + Sheet, + SheetContent, + SheetHeader, + SheetTitle, + SheetDescription, + SheetFooter, +} from "@/components/ui/sheet"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { formatDate } from "@/lib/utils"; + +// Filter chips: "all" plus one per concern group. +const CONCERN_FILTERS = ["all", ...PROMPT_GROUPS.map((g) => g.name)] as const; +/** Types this page manages — everything grouped under a concern. */ +const MANAGED_PROMPT_TYPES = new Set(PROMPT_TYPES.map((t) => t.value as string)); + +export default function PromptsPage() { + const queryClient = useQueryClient(); + const [concern, setConcern] = useState<(typeof CONCERN_FILTERS)[number]>("all"); + const [editorOpen, setEditorOpen] = useState(false); + const [editing, setEditing] = useState(null); + + const { data, isLoading, isError, error, refetch, isFetching } = useQuery({ + queryKey: ["prompts-library"], + queryFn: () => listAllPrompts(), + }); + const prompts = data ?? []; + // An API or database failure must never render as "there are no prompts": + // that reads as a successful empty library and invites an admin to recreate + // prompts that already exist. Cached data stays visible across a failed + // refetch, with the banner explaining that what's shown may be stale. + const loadFailed = isError && data === undefined; + + const setDefaultMut = useMutation({ + mutationFn: (id: string) => setPromptDefault(id), + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: ["prompts-library"] }); + // The preset editors read their picker options from their own query, so + // a promotion here must refresh them too or they keep showing the old + // default badge. + queryClient.invalidateQueries({ queryKey: ["prompts-for-presets"] }); + toast.success("Default prompt updated"); + }, + onError: (e) => toast.error(e.message), + }); + + const deleteMut = useMutation({ + mutationFn: (id: string) => deletePrompt(id), + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: ["prompts-library"] }); + queryClient.invalidateQueries({ queryKey: ["prompts-for-presets"] }); + toast.success("Prompt deleted"); + }, + onError: (e) => toast.error(e.message), + }); + + const openCreate = () => { + setEditing(null); + setEditorOpen(true); + }; + const openEdit = (p: PromptResponse) => { + setEditing(p); + setEditorOpen(true); + }; + + const visibleGroups = PROMPT_GROUPS.filter((g) => concern === "all" || g.name === concern); + // The API can return types this page deliberately doesn't surface yet (today: + // spoken_style_answer, driven by a chat metadata flag rather than by anything + // configurable here). Only the grouped types render, so count those too — + // otherwise the header advertises more prompts than there are cards. + const managed = prompts.filter((p) => MANAGED_PROMPT_TYPES.has(p.prompt_type)); + const customCount = managed.filter((p) => !p.is_default).length; + + return ( +
+ + New Prompt + + } + /> + +
+
+ {CONCERN_FILTERS.map((c) => ( + + ))} +
+ {!isLoading && !loadFailed && ( + + {managed.length} prompt{managed.length === 1 ? "" : "s"} · {customCount} custom + + )} +
+ + {isError && ( +
+ + {loadFailed ? "Could not load the prompt library." : "Could not refresh the prompt library; showing the last known data."}{" "} + {(error as Error)?.message} + + +
+ )} + + {isLoading ? ( +
+ {[1, 2, 3, 4, 5, 6].map((i) => ( + + ))} +
+ ) : loadFailed ? null : ( +
+ {visibleGroups.map((group) => { + const groupTypes = new Set(group.types.map((t) => t.value)); + const groupPrompts = prompts.filter((p) => groupTypes.has(p.prompt_type)); + return ( +
+
+

+ {group.name} +

+

{group.description}

+
+ {groupPrompts.length === 0 ? ( +

+ No {group.name.toLowerCase()} prompts yet. +

+ ) : ( +
+ {groupPrompts.map((prompt) => ( + openEdit(prompt)} + onSetDefault={() => setDefaultMut.mutate(prompt.id)} + onDelete={() => deleteMut.mutate(prompt.id)} + /> + ))} +
+ )} +
+ ); + })} +
+ )} + + +
+ ); +} + +/* ---------- Prompt card ---------- */ + +function PromptCard({ + prompt, + onEdit, + onSetDefault, + onDelete, +}: { + prompt: PromptResponse; + onEdit: () => void; + onSetDefault: () => void; + onDelete: () => void; +}) { + const used = prompt.used_by; + return ( + +
+ 0 + ? "text-xs bg-amber-50 text-amber-700 border-amber-200 dark:bg-amber-950/30 dark:text-amber-100 dark:border-amber-900/60" + : "text-xs bg-muted text-muted-foreground border-transparent" + } + > + + {used > 0 ? `${used} partition${used === 1 ? "" : "s"}` : "Unused"} + +
+ +

{prompt.prompt_type}

+
+ {prompt.name || "Untitled"} + {prompt.is_default && ( + Default + )} +
+ + +

+ {prompt.content.slice(0, 160) || "(empty)"} +

+
Updated {formatDate(prompt.updated_at)}
+
+ {!prompt.is_default && ( + + )} + + {prompt.is_default ? ( + + ) : ( + 0 + ? `"${prompt.name}" is selected by ${used} partition${used === 1 ? "" : "s"}. They will fall back to the default. Delete anyway?` + : `This will permanently delete "${prompt.name}".` + } + onConfirm={onDelete} + > + + + )} +
+ + + ); +} + +/* ---------- Editor drawer ---------- */ + +function PromptEditorSheet({ + open, + onOpenChange, + editing, +}: { + open: boolean; + onOpenChange: (v: boolean) => void; + editing: PromptResponse | null; +}) { + const queryClient = useQueryClient(); + const [promptType, setPromptType] = useState("sys_prompt"); + const [name, setName] = useState(""); + const [content, setContent] = useState(""); + + // Sync the form to the editing target each time the drawer opens (reset for + // "create"). Controlled-drawer reset pattern. + /* eslint-disable react-hooks/set-state-in-effect */ + useEffect(() => { + if (!open) return; + if (editing) { + setPromptType(editing.prompt_type); + setName(editing.name); + setContent(editing.content); + } else { + setPromptType("sys_prompt"); + setName(""); + setContent(""); + } + }, [open, editing]); + /* eslint-enable react-hooks/set-state-in-effect */ + + const createMut = useMutation({ + mutationFn: createPrompt, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: ["prompts-library"] }); + queryClient.invalidateQueries({ queryKey: ["prompts-for-presets"] }); + toast.success("Prompt created"); + onOpenChange(false); + }, + onError: (e) => toast.error(e.message), + }); + + const updateMut = useMutation({ + mutationFn: ({ id, ...data }: { id: string; name: string; content: string }) => + updatePrompt(id, data), + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: ["prompts-library"] }); + queryClient.invalidateQueries({ queryKey: ["prompts-for-presets"] }); + toast.success("Prompt updated"); + onOpenChange(false); + }, + onError: (e) => toast.error(e.message), + }); + + const loading = createMut.isPending || updateMut.isPending; + const effectiveType = editing?.prompt_type ?? promptType; + const templateCheck = validatePlaceholders(content, effectiveType); + // Presets and partitions reference a prompt by *name*, so a rename silently + // orphans every selection pointing at the old one — they fall back to the + // global default. Warn before that happens instead of letting the drawer's + // "changes apply everywhere" promise quietly become false. + const isRename = !!editing && name.trim() !== editing.name; + const renameBreaksRefs = isRename && (editing?.used_by ?? 0) > 0; + + const handleSubmit = (e: React.FormEvent) => { + e.preventDefault(); + if (!name.trim()) { + toast.error("Name is required"); + return; + } + if (!content.trim()) { + toast.error("Content is required"); + return; + } + // Mirror the API's template rules here so a malformed template is caught + // before the request instead of coming back as a 422. + if (templateCheck.malformed) { + toast.error(templateCheck.error ?? "The template has unbalanced braces."); + return; + } + if (templateCheck.unknown.length > 0) { + toast.error(`Unknown variable(s): ${templateCheck.unknown.map((v) => `{${v}}`).join(", ")}`); + return; + } + if ( + renameBreaksRefs && + !window.confirm( + `"${editing?.name}" is selected by ${editing?.used_by} partition(s). ` + + `Renaming it to "${name.trim()}" drops those selections — they fall back to the ` + + `default ${promptTypeLabel(effectiveType).toLowerCase()}. Rename anyway?`, + ) + ) { + return; + } + if (editing) { + updateMut.mutate({ id: editing.id, name, content }); + } else { + createMut.mutate({ prompt_type: promptType, name, content }); + } + }; + + return ( + + + + {editing ? "Edit prompt" : "New prompt"} + + Changes apply to every preset and partition that selects this prompt. + + + +
+
+ + setName(e.target.value)} + required + /> + {renameBreaksRefs && ( +

+ Selected by {editing?.used_by} partition(s) — renaming drops those selections. +

+ )} +
+
+ + {editing ? ( + + ) : ( + + )} +
+
+ + + + {editing && ( +

+ + {editing.used_by > 0 + ? `Selected by ${editing.used_by} partition${editing.used_by === 1 ? "" : "s"}.` + : "Not selected by any partition yet."} +

+ )} + + + + + + + + + ); +} + +/* ---------- Template editor (edit / preview + {var} helpers) ---------- */ + +function PromptTemplateEditor({ + promptType, + value, + onChange, +}: { + promptType: string; + value: string; + onChange: (v: string) => void; +}) { + const [tab, setTab] = useState<"edit" | "preview">("edit"); + const textareaRef = useRef(null); + const variables = PROMPT_TYPE_VARIABLES[promptType] ?? []; + const hasVariables = variables.length > 0; + + const validation = useMemo(() => validatePlaceholders(value, promptType), [value, promptType]); + const preview = useMemo(() => renderPreview(value, promptType), [value, promptType]); + + const insertVariable = (varName: string) => { + const ta = textareaRef.current; + if (!ta) return; + const start = ta.selectionStart; + const end = ta.selectionEnd; + const placeholder = `{${varName}}`; + onChange(value.slice(0, start) + placeholder + value.slice(end)); + requestAnimationFrame(() => { + ta.focus(); + ta.setSelectionRange(start + placeholder.length, start + placeholder.length); + }); + }; + + return ( +
+
+ +
+ + +
+
+ + {hasVariables && tab === "edit" && ( +
+ Variables: + {variables.map((v) => ( + + ))} +
+ )} + + {tab === "edit" ? ( + <> +