diff --git a/.env.example b/.env.example index 7239a69a6..a502bcf9f 100644 --- a/.env.example +++ b/.env.example @@ -1,6 +1,79 @@ +# ============================================================================ +# LAMB root Docker Compose environment (docker-compose.yaml) +# ============================================================================ +# Copy this file to `.env` in the repository root (next to docker-compose.yaml). +# Docker Compose loads it automatically. Per-service application config still +# lives in backend/.env, lamb-kb-server/backend/.env, etc. + +# Absolute path to this checkout on the host. REQUIRED — the dev compose +# bind-mounts it into every container, so it must be correct. LAMB_PROJECT_PATH=/opt/lamb + +# Open WebUI integration OWI_BASE_URL=http://openwebui:8080 OWI_PUBLIC_BASE_URL=http://localhost:8080 OWI_PATH=/opt/lamb/open-webui/backend/data + +# LAMB data + legacy KB server (port 9090) LAMB_DB_PATH=/opt/lamb -LAMB_KB_SERVER=http://kb:9090 \ No newline at end of file +LAMB_KB_SERVER=http://kb:9090 + +# New KB server (Knowledge Stores, port 9092) — backs the KG-RAG feature. +# The backend reaches it at these values; the token must match the +# LAMB_API_TOKEN the kb-v2 service starts with (compose default: 0p3n-w3bu!). +LAMB_KB_SERVER_V2=http://kb-v2:9092 +LAMB_KB_SERVER_V2_TOKEN=0p3n-w3bu! + +# Library Manager (port 9091) bearer token — must match on backend + service. +LIBRARY_MANAGER_TOKEN=change-me + +# ============================================================================ +# KG-RAG / Semantic Graph (optional feature on this branch) +# ============================================================================ +# KG-RAG adds a Neo4j-backed semantic graph on top of vector retrieval in the +# new KB server (kb-v2). It is OFF by default and fully isolated: with +# KG_RAG_ENABLED=false the /graph + /benchmarks endpoints are never mounted and +# every existing simple / hierarchical RAG flow behaves exactly as before. +# +# Neo4j only starts under the `kg-rag` compose profile, so the base stack stays +# unchanged unless you explicitly opt in: +# +# docker compose --profile kg-rag up -d +# +# See Documentation/kg-rag-deployment.md for the full walkthrough. + +# Master switch. Set to true (and use --profile kg-rag) to enable KG-RAG. +KG_RAG_ENABLED=false + +# Run LLM concept extraction + a Neo4j write on every ingestion job. When true, +# single-document ingestion grows from sub-second to several seconds. Set false +# to keep ingestion fast and back-populate the graph later via the migrate API. +KG_RAG_INDEX_ON_INGEST=false + +# OpenAI key used for concept/entity extraction. Per-request keys (threaded +# from the LAMB org config) are preferred; this is the ingestion-time fallback. +KG_RAG_OPENAI_API_KEY= + +# Models used for graph reasoning and concept extraction. +KG_RAG_CHAT_MODEL=gpt-4o-mini +# Leave empty to fall back to KG_RAG_CHAT_MODEL / gpt-4o-mini. +KG_RAG_EXTRACTION_MODEL= + +# Neo4j connection. The neo4j service is created from these values too, so the +# password set here is the password the database is initialised with. +KG_RAG_NEO4J_URI=bolt://neo4j:7687 +KG_RAG_NEO4J_USER=neo4j +KG_RAG_NEO4J_PASSWORD=lamb-kg-rag-password + +# Graph traversal tuning (clamped server-side: depth 1-4, limit_factor 1-20). +KG_RAG_GRAPH_DEPTH=2 +KG_RAG_LIMIT_FACTOR=4 +KG_RAG_EXTRACTION_MAX_WORKERS=4 + +# kb-v2 query-plugin governance (ENABLE | ADVANCED | DISABLE). +PLUGIN_KG_RAG_QUERY=ADVANCED + +# Neo4j JVM sizing (optional — raise for larger graphs). +# NEO4J_HEAP_INITIAL_SIZE=512m +# NEO4J_HEAP_MAX_SIZE=1G +# NEO4J_PAGECACHE_SIZE=512m \ No newline at end of file diff --git a/.env.next.example b/.env.next.example index 9ab134942..15906efe6 100644 --- a/.env.next.example +++ b/.env.next.example @@ -20,8 +20,10 @@ OWI_PATH=/data/openwebui OWI_BASE_URL=http://openwebui:8080 # Auth/secrets required by current backend config -LAMB_BEARER_TOKEN=change-me -SIGNUP_SECRET_KEY=change-me +# Set each to a strong, unique, random value (e.g. `openssl rand -hex 32`). +# Do NOT deploy with these placeholders. (#415) +LAMB_BEARER_TOKEN=CHANGE_ME_strong_random_secret +SIGNUP_SECRET_KEY=CHANGE_ME_strong_random_secret # OpenAI provider required by current backend config OPENAI_BASE_URL=https://api.openai.com/v1 @@ -30,7 +32,7 @@ OPENAI_MODEL=gpt-4o-mini # OpenWebUI bootstrap admin required by current backend config OWI_ADMIN_NAME=Admin User OWI_ADMIN_EMAIL=admin@owi.com -OWI_ADMIN_PASSWORD=admin +OWI_ADMIN_PASSWORD=CHANGE_ME_strong_admin_password # ============================================================================ # OPTIONAL VARIABLES (uncomment to override defaults) @@ -108,3 +110,41 @@ OWI_ADMIN_PASSWORD=admin # CADDY_EMAIL=admin@yourdomain.com # LAMB_PUBLIC_HOST=lamb.yourdomain.com # OWI_PUBLIC_HOST=owi.lamb.yourdomain.com + +# ============================================================================ +# KG-RAG / Semantic Graph (optional) +# ============================================================================ +# Enable Graph RAG / KG-RAG in the new KB server. Requires the ``kg-rag`` +# Compose profile (starts Neo4j): +# docker compose -f docker-compose.next.yaml --profile kg-rag up -d +# +# When KG_RAG_ENABLED=false (default), graph endpoints are not registered +# and existing simple / hierarchical / parent-child RAG flows are unchanged. +# +# Per-collection opt-in: even with the flag on, a Knowledge Store only +# uses graph features when its ``graph_enabled`` field is true (set in +# create/migrate flow). +# +# OpenAI credentials for concept extraction can be sent per-request +# (preferred) via the X-OpenAI-Api-Key header on /graph endpoints. The +# server-level KG_RAG_OPENAI_API_KEY is a fallback used during ingestion +# when no per-request key has been threaded through. + +# KG_RAG_ENABLED=false +# Note on KG_RAG_INDEX_ON_INGEST: when enabled, every ingestion job ALSO +# runs LLM concept extraction (1 OpenAI call per parent text) and a Neo4j +# write. Expect single-document ingestion latency to grow from sub-second +# to multiple seconds. Set to false to keep ingestion fast and run +# `/graph/.../migrate` manually instead. +# KG_RAG_INDEX_ON_INGEST=true +# KG_RAG_OPENAI_API_KEY= +# KG_RAG_CHAT_MODEL=gpt-4o-mini +# KG_RAG_EXTRACTION_MODEL=gpt-5-nano +# KG_RAG_NEO4J_URI=bolt://neo4j:7687 +# KG_RAG_NEO4J_USER=neo4j +# KG_RAG_NEO4J_PASSWORD=change-this-password +# KG_RAG_GRAPH_DEPTH=2 +# KG_RAG_LIMIT_FACTOR=4 +# KG_RAG_EXTRACTION_MAX_WORKERS=4 +# Per-request OpenAI timeout for extraction calls. Default 60s. +# KG_RAG_OPENAI_TIMEOUT_SECONDS=60 diff --git a/.gitignore b/.gitignore index d87e37816..1a1db64df 100644 --- a/.gitignore +++ b/.gitignore @@ -202,7 +202,6 @@ docs/ frontend-architecture-diagrams.md opencode.json package-lock.json - .playwright-mcp aac_logs/ @@ -211,4 +210,7 @@ aac_logs/ # Worktrees .worktrees/ .agents/ -skills-lock.json \ No newline at end of file +skills-lock.json +test-results/ +testing/load/document-cache-test/ +testing/load/kv-cache-classroom/ diff --git a/CLAUDE.md b/CLAUDE.md index f29d007c6..965c6614c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -64,12 +64,14 @@ pytest -m "not slow" # skip slow tests # Setup (first time only) python3 -m venv .venv && source .venv/bin/activate && pip install -e ".[all,dev]" -# Run tests -pytest tests/ -v # all 52 tests -pytest tests/ --cov=backend # with coverage report +# Run tests (3 tiers: unit / integration / e2e — see tests/README.md) +./scripts/run_tests.sh # all tiers + the 95% line+branch coverage gate +pytest tests/unit/ -q # a single tier (unit | integration | e2e) +pytest -m "not slow" -q # skip subprocess/Docker tests +pytest tests/unit/ tests/integration/ --cov=backend --cov-branch --cov-fail-under=95 ruff check backend/ tests/ --select E,F,W,I,UP,SIM --target-version py311 # lint ``` -Tests use an in-process ASGI client (no server needed). YouTube transcript test uses a cached response in `tests/.yt_cache/` to avoid rate limits. +Unit tests call modules directly; integration tests use an in-process ASGI client against the real DB + worker; e2e tests hit a `uvicorn` subprocess over real HTTP (plus a Docker smoke test that skips without a daemon). External SDKs (markitdown, Firecrawl, PyMuPDF, OpenAI) are mocked at their import boundary; YouTube uses a cached response in `tests/.yt_cache/`. Combined `backend/` line+branch coverage is enforced at ≥95%. ### Applying backend env changes in Docker ```bash @@ -131,7 +133,7 @@ These terms are used consistently throughout the codebase and must not be confus - Every imported document is stored in a common structured format: `metadata.json` + `source_ref.json` + `original/` + `content/full.md` + `content/pages/` + `content/images/`. - Five import plugins (simple, markitdown, markitdown_plus, url, youtube), each governed by a tri-state env var (`DISABLE|SIMPLIFIED|ADVANCED`). - Internal service: no CORS, no published ports, non-root Docker container, single-instance file lock. -- 52 tests, 80% coverage. Full README at `library-manager/README.md`. +- ~710 tests across unit/integration/e2e tiers; ≥95% line+branch coverage gate (currently ~98%). Full README at `library-manager/README.md`; tier contract at `library-manager/tests/README.md`. **LAMB integration:** Creator Interface endpoints (`/creator/libraries/...`) validate ACL and proxy to the Library Manager. The `lamb-cli` commands (`lamb library ...`) are in `lamb-cli/src/lamb_cli/commands/library.py`. Svelte frontend at `/libraries` route with components in `frontend/svelte-app/src/lib/components/libraries/`. diff --git a/Documentation/kg-rag-deployment.md b/Documentation/kg-rag-deployment.md new file mode 100644 index 000000000..c6f06c111 --- /dev/null +++ b/Documentation/kg-rag-deployment.md @@ -0,0 +1,157 @@ +# KG-RAG / Semantic-Graph Deployment Guide + +This guide explains how to bring up the LAMB stack with the **KG-RAG (Knowledge-Graph RAG)** feature enabled — the Neo4j-backed semantic-graph layer that augments vector retrieval in the new KB server. + +KG-RAG is **optional and off by default**. With `KG_RAG_ENABLED=false` the graph endpoints are never mounted and every existing RAG flow (simple, hierarchical, single-file, context-aware) behaves exactly as before. You only get the extra services and behaviour when you explicitly opt in as described below. + +> Scope: this targets the development `docker-compose.yaml` at the repository root, which is the compose file that wires the new KB server (`kb-v2`, port 9092) and Neo4j. It bind-mounts the checkout into every container and installs dependencies on first start. + +--- + +## What KG-RAG adds to the stack + +| Service | Port | Role with KG-RAG | Started by | +|---|---|---|---| +| `kb-v2` (new KB server) | 9092 | Concept extraction, graph write, graph-augmented query | always (base stack) | +| `neo4j` | 7474 (HTTP), 7687 (Bolt) | Semantic graph + curation audit log | **only `--profile kg-rag`** | + +`kb-v2` always runs, but it only mounts the `/graph` + `/benchmarks` routers and talks to Neo4j when `KG_RAG_ENABLED=true`. Neo4j itself sits behind the `kg-rag` Compose profile, so the base `docker compose up` never starts it. + +--- + +## Prerequisites + +- Docker Engine 24+ and the Docker Compose v2 plugin (`docker compose version`). +- This repository checked out on the **`feat/kg-rag-new-kbserver-arch`** branch. +- An OpenAI-compatible API key if you want concept extraction to actually run (extraction is an LLM call). Without a key the stack still starts; ingestion just skips graph writes / logs an extraction error. +- ~2 GB free RAM for Neo4j (default heap 1 GB + page cache). + +--- + +## Step 1 — Create the root `.env` + +From the repository root: + +```bash +cp .env.example .env +``` + +Edit `.env` and set, at minimum: + +```bash +# REQUIRED — absolute path to THIS checkout on the host +LAMB_PROJECT_PATH=/absolute/path/to/lamb + +# Turn the feature on +KG_RAG_ENABLED=true + +# Pick a Neo4j password (this initialises the database on first boot) +KG_RAG_NEO4J_PASSWORD=lamb-kg-rag-password + +# Key used for concept extraction at ingest time (optional but recommended) +KG_RAG_OPENAI_API_KEY=sk-... +``` + +Everything else has working defaults (see the comments in `.env.example`). The KG-RAG block controls graph behaviour: + +| Variable | Default | Meaning | +|---|---|---| +| `KG_RAG_ENABLED` | `false` | Master switch. Must be `true` for any graph behaviour. | +| `KG_RAG_INDEX_ON_INGEST` | `false` | If `true`, run extraction + graph write on every ingestion job (slower). If `false`, ingest fast and back-populate later via the migrate API. | +| `KG_RAG_OPENAI_API_KEY` | empty | Fallback key for extraction (per-request keys from LAMB are preferred). | +| `KG_RAG_CHAT_MODEL` | `gpt-4o-mini` | Model for graph reasoning. | +| `KG_RAG_EXTRACTION_MODEL` | empty → chat model | Model for concept/entity extraction. | +| `KG_RAG_NEO4J_URI` | `bolt://neo4j:7687` | Neo4j Bolt URI (service name inside the compose network). | +| `KG_RAG_NEO4J_USER` / `KG_RAG_NEO4J_PASSWORD` | `neo4j` / `lamb-kg-rag-password` | Neo4j credentials; the DB is initialised with these. | +| `KG_RAG_GRAPH_DEPTH` | `2` | Graph traversal depth (clamped 1–4). `.env.example` ships `2`; note the kb-v2 service's bare compose fallback is `1`, so keep this line set rather than deleting it. | +| `KG_RAG_LIMIT_FACTOR` | `4` | Candidate-expansion factor (clamped 1–20). | + +The same file also sets `LAMB_KB_SERVER_V2_TOKEN` and `LIBRARY_MANAGER_TOKEN`; leave the defaults unless you change them on the services too. + +## Step 2 — Create `backend/.env` + +The backend container loads `backend/.env` directly (`env_file:`), so it must exist. From the **repository root**: + +```bash +cp backend/.env.example backend/.env +``` + +Set a real `OPENAI_API_KEY` (and any other provider keys) inside `backend/.env`. No KG-RAG-specific edits are needed here — the backend learns the new KB server address from `LAMB_KB_SERVER_V2` / `LAMB_KB_SERVER_V2_TOKEN`, which the compose file injects from the root `.env`. + +> You do **not** need to create a `.env` for the KB servers by hand. `kb-v2` (the KG-RAG server) takes all of its config from the compose `environment:` block — it never reads a `.env` file. The legacy `kb` service auto-creates `lamb-kb-server-stable/backend/.env` from its `.env.example` on first start. Either way, the KG-RAG settings come from compose environment variables sourced from the root `.env`. + +## Step 3 — Start the stack with the `kg-rag` profile + +```bash +docker compose --profile kg-rag up -d +``` + +The `--profile kg-rag` flag is what adds the `neo4j` service. Without it, `kb-v2` starts but has no graph database to talk to. + +> This assumes a clean host: ports 9099, 9092, 9090, 9091, 8080, 5173, 7474 and 7687 must be free. If a LAMB stack is already running, `up` will fail on a port conflict — bring the existing one down first (`docker compose --profile kg-rag down`) or start the new one with a distinct `COMPOSE_PROJECT_NAME` and remapped ports. + +First boot is slow: the containers `pip install` / `npm build` against the bind-mounted source. Watch progress with: + +```bash +docker compose --profile kg-rag ps +docker compose logs -f kb-v2 +``` + +You can also bring up just the feature-critical services to validate KG-RAG quickly without building the frontend/OpenWebUI: + +```bash +docker compose --profile kg-rag up -d neo4j kb-v2 +``` + +## Step 4 — Verify KG-RAG is live + +1. **Neo4j is healthy:** + ```bash + docker compose ps neo4j # STATUS shows "healthy" (ps lists running containers; no profile flag needed) + ``` + The Neo4j browser is at (log in with `KG_RAG_NEO4J_USER` / `KG_RAG_NEO4J_PASSWORD`). + +2. **kb-v2 mounted the graph routers** — the startup log prints this line only when the feature is on: + ```bash + docker compose logs kb-v2 | grep "KG-RAG enabled" + # → KG-RAG enabled: mounted /graph and /benchmarks routers + ``` + +3. **kb-v2 is healthy and the graph surface exists:** + ```bash + curl -s http://localhost:9092/health # {"status":"ok",...} + curl -s http://localhost:9092/openapi.json | grep -o '"/graph[^"]*"' | head + # → graph paths appear ONLY when KG_RAG_ENABLED=true + ``` + With the feature off, `/graph` paths are absent from the OpenAPI spec (callers get a clean 404). + +## Step 5 — Use it from LAMB + +1. Open the frontend at (or the Svelte dev server at ). +2. Create a **Knowledge Store** via the "Create Knowledge" wizard. In the Setup step, enable the **graph / `graph_enabled`** toggle and pick the extraction vendor/model. This per-store opt-in is required *in addition* to `KG_RAG_ENABLED` — the flag enables the capability, the toggle enables it for that store. Store setup is locked after creation. +3. Link Library items into the store to ingest them. If `KG_RAG_INDEX_ON_INGEST=true`, concepts are extracted and written to Neo4j during ingestion; otherwise run the migrate/back-populate action from the store's graph view. +4. Inspect and curate the graph in the store's **Graph** view (Sigma.js): approve/reject/rename/merge concepts; every edit is recorded as an audit `ChangeEvent`. +5. Attach the store to an assistant whose RAG processor is `knowledge_store_rag` to use graph-augmented retrieval at query time. + +## Turning KG-RAG off again + +```bash +docker compose --profile kg-rag down # stop everything incl. Neo4j +``` + +Set `KG_RAG_ENABLED=false` in `.env` and start normally (`docker compose up -d`, no profile). The graph routers disappear, Neo4j is not started, and all vector RAG flows continue unchanged. Neo4j data persists in the `neo4j-data` volume for next time. + +--- + +## Troubleshooting + +| Symptom | Cause / fix | +|---|---| +| `kb-v2` log has no "KG-RAG enabled" line | `KG_RAG_ENABLED` not `true` in `.env`, or container not recreated — run `docker compose up -d --force-recreate kb-v2`. | +| `/graph` calls return 404 | Feature flag off, or you hit the legacy `kb` (9090) instead of `kb-v2` (9092). | +| kb-v2 logs Neo4j connection errors | Neo4j not started (missing `--profile kg-rag`), wrong `KG_RAG_NEO4J_URI`, or password mismatch. Confirm `docker compose ps neo4j` is healthy. | +| Ingestion is suddenly slow | `KG_RAG_INDEX_ON_INGEST=true` runs an LLM call per parent text. Set it `false` and back-populate via the migrate API. | +| Extraction silently does nothing | No OpenAI key reached the server — set `KG_RAG_OPENAI_API_KEY` or thread a per-request key from the LAMB org config. | +| `.env` changes not picked up | Compose only re-reads env on container (re)create: `docker compose up -d --force-recreate ` (a plain `restart` is not enough). | + +See `lamb-kb-server/Documentation/` for the architecture decision records (ADR-KS-1…17) behind this design. diff --git a/README.md b/README.md index ccf9252a2..91d9b793c 100644 --- a/README.md +++ b/README.md @@ -89,6 +89,12 @@ For development or custom deployments: 📘 **[Complete Installation Guide](Documentation/installationguide.md)** - Step-by-step manual setup for all components +### Optional: KG-RAG (Semantic Graph) + +To bring up the stack with the Neo4j-backed Knowledge-Graph RAG feature enabled: + +📘 **[KG-RAG Deployment Guide](Documentation/kg-rag-deployment.md)** - Enable graph-augmented retrieval with the `kg-rag` Compose profile + ### Quick Overview LAMB requires four main services: diff --git a/backend/.coveragerc b/backend/.coveragerc new file mode 100644 index 000000000..7a8b1d12b --- /dev/null +++ b/backend/.coveragerc @@ -0,0 +1,6 @@ +[report] +exclude_lines = + pragma: no cover + if TYPE_CHECKING: + raise NotImplementedError + if __name__ == .__main__.: diff --git a/backend/.env.example b/backend/.env.example index 219a26f9a..ba70d58db 100644 --- a/backend/.env.example +++ b/backend/.env.example @@ -30,6 +30,13 @@ LAMB_BACKEND_HOST=http://backend:9099 #PIPELINES_BEARER_TOKEN=0p3n-w3bu! LAMB_BEARER_TOKEN=0p3n-w3bu! +# Secret used to HMAC-sign the public citation permalinks students click from +# chat. Per-org keys are derived from it, so rotating it revokes every signed +# link at once. MUST stay stable across restarts (changing it breaks links in +# existing chats). Defaults to LAMB_BEARER_TOKEN if unset; set a dedicated +# value in production. +#LAMB_PERMALINK_SIGNING_SECRET=change-me-to-a-long-random-string + # Signup Configuration SIGNUP_ENABLED=true diff --git a/backend/config.py b/backend/config.py index f8c0224fe..ca6971346 100644 --- a/backend/config.py +++ b/backend/config.py @@ -38,6 +38,16 @@ def _env_bool(name: str, default: bool) -> bool: lamb_token = lamb_token.strip() LAMB_BEARER_TOKEN = lamb_token + +# Secret used to HMAC-sign public citation permalinks (the signed "capability" +# URLs students click from chat). Per-org keys are derived from this master +# secret, so rotating it revokes every signed link at once. MUST be stable +# across restarts — if it changes, links in existing chats stop resolving. +# Falls back to LAMB_BEARER_TOKEN so the feature works out of the box, but a +# dedicated value is recommended in production. +LAMB_PERMALINK_SIGNING_SECRET = ( + os.getenv('LAMB_PERMALINK_SIGNING_SECRET') or lamb_token +) PIPELINES_BEARER_TOKEN = lamb_token # Legacy alias for backward compatibility PIPELINES_DIR = os.getenv("PIPELINES_DIR", "./lamb_assistants") @@ -109,16 +119,26 @@ def _env_bool(name: str, default: bool) -> bool: # LAMB Native Authentication # JWT secret for signing LAMB-issued tokens. Resolution order: # 1. LAMB_JWT_SECRET (explicit, for full OWI decoupling) -# 2. WEBUI_SECRET_KEY / WEBUI_JWT_SECRET_KEY env vars -# 3. OWI's hardcoded default ("t0p-s3cr3t") — matches open_webui/env.py -# This ensures zero-config compatibility: LAMB signs tokens with the same -# secret OWI uses, so existing OWI tokens decode seamlessly. +# 2. WEBUI_SECRET_KEY / WEBUI_JWT_SECRET_KEY (legacy OWI env vars — zero-config +# compat: LAMB signs with the same secret OWI uses, so OWI tokens decode) +# 3. OWI's historical default ("t0p-s3cr3t") as a last resort. +# When the explicit var is absent we fall back to the legacy OWI variables rather +# than failing. The literal default is insecure (the auth layer trusts the JWT +# `role` claim), so warn loudly and require an override in production (#412). LAMB_JWT_SECRET = ( os.getenv('LAMB_JWT_SECRET') or os.getenv('WEBUI_SECRET_KEY') or os.getenv('WEBUI_JWT_SECRET_KEY') - or 't0p-s3cr3t' ) +if not LAMB_JWT_SECRET: + import sys + LAMB_JWT_SECRET = 't0p-s3cr3t' + print( + "WARNING [config]: no LAMB_JWT_SECRET / WEBUI_SECRET_KEY / WEBUI_JWT_SECRET_KEY " + "set; falling back to the insecure default JWT secret. Set a strong secret " + "before production — the auth layer trusts the JWT role claim (#412).", + file=sys.stderr, + ) # OWI Admin Configuration OWI_ADMIN_NAME = os.getenv('OWI_ADMIN_NAME') diff --git a/backend/creator_interface/assistant_router.py b/backend/creator_interface/assistant_router.py index 50af23045..22c961605 100644 --- a/backend/creator_interface/assistant_router.py +++ b/backend/creator_interface/assistant_router.py @@ -27,6 +27,7 @@ import config from utils.name_sanitizer import sanitize_assistant_name_with_prefix from lamb.logging_config import get_logger +from creator_interface.metadata_validators import validate_update_plugin_metadata as _validate_metadata # Configuration # Use LAMB_BACKEND_HOST for internal server-to-server requests @@ -316,61 +317,10 @@ def sanitize_filename(filename: str) -> str: return filename[:100] if filename else "assistant_export" -REQUIRED_PLUGIN_METADATA_KEYS = ( - "prompt_processor", - "connector", - "llm", - "rag_processor", -) - - def validate_update_plugin_metadata( original_body: Dict[str, Any] ) -> Tuple[Optional[str], Optional[str]]: - """ - Validate assistant plugin metadata for updates. - - Updates must provide complete plugin metadata so the backend never replaces a - valid stored configuration with partial or blank data. - - Returns: - Tuple[Optional[str], Optional[str]]: (normalized_metadata_json, error_message) - """ - raw_metadata = original_body.get("metadata", original_body.get("api_callback")) - - if raw_metadata is None: - return None, ( - "Assistant updates must include metadata with prompt_processor, " - "connector, llm, and rag_processor." - ) - - if isinstance(raw_metadata, dict): - metadata_dict = raw_metadata - elif isinstance(raw_metadata, str): - if not raw_metadata.strip(): - return None, "Assistant metadata cannot be empty on update." - try: - parsed = json.loads(raw_metadata) - except json.JSONDecodeError as e: - return None, f"Assistant metadata must be valid JSON: {str(e)}" - if not isinstance(parsed, dict): - return None, "Assistant metadata must be a JSON object." - metadata_dict = parsed - else: - return None, "Assistant metadata must be a JSON string or object." - - missing_keys = [ - key for key in REQUIRED_PLUGIN_METADATA_KEYS - if not isinstance(metadata_dict.get(key), str) or not metadata_dict.get(key).strip() - ] - if missing_keys: - return None, ( - "Assistant metadata is incomplete. Missing required plugin fields: " - + ", ".join(missing_keys) - ) - - normalized_metadata = json.dumps(metadata_dict) - return normalized_metadata, None + return _validate_metadata(original_body) def _ensure_metadata_defaults(metadata_raw) -> str: @@ -556,8 +506,8 @@ async def create_assistant_directly(request: Request, auth: AuthContext = Depend def check_assistant_exists(prefixed_name: str) -> bool: """Check if an assistant with this name exists for this owner""" try: - existing = db_manager.get_assistant_by_name_and_owner( - prefixed_name, + existing = db_manager.get_assistant_by_name( + prefixed_name, creator_user['email'] ) return existing is not None @@ -1218,6 +1168,11 @@ async def update_assistant_proxy(assistant_id: int, request: Request, auth: Auth creator_user = auth.user logger.info(f"User {creator_user.get('email')} attempting to update assistant {assistant_id}.") + # Enforce object-level authorization before any mutation (#408). + # 404 (not 403) when the caller has no access at all, to avoid + # leaking existence; owner or org-admin required to modify. + auth.require_assistant_access(assistant_id, level="owner_or_admin") + # Fetch current assistant to merge with partial updates (#328) assistant_service = AssistantService() current = assistant_service.get_assistant_by_id(assistant_id) @@ -1264,6 +1219,12 @@ async def update_assistant_proxy(assistant_id: int, request: Request, auth: Auth if error: logger.error(f"Error preparing update body for assistant {assistant_id}: {error}") raise HTTPException(status_code=400, detail=error) + + # Owner is immutable on update — never derive it from the caller (#408). + # prepare_assistant_body() sets owner to the calling user (correct for + # create, an ownership-takeover hole on update); preserve the existing owner. + new_body["owner"] = current.owner + logger.info(f"Prepared body for update (Assistant ID {assistant_id}): {new_body}") # Create a mock request object with the prepared body diff --git a/backend/creator_interface/kb_server_manager.py b/backend/creator_interface/kb_server_manager.py index 8f25595a5..213be2226 100644 --- a/backend/creator_interface/kb_server_manager.py +++ b/backend/creator_interface/kb_server_manager.py @@ -23,8 +23,12 @@ # Get environment variables _raw_kb_server = os.getenv('LAMB_KB_SERVER', None) -# In dev/test: if the configured URL is the Docker service name that isn't running, redirect to host -_KB_REDIRECTS = {'http://kb:9090': 'http://172.18.0.1:9090'} +# Redirect table for dev environments where the configured URL is unreachable. +# In the normal docker-compose topology, `kb` resolves correctly inside the +# `lamb-*` bridge network, so the table is empty. The runtime fallback in +# is_kb_server_available() still tries host.docker.internal as a backup when +# the primary `http://kb:9090` is unreachable. +_KB_REDIRECTS: dict[str, str] = {} LAMB_KB_SERVER = _KB_REDIRECTS.get(_raw_kb_server, _raw_kb_server) or 'http://172.17.0.1:9090' LAMB_KB_SERVER_TOKEN = os.getenv('LAMB_KB_SERVER_TOKEN') if not LAMB_KB_SERVER_TOKEN: @@ -43,7 +47,7 @@ def __init__(self): self.global_kb_server_token = LAMB_KB_SERVER_TOKEN self.kb_server_configured = KB_SERVER_CONFIGURED - def _get_kb_config_for_user(self, creator_user: Dict[str, Any]) -> Dict[str, str]: + def _get_kb_config_for_user(self, creator_user: Dict[str, Any]) -> Dict[str, Any]: """ Resolve KB server configuration based on user's organization. Uses organization-specific config if available, falls back to environment variables. @@ -52,7 +56,7 @@ def _get_kb_config_for_user(self, creator_user: Dict[str, Any]) -> Dict[str, str creator_user: Dict containing user information with 'email' and 'organization_id' Returns: - Dict with 'url' and 'token' keys for KB server connection + Dict with 'url', 'token', and optionally 'embedding_model' and 'collection_defaults' """ from lamb.completions.org_config_resolver import OrganizationConfigResolver @@ -78,11 +82,19 @@ def _get_kb_config_for_user(self, creator_user: Dict[str, Any]) -> Dict[str, str if not api_token: api_token = self.global_kb_server_token org_url = kb_config.get('server_url') - resolved_url = _KB_REDIRECTS.get(org_url, org_url) - return { - 'url': resolved_url, + result = { + 'url': _KB_REDIRECTS.get(org_url, org_url), 'token': api_token } + # Pass through embedding model and collection defaults if configured + embedding_model = kb_config.get('embedding_model') + if embedding_model: + result['embedding_model'] = embedding_model + logger.info(f"Using org embedding model: {embedding_model}") + collection_defaults = kb_config.get('collection_defaults') + if collection_defaults: + result['collection_defaults'] = collection_defaults + return result else: logger.info(f"No organization KB config for {user_email}, using global config") @@ -497,7 +509,7 @@ async def create_knowledge_base(self, kb_data: KnowledgeBaseCreate, creator_user # Update kb_data with sanitized name kb_data.name = sanitized_name - + # Create collection in KB server # Apply host redirect for dev/test environments where 'kb' DNS doesn't resolve _alt_url = _KB_REDIRECTS.get(kb_server_url) @@ -521,18 +533,33 @@ async def create_knowledge_base(self, kb_data: KnowledgeBaseCreate, creator_user except Exception as md_err: logger.warning(f"Error extracting metadata description: {str(md_err)}") + # Resolve embedding model: use org config if set, otherwise default + org_embedding_model = kb_config.get('embedding_model') + if org_embedding_model: + # Use org-specified model; let KB server resolve vendor/endpoint from its env + embeddings_model = { + "model": org_embedding_model, + "vendor": "default", + "api_endpoint": "default", + "apikey": "default" + } + logger.info(f"Using org-configured embedding model: {org_embedding_model}") + else: + # No org config; fully default (KB server resolves all from its env) + embeddings_model = { + "model": "default", + "vendor": "default", + "api_endpoint": "default", + "apikey": "default" + } + # Prepare collection data according to KB server API collection_data = { "name": kb_data.name, "description": description, "owner": str(creator_user.get('id')), # Use ID instead of email for privacy (as string) "visibility": kb_data.access_control or "private", - "embeddings_model": { - "model": "default", - "vendor": "default", - "api_endpoint": "default", - "apikey": "default" - } + "embeddings_model": embeddings_model } # Log the final data being sent to the KB server diff --git a/backend/creator_interface/knowledge_store_client.py b/backend/creator_interface/knowledge_store_client.py index fc5c45803..254065a53 100644 --- a/backend/creator_interface/knowledge_store_client.py +++ b/backend/creator_interface/knowledge_store_client.py @@ -260,6 +260,10 @@ async def create_collection( embedding_endpoint: str = "", embedding_params: Dict[str, Any] = None, vector_db_params: Dict[str, Any] = None, + graph_enabled: bool = False, + extraction_vendor: Optional[str] = None, + extraction_model: Optional[str] = None, + extraction_endpoint: Optional[str] = None, creator_user: Dict[str, Any] = None, ) -> Dict: """Create a collection on the KB Server. @@ -289,9 +293,24 @@ async def create_collection( "embedding_params": embedding_params or {}, "vector_db_backend": vector_db_backend, "vector_db_params": vector_db_params or {}, + "graph_enabled": bool(graph_enabled), } + # Extraction config is only persisted when graph is enabled. + if graph_enabled and (extraction_vendor or extraction_model): + payload["extraction"] = { + "vendor": extraction_vendor or None, + "model": extraction_model or None, + "api_endpoint": extraction_endpoint or None, + } return await self._request("POST", "/collections", config, json=payload) + async def get_llm_vendors( + self, creator_user: Dict[str, Any] = None + ) -> Dict: + """Fetch registered LLM extraction vendors from kb-v2.""" + config = self._get_ks_config(creator_user) + return await self._request("GET", "/llm-vendors", config) + async def get_collection(self, knowledge_store_id: str, creator_user: Dict[str, Any] = None) -> Dict: """Get a collection by ID.""" @@ -434,6 +453,167 @@ async def cancel_job(self, job_id: str, config = self._get_ks_config(creator_user) return await self._request("POST", f"/jobs/{job_id}/cancel", config) + async def retry_job(self, job_id: str, + creator_user: Dict[str, Any] = None) -> Dict: + """Retry a failed ingestion job, reusing its cached credentials. + + The KB Server keeps the document payload and embedding credentials in + memory across attempts, so the retry does not re-send them. Raises an + ``HTTPException`` with the KB Server's status code on 404/409/410 (job + missing, not retryable / exhausted, or credentials expired). + """ + config = self._get_ks_config(creator_user) + return await self._request("POST", f"/jobs/{job_id}/retry", config) + + async def get_job_retry_available(self, job_id: str, + creator_user: Dict[str, Any] = None) -> Dict: + """Report whether a failed job can still be retried in place.""" + config = self._get_ks_config(creator_user) + return await self._request("GET", f"/jobs/{job_id}/retry-available", config) + + # ------------------------------------------------------------------ + # KG-RAG / Semantic Graph proxy + # ------------------------------------------------------------------ + # These endpoints exist on the KB Server only when ``KG_RAG_ENABLED=true`` + # is set in its env. LAMB proxies through to keep the per-org token / + # URL resolution in one place. + + async def get_graph_status(self, creator_user: Dict[str, Any] = None) -> Dict: + config = self._get_ks_config(creator_user) + return await self._request("GET", "/graph/status", config) + + async def migrate_collection_to_graph( + self, + knowledge_store_id: str, + openai_api_key: str = "", + creator_user: Dict[str, Any] = None, + ) -> Dict: + config = self._get_ks_config(creator_user) + extra_headers = ( + {"X-OpenAI-Api-Key": openai_api_key} if openai_api_key else {} + ) + headers = {**self._headers(config["token"]), **extra_headers} + url = ( + f"{config['url'].rstrip('/')}" + f"/graph/collections/{knowledge_store_id}/migrate" + ) + try: + async with httpx.AsyncClient(timeout=600.0) as client: + response = await client.post(url, headers=headers) + if response.is_success: + return response.json() if response.content else {} + detail = response.json().get("detail", response.text) + raise HTTPException( + status_code=response.status_code, + detail=f"Knowledge Store server error: {detail}", + ) + except httpx.RequestError as exc: + logger.error("Knowledge Store connection error: %s", exc) + raise HTTPException( + status_code=503, + detail="Unable to connect to Knowledge Store server", + ) + + async def get_graph_snapshot( + self, + knowledge_store_id: str, + params: Dict[str, Any] = None, + creator_user: Dict[str, Any] = None, + ) -> Dict: + config = self._get_ks_config(creator_user) + return await self._request( + "GET", + f"/graph/collections/{knowledge_store_id}/snapshot", + config, + params=params or {}, + ) + + async def list_graph_changes( + self, + knowledge_store_id: str, + params: Dict[str, Any] = None, + creator_user: Dict[str, Any] = None, + ) -> Dict: + config = self._get_ks_config(creator_user) + return await self._request( + "GET", + f"/graph/collections/{knowledge_store_id}/changes", + config, + params=params or {}, + ) + + async def graph_concept_rename( + self, + knowledge_store_id: str, + concept: str, + body: Dict[str, Any], + creator_user: Dict[str, Any] = None, + ) -> Dict: + config = self._get_ks_config(creator_user) + return await self._request( + "PATCH", + f"/graph/collections/{knowledge_store_id}/concepts/{concept}/rename", + config, + json=body, + ) + + async def graph_concepts_merge( + self, + knowledge_store_id: str, + body: Dict[str, Any], + creator_user: Dict[str, Any] = None, + ) -> Dict: + config = self._get_ks_config(creator_user) + return await self._request( + "POST", + f"/graph/collections/{knowledge_store_id}/concepts/merge", + config, + json=body, + ) + + async def graph_concept_curation( + self, + knowledge_store_id: str, + concept: str, + body: Dict[str, Any], + creator_user: Dict[str, Any] = None, + ) -> Dict: + config = self._get_ks_config(creator_user) + return await self._request( + "PATCH", + f"/graph/collections/{knowledge_store_id}/concepts/{concept}/curation", + config, + json=body, + ) + + async def graph_relationship_edit( + self, + knowledge_store_id: str, + body: Dict[str, Any], + creator_user: Dict[str, Any] = None, + ) -> Dict: + config = self._get_ks_config(creator_user) + return await self._request( + "PATCH", + f"/graph/collections/{knowledge_store_id}/relationships", + config, + json=body, + ) + + async def graph_relationship_curation( + self, + knowledge_store_id: str, + body: Dict[str, Any], + creator_user: Dict[str, Any] = None, + ) -> Dict: + config = self._get_ks_config(creator_user) + return await self._request( + "PATCH", + f"/graph/collections/{knowledge_store_id}/relationships/curation", + config, + json=body, + ) + # ------------------------------------------------------------------ # Org-level discovery (for the UI options endpoint) # ------------------------------------------------------------------ diff --git a/backend/creator_interface/knowledge_store_graph_router.py b/backend/creator_interface/knowledge_store_graph_router.py new file mode 100644 index 000000000..a4a611bd9 --- /dev/null +++ b/backend/creator_interface/knowledge_store_graph_router.py @@ -0,0 +1,288 @@ +"""Creator-Interface proxy for KG-RAG / semantic-graph endpoints. + +Routes mount at ``/creator/knowledge-stores/{ks_id}/graph/...``. Each call +resolves the per-org KB Server URL/token through ``KnowledgeStoreClient`` +and proxies to the new KB Server's ``/graph`` router. + +Benchmark runs are no longer exposed through the UI — they live as +standalone scripts that hit the KB Server's ``/benchmarks`` router +directly (see ``memoria/run_*_bench.py``). The KB Server only exposes +the graph router when ``KG_RAG_ENABLED=true``; if the flag is off the +proxy will return 503. The frontend uses ``/graph/status`` to gate its +UI accordingly. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional + +from fastapi import APIRouter, Body, Depends, HTTPException, Query +from pydantic import BaseModel, Field + +from lamb.auth_context import AuthContext, get_auth_context +from lamb.database_manager import LambDatabaseManager + +from .knowledge_store_client import KnowledgeStoreClient + +logger = logging.getLogger(__name__) + +router = APIRouter() +_client = KnowledgeStoreClient() +_db = LambDatabaseManager() + + +def _assert_ks_access(ks_id: str, auth: AuthContext) -> Dict[str, Any]: + """Load the Knowledge Store row and check the caller can use it. + + Mirrors the access checks in ``knowledge_store_router`` — owner or a + shared store within the same org. Raises 404/403 on failure. + """ + ks = _db.get_knowledge_store(ks_id) + if not ks: + raise HTTPException(status_code=404, detail="Knowledge Store not found") + if ( + ks.get("owner_user_id") != auth.user.get("id") + and not ks.get("is_shared") + ): + raise HTTPException(status_code=403, detail="Forbidden") + if ks.get("organization_id") != auth.organization.get("id"): + raise HTTPException(status_code=403, detail="Forbidden") + return ks + + +# ---------------------------------------------------------------------- +# Status / migration +# ---------------------------------------------------------------------- + + +@router.get("/graph/status") +async def get_graph_status(auth: AuthContext = Depends(get_auth_context)): + """Get KG-RAG feature flag + Neo4j availability from the KB Server. + + The frontend reads this on the Knowledge Stores page so it knows + whether to show Graph / Benchmark tabs at all. + """ + return await _client.get_graph_status(creator_user=auth.user) + + +class MigrateRequest(BaseModel): + openai_api_key: str = Field( + default="", + description=( + "Per-request OpenAI key for concept extraction. Preferred over " + "the KB server's permanent KG_RAG_OPENAI_API_KEY." + ), + ) + + +@router.post("/{ks_id}/graph/migrate") +async def migrate_knowledge_store_to_graph( + ks_id: str, + body: MigrateRequest = Body(default_factory=MigrateRequest), + auth: AuthContext = Depends(get_auth_context), +): + """Run KG-RAG migration for an existing Knowledge Store. + + Resolves the OpenAI key from (1) the explicit per-request body, + (2) the org's ``providers.openai.api_key``, or fails to use the KB + server-level fallback if both are empty. + """ + _assert_ks_access(ks_id, auth) + + api_key = (body.openai_api_key or "").strip() + if not api_key: + from lamb.completions.org_config_resolver import OrganizationConfigResolver + resolver = OrganizationConfigResolver(auth.user.get("email")) + try: + api_key = resolver.get_provider_api_key("openai") or "" + except Exception: # noqa: BLE001 — org config may be missing + api_key = "" + + return await _client.migrate_collection_to_graph( + knowledge_store_id=ks_id, + openai_api_key=api_key, + creator_user=auth.user, + ) + + +# ---------------------------------------------------------------------- +# Read endpoints +# ---------------------------------------------------------------------- + + +@router.get("/{ks_id}/graph/snapshot") +async def get_graph_snapshot( + ks_id: str, + concept: Optional[str] = Query(default=None), + document_id: Optional[str] = Query(default=None), + chunk_id: Optional[str] = Query(default=None), + filename: Optional[str] = Query(default=None), + include_chunks: bool = Query(default=True), + limit: int = Query(default=60, ge=1, le=10000), + auth: AuthContext = Depends(get_auth_context), +): + _assert_ks_access(ks_id, auth) + params = { + "include_chunks": str(include_chunks).lower(), + "limit": limit, + } + if concept: + params["concept"] = concept + if document_id: + params["document_id"] = document_id + if chunk_id: + params["chunk_id"] = chunk_id + if filename: + params["filename"] = filename + return await _client.get_graph_snapshot( + knowledge_store_id=ks_id, params=params, creator_user=auth.user, + ) + + +@router.get("/{ks_id}/graph/changes") +async def list_graph_changes( + ks_id: str, + concept: Optional[str] = Query(default=None), + document_id: Optional[str] = Query(default=None), + filename: Optional[str] = Query(default=None), + operation: Optional[str] = Query(default=None), + limit: int = Query(default=25, ge=1, le=200), + auth: AuthContext = Depends(get_auth_context), +): + _assert_ks_access(ks_id, auth) + params: Dict[str, Any] = {"limit": limit} + if concept: + params["concept"] = concept + if document_id: + params["document_id"] = document_id + if filename: + params["filename"] = filename + if operation: + params["operation"] = operation + return await _client.list_graph_changes( + knowledge_store_id=ks_id, params=params, creator_user=auth.user, + ) + + +# ---------------------------------------------------------------------- +# Curation +# ---------------------------------------------------------------------- + + +class ConceptRenameRequest(BaseModel): + new_name: str = Field(..., min_length=1) + actor: str = "graph-curation-api" + reason: str = "" + + +class ConceptMergeRequest(BaseModel): + source_names: List[str] = Field(..., min_length=1) + target_name: str = Field(..., min_length=1) + actor: str = "graph-curation-api" + reason: str = "" + + +class ConceptCurationRequest(BaseModel): + notes: Optional[str] = None + tags: Optional[List[str]] = None + verification_state: Optional[str] = None + actor: str = "graph-curation-api" + reason: str = "" + + +class RelationshipEditRequest(BaseModel): + source_concept: str + target_concept: str + relation: str + new_relation: Optional[str] = None + weight: Optional[float] = None + description: Optional[str] = None + evidence: Optional[str] = None + notes: Optional[str] = None + tags: Optional[List[str]] = None + verification_state: Optional[str] = None + actor: str = "graph-curation-api" + reason: str = "" + + +class RelationshipCurationRequest(BaseModel): + source_concept: str + target_concept: str + relation: str + notes: Optional[str] = None + tags: Optional[List[str]] = None + verification_state: Optional[str] = None + actor: str = "graph-curation-api" + reason: str = "" + + +@router.patch("/{ks_id}/graph/concepts/{concept}/rename") +async def rename_concept( + ks_id: str, + concept: str, + body: ConceptRenameRequest, + auth: AuthContext = Depends(get_auth_context), +): + _assert_ks_access(ks_id, auth) + return await _client.graph_concept_rename( + knowledge_store_id=ks_id, concept=concept, + body=body.model_dump(), creator_user=auth.user, + ) + + +@router.post("/{ks_id}/graph/concepts/merge") +async def merge_concepts( + ks_id: str, + body: ConceptMergeRequest, + auth: AuthContext = Depends(get_auth_context), +): + _assert_ks_access(ks_id, auth) + return await _client.graph_concepts_merge( + knowledge_store_id=ks_id, body=body.model_dump(), + creator_user=auth.user, + ) + + +@router.patch("/{ks_id}/graph/concepts/{concept}/curation") +async def curate_concept( + ks_id: str, + concept: str, + body: ConceptCurationRequest, + auth: AuthContext = Depends(get_auth_context), +): + _assert_ks_access(ks_id, auth) + return await _client.graph_concept_curation( + knowledge_store_id=ks_id, concept=concept, + body=body.model_dump(), creator_user=auth.user, + ) + + +@router.patch("/{ks_id}/graph/relationships") +async def edit_relationship( + ks_id: str, + body: RelationshipEditRequest, + auth: AuthContext = Depends(get_auth_context), +): + _assert_ks_access(ks_id, auth) + return await _client.graph_relationship_edit( + knowledge_store_id=ks_id, body=body.model_dump(), + creator_user=auth.user, + ) + + +@router.patch("/{ks_id}/graph/relationships/curation") +async def curate_relationship( + ks_id: str, + body: RelationshipCurationRequest, + auth: AuthContext = Depends(get_auth_context), +): + _assert_ks_access(ks_id, auth) + return await _client.graph_relationship_curation( + knowledge_store_id=ks_id, body=body.model_dump(), + creator_user=auth.user, + ) + + + + diff --git a/backend/creator_interface/knowledge_store_router.py b/backend/creator_interface/knowledge_store_router.py index 946cf2583..f2e08db72 100644 --- a/backend/creator_interface/knowledge_store_router.py +++ b/backend/creator_interface/knowledge_store_router.py @@ -58,6 +58,18 @@ class KnowledgeStoreCreate(BaseModel): # Extra knobs declared by the vector-DB backend's plugin schema. # Empty for today's chromadb/qdrant backends. vector_db_params: Optional[Dict[str, Any]] = None + # Optional semantic-graph / KG-RAG opt-in. Forwarded to the KB Server + # so the collection is created with ``graph_enabled=true`` and + # ingestion-time concept extraction runs against it. Requires + # ``KG_RAG_ENABLED=true`` on the KB server. + graph_enabled: bool = False + # Locked-at-creation extraction config. Only meaningful when + # ``graph_enabled=true``. ``None`` everywhere means "fall back to the + # KB server's env defaults" (back-compat for the previous toggle-only + # surface). + extraction_vendor: Optional[str] = None + extraction_model: Optional[str] = None + extraction_endpoint: Optional[str] = None class KnowledgeStoreUpdate(BaseModel): @@ -167,6 +179,14 @@ async def get_options(auth: AuthContext = Depends(get_auth_context)): ) +@router.get("/llm-vendors") +async def get_llm_vendors(auth: AuthContext = Depends(get_auth_context)): + """Return registered LLM extraction vendors (for KG-RAG concept extraction) + along with their default-model lists so the create UI can render a picker. + """ + return await _client.get_llm_vendors(creator_user=auth.user) + + # ====================================================================== # Knowledge Store CRUD # ====================================================================== @@ -242,6 +262,10 @@ async def create_knowledge_store( embedding_endpoint=resolved_endpoint or "", embedding_params=body.embedding_params, vector_db_params=body.vector_db_params, + graph_enabled=bool(body.graph_enabled), + extraction_vendor=body.extraction_vendor, + extraction_model=body.extraction_model, + extraction_endpoint=body.extraction_endpoint, creator_user=auth.user, ) except Exception as e: @@ -289,11 +313,15 @@ async def get_knowledge_store( entry["server_status"] = server_data.get("status") entry["document_count"] = server_data.get("document_count", 0) entry["chunk_count"] = server_data.get("chunk_count", 0) + entry["graph_enabled"] = bool(server_data.get("graph_enabled", False)) + entry["extraction"] = server_data.get("extraction") except Exception as e: logger.warning(f"Could not fetch collection metadata from KB Server for {ks_id}: {e}") entry["server_status"] = None entry["document_count"] = None entry["chunk_count"] = None + entry["graph_enabled"] = False + entry["extraction"] = None entry["content"] = _db.get_kb_content_links_for_ks(ks_id) entry["is_owner"] = entry.get("owner_user_id") == auth.user.get("id") @@ -584,7 +612,6 @@ async def add_content( } - @router.get("/{ks_id}/content/{library_item_id}") async def get_content_link( ks_id: str, @@ -622,6 +649,56 @@ async def get_content_link( return link +@router.post("/{ks_id}/content/{library_item_id}/retry") +async def retry_content( + ks_id: str, + library_item_id: str, + auth: AuthContext = Depends(get_auth_context), +): + """Retry a failed indexing job for a single content link. + + The KB Server retains the document payload and embedding credentials in + memory across attempts, so the retry re-uses them — nothing is re-sent. + Propagates the KB Server's response: + + - 404 if the link or job is unknown, + - 409 if the job is not failed or has exhausted its retries, + - 410 if the retention window elapsed (the user must re-add the content). + """ + auth.require_knowledge_store_access(ks_id, level="owner") + link = _db.get_kb_content_link(ks_id, library_item_id) + if not link: + raise HTTPException(status_code=404, detail="Content link not found") + + job_id = link.get("kb_job_id") + if not job_id: + raise HTTPException( + status_code=409, + detail="This content has no indexing job to retry.", + ) + + job = await _client.retry_job(job_id, creator_user=auth.user) + + # The KB Server flipped the job back to pending; mirror that on the link so + # the UI shows the spinner again and resumes polling. + # Pass an empty string (not None) so the previous error is actually + # cleared — update_kb_content_link_status skips fields that are None. + _db.update_kb_content_link_status( + link_id=link["id"], + status="pending", + error_message="", + ) + _audit(auth, "knowledge_store.retry_content", "knowledge_store", ks_id, { + "library_item_id": library_item_id, + "kb_job_id": job_id, + }) + return { + "message": "Indexing retry queued.", + "job_id": job_id, + "status": job.get("status", "pending"), + } + + @router.delete("/{ks_id}/content/{library_item_id}") async def remove_content( ks_id: str, diff --git a/backend/creator_interface/library_router.py b/backend/creator_interface/library_router.py index 63a5a7f09..de09baaad 100644 --- a/backend/creator_interface/library_router.py +++ b/backend/creator_interface/library_router.py @@ -1072,6 +1072,118 @@ async def export_library( permalink_proxy_router = APIRouter() +# ---------------------------------------------------------------------- +# Public, HMAC-signed citation permalinks (student-clickable, no login) +# ---------------------------------------------------------------------- +# Defined BEFORE the authenticated catch-all below so `/docs/public/...` is not +# swallowed by `/docs/{org_id}/.../{subpath:path}` (FastAPI matches in order). +# A student clicking a citation in the OpenWebUI chat arrives with no LAMB +# credential, so the link authorizes itself via an org-scoped HMAC signature +# (see permalink_signing). Traffic still flows OpenWebUI -> LAMB -> Library +# Manager; OWI never contacts the Library Manager directly. + +_PUBLIC_VIEW_CSS = ( + "body{font-family:system-ui,-apple-system,Segoe UI,Roboto,sans-serif;" + "max-width:820px;margin:0 auto;padding:1.5rem;color:#1f2937;line-height:1.6}" + "header{display:flex;align-items:center;justify-content:space-between;gap:1rem;" + "border-bottom:1px solid #e5e7eb;padding-bottom:.75rem;margin-bottom:1.25rem}" + "h1{font-size:1.15rem;margin:0;word-break:break-word}" + "a.dl{flex:none;background:#2271b3;color:#fff;text-decoration:none;border-radius:6px;" + "padding:.45rem .8rem;font-size:.85rem}" + "article img{max-width:100%}article pre{overflow:auto;background:#f3f4f6;padding:.75rem;border-radius:6px}" + "article table{border-collapse:collapse}article td,article th{border:1px solid #e5e7eb;padding:.3rem .5rem}" +) + + +@permalink_proxy_router.get("/docs/public/{org_id}/{library_id}/{item_id}/view") +async def public_permalink_view( + org_id: str, + library_id: str, + item_id: str, + name: str = Query("", description="Display filename for the cited item"), + sig: str = Query(..., description="Org-scoped HMAC signature"), +): + """Render a cited item's content as an HTML page titled with its filename. + + Public (signature-authorized) — no login required. Shows the item's + extracted markdown rendered to HTML under the original filename, plus a + "Download original" button when an original file exists. + """ + from html import escape # noqa: PLC0415 + from urllib.parse import quote # noqa: PLC0415 + + from creator_interface.permalink_signing import sign, verify # noqa: PLC0415 + + if not verify(org_id, f"{org_id}/{library_id}/{item_id}/view", sig): + raise HTTPException(status_code=403, detail="Invalid or expired link") + + # Proxy the full markdown from the Library Manager using a system context + # (the signature is the authorization; there is no end-user identity here). + response = await _client.proxy_content( + library_id=library_id, + item_id=item_id, + subpath="content", + creator_user={}, + ) + markdown_text = response.content.decode("utf-8", errors="replace") + + try: + from markdown_it import MarkdownIt # noqa: PLC0415 + body_html = MarkdownIt("commonmark", {"html": False, "linkify": True}).render(markdown_text) + except Exception: # noqa: BLE001 — fall back to preformatted text + body_html = f"
{escape(markdown_text)}
" + + title = escape(name or "Source document") + download_html = "" + # Offer the original download only when the cited item is a real file + # (has an extension); URL/transcript sources have no original to download. + if name and "." in name: + dl_sig = sign(org_id, f"{org_id}/{library_id}/{item_id}/original/{name}") + dl_url = ( + f"/docs/public/{org_id}/{library_id}/{item_id}/original/{quote(name)}?sig={dl_sig}" + ) + download_html = f'Download original' + + page = ( + f"" + f"" + f"{title}" + f"

{title}

{download_html}
" + f"
{body_html}
" + ) + return Response(content=page, media_type="text/html; charset=utf-8") + + +@permalink_proxy_router.get( + "/docs/public/{org_id}/{library_id}/{item_id}/original/{filename:path}" +) +async def public_permalink_original( + org_id: str, + library_id: str, + item_id: str, + filename: str, + sig: str = Query(..., description="Org-scoped HMAC signature"), +): + """Stream the original uploaded file for a cited item (signature-authorized).""" + from creator_interface.permalink_signing import verify # noqa: PLC0415 + + if not verify(org_id, f"{org_id}/{library_id}/{item_id}/original/{filename}", sig): + raise HTTPException(status_code=403, detail="Invalid or expired link") + + response = await _client.proxy_content( + library_id=library_id, + item_id=item_id, + subpath=f"original/{filename}", + creator_user={}, + ) + content_type = response.headers.get("content-type", "application/octet-stream") + return Response( + content=response.content, + media_type=content_type, + headers={"Content-Disposition": f'attachment; filename="{filename}"'}, + ) + + @permalink_proxy_router.get("/docs/{org_id}/{library_id}/{item_id}/{subpath:path}") async def permalink_proxy( org_id: str, diff --git a/backend/creator_interface/main.py b/backend/creator_interface/main.py index 88b730cb8..3d3e6d9b5 100644 --- a/backend/creator_interface/main.py +++ b/backend/creator_interface/main.py @@ -1,6 +1,9 @@ from .chats_router import router as chats_router from .library_router import router as library_router -from .knowledge_store_router import router as knowledge_store_router +from .knowledge_store_router import router as knowledge_store_router # pragma: no cover - import-time only; creator_interface.main pulls in openai (uninstalled in test env), so unreachable under pytest +from .knowledge_store_graph_router import ( # pragma: no cover - see note above + router as knowledge_store_graph_router, +) from .analytics_router import router as analytics_router from .prompt_templates_router import router as prompt_templates_router from .evaluaitor_router import router as evaluaitor_router @@ -144,7 +147,12 @@ async def stop_news_cache_refresh_loop(): # Include the Knowledge Store router (new KB Server, port 9092). Distinct # from the legacy /knowledgebases routes which serve the stable KB Server. -router.include_router(knowledge_store_router, prefix="/knowledge-stores") +router.include_router(knowledge_store_router, prefix="/knowledge-stores") # pragma: no cover - import-time router wiring; module unimportable under pytest (see top) +# KG-RAG graph proxy. Mounted on the same prefix so the frontend +# can keep ``/creator/knowledge-stores/...`` as the only KB-related base. +router.include_router( # pragma: no cover - see note above + knowledge_store_graph_router, prefix="/knowledge-stores" +) # REMOVED: assistant_sharing_router - functionality moved to services, accessed via /creator/lamb/* proxy diff --git a/backend/creator_interface/metadata_validators.py b/backend/creator_interface/metadata_validators.py new file mode 100644 index 000000000..4d95a0a59 --- /dev/null +++ b/backend/creator_interface/metadata_validators.py @@ -0,0 +1,91 @@ +import json +import logging +from typing import Dict, Any, Tuple, Optional + +logger = logging.getLogger(__name__) + +REQUIRED_PLUGIN_METADATA_KEYS = ( + "prompt_processor", + "connector", + "llm", + "rag_processor", +) + + +def validate_update_plugin_metadata( + original_body: Dict[str, Any] +) -> Tuple[Optional[str], Optional[str]]: + """ + Validate assistant plugin metadata for updates. + + Updates must provide complete plugin metadata so the backend never replaces a + valid stored configuration with partial or blank data. + + Returns: + Tuple[Optional[str], Optional[str]]: (normalized_metadata_json, error_message) + """ + raw_metadata = original_body.get("metadata", original_body.get("api_callback")) + + if raw_metadata is None: + return None, ( + "Assistant updates must include metadata with prompt_processor, " + "connector, llm, and rag_processor." + ) + + if isinstance(raw_metadata, dict): + metadata_dict = raw_metadata + elif isinstance(raw_metadata, str): + if not raw_metadata.strip(): + return None, "Assistant metadata cannot be empty on update." + try: + parsed = json.loads(raw_metadata) + except json.JSONDecodeError as e: + return None, f"Assistant metadata must be valid JSON: {str(e)}" + if not isinstance(parsed, dict): + return None, "Assistant metadata must be a JSON object." + metadata_dict = parsed + else: + return None, "Assistant metadata must be a JSON string or object." + + missing_keys = [ + key for key in REQUIRED_PLUGIN_METADATA_KEYS + if not isinstance(metadata_dict.get(key), str) or not metadata_dict.get(key).strip() + ] + if missing_keys: + return None, ( + "Assistant metadata is incomplete. Missing required plugin fields: " + + ", ".join(missing_keys) + ) + + # single_file_rag specific validation + if metadata_dict.get("rag_processor") == "single_file_rag": + has_library_ref = ( + metadata_dict.get("library_id") and metadata_dict.get("item_id") + ) + has_file_ref = bool(metadata_dict.get("file_path")) + if not has_library_ref and not has_file_ref: + return None, ( + "single_file_rag requires either library_id + item_id " + "or file_path in metadata" + ) + if metadata_dict.get("library_id") and not metadata_dict.get("item_id"): + return None, "single_file_rag: library_id requires item_id" + if metadata_dict.get("item_id") and not metadata_dict.get("library_id"): + return None, "single_file_rag: item_id requires library_id" + + # document_rag specific validation + if metadata_dict.get("document_rag") == "library_file_rag": + has_library_ref = ( + metadata_dict.get("library_id") and metadata_dict.get("item_id") + ) + if not has_library_ref: + return None, ( + "document_rag=library_file_rag requires library_id + item_id in metadata" + ) + if metadata_dict.get("library_id") and not metadata_dict.get("item_id"): + return None, "document_rag: library_id requires item_id" + if metadata_dict.get("item_id") and not metadata_dict.get("library_id"): + return None, "document_rag: item_id requires library_id" + + normalized_metadata = json.dumps(metadata_dict) + return normalized_metadata, None diff --git a/backend/creator_interface/permalink_signing.py b/backend/creator_interface/permalink_signing.py new file mode 100644 index 000000000..0b6e720b0 --- /dev/null +++ b/backend/creator_interface/permalink_signing.py @@ -0,0 +1,52 @@ +"""HMAC signing for public citation permalinks. + +When a learning assistant cites a source in chat, the student clicking the +citation arrives as an unauthenticated, cross-origin browser navigation that +carries no LAMB credential. So the link must authorize itself: it embeds an +HMAC signature that a public serving route verifies. + +Design (decided with the project lead): +- **Per-organization isolation.** The signing key is derived from the master + secret AND the org id, so a signature minted for one org's item can never be + reused to forge another org's link. Rotating one org's effective key (by + rotating the master secret) is the kill-switch. +- **No expiry.** Links live as long as the chat history that contains them, so + old conversations stay consistent. Revocation is coarse by design: rotate + ``LAMB_PERMALINK_SIGNING_SECRET`` (revokes everything) or delete the library + item (its link then 404s at the proxy). +- **Capability semantics.** Within an org, anyone holding a specific link can + open that one item — there is no identity to check on the click. This is the + accepted trade-off for letting students reach cited sources without a login. +""" + +from __future__ import annotations + +import hashlib +import hmac + +from config import LAMB_PERMALINK_SIGNING_SECRET + +_MASTER = (LAMB_PERMALINK_SIGNING_SECRET or "").encode("utf-8") + + +def _org_key(org_id: str) -> bytes: + """Derive a per-organization signing key from the master secret.""" + return hmac.new(_MASTER, f"perm:{org_id}".encode(), hashlib.sha256).digest() + + +def sign(org_id: str, resource: str) -> str: + """Return the hex HMAC signature for ``resource`` scoped to ``org_id``. + + ``resource`` is a canonical, slash-joined string identifying exactly what + may be served, e.g. ``"3/lib-uuid/item-uuid/view"`` or + ``"3/lib-uuid/item-uuid/original/cv.pdf"``. + """ + return hmac.new(_org_key(str(org_id)), resource.encode("utf-8"), hashlib.sha256).hexdigest() + + +def verify(org_id: str, resource: str, signature: str) -> bool: + """Constant-time check that ``signature`` is valid for ``(org_id, resource)``.""" + if not signature: + return False + expected = sign(org_id, resource) + return hmac.compare_digest(expected, signature) diff --git a/backend/lamb/completions/citation_sources.py b/backend/lamb/completions/citation_sources.py new file mode 100644 index 000000000..f4d130a17 --- /dev/null +++ b/backend/lamb/completions/citation_sources.py @@ -0,0 +1,137 @@ +"""Render RAG citations so a student can click through to the cited item. + +Open WebUI's backend, on the normal (websocket) chat path, forwards only the +assistant's ``content`` from an external model's stream — it drops any +top-level ``sources`` field (it renders only sources it generates itself). So +to surface clickable citations through OWI we append a Markdown **Sources** +section to the answer content; OWI renders it like any other markdown. + +Each source links to a LAMB-served view page via an org-scoped, HMAC-signed +"capability" URL (see ``creator_interface.permalink_signing``) so a logged-out +student can open the cited item. The inline ``[N]`` markers in the body point +at the numbered list. + +``build_owi_sources`` (the OWI ``sources`` schema) is kept for non-streaming / +spec-compliant clients that *do* consume the field. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional +from urllib.parse import quote + +from config import LAMB_WEB_HOST +from creator_interface.permalink_signing import sign +from lamb.logging_config import get_logger + +logger = get_logger(__name__, component="RAG") + + +def _parse_permalink(source: Dict[str, Any]) -> Optional[tuple]: + """Extract ``(org, lib, item, original_filename)`` from a source's permalinks. + + Permalinks have the form ``/docs/{org}/{lib}/{item}/...``; the original + filename (when present) is the last segment of ``permalink_original``. + """ + for key in ("permalink_markdown", "permalink_page", "permalink_original"): + permalink = source.get(key) + if permalink and permalink.startswith("/docs/"): + parts = permalink.split("/") # ['', 'docs', org, lib, item, ...] + if len(parts) >= 5: + org, lib, item = parts[2], parts[3], parts[4] + filename = "" + original = source.get("permalink_original") + if original: + filename = original.rstrip("/").split("/")[-1] + return org, lib, item, filename + return None + + +def _signed_view_url(org: str, lib: str, item: str, filename: str) -> str: + """Mint the absolute, org-scoped signed URL for a cited item's view page.""" + sig = sign(org, f"{org}/{lib}/{item}/view") + base = LAMB_WEB_HOST.rstrip("/") + return f"{base}/docs/public/{org}/{lib}/{item}/view?name={quote(filename)}&sig={sig}" + + +def build_sources_markdown(rag_context: Any) -> str: + """Build a Markdown 'Sources' section appended to the answer content. + + Sources are grouped by library item (chunks are numbered individually by + the RAG processor, but several chunks often come from the same document). + Each line lists every citation number that points at the item followed by a + clickable link, e.g. ``[1][3][5] [cv.pdf](https://…/view?…)`` — so every + inline ``[N]`` marker in the answer resolves to a clickable entry. Returns + ``""`` when there are no sources. Lines deliberately avoid the ``[N]: url`` + reference-definition form, which OWI strips. + """ + if not isinstance(rag_context, dict): + return "" + groups: Dict[str, Dict[str, Any]] = {} + order: List[str] = [] + for src in rag_context.get("sources", []) or []: + n = src.get("n") + filename = src.get("title") or "Source" + url = "" + parsed = _parse_permalink(src) + item_key = None + if parsed: + org, lib, item, original_filename = parsed + item_key = f"{org}/{lib}/{item}" + filename = original_filename or filename + try: + url = _signed_view_url(org, lib, item, filename) + except Exception: # noqa: BLE001 — never let citation building break a chat + logger.warning("Failed to sign citation URL for source %s", n) + key = item_key or f"_{n}" + if key not in groups: + groups[key] = {"ns": [], "filename": filename, "url": url} + order.append(key) + if n is not None: + groups[key]["ns"].append(n) + + if not order: + return "" + lines: List[str] = [] + for key in order: + g = groups[key] + marks = "".join(f"[{n}]" for n in sorted(g["ns"])) or "-" + link = f"[{g['filename']}]({g['url']})" if g["url"] else g["filename"] + lines.append(f"{marks} {link}") + return "\n\n---\n\n**Sources**\n\n" + "\n\n".join(lines) + "\n" + + +def build_owi_sources(rag_context: Any) -> List[Dict[str, Any]]: + """Map LAMB RAG sources to Open WebUI's citations schema. + + Kept for non-streaming responses and any spec-compliant client that + consumes a top-level ``sources`` field. (OWI's websocket chat path drops + this for external models — the Markdown section above is what reaches the + student there.) + """ + if not isinstance(rag_context, dict): + return [] + owi: List[Dict[str, Any]] = [] + for src in rag_context.get("sources", []) or []: + n = src.get("n") + name = str(n) if n is not None else (src.get("title") or "?") + text = src.get("text") or "" + filename = src.get("title") or "Source" + url = "" + parsed = _parse_permalink(src) + if parsed: + org, lib, item, original_filename = parsed + filename = original_filename or filename + try: + url = _signed_view_url(org, lib, item, filename) + except Exception: # noqa: BLE001 + logger.warning("Failed to sign citation URL for source %s", name) + excerpt = f"**{filename}**\n\n{text}" if text else f"**{filename}**" + score = src.get("score") + owi.append({ + "source": {"name": name, "url": url}, + "document": [excerpt], + "metadata": [{"name": name, "filename": filename}], + "distances": [score] if isinstance(score, (int, float)) else [], + }) + return owi diff --git a/backend/lamb/completions/main.py b/backend/lamb/completions/main.py index 8cd712d31..1fcecf306 100644 --- a/backend/lamb/completions/main.py +++ b/backend/lamb/completions/main.py @@ -11,6 +11,7 @@ from lamb.logging_config import get_logger from lamb.auth_context import AuthContext, get_optional_auth_context from lamb.completions.task_routing import maybe_route_non_streaming_task +from lamb.completions.plugin_config import parse_plugin_config, process_completion_request from utils.langsmith_config import traceable_llm_call, add_trace_metadata, is_tracing_enabled import traceback import asyncio @@ -168,8 +169,9 @@ async def create_completion( pps, connectors, rag_processors = load_and_validate_plugins(plugin_config) logger.debug(f"Plugins loaded: {pps}, {connectors}, {rag_processors}") rag_context = await get_rag_context(request, rag_processors, plugin_config["rag_processor"], assistant_details) + document_context = await get_rag_context(request, rag_processors, plugin_config.get("document_rag", ""), assistant_details) logger.debug(f"RAG context: {rag_context}") - messages = process_completion_request(request, assistant_details, plugin_config, rag_context, pps) + messages = process_completion_request(request, assistant_details, plugin_config, rag_context, pps, document_context) logger.debug(f"Messages: {messages}") stream = request.get("stream", False) logger.debug(f"Stream mode: {stream}") @@ -255,6 +257,26 @@ def get_assistant_details(assistant: int) -> Any: return assistant_details +def _assistant_exposes_sources(assistant: Any) -> bool: + """Whether this assistant opts in to student-clickable cited-source links. + + Reads ``metadata.capabilities.expose_sources`` (mirroring the vision / + image_generation capability flags). Defaults to False, so existing + assistants — and any whose JSON lacks the key — never expose their cited + documents until the creator explicitly enables it. + """ + if not assistant: + return False + metadata_str = getattr(assistant, "metadata", None) or getattr(assistant, "api_callback", None) + if not metadata_str: + return False + try: + capabilities = json.loads(metadata_str).get("capabilities", {}) or {} + return bool(capabilities.get("expose_sources", False)) + except (json.JSONDecodeError, AttributeError, TypeError): + return False + + def _provider_for_connector(connector: str) -> str | None: """Map connector name to provider string used in model_pricing table.""" return {"openai": "openai", "anthropic": "anthropic"}.get(connector) @@ -303,41 +325,10 @@ def _check_quota(assistant_id: int, assistant_details) -> None: } ) -def parse_plugin_config(assistant_details) -> Dict[str, str]: - """ - Parse the metadata field from the assistant record. - Expects a JSON string with keys: prompt_processor, connector, llm, rag_processor. - """ - try: - # Handle empty string case by defaulting to an empty JSON object - if not assistant_details.metadata or assistant_details.metadata.strip() == '': - logger.warning(f"Empty metadata for assistant {assistant_details.id}, using default values") - callback = {} - else: - callback = json.loads(assistant_details.metadata) - except Exception as e: - logger.error(f"Failed to parse metadata for assistant {assistant_details.id}: {e}") - raise HTTPException(status_code=400, detail=f"Assistant metadata cannot be parsed: {e}") - - # Set default values if keys are missing - defaults = { - "prompt_processor": "default", - "connector": "openai", - "llm": "gpt-4", - "rag_processor": "" - } - - # Apply defaults for missing keys - for key in defaults: - if key not in callback: - callback[key] = defaults[key] - logger.info(f"Using default {key}={defaults[key]} for assistant {assistant_details.id}") - - return callback - def load_and_validate_plugins(plugin_config: Dict[str, str]) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]: """ Load plugin modules and verify that the requested plugins exist. + Also validates COMPATIBLE_RAG declared by the prompt processor. """ pps = load_plugins('pps') connectors = load_plugins('connectors') @@ -352,6 +343,31 @@ def load_and_validate_plugins(plugin_config: Dict[str, str]) -> Tuple[Dict[str, if plugin_config["rag_processor"] and plugin_config["rag_processor"] not in rag_processors: logger.error(f"RAG processor '{plugin_config['rag_processor']}' not found") raise HTTPException(status_code=400, detail=f"RAG processor '{plugin_config['rag_processor']}' not found") + if plugin_config.get("document_rag") and plugin_config["document_rag"] not in rag_processors: + logger.error(f"Document RAG processor '{plugin_config['document_rag']}' not found") + raise HTTPException(status_code=400, detail=f"Document RAG processor '{plugin_config['document_rag']}' not found") + + pps_name = plugin_config["prompt_processor"] + pps_module = importlib.import_module(f"lamb.completions.pps.{pps_name}") + compatible_rag = getattr(pps_module, "COMPATIBLE_RAG", None) + + if compatible_rag is not None: + if plugin_config["rag_processor"] and plugin_config["rag_processor"] not in compatible_rag: + raise HTTPException( + status_code=400, + detail=( + f"rag_processor '{plugin_config['rag_processor']}' not compatible with " + f"prompt_processor '{pps_name}'. Compatible: {compatible_rag}" + ), + ) + if plugin_config.get("document_rag") and plugin_config["document_rag"] not in compatible_rag: + raise HTTPException( + status_code=400, + detail=( + f"document_rag '{plugin_config['document_rag']}' not compatible with " + f"prompt_processor '{pps_name}'. Compatible: {compatible_rag}" + ), + ) return pps, connectors, rag_processors @@ -378,15 +394,6 @@ async def get_rag_context(request: Dict[str, Any], rag_processors: Dict[str, Any logger.debug("No RAG processor requested") return None -def process_completion_request(request: Dict[str, Any], assistant_details: Any, plugin_config: Dict[str, str], rag_context: Any, pps: Dict[str, Any]) -> Any: - """ - Process the prompt using the specified prompt processor and return prepared messages. - """ - logger.info("Processing completion request") - messages = pps[plugin_config["prompt_processor"]](request=request, assistant=assistant_details, rag_context=rag_context) - logger.debug(f"Processed messages: {messages}") - return messages - def load_plugins(plugin_type: str) -> Dict[str, Any]: """ Dynamically load plugins from the specified directory @@ -455,7 +462,8 @@ def _is_plugin_enabled(module_name: str) -> bool: async def run_lamb_assistant( request: Dict[str, Any], assistant: int, # now expect only an assistant id as int - headers: Optional[Dict[str, str]] = None # Add optional headers argument + headers: Optional[Dict[str, str]] = None, # Add optional headers argument + include_eval_metadata: bool = False # opt-in: attach retrieved RAG context to the response ): """ Implements a non WS version of create completion. @@ -494,7 +502,11 @@ async def run_lamb_assistant( ) pps, connectors, rag_processors = load_and_validate_plugins(plugin_config) rag_context = await get_rag_context(request, rag_processors, plugin_config["rag_processor"], assistant_details) - messages = process_completion_request(request, assistant_details, plugin_config, rag_context, pps) + document_context = await get_rag_context(request, rag_processors, plugin_config.get("document_rag", ""), assistant_details) + messages = process_completion_request(request, assistant_details, plugin_config, rag_context, pps, document_context) + # Clickable source links are opt-in per assistant (default off): only + # expose cited documents to students when the creator enabled it. + expose_sources = _assistant_exposes_sources(assistant_details) stream = request.get("stream", False) llm = plugin_config.get("llm") # Get LLM from config @@ -521,8 +533,32 @@ async def run_lamb_assistant( generator, usage_out = llm_response, None async def _tracked_stream(): + # Append a Markdown "Sources" section (clickable signed links) + # to the answer content just before [DONE]. We deliver it as + # answer CONTENT — not a top-level `sources` field — because + # Open WebUI's chat path forwards only `choices[].delta.content` + # from an external model and drops any `sources` field. + from lamb.completions.citation_sources import build_sources_markdown # noqa: PLC0415 + sources_md = build_sources_markdown(rag_context) if expose_sources else "" + sources_chunk = None + if sources_md: + _delta = {"choices": [{ + "index": 0, + "delta": {"content": sources_md}, + "finish_reason": None, + }]} + sources_chunk = f"data: {json.dumps(_delta)}\n\n" + sources_sent = False async for chunk in generator: + if ( + sources_chunk and not sources_sent + and isinstance(chunk, str) and "data: [DONE]" in chunk + ): + yield sources_chunk + sources_sent = True yield chunk + if sources_chunk and not sources_sent: + yield sources_chunk # Log usage when stream completes for tracked connectors if connector != "ollama" and usage_out and provider and assistant_details.organization_id is not None: db_manager.log_token_usage( @@ -559,6 +595,26 @@ async def _tracked_stream(): usage_data=llm_response["usage"] ) + # Opt-in: surface the retrieved RAG context so the evaluation + # framework can score answers against what was actually retrieved. + # Only attached when the caller set include_eval_metadata; the extra + # top-level key is ignored by standard OpenAI-compatible clients, so + # default callers get an unchanged response. Streaming is left as-is. + if include_eval_metadata and isinstance(rag_context, dict): + llm_response["eval_metadata"] = { + "rag_context": { + "context": rag_context.get("context", ""), + "sources": rag_context.get("sources", []), + } + } + + # Attach OpenWebUI-shaped citations so the panel renders on the + # non-streaming path too (standard OpenAI clients ignore the key). + from lamb.completions.citation_sources import build_owi_sources # noqa: PLC0415 + owi_sources = build_owi_sources(rag_context) if expose_sources else [] + if owi_sources: + llm_response["sources"] = owi_sources + return Response( content=json.dumps(llm_response, indent=2), # Ensure pretty printing if desired media_type="application/json", diff --git a/backend/lamb/completions/org_config_resolver.py b/backend/lamb/completions/org_config_resolver.py index 2767e8e52..a3241cefa 100644 --- a/backend/lamb/completions/org_config_resolver.py +++ b/backend/lamb/completions/org_config_resolver.py @@ -94,7 +94,8 @@ def get_knowledge_base_config(self) -> Dict[str, Any]: # Support both current and legacy shapes. # - Newer docs/configs: org_config.kb_server.{url, api_key} # - Older configs: setups[setup_name].knowledge_base.{server_url, api_token} - kb_config = setup.get("knowledge_base", {}) or org_config.get("kb_server", {}) + # Merge both, with kb_server (newer) taking priority for overlapping keys. + kb_config = {**setup.get("knowledge_base", {}), **org_config.get("kb_server", {})} # Fallback to env vars for system org if not kb_config and self.organization.get('is_system', False): @@ -107,10 +108,18 @@ def get_knowledge_base_config(self) -> Dict[str, Any]: if kb_config: server_url = kb_config.get("server_url") or kb_config.get("url") api_token = kb_config.get("api_token") or kb_config.get("api_key") or kb_config.get("token") - return { + result = { "server_url": server_url, "api_token": api_token, } + # Include embedding model and collection defaults if configured + embedding_model = kb_config.get("embedding_model") + if embedding_model: + result["embedding_model"] = embedding_model + collection_defaults = kb_config.get("collection_defaults") + if collection_defaults: + result["collection_defaults"] = collection_defaults + return result return {} @@ -175,11 +184,20 @@ def get_knowledge_store_config(self) -> Dict[str, Any]: if ks_config: return { - "server_url": ks_config.get("server_url") or ks_config.get("url"), + # Fall back to env when the org block carries only allow-lists + # (e.g. allowed_embedding_models) and omits the server URL/token, + # so a partial knowledge_store config still resolves the server. + "server_url": ( + ks_config.get("server_url") + or ks_config.get("url") + or os.getenv("LAMB_KB_SERVER_V2", "http://kb-server:9092") + ), "api_token": ( ks_config.get("api_token") or ks_config.get("api_key") or ks_config.get("token") + or os.getenv("LAMB_KB_SERVER_V2_TOKEN") + or None ), "allowed_vector_db_backends": ks_config.get("allowed_vector_db_backends", []), "allowed_chunking_strategies": ks_config.get("allowed_chunking_strategies", []), diff --git a/backend/lamb/completions/plugin_config.py b/backend/lamb/completions/plugin_config.py new file mode 100644 index 000000000..db937f5e7 --- /dev/null +++ b/backend/lamb/completions/plugin_config.py @@ -0,0 +1,51 @@ +import json +import logging +from typing import Dict, Any + +from fastapi import HTTPException + +logger = logging.getLogger(__name__) + + +def parse_plugin_config(assistant_details) -> Dict[str, str]: + """ + Parse the metadata field from the assistant record. + Expects a JSON string with keys: prompt_processor, connector, llm, rag_processor. + """ + try: + # Handle empty string case by defaulting to an empty JSON object + if not assistant_details.metadata or assistant_details.metadata.strip() == '': + logger.warning(f"Empty metadata for assistant {assistant_details.id}, using default values") + callback = {} + else: + callback = json.loads(assistant_details.metadata) + except Exception as e: + logger.error(f"Failed to parse metadata for assistant {assistant_details.id}: {e}") + raise HTTPException(status_code=400, detail=f"Assistant metadata cannot be parsed: {e}") + + # Set default values if keys are missing + defaults = { + "prompt_processor": "default", + "connector": "openai", + "llm": "gpt-4", + "rag_processor": "", + "document_rag": "" + } + + # Apply defaults for missing keys + for key in defaults: + if key not in callback: + callback[key] = defaults[key] + logger.info(f"Using default {key}={defaults[key]} for assistant {assistant_details.id}") + + return callback + + +def process_completion_request(request: Dict[str, Any], assistant_details: Any, plugin_config: Dict[str, str], rag_context: Any, pps: Dict[str, Any], document_context=None) -> Any: + """ + Process the prompt using the specified prompt processor and return prepared messages. + """ + logger.info("Processing completion request") + messages = pps[plugin_config["prompt_processor"]](request=request, assistant=assistant_details, rag_context=rag_context, document_context=document_context) + logger.debug(f"Processed messages: {messages}") + return messages diff --git a/backend/lamb/completions/pps/kvcache_augment.py b/backend/lamb/completions/pps/kvcache_augment.py new file mode 100644 index 000000000..0d2595353 --- /dev/null +++ b/backend/lamb/completions/pps/kvcache_augment.py @@ -0,0 +1,231 @@ +from typing import Dict, Any, List, Optional +from lamb.lamb_classes import Assistant +import json +import re +from lamb.logging_config import get_logger + +logger = get_logger(__name__, component="MAIN") + +COMPATIBLE_RAG = [ + "library_file_rag", + "knowledge_store_rag", + "query_rewriting_ks_rag", + "rubric_rag", + "no_rag", +] + +DEFAULT_RAG_PROMPT_TEMPLATE = ( + "Use the following context to answer the question. " + "If the context does not contain the answer, say you do not know.\n\n" + "Context:\n{context}\n\nQuestion: {user_input}" +) + +CITATION_INSTRUCTION = ( + "When you use information from the context above, cite the supporting " + "source inline using its bracketed number, e.g. [1] or [2][3]. Place the " + "citation immediately after the statement it supports. Only cite numbers " + "that appear in the context; do not invent citations." +) + + +def _build_full_context(rag_context: Any, cite: bool = True) -> str: + """Build the text that replaces ``{context}``. + + The retrieved context is already prefixed per chunk with ``[N]`` markers by + the KS RAG processor. When ``cite`` is True we keep those markers and append + the inline citation instruction so the model cites — the clickable source + list is then rendered separately (see ``lamb.completions.citation_sources``). + + When ``cite`` is False (the assistant did NOT opt in to exposing sources), + we strip the ``[N]`` prefixes and omit the instruction, so the answer has no + citation markers at all — students never see ``[1]`` pointing at a source + they cannot open. + """ + if not isinstance(rag_context, dict): + return str(rag_context) if rag_context else "" + context = rag_context.get("context", "") or "" + sources = rag_context.get("sources", []) or [] + if sources and cite: + return context + "\n\n" + CITATION_INSTRUCTION + if not cite and context: + # Drop the per-chunk "[N] " citation prefixes so the model has no cue + # to emit citation markers it cannot link to. + context = re.sub(r"(?m)^\[\d+\]\s+", "", context) + return context + + +def _exposes_sources(assistant: Assistant) -> bool: + """Whether this assistant opts in to student-clickable cited sources. + + Mirrors ``_has_vision_capability``; defaults to False so citations are only + produced when the creator enabled ``capabilities.expose_sources``. + """ + if not assistant: + return False + metadata_str = getattr(assistant, 'metadata', None) or getattr(assistant, 'api_callback', None) + if not metadata_str: + return False + try: + capabilities = json.loads(metadata_str).get('capabilities', {}) or {} + return bool(capabilities.get('expose_sources', False)) + except (json.JSONDecodeError, AttributeError): + return False + + +def _has_vision_capability(assistant: Assistant) -> bool: + if not assistant: + return False + metadata_str = getattr(assistant, 'metadata', None) or getattr(assistant, 'api_callback', None) + if not metadata_str: + return False + try: + metadata = json.loads(metadata_str) + capabilities = metadata.get('capabilities', {}) + return capabilities.get('vision', False) + except (json.JSONDecodeError, AttributeError): + return False + + +def _has_image_generation_capability(assistant: Assistant) -> bool: + if not assistant: + return False + metadata_str = getattr(assistant, 'metadata', None) or getattr(assistant, 'api_callback', None) + if not metadata_str: + return False + try: + metadata = json.loads(metadata_str) + capabilities = metadata.get('capabilities', {}) + return capabilities.get('image_generation', False) + except (json.JSONDecodeError, AttributeError): + return False + + +def prompt_processor( + request: Dict[str, Any], + assistant: Optional[Assistant] = None, + rag_context: Optional[Dict[str, Any]] = None, + document_context: Optional[Dict[str, Any]] = None, +) -> List[Dict[str, str]]: + messages = request.get('messages', []) + if not messages: + return messages + + # Only cite (inline [N] markers + clickable sources) when the assistant + # opted in to exposing sources; otherwise produce no citation markers. + cite = _exposes_sources(assistant) + + last_message = messages[-1]['content'] + processed_messages = [] + + if assistant: + system_content = assistant.system_prompt or "" + if document_context and isinstance(document_context, dict): + doc_text = document_context.get("context", "") + if doc_text: + labeled_doc = ( + "\n\n## REFERENCE DOCUMENT\n\n" + "This document has been selected by the assistant creator as a reference " + "that will likely be useful for many queries, as it is generally a helpful " + "document. Use it as context when answering questions.\n\n" + f"{doc_text}" + "\n\nIMPORTANT: The reference document above is available for this entire " + "conversation. Always consider it alongside any retrieved context when " + "answering questions. If the user's question relates to the document's " + "content, use it." + ) + system_content = (labeled_doc + "\n\n" + system_content) if system_content else labeled_doc + if system_content: + processed_messages.append({ + "role": "system", + "content": system_content + }) + + processed_messages.extend(messages[:-1]) + + if assistant.prompt_template: + has_vision = _has_vision_capability(assistant) + + if isinstance(last_message, list) and has_vision: + augmented_content = [] + text_parts = [] + for item in last_message: + if item.get('type') == 'text': + text_parts.append(item.get('text', '')) + user_input_text = ' '.join(text_parts) + + logger.debug(f"User message: {user_input_text}") + augmented_text = assistant.prompt_template.replace("{user_input}", "\n\n" + user_input_text + "\n\n") + + if rag_context: + full_context = _build_full_context(rag_context, cite) + augmented_text = augmented_text.replace("{context}", "\n\n" + full_context + "\n\n") + else: + augmented_text = augmented_text.replace("{context}", "") + + augmented_content.append({"type": "text", "text": augmented_text}) + for item in last_message: + if item.get('type') != 'text': + augmented_content.append(item) + + processed_messages.append({ + "role": messages[-1]['role'], + "content": augmented_content + }) + else: + if isinstance(last_message, list): + text_parts = [] + for item in last_message: + if item.get('type') == 'text': + text_parts.append(item.get('text', '')) + user_input_text = ' '.join(text_parts) + else: + user_input_text = str(last_message) + + logger.debug(f"User message: {user_input_text}") + prompt = assistant.prompt_template.replace("{user_input}", "\n\n" + user_input_text + "\n\n") + + if rag_context: + full_context = _build_full_context(rag_context, cite) + prompt = prompt.replace("{context}", "\n\n" + full_context + "\n\n") + else: + prompt = prompt.replace("{context}", "") + + processed_messages.append({ + "role": messages[-1]['role'], + "content": prompt + }) + else: + effective_template = None + if rag_context: + context_text = ( + rag_context.get("context", "") + if isinstance(rag_context, dict) + else str(rag_context) + ) + if context_text: + effective_template = DEFAULT_RAG_PROMPT_TEMPLATE + + if effective_template: + if isinstance(last_message, list): + text_parts = [] + for item in last_message: + if item.get('type') == 'text': + text_parts.append(item.get('text', '')) + user_input_text = ' '.join(text_parts) + else: + user_input_text = str(last_message) + + prompt = effective_template.replace("{user_input}", "\n\n" + user_input_text + "\n\n") + full_context = _build_full_context(rag_context, cite) + prompt = prompt.replace("{context}", "\n\n" + full_context + "\n\n") + + processed_messages.append({ + "role": messages[-1]['role'], + "content": prompt + }) + else: + processed_messages.append(messages[-1]) + + return processed_messages + + return messages diff --git a/backend/lamb/completions/pps/simple_augment.py b/backend/lamb/completions/pps/simple_augment.py index 8239a00eb..f3b8c5594 100644 --- a/backend/lamb/completions/pps/simple_augment.py +++ b/backend/lamb/completions/pps/simple_augment.py @@ -5,6 +5,15 @@ logger = get_logger(__name__, component="MAIN") +COMPATIBLE_RAG = [ + "simple_rag", + "context_aware_rag", + "hierarchical_rag", + "single_file_rag", + "rubric_rag", + "no_rag", +] + def _inject_rag_context(template: str, rag_context: Optional[Dict[str, Any]]) -> str: """Replace {context} placeholder with RAG context and formatted sources. @@ -100,7 +109,8 @@ def _has_image_generation_capability(assistant: Assistant) -> bool: def prompt_processor( request: Dict[str, Any], assistant: Optional[Assistant] = None, - rag_context: Optional[Dict[str, Any]] = None + rag_context: Optional[Dict[str, Any]] = None, + document_context: Optional[Dict[str, Any]] = None ) -> List[Dict[str, str]]: """ Simple augment prompt processor that: @@ -130,27 +140,8 @@ def prompt_processor( # Add previous messages except the last one processed_messages.extend(messages[:-1]) - # If RAG context was produced but the assistant has an empty / missing - # prompt template, substitute the default template so the retrieved - # chunks actually reach the LLM. Otherwise the {context} substitution - # below would silently drop them. (Defect D3 — lifecycle 2026-05-03.) - effective_template = assistant.prompt_template - if (not effective_template) and rag_context: - context_text = ( - rag_context.get("context", "") - if isinstance(rag_context, dict) - else str(rag_context) - ) - if context_text: - logger.info( - "simple_augment: applying DEFAULT_RAG_PROMPT_TEMPLATE " - "because assistant has empty prompt_template but " - "rag_context is present (defect D3 fallback)." - ) - effective_template = DEFAULT_RAG_PROMPT_TEMPLATE - # Process the last message using the prompt template - if effective_template: + if assistant.prompt_template: # Check if assistant has vision capabilities has_vision = _has_vision_capability(assistant) @@ -169,7 +160,7 @@ def prompt_processor( # Create augmented text content with template logger.debug(f"User message: {user_input_text}") - augmented_text = effective_template.replace("{user_input}", "\n\n" + user_input_text + "\n\n") + augmented_text = assistant.prompt_template.replace("{user_input}", "\n\n" + user_input_text + "\n\n") augmented_text = _inject_rag_context(augmented_text, rag_context) @@ -205,7 +196,7 @@ def prompt_processor( # Replace placeholders in template logger.debug(f"User message: {user_input_text}") - prompt = effective_template.replace("{user_input}", "\n\n" + user_input_text + "\n\n") + prompt = assistant.prompt_template.replace("{user_input}", "\n\n" + user_input_text + "\n\n") prompt = _inject_rag_context(prompt, rag_context) diff --git a/backend/lamb/completions/rag/_ks_query_helpers.py b/backend/lamb/completions/rag/_ks_query_helpers.py new file mode 100644 index 000000000..a0b7e0618 --- /dev/null +++ b/backend/lamb/completions/rag/_ks_query_helpers.py @@ -0,0 +1,148 @@ +"""Shared helpers for Knowledge Store RAG processors. + +Extracted from ``knowledge_store_rag.py`` so that sibling processors +(e.g. ``query_rewriting_ks_rag``) can reuse the same query and +source-extraction logic without code duplication. +""" + +from __future__ import annotations + +import json +from typing import Any, Dict, List, Optional + +from creator_interface.knowledge_store_client import KnowledgeStoreClient +from lamb.database_manager import LambDatabaseManager +from lamb.logging_config import get_logger + +logger = get_logger(__name__, component="RAG") + +_client = KnowledgeStoreClient() +_db = LambDatabaseManager() + + +def serialize_assistant(assistant) -> Dict[str, Any]: + """Best-effort JSON-safe view of the assistant for logging / response.""" + if not assistant: + return {} + out: Dict[str, Any] = {} + for key in ( + "id", "name", "description", "system_prompt", "prompt_template", + "RAG_collections", "RAG_Top_k", "published", "published_at", "owner", + ): + if hasattr(assistant, key): + try: + value = getattr(assistant, key) + json.dumps({key: value}) + out[key] = value + except (TypeError, OverflowError, Exception): + out[key] = str(value) + return out + + +def build_user_dict_from_owner(owner_email: str) -> Dict[str, Any]: + """Build the minimal ``creator_user`` dict the client expects.""" + return {"email": owner_email} + + +async def query_one_ks( + ks_id: str, + query_text: str, + top_k: int, + owner_email: str, +) -> Dict[str, Any]: + """Query a single Knowledge Store and return the raw KB Server response.""" + ks_entry = _db.get_knowledge_store(ks_id) + if not ks_entry: + return {"status": "error", "error": f"Knowledge Store {ks_id} not found in LAMB DB."} + + creator_user = build_user_dict_from_owner(owner_email) + embedding_api_key = _client.resolve_embedding_api_key( + creator_user=creator_user, + vendor=ks_entry["embedding_vendor"], + ) + embedding_endpoint = ks_entry.get("embedding_endpoint") or "" + + try: + response = await _client.query( + knowledge_store_id=ks_id, + query_text=query_text, + top_k=top_k, + embedding_api_key=embedding_api_key, + embedding_api_endpoint=embedding_endpoint, + creator_user=creator_user, + ) + return {"status": "success", "data": response} + except Exception as e: + logger.error(f"Error querying Knowledge Store {ks_id}: {e}") + return {"status": "error", "error": str(e)} + + +def _build_source(ks_id: str, chunk: Dict[str, Any]) -> Dict[str, Any]: + """Build a single LAMB-side ``source`` dict from one KB Server chunk.""" + meta = chunk.get("metadata", {}) or {} + source: Dict[str, Any] = { + "knowledge_store_id": ks_id, + "title": meta.get("source_title") + or meta.get("title") + or meta.get("library_name") + or "Source", + "score": chunk.get("score"), + "source_item_id": meta.get("source_item_id"), + } + for key in ("permalink_original", "permalink_markdown", "permalink_page"): + if meta.get(key): + source[key] = meta[key] + if meta.get("library_id"): + source["library_id"] = meta["library_id"] + if meta.get("library_name"): + source["library_name"] = meta["library_name"] + primary = ( + meta.get("permalink_page") + or meta.get("permalink_markdown") + or meta.get("permalink_original") + ) + if primary: + source["url"] = primary + return source + + +def _extract_sources( + ks_id: str, + results: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Build the LAMB-side ``sources`` list from KB Server query results.""" + return [_build_source(ks_id, chunk) for chunk in results] + + +def build_context_and_sources( + ks_results: List[tuple], +) -> tuple: + """Build the numbered RAG context and an aligned ``sources`` list. + + ``ks_results`` is a list of ``(ks_id, chunks)`` pairs for the Knowledge + Stores that returned successfully. Every chunk that carries text is given + a stable 1-based citation number ``n``: its context block is prefixed with + ``[n]`` and the matching source dict carries the same ``n``. This lets the + LLM cite inline with ``[N]`` markers while the caller renders a numbered + sources list that lines up exactly with those markers. + + Returns ``(combined_context, sources)``. + """ + context_parts: List[str] = [] + sources: List[Dict[str, Any]] = [] + n = 0 + for ks_id, chunks in ks_results: + for chunk in chunks or []: + text = (chunk.get("text") or "").strip() + if not text: + continue + n += 1 + context_parts.append(f"[{n}] {text}") + source = _build_source(ks_id, chunk) + source["n"] = n + # Keep the chunk text on the source so the citation UI can show the + # supporting excerpt (e.g. OWI's citations panel). + source["text"] = text + sources.append(source) + combined_context = "\n\n".join(context_parts) if context_parts else "" + return combined_context, sources diff --git a/backend/lamb/completions/rag/_query_rewriting_helper.py b/backend/lamb/completions/rag/_query_rewriting_helper.py new file mode 100644 index 000000000..08e62be6d --- /dev/null +++ b/backend/lamb/completions/rag/_query_rewriting_helper.py @@ -0,0 +1,116 @@ +"""Shared query rewriting logic for RAG processors. + +Extracted from ``context_aware_rag.py`` so that both the legacy +context-aware processor and the new ``query_rewriting_ks_rag`` can +reuse the same small-fast-model query optimization without duplication. +""" + +from __future__ import annotations + +from typing import Any, Dict, List + +from lamb.completions.small_fast_model_helper import ( + invoke_small_fast_model, + is_small_fast_model_configured, +) +from lamb.logging_config import get_logger + +logger = get_logger(__name__, component="RAG") + + +def get_last_user_message_text(messages: List[Dict[str, Any]]) -> str: + """Extract text from the last user message, handling multimodal content.""" + for msg in reversed(messages or []): + if msg.get("role") == "user": + content = msg.get("content", "") + if isinstance(content, list): + text_parts = [item.get("text", "") for item in content if item.get("type") == "text"] + return " ".join(text_parts) + return content or "" + return "" + + +async def generate_optimal_query( + messages: List[Dict[str, Any]], + assistant_owner: str, +) -> str: + """Use the small-fast-model to generate an optimized RAG query. + + Falls back silently to the last user message when: + - small-fast-model is not configured + - the model call fails + - the model returns an empty response + """ + if not is_small_fast_model_configured(assistant_owner): + logger.info("Small-fast-model not configured, using last user message as query") + return get_last_user_message_text(messages) + + recent_messages = messages[-10:] if len(messages) > 10 else messages + conversation_summary = [] + for msg in recent_messages: + role = msg.get("role", "") + content = msg.get("content", "") + if isinstance(content, list): + text_parts = [item.get("text", "") for item in content if item.get("type") == "text"] + content = " ".join(text_parts) + if len(content) > 500: + content = content[:500] + "..." + conversation_summary.append(f"{role.upper()}: {content}") + + conversation_text = "\n".join(conversation_summary) + + system_prompt = ( + "You are a query optimization assistant for a RAG " + "(Retrieval-Augmented Generation) system.\n\n" + "Your task is to analyze the conversation history and generate an " + "optimal search query that will retrieve the most relevant documents " + "from a knowledge base.\n\n" + "Guidelines:\n" + "1. Consider the full conversation context, not just the last message\n" + "2. Identify the core information need\n" + "3. Include relevant keywords and concepts\n" + "4. If the conversation references previous topics, incorporate them\n" + "5. Make the query specific and focused\n" + "6. Keep the query concise (1-3 sentences max)\n" + "7. Return ONLY the optimized query, nothing else\n\n" + "Example:\n" + "CONVERSATION:\n" + "USER: What is photosynthesis?\n" + "ASSISTANT: Photosynthesis is the process by which plants convert light energy into chemical energy.\n" + "USER: How does it work in detail?\n\n" + "OPTIMAL QUERY: detailed explanation of photosynthesis process mechanism light energy conversion chlorophyll\n\n" + "Now generate the optimal query for the following conversation:" + ) + + user_prompt = f"CONVERSATION:\n{conversation_text}\n\nOPTIMAL QUERY:" + + enhancement_messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ] + + try: + logger.info("Generating optimal query using small-fast-model...") + response = await invoke_small_fast_model( + messages=enhancement_messages, + assistant_owner=assistant_owner, + stream=False, + ) + + optimized_query = "" + if isinstance(response, dict): + if "choices" in response and len(response["choices"]) > 0: + optimized_query = response["choices"][0]["message"]["content"].strip() + elif "message" in response: + optimized_query = response["message"].get("content", "").strip() + + if optimized_query: + logger.info(f"Optimized query generated: {optimized_query[:100]}...") + return optimized_query + + logger.warning("Empty response from small-fast-model, falling back to last user message") + return get_last_user_message_text(messages) + + except Exception as e: + logger.error(f"Error generating optimal query: {e}", exc_info=True) + return get_last_user_message_text(messages) diff --git a/backend/lamb/completions/rag/context_aware_rag.py b/backend/lamb/completions/rag/context_aware_rag.py index c3e77ab8a..9f72f5b8d 100644 --- a/backend/lamb/completions/rag/context_aware_rag.py +++ b/backend/lamb/completions/rag/context_aware_rag.py @@ -4,152 +4,12 @@ from typing import Dict, Any, List from lamb.lamb_classes import Assistant from lamb.completions.org_config_resolver import OrganizationConfigResolver +from lamb.completions.rag._query_rewriting_helper import generate_optimal_query from lamb.logging_config import get_logger logger = get_logger(__name__, component="RAG") -async def _generate_optimal_query(messages: List[Dict[str, Any]], assistant: Assistant) -> str: - """ - Use the small-fast-model to analyze the full conversation and generate - an optimal query for RAG retrieval. - - Args: - messages: Full conversation history - assistant: Assistant object with owner information - - Returns: - Optimized query string for RAG retrieval - """ - try: - from lamb.completions.small_fast_model_helper import invoke_small_fast_model, is_small_fast_model_configured - - # Check if small-fast-model is configured - if not is_small_fast_model_configured(assistant.owner): - logger.info( - "Small-fast-model not configured, using last user message as query") - # Fallback: return last user message - for msg in reversed(messages): - if msg.get("role") == "user": - content = msg.get("content", "") - if isinstance(content, list): - # Extract text from multimodal content - text_parts = [ - item.get('text', '') for item in content if item.get('type') == 'text'] - return ' '.join(text_parts) - return content - return "" - - # Build a condensed conversation history (last 5 turns max to save tokens) - conversation_summary = [] - recent_messages = messages[-10:] if len(messages) > 10 else messages - - for msg in recent_messages: - role = msg.get("role", "") - content = msg.get("content", "") - - # Handle multimodal content - if isinstance(content, list): - text_parts = [item.get('text', '') - for item in content if item.get('type') == 'text'] - content = ' '.join(text_parts) - - # Truncate very long messages - if len(content) > 500: - content = content[:500] + "..." - - conversation_summary.append(f"{role.upper()}: {content}") - - conversation_text = "\n".join(conversation_summary) - - # Create prompt for query optimization - system_prompt = """You are a query optimization assistant for a RAG (Retrieval-Augmented Generation) system. - -Your task is to analyze the conversation history and generate an optimal search query that will retrieve the most relevant documents from a knowledge base. - -Guidelines: -1. Consider the full conversation context, not just the last message -2. Identify the core information need -3. Include relevant keywords and concepts -4. If the conversation references previous topics, incorporate them -5. Make the query specific and focused -6. Keep the query concise (1-3 sentences max) -7. Return ONLY the optimized query, nothing else - -Example: -CONVERSATION: -USER: What is photosynthesis? -ASSISTANT: Photosynthesis is the process by which plants convert light energy into chemical energy. -USER: How does it work in detail? - -OPTIMAL QUERY: detailed explanation of photosynthesis process mechanism light energy conversion chlorophyll - -Now generate the optimal query for the following conversation:""" - - user_prompt = f"""CONVERSATION: -{conversation_text} - -OPTIMAL QUERY:""" - - # Invoke small-fast-model - enhancement_messages = [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_prompt} - ] - - logger.info("🔍 Generating optimal query using small-fast-model...") - logger.debug( - f"Query optimization context: {len(recent_messages)} messages") - - response = await invoke_small_fast_model( - messages=enhancement_messages, - assistant_owner=assistant.owner, - stream=False - ) - - # Extract the optimized query from response - optimized_query = "" - if isinstance(response, dict): - if 'choices' in response and len(response['choices']) > 0: - optimized_query = response['choices'][0]['message']['content'].strip( - ) - elif 'message' in response: - optimized_query = response['message'].get( - 'content', '').strip() - - if optimized_query: - logger.info( - f"✅ Optimized query generated: {optimized_query[:100]}...") - logger.debug(f"Full optimized query: {optimized_query}") - return optimized_query - else: - logger.warning( - "Empty response from small-fast-model, falling back to last user message") - # Fallback to last user message - for msg in reversed(messages): - if msg.get("role") == "user": - content = msg.get("content", "") - if isinstance(content, list): - text_parts = [ - item.get('text', '') for item in content if item.get('type') == 'text'] - return ' '.join(text_parts) - return content - return "" - - except Exception as e: - logger.error(f"Error generating optimal query: {e}", exc_info=True) - # Fallback: return last user message - for msg in reversed(messages): - if msg.get("role") == "user": - content = msg.get("content", "") - if isinstance(content, list): - text_parts = [ - item.get('text', '') for item in content if item.get('type') == 'text'] - return ' '.join(text_parts) - return content - return "" - - async def rag_processor(messages: List[Dict[str, Any]], assistant: Assistant = None, request: Dict[str, Any] = None) -> Dict[str, Any]: """ Context-aware RAG processor that returns context from the knowledge base server @@ -197,7 +57,8 @@ async def rag_processor(messages: List[Dict[str, Any]], assistant: Assistant = N # Generate optimal query from full conversation context logger.info("Analyzing conversation context for optimal query generation...") - optimal_query = await _generate_optimal_query(messages, assistant) + owner_email = getattr(assistant, "owner", "") or "" + optimal_query = await generate_optimal_query(messages, owner_email) if not optimal_query: # Fallback: extract last user message @@ -362,7 +223,7 @@ async def rag_processor(messages: List[Dict[str, Any]], assistant: Assistant = N print("===========================================\n") - # Print a summary of all responses + # Print a summary of all queries print("\n===== SUMMARY OF ALL QUERIES =====") sources = [] contexts = [] diff --git a/backend/lamb/completions/rag/hierarchical_rag.py b/backend/lamb/completions/rag/hierarchical_rag.py index ad82e1d84..ddb8d893a 100644 --- a/backend/lamb/completions/rag/hierarchical_rag.py +++ b/backend/lamb/completions/rag/hierarchical_rag.py @@ -316,7 +316,7 @@ async def rag_processor(messages: List[Dict[str, Any]], assistant: Assistant = N # Parse the JSON response raw_response = response.json() # Print the entire raw response - print(f"Response Summary: {len(raw_response.get('documents', []))} documents returned") + print(f"Response Summary: {len(raw_response.get('results', raw_response.get('documents', [])))} documents returned") print(f"Raw Response:\n{json.dumps(raw_response, indent=2)}") # Store the response @@ -350,7 +350,8 @@ async def rag_processor(messages: List[Dict[str, Any]], assistant: Assistant = N for cid, result in all_responses.items(): status = result["status"] if status == "success": - documents = result["data"].get("documents", []) + # KB server returns hits under 'results'; fall back to legacy 'documents' (#330) + documents = result["data"].get("results", result["data"].get("documents", [])) doc_count = len(documents) print(f"Collection {cid}: {status} - {doc_count} documents") diff --git a/backend/lamb/completions/rag/knowledge_store_rag.py b/backend/lamb/completions/rag/knowledge_store_rag.py index 4ec066d51..0d21da82d 100644 --- a/backend/lamb/completions/rag/knowledge_store_rag.py +++ b/backend/lamb/completions/rag/knowledge_store_rag.py @@ -21,146 +21,32 @@ from __future__ import annotations import asyncio -import json from typing import Any, Dict, List, Optional -from creator_interface.knowledge_store_client import KnowledgeStoreClient -from lamb.database_manager import LambDatabaseManager from lamb.lamb_classes import Assistant from lamb.logging_config import get_logger +from lamb.completions.rag._ks_query_helpers import ( + serialize_assistant, + query_one_ks, + build_context_and_sources, +) logger = get_logger(__name__, component="RAG") -_client = KnowledgeStoreClient() -_db = LambDatabaseManager() - - -def _serialize_assistant(assistant: Optional[Assistant]) -> Dict[str, Any]: - """Best-effort JSON-safe view of the assistant for logging / response.""" - if not assistant: - return {} - out: Dict[str, Any] = {} - for key in ( - "id", "name", "description", "system_prompt", "prompt_template", - "RAG_collections", "RAG_Top_k", "published", "published_at", "owner", - ): - if hasattr(assistant, key): - try: - value = getattr(assistant, key) - json.dumps({key: value}) - out[key] = value - except (TypeError, OverflowError, Exception): - out[key] = str(value) - return out - - -def _build_user_dict_from_owner(owner_email: str) -> Dict[str, Any]: - """Build the minimal ``creator_user`` dict the client expects. - - The client only uses ``email`` for org-config resolution, so we keep - this lightweight rather than fetching the full row. - """ - return {"email": owner_email} - - -async def _query_one_ks( - ks_id: str, - query_text: str, - top_k: int, - owner_email: str, -) -> Dict[str, Any]: - """Query a single Knowledge Store and return the raw KB Server response. - - Resolves the embedding API key per-KS by looking up the locked vendor - in LAMB DB and reading the org's provider key. - """ - ks_entry = _db.get_knowledge_store(ks_id) - if not ks_entry: - return {"status": "error", "error": f"Knowledge Store {ks_id} not found in LAMB DB."} - - creator_user = _build_user_dict_from_owner(owner_email) - embedding_api_key = _client.resolve_embedding_api_key( - creator_user=creator_user, - vendor=ks_entry["embedding_vendor"], - ) - embedding_endpoint = ks_entry.get("embedding_endpoint") or "" - - try: - response = await _client.query( - knowledge_store_id=ks_id, - query_text=query_text, - top_k=top_k, - embedding_api_key=embedding_api_key, - embedding_api_endpoint=embedding_endpoint, - creator_user=creator_user, - ) - return {"status": "success", "data": response} - except Exception as e: - logger.error(f"Error querying Knowledge Store {ks_id}: {e}") - return {"status": "error", "error": str(e)} - - -def _extract_sources( - ks_id: str, - results: List[Dict[str, Any]], -) -> List[Dict[str, Any]]: - """Build the LAMB-side ``sources`` list from KB Server query results. - - The new KB Server attaches permalink metadata to every chunk (set at - ingestion time by ``knowledge_store_router.add_content``). Permalinks - are LAMB-relative URLs into ``/docs/{org}/{lib}/{item}/...``, so they - resolve through LAMB's ACL-enforced proxy. - """ - sources: List[Dict[str, Any]] = [] - for chunk in results: - meta = chunk.get("metadata", {}) or {} - source: Dict[str, Any] = { - "knowledge_store_id": ks_id, - "title": meta.get("source_title") - or meta.get("title") - or meta.get("library_name") - or "Source", - "score": chunk.get("score"), - "source_item_id": meta.get("source_item_id"), - } - # Permalink URLs are sent into the KB Server as LAMB-relative paths - # at ingestion time; surface whichever variants are present. - for key in ("permalink_original", "permalink_markdown", "permalink_page"): - if meta.get(key): - source[key] = meta[key] - if meta.get("library_id"): - source["library_id"] = meta["library_id"] - if meta.get("library_name"): - source["library_name"] = meta["library_name"] - # Pick a primary URL for renderers that use a single `url` field - # (matches simple_rag's shape so the chat UI's citation renderer - # works without changes). - primary = ( - meta.get("permalink_page") - or meta.get("permalink_markdown") - or meta.get("permalink_original") - ) - if primary: - source["url"] = primary - sources.append(source) - return sources - async def _run( messages: List[Dict[str, Any]], assistant: Optional[Assistant], request: Optional[Dict[str, Any]], ) -> Dict[str, Any]: - """Async core. The exported ``rag_processor`` wraps this in a sync runner - if needed so it works under either the sync or async dispatch path.""" + """Async core.""" logger.info( "Using knowledge_store_rag processor with assistant: %s", assistant.name if assistant else "None", ) - assistant_dict = _serialize_assistant(assistant) + assistant_dict = serialize_assistant(assistant) - # Last user message is the query. last_user_message = "" for msg in reversed(messages or []): if msg.get("role") == "user": @@ -201,33 +87,25 @@ async def _run( f"{getattr(assistant, 'id', '?')} (top_k={top_k})" ) - # Run all KS queries concurrently — each is its own httpx call. raw_responses = await asyncio.gather( - *[_query_one_ks(ks_id, last_user_message, top_k, owner_email) for ks_id in ks_ids], + *[query_one_ks(ks_id, last_user_message, top_k, owner_email) for ks_id in ks_ids], return_exceptions=False, ) all_responses: Dict[str, Any] = {} - sources: List[Dict[str, Any]] = [] - contexts: List[str] = [] - any_success = False + successful: List[tuple] = [] for ks_id, result in zip(ks_ids, raw_responses): all_responses[ks_id] = result if result.get("status") == "success": - any_success = True data = result.get("data", {}) or {} chunks = data.get("results", []) or [] logger.info(f"KS {ks_id}: {len(chunks)} chunks returned") - sources.extend(_extract_sources(ks_id, chunks)) - for chunk in chunks: - text = chunk.get("text") or "" - if text: - contexts.append(text) + successful.append((ks_id, chunks)) else: logger.warning(f"KS {ks_id}: {result.get('error', 'unknown error')}") - combined_context = "\n\n".join(contexts) if contexts else "" + combined_context, sources = build_context_and_sources(successful) return { "context": combined_context, @@ -242,11 +120,5 @@ async def rag_processor( assistant: Assistant = None, request: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: - """RAG processor entry point — discovered automatically by the plugin - loader (``backend/lamb/completions/main.py:load_plugins('rag')``). - - Async by design: the underlying KB Server client uses ``httpx.AsyncClient`` - and the dispatcher in ``get_rag_context`` already handles both sync and - async processors. - """ + """RAG processor entry point.""" return await _run(messages, assistant, request) diff --git a/backend/lamb/completions/rag/library_file_rag.py b/backend/lamb/completions/rag/library_file_rag.py new file mode 100644 index 000000000..995e04d44 --- /dev/null +++ b/backend/lamb/completions/rag/library_file_rag.py @@ -0,0 +1,72 @@ +import os +from typing import Dict, Any, List, Optional +from lamb.lamb_classes import Assistant +import json +import logging +import httpx + +logger = logging.getLogger('lamb.completions.rag.library_file_rag') +logger.setLevel(logging.WARNING) + + +def rag_processor( + messages: List[Dict[str, Any]], + assistant: Assistant = None, + request: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + empty_result = {"context": "", "sources": []} + + if assistant is None: + logger.warning("No assistant provided") + return empty_result + + try: + metadata = json.loads(assistant.metadata) if assistant.metadata else {} + except (json.JSONDecodeError, TypeError): + logger.warning("Invalid metadata JSON") + return empty_result + + library_id = metadata.get("library_id") + item_id = metadata.get("item_id") + + if not library_id or not item_id: + logger.warning("library_id and item_id are required for library_file_rag") + return empty_result + + return _fetch_from_library_manager(library_id, item_id) + + +def _fetch_from_library_manager(library_id: str, item_id: str) -> Dict[str, Any]: + empty_result = {"context": "", "sources": []} + + lm_url = os.environ.get("LAMB_LIBRARY_SERVER", "").rstrip("/") + lm_token = os.environ.get("LAMB_LIBRARY_TOKEN", "") + + if not lm_url or not lm_token: + logger.warning("LAMB_LIBRARY_SERVER or LAMB_LIBRARY_TOKEN not configured") + return empty_result + + url = f"{lm_url}/libraries/{library_id}/items/{item_id}/content" + headers = {"Authorization": f"Bearer {lm_token}"} + + try: + response = httpx.get(url, params={"format": "markdown"}, headers=headers, timeout=30.0) + if response.status_code == 200: + content = response.text + return { + "context": content, + "sources": [{ + "title": "Library Document", + "url": f"/docs/{library_id}/{item_id}", + "similarity": 1.0, + }], + } + else: + logger.warning( + f"Library Manager returned {response.status_code} for " + f"library={library_id} item={item_id}" + ) + return empty_result + except httpx.HTTPError as e: + logger.warning(f"Failed to fetch from Library Manager: {e}") + return empty_result diff --git a/backend/lamb/completions/rag/query_rewriting_ks_rag.py b/backend/lamb/completions/rag/query_rewriting_ks_rag.py new file mode 100644 index 000000000..ce3e67507 --- /dev/null +++ b/backend/lamb/completions/rag/query_rewriting_ks_rag.py @@ -0,0 +1,117 @@ +"""RAG processor with query rewriting backed by Knowledge Stores. + +Combines the query-optimization logic from ``_query_rewriting_helper.py`` +(small-fast-model generates an optimal search query from the full +conversation) with the Knowledge Store query pipeline from +``_ks_query_helpers.py`` (async httpx, per-request embedding +credentials, permalink-based citations). + +When the small-fast-model is not configured for the organization, the +processor silently falls back to using the last user message as the +query — identical to ``knowledge_store_rag.py`` behavior. + +Assistants opt into this by setting +``rag_processor='query_rewriting_ks_rag'`` in their plugin config. +""" + +from __future__ import annotations + +import asyncio +from typing import Any, Dict, List, Optional + +from lamb.lamb_classes import Assistant +from lamb.logging_config import get_logger +from lamb.completions.rag._ks_query_helpers import ( + serialize_assistant, + query_one_ks, + build_context_and_sources, +) +from lamb.completions.rag._query_rewriting_helper import generate_optimal_query + +logger = get_logger(__name__, component="RAG") + + +async def _run( + messages: List[Dict[str, Any]], + assistant: Optional[Assistant], + request: Optional[Dict[str, Any]], +) -> Dict[str, Any]: + """Async core.""" + logger.info( + "Using query_rewriting_ks_rag processor with assistant: %s", + assistant.name if assistant else "None", + ) + + assistant_dict = serialize_assistant(assistant) + + if not assistant or not getattr(assistant, "RAG_collections", None): + return { + "context": "No Knowledge Stores specified in the assistant configuration", + "sources": [], + "assistant_data": assistant_dict, + } + + owner_email = getattr(assistant, "owner", "") or "" + optimal_query = await generate_optimal_query(messages, owner_email) + + if not optimal_query: + return { + "context": "No user message found to use for the query", + "sources": [], + "assistant_data": assistant_dict, + } + + ks_ids = [ + cid.strip() + for cid in (assistant.RAG_collections or "").split(",") + if cid.strip() + ] + if not ks_ids: + return { + "context": "RAG_collections is empty or improperly formatted", + "sources": [], + "assistant_data": assistant_dict, + } + + top_k = getattr(assistant, "RAG_Top_k", 3) or 3 + + logger.info( + f"query_rewriting_ks_rag: querying {len(ks_ids)} KS(es) with query " + f"'{optimal_query[:80]}...' (top_k={top_k})" + ) + + raw_responses = await asyncio.gather( + *[query_one_ks(ks_id, optimal_query, top_k, owner_email) for ks_id in ks_ids], + return_exceptions=False, + ) + + all_responses: Dict[str, Any] = {} + successful: List[tuple] = [] + + for ks_id, result in zip(ks_ids, raw_responses): + all_responses[ks_id] = result + if result.get("status") == "success": + data = result.get("data", {}) or {} + chunks = data.get("results", []) or [] + logger.info(f"KS {ks_id}: {len(chunks)} chunks returned") + successful.append((ks_id, chunks)) + else: + logger.warning(f"KS {ks_id}: {result.get('error', 'unknown error')}") + + combined_context, sources = build_context_and_sources(successful) + + return { + "context": combined_context, + "sources": sources, + "assistant_data": assistant_dict, + "raw_responses": all_responses, + } + + +async def rag_processor( + messages: List[Dict[str, Any]], + assistant: Assistant = None, + request: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + """RAG processor entry point.""" + return await _run(messages, assistant, request) diff --git a/backend/lamb/database_manager.py b/backend/lamb/database_manager.py index dbc4b559d..05d1bf992 100644 --- a/backend/lamb/database_manager.py +++ b/backend/lamb/database_manager.py @@ -620,7 +620,15 @@ def initialize_system_organization(self): self.ensure_system_admin(system_org_id) def ensure_system_admin(self, system_org_id: int): - """Ensure the system admin exists and has proper roles in both OWI and LAMB""" + """Ensure the system admin exists with system role 'admin' AND an 'admin' + organization role in the system org. + + Idempotent and self-healing: runs on every startup, verifies the end state, + and repairs anything missing. Return values are checked and any incomplete + bootstrap is logged at WARNING (not swallowed), so a half-provisioned admin + — which silently hides the Admin menu — becomes visible and is fixed on the + next boot rather than persisting forever. (#405) + """ # First, ensure OWI admin exists self.create_admin_user() @@ -635,40 +643,56 @@ def ensure_system_admin(self, system_org_id: int): password=config.OWI_ADMIN_PASSWORD, organization_id=system_org_id ) - if admin_user_id: - logger.info( - f"Created LAMB admin user: {config.OWI_ADMIN_EMAIL}") - # Assign admin role in organization - self.assign_organization_role( - organization_id=system_org_id, - user_id=admin_user_id, - role="admin" - ) - logger.info( - f"Assigned admin role to user {admin_user_id} in system organization") + if not admin_user_id: + # First-boot races (e.g. migration/column visibility) can land here. + # Bail loudly; the next startup re-runs this and self-heals. + logger.warning( + f"ensure_system_admin: could not create LAMB admin user " + f"{config.OWI_ADMIN_EMAIL}; will retry on next startup") + return + logger.info(f"Created LAMB admin user: {config.OWI_ADMIN_EMAIL}") + # Re-fetch so we self-heal against the persisted row below. + admin_user = self.get_creator_user_by_email(config.OWI_ADMIN_EMAIL) or {} + admin_user_id = admin_user.get('id', admin_user_id) else: - # User exists, ensure they have correct organization and role admin_user_id = admin_user['id'] - - # Check and update organization if needed + # Ensure the admin sits in the system organization if admin_user.get('organization_id') != system_org_id: self.update_user_organization(admin_user_id, system_org_id) - logger.info(f"Updated admin user organization to system org") - - # Check and assign admin role if needed - current_role = self.get_user_organization_role( - admin_user_id, system_org_id) - if current_role != "admin": - self.assign_organization_role( - organization_id=system_org_id, - user_id=admin_user_id, - role="admin" - ) - logger.info( - f"Updated admin user role to 'admin' in system organization") + logger.info("ensure_system_admin: moved admin user to system org") - # Ensure LAMB-side system role matches OWI (bootstrap admin should be system admin). - self.update_creator_user_role(config.OWI_ADMIN_EMAIL, 'admin') + # Self-heal the organization role (the row found empty in #405). + if self.get_user_organization_role(admin_user_id, system_org_id) != "admin": + if self.assign_organization_role( + organization_id=system_org_id, user_id=admin_user_id, role="admin"): + logger.info("ensure_system_admin: assigned admin org-role in system org") + else: + logger.warning( + f"ensure_system_admin: FAILED to assign admin org-role to user " + f"{admin_user_id} in org {system_org_id}") + + # Self-heal the LAMB-side system role (bootstrap admin must be system admin). + if (admin_user or {}).get('role') != 'admin': + if not self.update_creator_user_role(config.OWI_ADMIN_EMAIL, 'admin'): + logger.warning( + f"ensure_system_admin: FAILED to set system role 'admin' for " + f"{config.OWI_ADMIN_EMAIL}") + + # Final verification — loud if still incomplete, so the #405 silent-failure + # chain can't hide a non-admin admin (no Admin menu) again. + final = self.get_creator_user_by_email(config.OWI_ADMIN_EMAIL) or {} + final_org_role = ( + self.get_user_organization_role(admin_user_id, system_org_id) + if admin_user_id else None) + if final.get('role') != 'admin' or final_org_role != 'admin': + logger.warning( + f"ensure_system_admin INCOMPLETE for {config.OWI_ADMIN_EMAIL}: " + f"system_role={final.get('role')!r}, org_role={final_org_role!r}. " + f"Admin menu will not appear; will retry on next startup.") + else: + logger.info( + f"ensure_system_admin: verified {config.OWI_ADMIN_EMAIL} has " + f"system role + org role = admin") def run_migrations(self): """Run database migrations for schema updates""" @@ -2983,6 +3007,25 @@ def get_user_organizations(self, user_id: int) -> List[Dict[str, Any]]: finally: connection.close() + def get_user_organization(self, user_id: int) -> Optional[Dict[str, Any]]: + """The user's primary organization as a full dict (incl. parsed config), or None. + + Resolves the first organization the user belongs to (system org sorts first) + and returns the complete record via get_organization_by_id so callers get + `config`. Added for #325 (rubric visibility/access paths called this). + """ + orgs = self.get_user_organizations(user_id) + if not orgs: + return None + return self.get_organization_by_id(orgs[0]['id']) + + def get_user_organization_by_email(self, email: str) -> Optional[Dict[str, Any]]: + """The user's primary organization looked up by email, or None (#325).""" + user = self.get_creator_user_by_email(email) + if not user or not user.get('id'): + return None + return self.get_user_organization(user['id']) + def get_user_organization_role(self, user_id: int, organization_id: int) -> Optional[str]: """ Get the user's role in a specific organization @@ -5231,6 +5274,41 @@ def publish_assistant(self, assistant_id: int, assistant_name: str, assistant_ow connection.close() return False + def update_assistant_publication(self, assistant_id: int, group_id: str, + group_name: str, oauth_consumer_name: Optional[str]) -> bool: + """Update the publication record for an already-published assistant (#397). + + Updates group/oauth fields in place (name/owner are unchanged on update). + Returns True if a publication row was updated. + """ + connection = self.get_connection() + if connection: + try: + with connection: + cursor = connection.cursor() + cursor.execute(f""" + UPDATE {self.table_prefix}assistant_publish + SET group_id = ?, group_name = ?, oauth_consumer_name = ? + WHERE assistant_id = ? + """, (group_id, group_name, oauth_consumer_name, assistant_id)) + updated = cursor.rowcount + if updated > 0: + logger.info(f"Updated publication for assistant {assistant_id}") + else: + logger.warning( + f"update_assistant_publication: no publication row for assistant {assistant_id}") + return updated > 0 + except sqlite3.Error as e: + if "UNIQUE constraint failed" in str(e) and "oauth_consumer_name" in str(e): + logger.error( + f"Error updating publication for {assistant_id}: oauth_consumer_name '{oauth_consumer_name}' already in use.") + return False + logger.error(f"Error updating publication for assistant {assistant_id}: {e}") + return False + finally: + connection.close() + return False + def get_published_assistants(self) -> list: """Get list of published assistants, optionally filtered by owner""" connection = self.get_connection() diff --git a/backend/lamb/evaluaitor/rubric_database.py b/backend/lamb/evaluaitor/rubric_database.py index 8b0e4fdcf..c424c1e41 100644 --- a/backend/lamb/evaluaitor/rubric_database.py +++ b/backend/lamb/evaluaitor/rubric_database.py @@ -476,7 +476,13 @@ def set_showcase_status(self, rubric_id: str, is_showcase: bool, admin_email: st if not system_org: return False - admin_check = self.db_manager.get_user_organization_role(system_org['id'], admin_email) + # get_user_organization_role(user_id, organization_id) — resolve the + # admin's user_id from their email; args were previously swapped (#325). + admin_user = self.db_manager.get_creator_user_by_email(admin_email) + admin_check = ( + self.db_manager.get_user_organization_role(admin_user['id'], system_org['id']) + if admin_user and admin_user.get('id') else None + ) if admin_check != 'admin': return False diff --git a/backend/lamb/lti_activity_manager.py b/backend/lamb/lti_activity_manager.py index 8a087883c..18f8a1b37 100644 --- a/backend/lamb/lti_activity_manager.py +++ b/backend/lamb/lti_activity_manager.py @@ -9,6 +9,7 @@ import os import re import time +import secrets import hmac import hashlib import base64 @@ -303,9 +304,12 @@ def reconfigure_activity( for aid in to_add: owi_model.add_group_to_model(f"lamb_assistant.{aid}", owi_group_id, "read") - # Remove group from old models + # Remove group from old models (check the result — it used to silently + # fail because remove_group_from_model was broken, #399) for aid in to_remove: - owi_model.remove_group_from_model(f"lamb_assistant.{aid}", owi_group_id, "read") + ok = owi_model.remove_group_from_model(f"lamb_assistant.{aid}", owi_group_id, "read") + if not ok: + logger.warning(f"Failed to remove group {owi_group_id} from model lamb_assistant.{aid} (access may persist)") # Update DB if to_remove: @@ -358,7 +362,7 @@ def handle_student_launch( owi_user = self.owi_user_manager.create_user( name=display_name, email=email, - password=f"lti_activity_{activity['id']}", + password=secrets.token_urlsafe(32), # Random; header-trust signin, not this password (#411) role="user" ) if not owi_user: diff --git a/backend/lamb/lti_users_router.py b/backend/lamb/lti_users_router.py index fdff6cd28..4d2932ec1 100644 --- a/backend/lamb/lti_users_router.py +++ b/backend/lamb/lti_users_router.py @@ -9,6 +9,7 @@ from lamb.logging_config import get_logger from urllib.parse import unquote import os +import secrets import hmac import hashlib import base64 @@ -61,7 +62,7 @@ async def create_lti_user(request: Request, current_user: str = Depends(get_curr owi_user = owi_user_manager.create_user( name=lti_user.user_display_name, email=lti_user.user_email, - password=lti_user.assistant_id, # Using assistant_id as password + password=secrets.token_urlsafe(32), # Random; mirror users sign in via header-trust, not this password (#411) role="user" ) @@ -209,7 +210,7 @@ async def sign_in_lti_user(request: Request, current_user: str = Depends(get_cur owi_user = owi_user_manager.create_user( name=username, email=email, - password=str(published_assistant['assistant_id']), + password=secrets.token_urlsafe(32), # Random; header-trust signin, not this password (#411) role="user" ) diff --git a/backend/lamb/owi_bridge/owi_group.py b/backend/lamb/owi_bridge/owi_group.py index 1b587ffca..e25dc1e57 100644 --- a/backend/lamb/owi_bridge/owi_group.py +++ b/backend/lamb/owi_bridge/owi_group.py @@ -456,8 +456,9 @@ def get_group_users(self, group_id: str) -> List[Dict]: List[Dict]: List of users with their details """ try: - # First verify the group exists - group = self.db.get_group_by_id(group_id) + # First verify the group exists (get_group_by_id is a method of this + # class, not of self.db / OwiDatabaseManager). + group = self.get_group_by_id(group_id) if not group: return None diff --git a/backend/lamb/owi_bridge/owi_model.py b/backend/lamb/owi_bridge/owi_model.py index a355ed533..af8772491 100644 --- a/backend/lamb/owi_bridge/owi_model.py +++ b/backend/lamb/owi_bridge/owi_model.py @@ -231,34 +231,41 @@ def remove_group_from_model( Remove a group from model's access control list """ try: + # Use execute_query (OwiDatabaseManager has no .execute); mirror + # add_group_to_model. Previously self.db.execute(...) raised + # AttributeError, was swallowed, and every removal silently no-op'd — + # so unsharing never propagated to OWI (#399). query = "SELECT access_control FROM model WHERE id = ?" - result = self.db.execute(query, (model_id,)).fetchone() - + result = self.db.execute_query(query, (model_id,), fetch_one=True) + if not result: return False - - access_control = json.loads(result[0]) - + + access_control = json.loads(result[0]) if result[0] else { + "read": {"group_ids": [], "user_ids": []}, + "write": {"group_ids": [], "user_ids": []} + } + if permission_type not in ["read", "write"]: raise ValueError("Permission type must be 'read' or 'write'") - + if group_id in access_control[permission_type]["group_ids"]: access_control[permission_type]["group_ids"].remove(group_id) - + update_query = """ - UPDATE model + UPDATE model SET access_control = ?, updated_at = ? WHERE id = ? """ - + current_time = int(datetime.now().timestamp()) - self.db.execute( + success = self.db.execute_query( update_query, (json.dumps(access_control), current_time, model_id) ) - - return True + + return success is not None except Exception as e: print(f"Error removing group from model: {e}") @@ -271,11 +278,11 @@ def get_model_groups(self, model_id: str) -> Dict[str, List[str]]: """ try: query = "SELECT access_control FROM model WHERE id = ?" - result = self.db.execute(query, (model_id,)).fetchone() - - if not result: + result = self.db.execute_query(query, (model_id,), fetch_one=True) + + if not result or not result[0]: return {"read": [], "write": []} - + access_control = json.loads(result[0]) return { "read": access_control["read"]["group_ids"], diff --git a/backend/lamb/owi_bridge/owi_router.py b/backend/lamb/owi_bridge/owi_router.py index 4ab3300ff..b570883d0 100644 --- a/backend/lamb/owi_bridge/owi_router.py +++ b/backend/lamb/owi_bridge/owi_router.py @@ -791,18 +791,19 @@ async def update_user_role(request: Request): logger.error(f"[ULTRA_BASIC] update_user_role endpoint was called with method {request.method}") try: - # Check API key manually to ensure we're authenticated + # Validate the bearer token against the system API key (#409). + # Previously this checked only the "Bearer " prefix and never + # compared the token, so any bearer-shaped header authorized a + # role change (admin privilege escalation). Do not log the header. auth_header = request.headers.get("Authorization", "") - logger.error(f"[ULTRA_BASIC] Auth header: {auth_header}") - - # Very basic auth validation - if not auth_header.startswith("Bearer "): - logger.error("[ULTRA_BASIC] Missing or invalid authorization header") + token = auth_header[len("Bearer "):] if auth_header.startswith("Bearer ") else "" + if not token or token != API_KEY: + logger.error("[ULTRA_BASIC] Missing or invalid authorization token") return JSONResponse( status_code=401, content={"error": "Missing or invalid authorization"} ) - + # Manually parse the request body - no FastAPI magic body = await request.body() logger.error(f"[ULTRA_BASIC] Raw request body: {body}") diff --git a/backend/lamb/services/assistant_service.py b/backend/lamb/services/assistant_service.py index afce2e265..177fa348b 100644 --- a/backend/lamb/services/assistant_service.py +++ b/backend/lamb/services/assistant_service.py @@ -238,9 +238,13 @@ def update_assistant_publication( oauth_consumer_name=oauth_consumer_name ) - def unpublish_assistant(self, assistant_id: int, group_id: str) -> bool: - """Unpublish an assistant""" - return self.db_manager.unpublish_assistant(assistant_id, group_id) + def unpublish_assistant(self, assistant_id: int, group_id: str = None) -> bool: + """Unpublish an assistant. + + The publication row is keyed by assistant_id, so group_id is accepted for + call-site compatibility but not needed by the DB layer (#397). + """ + return self.db_manager.unpublish_assistant(assistant_id) def validate_assistant_name(self, name: str) -> Tuple[bool, Optional[str]]: """ diff --git a/backend/lamb/services/assistant_sharing_service.py b/backend/lamb/services/assistant_sharing_service.py index d4bced6b3..fafee2bf0 100644 --- a/backend/lamb/services/assistant_sharing_service.py +++ b/backend/lamb/services/assistant_sharing_service.py @@ -349,9 +349,40 @@ def _sync_assistant_to_owi_group(self, assistant_id: int): ) group_id = result.get('id') - # Add all users to the assistant_X group + # Add all entitled users to the assistant_X group self._add_users_to_owi_group(group_id, user_emails) + # Reconcile: remove members who are no longer entitled (#399). Previously + # the sync only ever added, so unsharing left the revoked user in the OWI + # group and they kept runtime access. Diff current members (by OWI user_id) + # against the desired set and remove the difference. We reconcile by id and + # use only the working group primitives (get_group_by_id + update_group via + # remove_user_from_group); get_group_users_by_emails relies on a broken + # OwiDatabaseManager.get_group_by_id call and silently returns nothing. + desired_ids = set() + for email in user_emails: + if not email: + continue + u = self.user_manager.get_user_by_email(email) + if u and u.get('id'): + desired_ids.add(u['id']) + + import json as json_lib + group = self.group_manager.get_group_by_id(group_id) + current_ids = (group or {}).get('user_ids', []) or [] + if isinstance(current_ids, str): + try: + current_ids = json_lib.loads(current_ids) + except Exception: + current_ids = [] + to_remove_ids = [uid for uid in current_ids if uid not in desired_ids] + # Safety: never reconcile against an empty desired set (e.g. owner OWI + # lookup failed) — that would strip every member. + if to_remove_ids and desired_ids: + logger.info(f"Unshare reconcile: removing {len(to_remove_ids)} user_id(s) from group {group_id}") + for uid in to_remove_ids: + self.group_manager.remove_user_from_group(group_id, uid) + def _add_users_to_owi_group(self, group_id: str, user_emails: List[str]): """Add users to OWI group by email""" for email in user_emails: diff --git a/backend/main.py b/backend/main.py index 458adb565..b11361f46 100644 --- a/backend/main.py +++ b/backend/main.py @@ -753,6 +753,11 @@ async def generate_openai_chat_completion(request: Request): model = form_data.get('model') messages = form_data.get('messages', []) stream = form_data.get('stream', False) + # Opt-in eval flag: when true, the response carries the retrieved RAG + # context under `eval_metadata` (consumed by the evaluation framework). + # Captured here because the DummyFormData wrapper below only preserves + # model/messages/stream; absent by default so normal callers are unaffected. + include_eval_metadata = bool(form_data.get('include_eval_metadata', False)) multimodal_logger.info("Final parsed data:") multimodal_logger.info(f"Model: {model}") @@ -822,10 +827,14 @@ def model_dump(self): request_data = form_data.model_dump() multimodal_logger.debug(f"Request data being sent: {json.dumps(request_data, indent=2)[:1000]}...") + # Pass the opt-in eval flag as an explicit arg, NOT inside request_data: + # the connector forwards request body fields to the provider SDK, so an + # unknown key in the body would break the upstream API call. response = await run_lamb_assistant( request=request_data, assistant=assistant_id, - headers=common_headers # Pass headers to the assistant runner + headers=common_headers, # Pass headers to the assistant runner + include_eval_metadata=include_eval_metadata ) multimodal_logger.info(f"Assistant returned response, type: {type(response)}") diff --git a/backend/pytest.ini b/backend/pytest.ini index 0e42f9c6e..376a3ba78 100644 --- a/backend/pytest.ini +++ b/backend/pytest.ini @@ -1,6 +1,7 @@ [pytest] testpaths = tests python_files = test_*.py +asyncio_mode = auto filterwarnings = ignore::DeprecationWarning ignore::PendingDeprecationWarning diff --git a/backend/tests/test_auth_context_diffcov.py b/backend/tests/test_auth_context_diffcov.py new file mode 100644 index 000000000..953c1b256 --- /dev/null +++ b/backend/tests/test_auth_context_diffcov.py @@ -0,0 +1,175 @@ +""" +Diff-coverage tests for lamb.auth_context knowledge-store access methods. + +Covers can_access_knowledge_store and require_knowledge_store_access +(feature-added lines). Mirrors the construction style of test_auth_context.py. + +Run with: pytest backend/tests/test_auth_context_diffcov.py -v +""" + +from unittest.mock import patch + +import pytest +from fastapi import HTTPException + +from lamb.auth_context import AuthContext + + +# --------------------------------------------------------------------------- +# Helpers (mirrors test_auth_context.py construction style) +# --------------------------------------------------------------------------- + +def _make_user( + user_id=1, + email="user@example.com", + name="Test User", + organization_id=10, + user_type="creator", + role="user", +): + return { + "id": user_id, + "email": email, + "name": name, + "organization_id": organization_id, + "user_type": user_type, + "role": role, + "user_config": {}, + "enabled": True, + "lti_user_id": None, + "auth_provider": "password", + "password_hash": None, + } + + +def _make_organization(org_id=10): + config = { + "features": { + "rag_enabled": True, + "mcp_enabled": True, + "lti_publishing": True, + "signup_enabled": False, + "sharing_enabled": True, + } + } + return { + "id": org_id, + "name": "Test Org", + "slug": "test-org", + "is_system": False, + "status": "active", + "config": config, + "created_at": "2024-01-01", + "updated_at": "2024-01-01", + } + + +def _make_auth_context(user=None, org=None, org_role="member", is_admin=False): + user = user if user is not None else _make_user() + payload = {"email": user.get("email"), "role": "admin" if is_admin else "user", "sub": "1"} + org = org if org is not None else _make_organization() + features = org.get("config", {}).get("features", {}) + + return AuthContext( + user=user, + token_payload=payload, + is_system_admin=is_admin, + organization_role=org_role, + is_org_admin=org_role in ("owner", "admin"), + organization=org, + features=features, + ) + + +# --------------------------------------------------------------------------- +# can_access_knowledge_store +# --------------------------------------------------------------------------- + +class TestCanAccessKnowledgeStore: + def test_no_user_id_returns_none(self): + # Lines 269-271: user has no id + ctx = _make_auth_context(user=_make_user(user_id=None)) + assert ctx.can_access_knowledge_store("ks-1") == "none" + + def test_db_grants_owner(self): + # Lines 273-275: db says can_access True with access_type + ctx = _make_auth_context() + with patch("lamb.auth_context._db") as mock_db: + mock_db.user_can_access_knowledge_store.return_value = (True, "owner") + assert ctx.can_access_knowledge_store("ks-1") == "owner" + + def test_db_grants_shared(self): + ctx = _make_auth_context() + with patch("lamb.auth_context._db") as mock_db: + mock_db.user_can_access_knowledge_store.return_value = (True, "shared") + assert ctx.can_access_knowledge_store("ks-1") == "shared" + + def test_system_admin_fallback(self): + # Lines 277-278: system admin gets owner even when db denies + ctx = _make_auth_context(is_admin=True) + with patch("lamb.auth_context._db") as mock_db: + mock_db.user_can_access_knowledge_store.return_value = (False, "none") + assert ctx.can_access_knowledge_store("ks-1") == "owner" + + def test_org_admin_same_org(self): + # Lines 280-283: org admin, KS in same org -> owner + ctx = _make_auth_context(org_role="admin") + with patch("lamb.auth_context._db") as mock_db: + mock_db.user_can_access_knowledge_store.return_value = (False, "none") + mock_db.get_knowledge_store.return_value = {"id": "ks-1", "organization_id": 10} + assert ctx.can_access_knowledge_store("ks-1") == "owner" + + def test_org_admin_different_org_returns_none(self): + # Lines 280-283 (false branch) + 285: org admin, KS in other org -> none + ctx = _make_auth_context(org_role="admin") + with patch("lamb.auth_context._db") as mock_db: + mock_db.user_can_access_knowledge_store.return_value = (False, "none") + mock_db.get_knowledge_store.return_value = {"id": "ks-1", "organization_id": 99} + assert ctx.can_access_knowledge_store("ks-1") == "none" + + def test_org_admin_missing_entry_returns_none(self): + # Line 285: org admin, KS not found -> none + ctx = _make_auth_context(org_role="admin") + with patch("lamb.auth_context._db") as mock_db: + mock_db.user_can_access_knowledge_store.return_value = (False, "none") + mock_db.get_knowledge_store.return_value = None + assert ctx.can_access_knowledge_store("ks-1") == "none" + + def test_member_denied_returns_none(self): + # Line 285: plain member, db denies -> none + ctx = _make_auth_context(org_role="member") + with patch("lamb.auth_context._db") as mock_db: + mock_db.user_can_access_knowledge_store.return_value = (False, "none") + assert ctx.can_access_knowledge_store("ks-1") == "none" + + +# --------------------------------------------------------------------------- +# require_knowledge_store_access +# --------------------------------------------------------------------------- + +class TestRequireKnowledgeStoreAccess: + def test_any_level_passes(self): + ctx = _make_auth_context() + with patch.object(ctx, "can_access_knowledge_store", return_value="shared"): + assert ctx.require_knowledge_store_access("ks-1", level="any") == "shared" + + def test_none_raises_404(self): + # Lines 300-301 + ctx = _make_auth_context() + with patch.object(ctx, "can_access_knowledge_store", return_value="none"): + with pytest.raises(HTTPException) as exc_info: + ctx.require_knowledge_store_access("ks-1") + assert exc_info.value.status_code == 404 + + def test_owner_level_denied_raises_403(self): + # Lines 303-304 + ctx = _make_auth_context() + with patch.object(ctx, "can_access_knowledge_store", return_value="shared"): + with pytest.raises(HTTPException) as exc_info: + ctx.require_knowledge_store_access("ks-1", level="owner") + assert exc_info.value.status_code == 403 + + def test_owner_level_passes(self): + ctx = _make_auth_context() + with patch.object(ctx, "can_access_knowledge_store", return_value="owner"): + assert ctx.require_knowledge_store_access("ks-1", level="owner") == "owner" diff --git a/backend/tests/test_citation_sources.py b/backend/tests/test_citation_sources.py new file mode 100644 index 000000000..c436f5f93 --- /dev/null +++ b/backend/tests/test_citation_sources.py @@ -0,0 +1,98 @@ +"""Tests for mapping LAMB RAG sources to the OpenWebUI citations schema.""" + +from urllib.parse import parse_qs, urlparse + +from lamb.completions.citation_sources import build_owi_sources, build_sources_markdown +from creator_interface.permalink_signing import verify + + +def _source(n, **over): + src = { + "n": n, + "title": "Doc", + "score": 0.9, + "text": "supporting excerpt", + "permalink_markdown": "/docs/3/lib-1/item-1/content", + "permalink_original": "/docs/3/lib-1/item-1/original/noa-ventura-cv.pdf", + } + src.update(over) + return src + + +class TestBuildOwiSources: + def test_empty_when_no_sources(self): + assert build_owi_sources({"sources": []}) == [] + assert build_owi_sources(None) == [] + + def test_shape_matches_owi_contract(self): + out = build_owi_sources({"sources": [_source(1)]}) + assert len(out) == 1 + entry = out[0] + # name == "1" so OWI auto-links the inline [1] marker. + assert entry["source"]["name"] == "1" + assert entry["metadata"][0]["name"] == "1" + # document[] holds the excerpt under the real filename. + assert entry["document"][0].startswith("**noa-ventura-cv.pdf**") + assert "supporting excerpt" in entry["document"][0] + assert entry["metadata"][0]["filename"] == "noa-ventura-cv.pdf" + assert entry["distances"] == [0.9] + + def test_url_is_absolute_signed_and_verifiable(self): + out = build_owi_sources({"sources": [_source(1)]}) + url = out[0]["source"]["url"] + assert url.startswith("http") # OWI only links http(s) urls + parsed = urlparse(url) + assert parsed.path == "/docs/public/3/lib-1/item-1/view" + sig = parse_qs(parsed.query)["sig"][0] + assert verify("3", "3/lib-1/item-1/view", sig) is True + # filename carried for the page title / download + assert parse_qs(parsed.query)["name"][0] == "noa-ventura-cv.pdf" + + def test_numbering_preserved_across_multiple_sources(self): + out = build_owi_sources({"sources": [_source(1), _source(2, title="Doc2")]}) + assert [e["source"]["name"] for e in out] == ["1", "2"] + + def test_missing_score_yields_empty_distances(self): + out = build_owi_sources({"sources": [_source(1, score=None)]}) + assert out[0]["distances"] == [] + + def test_no_permalink_leaves_url_empty(self): + src = {"n": 1, "title": "Doc", "text": "x"} + out = build_owi_sources({"sources": [src]}) + # No permalink → no signed URL (OWI renders it non-clickable). + assert out[0]["source"]["url"] == "" + assert out[0]["document"][0].startswith("**Doc**") + + +class TestBuildSourcesMarkdown: + def test_empty_when_no_sources(self): + assert build_sources_markdown({"sources": []}) == "" + assert build_sources_markdown(None) == "" + + def test_renders_clickable_numbered_list(self): + md = build_sources_markdown({"sources": [_source(1)]}) + assert "**Sources**" in md + # [N] marker plus a clickable markdown link to the signed view page. + assert "[1] [noa-ventura-cv.pdf](http" in md + # Not the reference-definition form (OWI strips "[1]: url"). + assert "[1]:" not in md + + def test_groups_chunks_from_same_item_into_one_line(self): + # Two chunks from the same item → one line listing BOTH numbers so each + # inline marker resolves: "[1][2] [filename](url)". + a = _source(1) + b = _source(2) # same permalinks/item as _source default + md = build_sources_markdown({"sources": [a, b]}) + assert md.count("[noa-ventura-cv.pdf](") == 1 # one clickable link + assert "[1][2] [noa-ventura-cv.pdf](http" in md + + def test_distinct_items_each_listed(self): + a = _source(1) + b = _source( + 2, + permalink_markdown="/docs/3/lib-2/item-2/content", + permalink_original="/docs/3/lib-2/item-2/original/other.pdf", + ) + md = build_sources_markdown({"sources": [a, b]}) + assert "noa-ventura-cv.pdf" in md and "other.pdf" in md + assert "[1]" in md and "[2]" in md diff --git a/backend/tests/test_database_manager_ks_diffcov.py b/backend/tests/test_database_manager_ks_diffcov.py new file mode 100644 index 000000000..5c46c1e16 --- /dev/null +++ b/backend/tests/test_database_manager_ks_diffcov.py @@ -0,0 +1,640 @@ +"""Diff-coverage tests for Knowledge-Store / kb_content_links DB methods. + +These exercise the real ``LambDatabaseManager`` against a temporary SQLite +database (built by the manager's own migrations under a tmp dir). Both the +happy paths and the error / None / not-found branches of the methods in the +``lamb.database_manager`` target range (7715-8364) are covered. +""" + +from __future__ import annotations + +import time +from unittest.mock import patch + +import pytest + +import config + + +@pytest.fixture +def dbm(tmp_path, monkeypatch): + """Construct a real manager backed by a fresh temp lamb_v4.db. + + Seeds one organization and two creator users so the FK constraints on + ``knowledge_stores`` / ``kb_content_links`` are satisfiable. + """ + monkeypatch.setattr(config, "LAMB_DB_PATH", str(tmp_path)) + from lamb.database_manager import LambDatabaseManager + + manager = LambDatabaseManager() + + # __init__ auto-creates the system organization with id=1; reuse it. + conn = manager.get_connection() + now = int(time.time()) + with conn: + cur = conn.cursor() + cur.execute( + f"INSERT INTO {manager.table_prefix}Creator_users " + f"(id, organization_id, user_email, user_name, user_type, user_config, created_at, updated_at) " + f"VALUES (10, 1, 'owner@example.com', 'Owner', 'creator', '{{}}', ?, ?)", + (now, now), + ) + cur.execute( + f"INSERT INTO {manager.table_prefix}Creator_users " + f"(id, organization_id, user_email, user_name, user_type, user_config, created_at, updated_at) " + f"VALUES (11, 1, 'other@example.com', 'Other', 'creator', '{{}}', ?, ?)", + (now, now), + ) + conn.close() + return manager + + +def _make_ks(dbm, ks_id="ks-1", name="KS One", owner=10, org=1, + chunking_params=None, status="active"): + return dbm.create_knowledge_store( + knowledge_store_id=ks_id, + name=name, + owner_user_id=owner, + organization_id=org, + chunking_strategy="simple", + embedding_vendor="openai", + embedding_model="text-embedding-3-small", + vector_db_backend="chromadb", + description="desc", + chunking_params=chunking_params, + status=status, + ) + + +# --------------------------------------------------------------------------- +# create_knowledge_store +# --------------------------------------------------------------------------- + + +def test_create_knowledge_store_success(dbm): + result = _make_ks(dbm, chunking_params={"size": 500}) + assert result == "ks-1" + + +def test_create_knowledge_store_duplicate_name(dbm): + assert _make_ks(dbm, ks_id="ks-a", name="Dup") == "ks-a" + # Same org + name violates UNIQUE -> IntegrityError -> None + assert _make_ks(dbm, ks_id="ks-b", name="Dup") is None + + +def test_create_knowledge_store_db_error(dbm): + # Invalid SQL via bad table prefix triggers sqlite3.Error branch + dbm.table_prefix = "nonexistent_" + assert _make_ks(dbm, ks_id="ks-err") is None + + +def test_create_knowledge_store_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert _make_ks(dbm, ks_id="ks-nc") is None + + +# --------------------------------------------------------------------------- +# update_knowledge_store_status +# --------------------------------------------------------------------------- + + +def test_update_status_success_and_missing(dbm): + _make_ks(dbm, ks_id="ks-s", status="provisional") + assert dbm.update_knowledge_store_status("ks-s", "active") is True + assert dbm.update_knowledge_store_status("missing", "active") is False + + +def test_update_status_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.update_knowledge_store_status("ks-s", "active") is False + + +def test_update_status_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.update_knowledge_store_status("ks-s", "active") is False + + +# --------------------------------------------------------------------------- +# get_knowledge_store +# --------------------------------------------------------------------------- + + +def test_get_knowledge_store_success(dbm): + _make_ks(dbm, ks_id="ks-g", chunking_params={"size": 100}) + ks = dbm.get_knowledge_store("ks-g") + assert ks is not None + assert ks["id"] == "ks-g" + assert ks["is_shared"] is False + assert ks["chunking_params"] == {"size": 100} + assert ks["owner_email"] == "owner@example.com" + + +def test_get_knowledge_store_bad_json_params(dbm): + _make_ks(dbm, ks_id="ks-bad") + # Corrupt chunking_params to a non-JSON string -> except -> {} + conn = dbm.get_connection() + with conn: + conn.execute( + f"UPDATE {dbm.table_prefix}knowledge_stores SET chunking_params = 'not json' WHERE id = ?", + ("ks-bad",), + ) + conn.close() + ks = dbm.get_knowledge_store("ks-bad") + assert ks["chunking_params"] == {} + + +def test_get_knowledge_store_not_found(dbm): + assert dbm.get_knowledge_store("nope") is None + + +def test_get_knowledge_store_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.get_knowledge_store("ks-g") is None + + +def test_get_knowledge_store_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.get_knowledge_store("ks-g") is None + + +# --------------------------------------------------------------------------- +# get_accessible_knowledge_stores +# --------------------------------------------------------------------------- + + +def test_get_accessible_knowledge_stores(dbm): + _make_ks(dbm, ks_id="ks-own", name="Owned", owner=10) + _make_ks(dbm, ks_id="ks-shared", name="Shared", owner=11) + dbm.toggle_knowledge_store_sharing("ks-shared", True) + # Also a non-shared KS owned by other user -> not visible + _make_ks(dbm, ks_id="ks-hidden", name="Hidden", owner=11) + # corrupt the params on one to hit the json-decode except branch + conn = dbm.get_connection() + with conn: + conn.execute( + f"UPDATE {dbm.table_prefix}knowledge_stores SET chunking_params = 'bad' WHERE id = ?", + ("ks-own",), + ) + conn.close() + + rows = dbm.get_accessible_knowledge_stores(user_id=10, organization_id=1) + ids = {r["id"] for r in rows} + assert ids == {"ks-own", "ks-shared"} + own = next(r for r in rows if r["id"] == "ks-own") + assert own["chunking_params"] == {} + assert own["is_shared"] is False + + +def test_get_accessible_knowledge_stores_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.get_accessible_knowledge_stores(10, 1) == [] + + +def test_get_accessible_knowledge_stores_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.get_accessible_knowledge_stores(10, 1) == [] + + +# --------------------------------------------------------------------------- +# user_can_access_knowledge_store +# --------------------------------------------------------------------------- + + +def test_user_can_access_owner(dbm): + _make_ks(dbm, ks_id="ks-acc", owner=10) + assert dbm.user_can_access_knowledge_store("ks-acc", 10) == (True, "owner") + + +def test_user_can_access_shared(dbm): + _make_ks(dbm, ks_id="ks-acc2", owner=10) + dbm.toggle_knowledge_store_sharing("ks-acc2", True) + assert dbm.user_can_access_knowledge_store("ks-acc2", 11) == (True, "shared") + + +def test_user_can_access_none_not_shared(dbm): + _make_ks(dbm, ks_id="ks-acc3", owner=10) + assert dbm.user_can_access_knowledge_store("ks-acc3", 11) == (False, "none") + + +def test_user_can_access_missing_ks(dbm): + assert dbm.user_can_access_knowledge_store("missing", 10) == (False, "none") + + +# --------------------------------------------------------------------------- +# toggle_knowledge_store_sharing +# --------------------------------------------------------------------------- + + +def test_toggle_sharing_success_and_missing(dbm): + _make_ks(dbm, ks_id="ks-t") + assert dbm.toggle_knowledge_store_sharing("ks-t", True) is True + assert dbm.toggle_knowledge_store_sharing("missing", True) is False + + +def test_toggle_sharing_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.toggle_knowledge_store_sharing("ks-t", True) is False + + +def test_toggle_sharing_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.toggle_knowledge_store_sharing("ks-t", True) is False + + +# --------------------------------------------------------------------------- +# update_knowledge_store +# --------------------------------------------------------------------------- + + +def test_update_knowledge_store_all_fields(dbm): + _make_ks(dbm, ks_id="ks-u") + assert dbm.update_knowledge_store( + "ks-u", name="New", description="New desc", + chunking_params={"size": 999}, + ) is True + ks = dbm.get_knowledge_store("ks-u") + assert ks["name"] == "New" + assert ks["description"] == "New desc" + assert ks["chunking_params"] == {"size": 999} + + +def test_update_knowledge_store_no_optional_fields(dbm): + _make_ks(dbm, ks_id="ks-u2") + # Only updated_at set -> still updates the existing row + assert dbm.update_knowledge_store("ks-u2") is True + # Missing id -> no rows + assert dbm.update_knowledge_store("missing") is False + + +def test_update_knowledge_store_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.update_knowledge_store("ks-u", name="x") is False + + +def test_update_knowledge_store_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.update_knowledge_store("ks-u", name="x") is False + + +# --------------------------------------------------------------------------- +# delete_knowledge_store +# --------------------------------------------------------------------------- + + +def test_delete_knowledge_store_success_and_missing(dbm): + _make_ks(dbm, ks_id="ks-d") + assert dbm.delete_knowledge_store("ks-d") is True + assert dbm.delete_knowledge_store("ks-d") is False + + +def test_delete_knowledge_store_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.delete_knowledge_store("ks-d") is False + + +def test_delete_knowledge_store_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.delete_knowledge_store("ks-d") is False + + +# --------------------------------------------------------------------------- +# register_kb_content_link +# --------------------------------------------------------------------------- + + +def _seed_library_item(dbm, lib_id="lib-1", item_id="item-1", + title="Item Title", org=1, owner=10): + now = int(time.time()) + conn = dbm.get_connection() + with conn: + conn.execute( + f"INSERT OR IGNORE INTO {dbm.table_prefix}libraries " + f"(id, organization_id, name, description, owner_user_id, status, created_at, updated_at) " + f"VALUES (?, ?, ?, '', ?, 'active', ?, ?)", + (lib_id, org, f"Library {lib_id}", owner, now, now), + ) + conn.execute( + f"INSERT INTO {dbm.table_prefix}library_items " + f"(id, library_id, organization_id, title, source_type, import_plugin, status, uploader_user_id, created_at, updated_at) " + f"VALUES (?, ?, ?, ?, 'file', 'simple', 'ready', ?, ?, ?)", + (item_id, lib_id, org, title, owner, now, now), + ) + conn.close() + + +def test_register_kb_content_link_success(dbm): + _make_ks(dbm, ks_id="ks-link") + _seed_library_item(dbm) + link_id = dbm.register_kb_content_link( + knowledge_store_id="ks-link", library_id="lib-1", + library_item_id="item-1", organization_id=1, + created_by_user_id=10, kb_job_id="job-1", status="pending", + ) + assert isinstance(link_id, int) + + +def test_register_kb_content_link_duplicate(dbm): + _make_ks(dbm, ks_id="ks-link2") + _seed_library_item(dbm) + assert dbm.register_kb_content_link( + "ks-link2", "lib-1", "item-1", 1, 10, + ) is not None + # Duplicate (ks, item) -> UNIQUE violation -> None + assert dbm.register_kb_content_link( + "ks-link2", "lib-1", "item-1", 1, 10, + ) is None + + +def test_register_kb_content_link_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.register_kb_content_link("ks", "lib", "item", 1, 10) is None + + +def test_register_kb_content_link_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.register_kb_content_link("ks", "lib", "item", 1, 10) is None + + +# --------------------------------------------------------------------------- +# update_kb_content_link_status +# --------------------------------------------------------------------------- + + +def test_update_link_status_by_id(dbm): + _make_ks(dbm, ks_id="ks-up") + _seed_library_item(dbm) + link_id = dbm.register_kb_content_link("ks-up", "lib-1", "item-1", 1, 10) + assert dbm.update_kb_content_link_status( + link_id=link_id, status="ready", kb_job_id="j2", + chunks_created=7, error_message="none", + ) is True + row = dbm.get_kb_content_link("ks-up", "item-1") + assert row["status"] == "ready" + assert row["chunks_created"] == 7 + + +def test_update_link_status_by_ks_item_pair(dbm): + _make_ks(dbm, ks_id="ks-up2") + _seed_library_item(dbm) + dbm.register_kb_content_link("ks-up2", "lib-1", "item-1", 1, 10) + assert dbm.update_kb_content_link_status( + knowledge_store_id="ks-up2", library_item_id="item-1", status="failed", + ) is True + + +def test_update_link_status_no_lookup_keys(dbm): + # Neither link_id nor (ks, item) pair -> returns False without query + assert dbm.update_kb_content_link_status(status="ready") is False + + +def test_update_link_status_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.update_kb_content_link_status(link_id=1, status="x") is False + + +def test_update_link_status_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.update_kb_content_link_status(link_id=1, status="x") is False + + +# --------------------------------------------------------------------------- +# delete_kb_content_link +# --------------------------------------------------------------------------- + + +def test_delete_link_success_and_missing(dbm): + _make_ks(dbm, ks_id="ks-dl") + _seed_library_item(dbm) + dbm.register_kb_content_link("ks-dl", "lib-1", "item-1", 1, 10) + assert dbm.delete_kb_content_link("ks-dl", "item-1") is True + assert dbm.delete_kb_content_link("ks-dl", "item-1") is False + + +def test_delete_link_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.delete_kb_content_link("ks", "item") is False + + +def test_delete_link_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.delete_kb_content_link("ks", "item") is False + + +# --------------------------------------------------------------------------- +# get_kb_content_links_for_ks +# --------------------------------------------------------------------------- + + +def test_get_links_for_ks_with_item(dbm): + _make_ks(dbm, ks_id="ks-list") + _seed_library_item(dbm) + dbm.register_kb_content_link("ks-list", "lib-1", "item-1", 1, 10) + rows = dbm.get_kb_content_links_for_ks("ks-list") + assert len(rows) == 1 + assert rows[0]["item_title"] == "Item Title" + assert rows[0]["library_name"] == "Library lib-1" + assert rows[0]["library_deleted"] is False + assert rows[0]["item_deleted"] is False + + +def test_get_links_for_ks_orphan_uses_audit_log(dbm): + """Link whose library/item rows are gone falls back to audit_log.""" + _make_ks(dbm, ks_id="ks-orphan") + now = int(time.time()) + conn = dbm.get_connection() + with conn: + # audit log entries for the recovered names + conn.execute( + f"INSERT INTO {dbm.table_prefix}audit_log " + f"(organization_id, actor_user_id, action, target_type, target_id, details, created_at) " + f"VALUES (1, 10, 'library.create', 'library', 'gone-lib', ?, ?)", + ('{"name": "Gone Library"}', now), + ) + conn.execute( + f"INSERT INTO {dbm.table_prefix}audit_log " + f"(organization_id, actor_user_id, action, target_type, target_id, details, created_at) " + f"VALUES (1, 10, 'library.upload', 'library_item', 'gone-item', ?, ?)", + ('{"filename": "gone.pdf"}', now), + ) + # content link pointing at non-existent library/item + conn.execute( + f"INSERT INTO {dbm.table_prefix}kb_content_links " + f"(knowledge_store_id, library_id, library_item_id, organization_id, status, chunks_created, created_by_user_id, created_at, updated_at) " + f"VALUES ('ks-orphan', 'gone-lib', 'gone-item', 1, 'ready', 0, 10, ?, ?)", + (now, now), + ) + conn.close() + rows = dbm.get_kb_content_links_for_ks("ks-orphan") + assert len(rows) == 1 + assert rows[0]["item_title"] == "gone.pdf" + assert rows[0]["library_name"] == "Gone Library" + assert rows[0]["library_deleted"] is True + assert rows[0]["item_deleted"] is True + + +def test_get_links_for_ks_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.get_kb_content_links_for_ks("ks") == [] + + +def test_get_links_for_ks_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.get_kb_content_links_for_ks("ks") == [] + + +# --------------------------------------------------------------------------- +# get_kb_content_link +# --------------------------------------------------------------------------- + + +def test_get_kb_content_link_found_and_missing(dbm): + _make_ks(dbm, ks_id="ks-one") + _seed_library_item(dbm) + dbm.register_kb_content_link("ks-one", "lib-1", "item-1", 1, 10) + row = dbm.get_kb_content_link("ks-one", "item-1") + assert row["knowledge_store_id"] == "ks-one" + assert dbm.get_kb_content_link("ks-one", "nope") is None + + +def test_get_kb_content_link_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.get_kb_content_link("ks", "item") is None + + +def test_get_kb_content_link_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.get_kb_content_link("ks", "item") is None + + +# --------------------------------------------------------------------------- +# get_kb_content_links_for_item +# --------------------------------------------------------------------------- + + +def test_get_links_for_item(dbm): + _make_ks(dbm, ks_id="ks-fi", name="KS FI") + _seed_library_item(dbm) + dbm.register_kb_content_link("ks-fi", "lib-1", "item-1", 1, 10) + rows = dbm.get_kb_content_links_for_item("item-1") + assert len(rows) == 1 + assert rows[0]["knowledge_store_name"] == "KS FI" + + +def test_get_links_for_item_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.get_kb_content_links_for_item("item") == [] + + +def test_get_links_for_item_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.get_kb_content_links_for_item("item") == [] + + +# --------------------------------------------------------------------------- +# get_kb_content_links_for_library +# --------------------------------------------------------------------------- + + +def test_get_links_for_library(dbm): + _make_ks(dbm, ks_id="ks-fl", name="KS FL") + _seed_library_item(dbm) + dbm.register_kb_content_link("ks-fl", "lib-1", "item-1", 1, 10) + rows = dbm.get_kb_content_links_for_library("lib-1") + assert len(rows) == 1 + assert rows[0]["knowledge_store_name"] == "KS FL" + assert rows[0]["item_title"] == "Item Title" + + +def test_get_links_for_library_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.get_kb_content_links_for_library("lib") == [] + + +def test_get_links_for_library_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.get_kb_content_links_for_library("lib") == [] + + +# --------------------------------------------------------------------------- +# get_knowledge_stores_for_library +# --------------------------------------------------------------------------- + + +def test_get_knowledge_stores_for_library(dbm): + _make_ks(dbm, ks_id="ks-gsl", name="KS GSL") + _seed_library_item(dbm, lib_id="lib-2", item_id="item-2") + _seed_library_item(dbm, lib_id="lib-2", item_id="item-3", title="Item Three") + # two links: one ready, one failed -> exercises the count branches + dbm.register_kb_content_link("ks-gsl", "lib-2", "item-2", 1, 10, status="ready") + dbm.register_kb_content_link("ks-gsl", "lib-2", "item-3", 1, 10, status="failed") + rows = dbm.get_knowledge_stores_for_library("lib-2") + assert len(rows) == 1 + r = rows[0] + assert r["id"] == "ks-gsl" + assert r["is_shared"] is False + assert r["item_count"] == 2 + assert r["ready_count"] == 1 + assert r["failed_count"] == 1 + + +def test_get_knowledge_stores_for_library_no_links(dbm): + # No links at all -> empty result (counts-None branch not triggered, but + # the SUM/COUNT-None normalization path is covered when a GROUP has nulls). + assert dbm.get_knowledge_stores_for_library("no-such-lib") == [] + + +def test_get_knowledge_stores_for_library_null_counts(dbm): + """A row whose count columns come back NULL hits the ``d[k] = 0`` branch. + + The real SQL can't produce NULL counts for a matched group, so a fake + connection returns a crafted row to exercise that normalization line. + """ + from contextlib import contextmanager + + class _FakeCursor: + description = [ + ("id",), ("name",), ("description",), ("chunking_strategy",), + ("embedding_vendor",), ("embedding_model",), ("vector_db_backend",), + ("is_shared",), ("organization_id",), ("owner_user_id",), + ("created_at",), ("updated_at",), ("item_count",), + ("ready_count",), ("failed_count",), + ] + + def execute(self, *a, **k): + return self + + def fetchall(self): + return [( + "ks-x", "X", "", "simple", "openai", "m", "chromadb", + 1, 1, 10, 0, 0, None, None, None, + )] + + class _FakeConn: + def cursor(self): + return _FakeCursor() + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def close(self): + pass + + with patch.object(dbm, "get_connection", return_value=_FakeConn()): + rows = dbm.get_knowledge_stores_for_library("lib-x") + assert rows[0]["item_count"] == 0 + assert rows[0]["ready_count"] == 0 + assert rows[0]["failed_count"] == 0 + assert rows[0]["is_shared"] is True + + +def test_get_knowledge_stores_for_library_no_connection(dbm): + with patch.object(dbm, "get_connection", return_value=None): + assert dbm.get_knowledge_stores_for_library("lib") == [] + + +def test_get_knowledge_stores_for_library_db_error(dbm): + dbm.table_prefix = "nonexistent_" + assert dbm.get_knowledge_stores_for_library("lib") == [] diff --git a/backend/tests/test_eval_metadata.py b/backend/tests/test_eval_metadata.py new file mode 100644 index 000000000..10ca2d2dd --- /dev/null +++ b/backend/tests/test_eval_metadata.py @@ -0,0 +1,91 @@ +"""Tests for the opt-in `include_eval_metadata` response field. + +When a caller sets `include_eval_metadata=True`, `run_lamb_assistant` attaches +the retrieved RAG context to the non-streaming response under +`eval_metadata.rag_context.context` (consumed by the lamb-eval framework). When +the flag is absent the response is unchanged. These tests exercise that logic +directly with the heavy dependencies (DB, plugin loading, connector, RAG) mocked. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +_BACKEND_ROOT = Path(__file__).parent.parent +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +from lamb.completions import main as completions_main # noqa: E402 + +_RETRIEVED = "RETRIEVED CHUNK TEXT" +_RAG_CONTEXT = {"context": _RETRIEVED, "sources": [{"title": "doc", "score": 0.9}]} + + +def _patches(rag_context): + """Patch every external dependency of run_lamb_assistant. + + connector is "ollama" so the token-usage logging branch is skipped. + """ + assistant_details = MagicMock(owner="creator@example.com", organization_id=None) + plugin_config = { + "connector": "ollama", + "rag_processor": "knowledge_store_rag", + "llm": "test-llm", + "document_rag": "", + } + connector_func = AsyncMock(return_value={"id": "chatcmpl-x", "choices": [{"message": {"content": "hi"}}]}) + return [ + patch.object(completions_main, "get_assistant_details", return_value=assistant_details), + patch.object(completions_main, "parse_plugin_config", return_value=plugin_config), + patch.object(completions_main, "_provider_for_connector", return_value=None), + patch.object(completions_main, "maybe_route_non_streaming_task", AsyncMock(return_value=None)), + patch.object( + completions_main, + "load_and_validate_plugins", + return_value=(MagicMock(), {"ollama": connector_func}, MagicMock()), + ), + patch.object(completions_main, "get_rag_context", AsyncMock(return_value=rag_context)), + patch.object(completions_main, "process_completion_request", return_value=[{"role": "user", "content": "q"}]), + ] + + +async def _run(include_eval_metadata, rag_context=_RAG_CONTEXT): + request = {"model": "lamb_assistant.1", "messages": [{"role": "user", "content": "q"}], "stream": False} + ctxs = _patches(rag_context) + for c in ctxs: + c.start() + try: + resp = await completions_main.run_lamb_assistant( + request=request, assistant=1, headers={}, include_eval_metadata=include_eval_metadata + ) + finally: + for c in ctxs: + c.stop() + return json.loads(resp.body) + + +@pytest.mark.asyncio +async def test_flag_true_attaches_retrieved_context(): + body = await _run(include_eval_metadata=True) + assert body["eval_metadata"]["rag_context"]["context"] == _RETRIEVED + assert body["eval_metadata"]["rag_context"]["sources"] + # The base OpenAI response is preserved. + assert body["choices"][0]["message"]["content"] == "hi" + + +@pytest.mark.asyncio +async def test_flag_absent_leaves_response_unchanged(): + body = await _run(include_eval_metadata=False) + assert "eval_metadata" not in body + + +@pytest.mark.asyncio +async def test_flag_true_but_no_rag_context_omits_eval_metadata(): + # A non-dict rag_context (e.g. no_rag returns None upstream) must not attach. + body = await _run(include_eval_metadata=True, rag_context=None) + assert "eval_metadata" not in body diff --git a/backend/tests/test_kb_server_manager_diffcov.py b/backend/tests/test_kb_server_manager_diffcov.py new file mode 100644 index 000000000..b3314873f --- /dev/null +++ b/backend/tests/test_kb_server_manager_diffcov.py @@ -0,0 +1,225 @@ +"""Diff-coverage unit tests for ``creator_interface.kb_server_manager``. + +Targets the branch-added lines: + - import-time bootstrap (15, 23, 29-30) — covered by importing the module + with ``LAMB_KB_SERVER_TOKEN`` set so the top-level guard does not raise. + - ``_get_kb_config_for_user`` org-config path + global fallback (82-83, 97) + - ``is_kb_server_available`` URL-try loop (127-139, 141, 144-146) + - ``create_knowledge_base`` host-redirect block (505-507) + +All subprocess/httpx/filesystem/env interaction is mocked; no network needed. +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +# Make ``backend/`` importable and ensure the import-time token guard passes. +_BACKEND_ROOT = Path(__file__).parent.parent +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +os.environ.setdefault("LAMB_KB_SERVER_TOKEN", "test-token") + +from creator_interface import kb_server_manager as ksm # noqa: E402 + + +@pytest.fixture +def manager(): + """A KBServerManager with a known global URL/token regardless of env.""" + mgr = ksm.KBServerManager() + mgr.global_kb_server_url = "http://kb:9090" + mgr.global_kb_server_token = "global-token" + mgr.kb_server_configured = True + return mgr + + +# --------------------------------------------------------------------------- +# _get_kb_config_for_user — org-config path (lines 82-83) and fallback (97) +# --------------------------------------------------------------------------- + + +def test_get_kb_config_org_config_without_api_token(manager): + """Org config present but missing api_token -> falls back to global token, + resolves url via redirects (lines 82-83).""" + fake_resolver = MagicMock() + fake_resolver.organization = {"name": "Acme"} + fake_resolver.get_knowledge_base_config.return_value = { + "server_url": "http://org-kb:9090", + # no api_token + } + + with patch( + "lamb.completions.org_config_resolver.OrganizationConfigResolver", + return_value=fake_resolver, + ): + cfg = manager._get_kb_config_for_user({"email": "u@example.com"}) + + assert cfg["url"] == "http://org-kb:9090" + assert cfg["token"] == "global-token" # fell back to global token (line 81-82) + + +def test_get_kb_config_org_config_with_redirect(manager): + """server_url matches a redirect entry -> resolved_url uses the redirect.""" + fake_resolver = MagicMock() + fake_resolver.organization = {"name": "Acme"} + fake_resolver.get_knowledge_base_config.return_value = { + "server_url": "http://org-kb:9090", + "api_token": "org-token", + } + + with patch.dict(ksm._KB_REDIRECTS, {"http://org-kb:9090": "http://redirected:9090"}, clear=False), \ + patch( + "lamb.completions.org_config_resolver.OrganizationConfigResolver", + return_value=fake_resolver, + ): + cfg = manager._get_kb_config_for_user({"email": "u@example.com"}) + + assert cfg["url"] == "http://redirected:9090" + assert cfg["token"] == "org-token" + + +def test_get_kb_config_fallback_to_global_with_redirect(manager): + """No org config -> global fallback path; resolved via redirects (line 97).""" + fake_resolver = MagicMock() + fake_resolver.organization = {"name": "Acme"} + fake_resolver.get_knowledge_base_config.return_value = None # no org config + + with patch.dict(ksm._KB_REDIRECTS, {"http://kb:9090": "http://redirected:9090"}, clear=False), \ + patch( + "lamb.completions.org_config_resolver.OrganizationConfigResolver", + return_value=fake_resolver, + ): + cfg = manager._get_kb_config_for_user({"email": "u@example.com"}) + + assert cfg["url"] == "http://redirected:9090" + assert cfg["token"] == "global-token" + + +# --------------------------------------------------------------------------- +# is_kb_server_available — URL-try loop (127-139, 141, 144-146) +# --------------------------------------------------------------------------- + + +def _make_async_client(get_mock): + """Build a context-manager mock standing in for httpx.AsyncClient.""" + client_instance = MagicMock() + client_instance.get = get_mock + cm = MagicMock() + cm.__aenter__ = AsyncMock(return_value=client_instance) + cm.__aexit__ = AsyncMock(return_value=False) + return MagicMock(return_value=cm) + + +@pytest.mark.asyncio +async def test_is_available_primary_ok(manager): + """Primary URL returns 200 -> True; covers urls_to_try build + 200 branch.""" + manager.global_kb_server_url = "http://kb:9090" + + resp = MagicMock() + resp.status_code = 200 + get_mock = AsyncMock(return_value=resp) + + with patch.object(ksm.httpx, "AsyncClient", _make_async_client(get_mock)): + result = await manager.is_kb_server_available() + + assert result is True + # Only the primary URL should have been hit. + get_mock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_is_available_fallback_to_docker_internal(manager): + """Primary kb:9090 fails, host.docker.internal succeeds -> True and + global_kb_server_url switched (lines 128-129, 137-139).""" + manager.global_kb_server_url = "http://kb:9090" + + primary_resp = MagicMock() + primary_resp.status_code = 500 # non-200 -> warning branch (141-143) + fallback_resp = MagicMock() + fallback_resp.status_code = 200 + + get_mock = AsyncMock(side_effect=[primary_resp, fallback_resp]) + + with patch.object(ksm.httpx, "AsyncClient", _make_async_client(get_mock)): + result = await manager.is_kb_server_available() + + assert result is True + assert manager.global_kb_server_url == "http://host.docker.internal:9090" + assert get_mock.await_count == 2 + + +@pytest.mark.asyncio +async def test_is_available_all_fail_raises(manager): + """All attempts raise -> exception branch (144-145) then return False (146).""" + manager.global_kb_server_url = "http://kb:9090" + + get_mock = AsyncMock(side_effect=RuntimeError("boom")) + + with patch.object(ksm.httpx, "AsyncClient", _make_async_client(get_mock)): + result = await manager.is_kb_server_available() + + assert result is False + assert get_mock.await_count == 2 # tried both primary + fallback + + +@pytest.mark.asyncio +async def test_is_available_config_resolution_failure(manager): + """creator_user given but config resolution raises -> early False.""" + with patch.object(manager, "_get_kb_config_for_user", side_effect=ValueError("nope")): + result = await manager.is_kb_server_available(creator_user={"email": "x@y.z"}) + assert result is False + + +@pytest.mark.asyncio +async def test_is_available_no_url_configured(manager): + """URL blank -> warning + False (covers the not-configured guard).""" + manager.global_kb_server_url = " " + result = await manager.is_kb_server_available() + assert result is False + + +# --------------------------------------------------------------------------- +# create_knowledge_base — host-redirect block (lines 505-507) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_create_kb_applies_host_redirect(manager): + """When kb_server_url is in _KB_REDIRECTS, it is swapped before the POST + (lines 505-507). Verifies the redirected URL is used in the request.""" + kb_data = ksm.KnowledgeBaseCreate(name="My KB", description="d", access_control="private") + + # Stub config resolution to a URL that has a redirect entry. + with patch.object( + manager, + "_get_kb_config_for_user", + return_value={"url": "http://kb:9090", "token": "tok"}, + ), patch.dict( + ksm._KB_REDIRECTS, {"http://kb:9090": "http://redirected:9090"}, clear=False + ): + post_resp = MagicMock() + post_resp.status_code = 201 + post_resp.json.return_value = {"id": "abc123"} + post_mock = AsyncMock(return_value=post_resp) + + client_instance = MagicMock() + client_instance.post = post_mock + cm = MagicMock() + cm.__aenter__ = AsyncMock(return_value=client_instance) + cm.__aexit__ = AsyncMock(return_value=False) + + with patch.object(ksm.httpx, "AsyncClient", MagicMock(return_value=cm)), \ + patch("lamb.database_manager.LambDatabaseManager") as DBM: + DBM.return_value = MagicMock() + result = await manager.create_knowledge_base(kb_data, {"id": 7, "organization_id": 1}) + + assert result["id"] == "abc123" + # The POST must have targeted the redirected host. + called_url = post_mock.await_args.args[0] + assert called_url == "http://redirected:9090/collections" diff --git a/backend/tests/test_knowledge_store_client.py b/backend/tests/test_knowledge_store_client.py new file mode 100644 index 000000000..da216aaf6 --- /dev/null +++ b/backend/tests/test_knowledge_store_client.py @@ -0,0 +1,545 @@ +"""Unit tests for ``creator_interface.knowledge_store_client.KnowledgeStoreClient``. + +``httpx`` and ``OrganizationConfigResolver`` are mocked so nothing leaves the +process. Covers the request/header/config internals, every thin proxy method +(via a mocked ``_request``), discovery error-wrapping, the org-options +aggregation, and allow-list validation. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import HTTPException + +import creator_interface.knowledge_store_client as ksc +from creator_interface.knowledge_store_client import ( + KnowledgeStoreClient, + KnowledgeStoreUnavailable, +) + + +def _cfg(**over): + base = { + "url": "http://kb", "token": "tok", + "allowed_vector_db_backends": [], "allowed_chunking_strategies": [], + "allowed_embedding_vendors": [], "allowed_embedding_models": {}, + } + base.update(over) + return base + + +@pytest.fixture +def client(monkeypatch): + c = KnowledgeStoreClient() + monkeypatch.setattr(c, "_get_ks_config", lambda creator_user=None: _cfg()) + return c + + +def run(coro): + return asyncio.run(coro) + + +# --------------------------------------------------------------------------- +# _headers +# --------------------------------------------------------------------------- + + +def test_headers_requires_token(): + c = KnowledgeStoreClient() + with pytest.raises(ValueError): + c._headers("") + assert c._headers("tok") == {"Authorization": "Bearer tok"} + + +# --------------------------------------------------------------------------- +# _get_ks_config +# --------------------------------------------------------------------------- + + +def test_get_ks_config_from_org_resolver(monkeypatch): + class _Resolver: + def __init__(self, email): + pass + + def get_knowledge_store_config(self): + return { + "server_url": "http://org-kb", + "api_token": "org-tok", + "allowed_embedding_vendors": ["openai"], + } + + monkeypatch.setattr(ksc, "OrganizationConfigResolver", _Resolver) + cfg = KnowledgeStoreClient()._get_ks_config({"email": "u@x.com"}) + assert cfg["url"] == "http://org-kb" + assert cfg["token"] == "org-tok" + assert cfg["allowed_embedding_vendors"] == ["openai"] + + +def test_get_ks_config_resolver_error_falls_back_to_global(monkeypatch): + class _Resolver: + def __init__(self, email): + raise RuntimeError("no config") + + monkeypatch.setattr(ksc, "OrganizationConfigResolver", _Resolver) + c = KnowledgeStoreClient() + c.global_server_url = "http://global-kb" + cfg = c._get_ks_config({"email": "u@x.com"}) + assert cfg["url"] == "http://global-kb" + + +def test_get_ks_config_no_global_raises(monkeypatch): + c = KnowledgeStoreClient() + c.global_server_url = "" + with pytest.raises(KnowledgeStoreUnavailable): + c._get_ks_config(None) + + +# --------------------------------------------------------------------------- +# resolve_embedding_api_key +# --------------------------------------------------------------------------- + + +def test_resolve_embedding_api_key_no_email(): + assert KnowledgeStoreClient().resolve_embedding_api_key(None, "openai") == "" + + +def test_resolve_embedding_api_key_success(monkeypatch): + class _Resolver: + def __init__(self, email): + pass + + def get_provider_api_key(self, vendor): + return "sk-org" + + monkeypatch.setattr(ksc, "OrganizationConfigResolver", _Resolver) + out = KnowledgeStoreClient().resolve_embedding_api_key({"email": "u@x.com"}, "openai") + assert out == "sk-org" + + +def test_resolve_embedding_api_key_resolver_error(monkeypatch): + class _Resolver: + def __init__(self, email): + raise RuntimeError("boom") + + monkeypatch.setattr(ksc, "OrganizationConfigResolver", _Resolver) + assert KnowledgeStoreClient().resolve_embedding_api_key({"email": "u@x.com"}, "x") == "" + + +# --------------------------------------------------------------------------- +# _wrap_discovery_error +# --------------------------------------------------------------------------- + + +def test_wrap_discovery_error(): + assert isinstance( + KnowledgeStoreClient._wrap_discovery_error(HTTPException(503, "down")), + KnowledgeStoreUnavailable, + ) + assert KnowledgeStoreClient._wrap_discovery_error(HTTPException(400, "bad")) is None + + +# --------------------------------------------------------------------------- +# _request (httpx mocked) +# --------------------------------------------------------------------------- + + +class _Resp: + def __init__(self, *, success=True, status=200, content=b"{}", payload=None, + text="", raise_json=False): + self.is_success = success + self.status_code = status + self.content = content + self._payload = payload if payload is not None else {} + self.text = text + self._raise_json = raise_json + + def json(self): + if self._raise_json: + raise ValueError("not json") + return self._payload + + +def _install_httpx(monkeypatch, *, response=None, request_error=None): + class FakeAsyncClient: + def __init__(self, *a, **k): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def request(self, method, url, headers=None, **kwargs): + if request_error is not None: + raise request_error + return response + + async def post(self, url, headers=None, **kwargs): + if request_error is not None: + raise request_error + return response + + monkeypatch.setattr(ksc.httpx, "AsyncClient", FakeAsyncClient) + monkeypatch.setattr( + ksc.httpx, "RequestError", ksc.httpx.RequestError, raising=False + ) + + +def test_request_success_json(client, monkeypatch): + _install_httpx(monkeypatch, response=_Resp(payload={"k": "v"}, content=b"{}")) + out = run(client._request("GET", "/x", _cfg())) + assert out == {"k": "v"} + + +def test_request_204_empty(client, monkeypatch): + _install_httpx(monkeypatch, response=_Resp(content=b"")) + assert run(client._request("DELETE", "/x", _cfg(), expect_204=True)) == {} + + +def test_request_error_with_json_detail(client, monkeypatch): + _install_httpx( + monkeypatch, + response=_Resp(success=False, status=409, payload={"detail": "conflict"}), + ) + with pytest.raises(HTTPException) as exc: + run(client._request("POST", "/x", _cfg())) + assert exc.value.status_code == 409 + assert "conflict" in exc.value.detail + + +def test_request_error_non_json_text(client, monkeypatch): + _install_httpx( + monkeypatch, + response=_Resp(success=False, status=500, text="boom", raise_json=True), + ) + with pytest.raises(HTTPException) as exc: + run(client._request("POST", "/x", _cfg())) + assert exc.value.status_code == 500 + + +def test_request_connection_error_503(client, monkeypatch): + _install_httpx(monkeypatch, request_error=ksc.httpx.ConnectError("refused")) + with pytest.raises(HTTPException) as exc: + run(client._request("GET", "/x", _cfg())) + assert exc.value.status_code == 503 + + +# --------------------------------------------------------------------------- +# proxy methods (mock _request) +# --------------------------------------------------------------------------- + + +@pytest.fixture +def proxied(client): + client._request = AsyncMock(return_value={"ok": True}) + return client + + +def _last_call(c): + args, kwargs = c._request.call_args + return args, kwargs + + +def test_create_collection_payload(proxied): + run(proxied.create_collection( + knowledge_store_id="ks1", organization_id=5, name="N", + chunking_strategy="simple", embedding_vendor="openai", + embedding_model="m", vector_db_backend="chromadb", + graph_enabled=True, extraction_vendor="openai", extraction_model="gpt", + )) + args, kwargs = _last_call(proxied) + assert args[0] == "POST" and args[1] == "/collections" + payload = kwargs["json"] + assert payload["id"] == "ks1" + assert payload["graph_enabled"] is True + assert payload["extraction"]["vendor"] == "openai" + + +def test_create_collection_no_extraction_when_graph_disabled(proxied): + run(proxied.create_collection( + knowledge_store_id="ks1", organization_id=5, name="N", + chunking_strategy="simple", embedding_vendor="openai", + embedding_model="m", vector_db_backend="chromadb", + graph_enabled=False, extraction_vendor="openai", + )) + _, kwargs = _last_call(proxied) + assert "extraction" not in kwargs["json"] + + +def test_simple_proxy_methods(proxied): + cases = [ + (proxied.get_llm_vendors(), "GET", "/llm-vendors"), + (proxied.get_collection("ks1"), "GET", "/collections/ks1"), + (proxied.delete_collection("ks1"), "DELETE", "/collections/ks1"), + (proxied.delete_content_by_source("ks1", "src1"), "DELETE", + "/collections/ks1/content/src1"), + (proxied.get_job_status("j1"), "GET", "/jobs/j1"), + (proxied.cancel_job("j1"), "POST", "/jobs/j1/cancel"), + (proxied.get_graph_status(), "GET", "/graph/status"), + ] + for coro, method, path in cases: + proxied._request.reset_mock() + run(coro) + args, _ = _last_call(proxied) + assert (args[0], args[1]) == (method, path) + + +def test_update_collection_builds_partial_body(proxied): + run(proxied.update_collection("ks1", name="New", chunking_params={"x": 1})) + args, kwargs = _last_call(proxied) + assert args[0] == "PUT" + assert kwargs["json"] == {"name": "New", "chunking_params": {"x": 1}} + + +def test_add_content_and_query_payloads(proxied): + run(proxied.add_content("ks1", [{"d": 1}], embedding_api_key="sk")) + _, kwargs = _last_call(proxied) + assert kwargs["json"]["embedding_credentials"]["api_key"] == "sk" + + proxied._request.reset_mock() + run(proxied.query("ks1", "hello", embedding_api_key="sk", top_k=7)) + args, kwargs = _last_call(proxied) + assert args[1] == "/collections/ks1/query" + assert kwargs["json"]["top_k"] == 7 + + +def test_graph_proxy_methods(proxied): + cases = [ + (proxied.get_graph_snapshot("ks1", params={"limit": 1}), "GET", + "/graph/collections/ks1/snapshot"), + (proxied.list_graph_changes("ks1", params={}), "GET", + "/graph/collections/ks1/changes"), + (proxied.graph_concept_rename("ks1", "alpha", {"new_name": "b"}), "PATCH", + "/graph/collections/ks1/concepts/alpha/rename"), + (proxied.graph_concepts_merge("ks1", {"target_name": "t"}), "POST", + "/graph/collections/ks1/concepts/merge"), + (proxied.graph_concept_curation("ks1", "alpha", {}), "PATCH", + "/graph/collections/ks1/concepts/alpha/curation"), + (proxied.graph_relationship_edit("ks1", {}), "PATCH", + "/graph/collections/ks1/relationships"), + (proxied.graph_relationship_curation("ks1", {}), "PATCH", + "/graph/collections/ks1/relationships/curation"), + ] + for coro, method, path in cases: + proxied._request.reset_mock() + run(coro) + args, _ = _last_call(proxied) + assert (args[0], args[1]) == (method, path) + + +# --------------------------------------------------------------------------- +# discovery error wrapping +# --------------------------------------------------------------------------- + + +def test_get_backends_success(client): + client._request = AsyncMock(return_value={"backends": [{"name": "chromadb"}]}) + assert run(client.get_backends())["backends"][0]["name"] == "chromadb" + + +def test_get_backends_5xx_wrapped(client): + client._request = AsyncMock(side_effect=HTTPException(503, "down")) + with pytest.raises(KnowledgeStoreUnavailable): + run(client.get_backends()) + + +def test_get_chunking_strategies_4xx_reraised(client): + client._request = AsyncMock(side_effect=HTTPException(400, "bad")) + with pytest.raises(HTTPException) as exc: + run(client.get_chunking_strategies()) + assert exc.value.status_code == 400 + + +def test_get_chunking_strategies_5xx_wrapped(client): + client._request = AsyncMock(side_effect=HTTPException(502, "down")) + with pytest.raises(KnowledgeStoreUnavailable): + run(client.get_chunking_strategies()) + + +def test_get_embedding_vendors_5xx_wrapped(client): + client._request = AsyncMock(side_effect=HTTPException(500, "boom")) + with pytest.raises(KnowledgeStoreUnavailable): + run(client.get_embedding_vendors()) + + +def test_get_backends_4xx_reraised(client): + client._request = AsyncMock(side_effect=HTTPException(403, "forbidden")) + with pytest.raises(HTTPException) as exc: + run(client.get_backends()) + assert exc.value.status_code == 403 + + +def test_get_embedding_vendors_4xx_reraised(client): + client._request = AsyncMock(side_effect=HTTPException(404, "missing")) + with pytest.raises(HTTPException) as exc: + run(client.get_embedding_vendors()) + assert exc.value.status_code == 404 + + +def test_update_collection_description_only(proxied): + run(proxied.update_collection("ks1", description="desc only")) + _, kwargs = _last_call(proxied) + assert kwargs["json"] == {"description": "desc only"} + + +# --------------------------------------------------------------------------- +# migrate_collection_to_graph (httpx mocked) +# --------------------------------------------------------------------------- + + +def test_migrate_success_with_key(client, monkeypatch): + _install_httpx(monkeypatch, response=_Resp(payload={"status": "ok"}, content=b"{}")) + out = run(client.migrate_collection_to_graph("ks1", openai_api_key="sk")) + assert out == {"status": "ok"} + + +def test_migrate_error(client, monkeypatch): + _install_httpx( + monkeypatch, + response=_Resp(success=False, status=503, payload={"detail": "no neo4j"}), + ) + with pytest.raises(HTTPException) as exc: + run(client.migrate_collection_to_graph("ks1")) + assert exc.value.status_code == 503 + + +def test_migrate_connection_error(client, monkeypatch): + _install_httpx(monkeypatch, request_error=ksc.httpx.ConnectError("refused")) + with pytest.raises(HTTPException) as exc: + run(client.migrate_collection_to_graph("ks1")) + assert exc.value.status_code == 503 + + +# --------------------------------------------------------------------------- +# get_org_options +# --------------------------------------------------------------------------- + + +def _vendor(name, *, model_default="m", endpoint_default="http://static"): + return { + "name": name, + "parameters": [ + {"name": "model", "default": model_default}, + {"name": "api_endpoint", "default": endpoint_default}, + ], + } + + +def test_get_org_options_no_user(monkeypatch): + c = KnowledgeStoreClient() + monkeypatch.setattr(c, "_get_ks_config", lambda cu=None: _cfg()) + c.get_backends = AsyncMock(return_value={"backends": [{"name": "chromadb"}]}) + c.get_chunking_strategies = AsyncMock( + return_value={"strategies": [{"name": "simple"}]}) + c.get_embedding_vendors = AsyncMock( + return_value={"vendors": [_vendor("openai")]}) + out = run(c.get_org_options(None)) + assert out["vector_db_backends"][0]["name"] == "chromadb" + # No user -> api_key_configured defaults True; model fallback from plugin. + assert out["embedding_vendors"][0]["api_key_configured"] is True + assert out["embedding_models"]["openai"] == ["m"] + + +def test_get_org_options_with_org_config(monkeypatch): + c = KnowledgeStoreClient() + monkeypatch.setattr( + c, "_get_ks_config", + lambda cu=None: _cfg(allowed_embedding_vendors=["openai"], + allowed_embedding_models={"openai": ["text-embed-3"]}), + ) + c.get_backends = AsyncMock(return_value={"backends": [{"name": "chromadb"}]}) + c.get_chunking_strategies = AsyncMock(return_value={"strategies": []}) + c.get_embedding_vendors = AsyncMock( + return_value={"vendors": [_vendor("openai"), _vendor("cohere")]}) + + class _Resolver: + def __init__(self, email): + pass + + def get_provider_config(self, vendor): + return {"api_key": "sk"} if vendor == "openai" else {} + + def get_provider_endpoint(self, vendor): + return "http://org-endpoint" + + monkeypatch.setattr(ksc, "OrganizationConfigResolver", _Resolver) + out = run(c.get_org_options({"email": "u@x.com"})) + # allow-list trims to openai only. + assert [v["name"] for v in out["embedding_vendors"]] == ["openai"] + vendor = out["embedding_vendors"][0] + assert vendor["api_key_configured"] is True + # org endpoint overrode the static default. + ep = next(p for p in vendor["parameters"] if p["name"] == "api_endpoint") + assert ep["default"] == "http://org-endpoint" + # explicit allowed-models list wins over plugin fallback. + assert out["embedding_models"]["openai"] == ["text-embed-3"] + + +def test_get_org_options_resolver_edge_branches(monkeypatch): + c = KnowledgeStoreClient() + monkeypatch.setattr(c, "_get_ks_config", lambda cu=None: _cfg()) + c.get_backends = AsyncMock(return_value={"backends": []}) + c.get_chunking_strategies = AsyncMock(return_value={"strategies": []}) + # One nameless vendor (skipped) + one whose provider lookups raise. + c.get_embedding_vendors = AsyncMock(return_value={ + "vendors": [{"parameters": []}, _vendor("openai")], + }) + + class _Resolver: + def __init__(self, email): + pass + + def get_provider_config(self, vendor): + raise RuntimeError("no provider cfg") + + def get_provider_endpoint(self, vendor): + raise ValueError("no endpoint") + + monkeypatch.setattr(ksc, "OrganizationConfigResolver", _Resolver) + out = run(c.get_org_options({"email": "u@x.com"})) + # The nameless vendor is preserved in the list but skipped for tagging; + # openai gets api_key_configured=False (provider cfg lookup failed). + openai = next(v for v in out["embedding_vendors"] if v.get("name") == "openai") + assert openai["api_key_configured"] is False + + +# --------------------------------------------------------------------------- +# validate_against_allow_list +# --------------------------------------------------------------------------- + + +def test_validate_all_allowed_returns_none(monkeypatch): + c = KnowledgeStoreClient() + monkeypatch.setattr(c, "_get_ks_config", lambda cu=None: _cfg()) + assert c.validate_against_allow_list( + {}, "simple", "openai", "m", "chromadb") is None + + +@pytest.mark.parametrize( + "field,cfg_key,bad", + [ + ("chunking", "allowed_chunking_strategies", "Chunking strategy"), + ("vendor", "allowed_embedding_vendors", "Embedding vendor"), + ("backend", "allowed_vector_db_backends", "Vector DB backend"), + ], +) +def test_validate_rejects_disallowed(monkeypatch, field, cfg_key, bad): + c = KnowledgeStoreClient() + cfg = _cfg(**{cfg_key: ["allowed-only"]}) + monkeypatch.setattr(c, "_get_ks_config", lambda cu=None: cfg) + msg = c.validate_against_allow_list( + {}, "simple", "openai", "m", "chromadb") + assert bad in msg + + +def test_validate_rejects_disallowed_model(monkeypatch): + c = KnowledgeStoreClient() + cfg = _cfg(allowed_embedding_models={"openai": ["good-model"]}) + monkeypatch.setattr(c, "_get_ks_config", lambda cu=None: cfg) + msg = c.validate_against_allow_list({}, "simple", "openai", "bad-model", "chromadb") + assert "Embedding model" in msg diff --git a/backend/tests/test_knowledge_store_graph_router.py b/backend/tests/test_knowledge_store_graph_router.py new file mode 100644 index 000000000..20a1943fd --- /dev/null +++ b/backend/tests/test_knowledge_store_graph_router.py @@ -0,0 +1,256 @@ +"""Tests for ``creator_interface.knowledge_store_graph_router``. + +The router is a thin ACL-checked proxy to the KB Server's ``/graph`` API. +We mount it on a bare app, override the auth dependency, and replace the +module-level ``_client`` (async KB Server client) and ``_db`` with fakes. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import creator_interface.knowledge_store_graph_router as gr +from lamb.auth_context import get_auth_context + + +class _FakeClient: + """Async KB Server client stub recording calls; every method echoes.""" + + def __init__(self): + self.calls = [] + + def _record(self, name, **kw): + self.calls.append((name, kw)) + return {"ok": True, "method": name, "args": kw} + + async def get_graph_status(self, **kw): + return self._record("status", **kw) + + async def migrate_collection_to_graph(self, **kw): + return self._record("migrate", **kw) + + async def get_graph_snapshot(self, **kw): + return self._record("snapshot", **kw) + + async def list_graph_changes(self, **kw): + return self._record("changes", **kw) + + async def graph_concept_rename(self, **kw): + return self._record("rename", **kw) + + async def graph_concepts_merge(self, **kw): + return self._record("merge", **kw) + + async def graph_concept_curation(self, **kw): + return self._record("curate_concept", **kw) + + async def graph_relationship_edit(self, **kw): + return self._record("edit_rel", **kw) + + async def graph_relationship_curation(self, **kw): + return self._record("curate_rel", **kw) + + +def _auth(user_id="u1", org_id="org1"): + return SimpleNamespace( + user={"id": user_id, "email": "u@x.com"}, + organization={"id": org_id}, + ) + + +def _ks(owner="u1", org="org1", shared=False): + return {"owner_user_id": owner, "organization_id": org, "is_shared": shared} + + +@pytest.fixture +def env(monkeypatch): + client = _FakeClient() + monkeypatch.setattr(gr, "_client", client) + + store_box = {"ks": _ks()} + + class _DB: + def get_knowledge_store(self, ks_id): + return store_box["ks"] + + monkeypatch.setattr(gr, "_db", _DB()) + + app = FastAPI() + app.include_router(gr.router) + auth_box = {"auth": _auth()} + app.dependency_overrides[get_auth_context] = lambda: auth_box["auth"] + tc = TestClient(app) + return SimpleNamespace(client=client, tc=tc, store_box=store_box, auth_box=auth_box) + + +# --------------------------------------------------------------------------- +# access control (_assert_ks_access) +# --------------------------------------------------------------------------- + + +def test_snapshot_404_when_ks_missing(env): + env.store_box["ks"] = None + assert env.tc.get("/ks-1/graph/snapshot").status_code == 404 + + +def test_snapshot_403_when_not_owner_and_not_shared(env): + env.store_box["ks"] = _ks(owner="someone-else", shared=False) + assert env.tc.get("/ks-1/graph/snapshot").status_code == 403 + + +def test_snapshot_403_when_org_mismatch(env): + env.store_box["ks"] = _ks(owner="u1", org="other-org") + assert env.tc.get("/ks-1/graph/snapshot").status_code == 403 + + +def test_shared_store_in_same_org_is_allowed(env): + env.store_box["ks"] = _ks(owner="someone-else", org="org1", shared=True) + assert env.tc.get("/ks-1/graph/snapshot").status_code == 200 + + +# --------------------------------------------------------------------------- +# status + read endpoints +# --------------------------------------------------------------------------- + + +def test_graph_status(env): + resp = env.tc.get("/graph/status") + assert resp.status_code == 200 + assert env.client.calls[0][0] == "status" + + +def test_snapshot_forwards_all_filters(env): + resp = env.tc.get( + "/ks-1/graph/snapshot", + params={ + "concept": "alpha", "document_id": "d1", "chunk_id": "c1", + "filename": "f.md", "include_chunks": "false", "limit": 10, + }, + ) + assert resp.status_code == 200 + params = env.client.calls[0][1]["params"] + assert params["concept"] == "alpha" + assert params["document_id"] == "d1" + assert params["chunk_id"] == "c1" + assert params["filename"] == "f.md" + assert params["include_chunks"] == "false" + assert params["limit"] == 10 + + +def test_changes_forwards_filters(env): + resp = env.tc.get( + "/ks-1/graph/changes", + params={"concept": "alpha", "document_id": "d1", "filename": "f", + "operation": "manual_rename_concept", "limit": 5}, + ) + assert resp.status_code == 200 + params = env.client.calls[0][1]["params"] + assert params["operation"] == "manual_rename_concept" + assert params["limit"] == 5 + + +# --------------------------------------------------------------------------- +# migrate (org-config fallback) +# --------------------------------------------------------------------------- + + +def test_migrate_with_explicit_key(env): + resp = env.tc.post("/ks-1/graph/migrate", json={"openai_api_key": "sk-explicit"}) + assert resp.status_code == 200 + assert env.client.calls[0][1]["openai_api_key"] == "sk-explicit" + + +def test_migrate_resolves_org_key(env, monkeypatch): + import lamb.completions.org_config_resolver as ocr + + class _Resolver: + def __init__(self, email): + pass + + def get_provider_api_key(self, vendor): + return "sk-from-org" + + monkeypatch.setattr(ocr, "OrganizationConfigResolver", _Resolver) + resp = env.tc.post("/ks-1/graph/migrate", json={}) + assert resp.status_code == 200 + assert env.client.calls[0][1]["openai_api_key"] == "sk-from-org" + + +def test_migrate_resolver_failure_falls_back_to_empty(env, monkeypatch): + import lamb.completions.org_config_resolver as ocr + + class _Resolver: + def __init__(self, email): + pass + + def get_provider_api_key(self, vendor): + raise RuntimeError("no org config") + + monkeypatch.setattr(ocr, "OrganizationConfigResolver", _Resolver) + resp = env.tc.post("/ks-1/graph/migrate", json={}) + assert resp.status_code == 200 + assert env.client.calls[0][1]["openai_api_key"] == "" + + +# --------------------------------------------------------------------------- +# curation endpoints +# --------------------------------------------------------------------------- + + +def test_rename_concept(env): + resp = env.tc.patch( + "/ks-1/graph/concepts/alpha/rename", json={"new_name": "beta"} + ) + assert resp.status_code == 200 + name, kw = env.client.calls[0] + assert name == "rename" + assert kw["concept"] == "alpha" + assert kw["body"]["new_name"] == "beta" + + +def test_merge_concepts(env): + resp = env.tc.post( + "/ks-1/graph/concepts/merge", + json={"source_names": ["a", "b"], "target_name": "c"}, + ) + assert resp.status_code == 200 + assert env.client.calls[0][0] == "merge" + + +def test_curate_concept(env): + resp = env.tc.patch( + "/ks-1/graph/concepts/alpha/curation", + json={"verification_state": "verified", "tags": ["t"]}, + ) + assert resp.status_code == 200 + assert env.client.calls[0][0] == "curate_concept" + + +def test_edit_relationship(env): + resp = env.tc.patch( + "/ks-1/graph/relationships", + json={"source_concept": "a", "target_concept": "b", "relation": "uses", + "new_relation": "depends_on"}, + ) + assert resp.status_code == 200 + assert env.client.calls[0][1]["body"]["new_relation"] == "depends_on" + + +def test_curate_relationship(env): + resp = env.tc.patch( + "/ks-1/graph/relationships/curation", + json={"source_concept": "a", "target_concept": "b", "relation": "uses", + "verification_state": "verified"}, + ) + assert resp.status_code == 200 + assert env.client.calls[0][0] == "curate_rel" + + +def test_merge_validation_error(env): + # missing target_name -> 422 from pydantic before the proxy is hit. + resp = env.tc.post("/ks-1/graph/concepts/merge", json={"source_names": ["a"]}) + assert resp.status_code == 422 diff --git a/backend/tests/test_knowledge_store_rag.py b/backend/tests/test_knowledge_store_rag.py new file mode 100644 index 000000000..eaadd78d8 --- /dev/null +++ b/backend/tests/test_knowledge_store_rag.py @@ -0,0 +1,162 @@ +"""Direct tests for the knowledge_store_rag RAG processor (_run). + +Covers the multi-KS fan-out, the early-return guards, error handling per +Knowledge Store, and the [N]-numbered context / aligned sources produced via +the shared _ks_query_helpers.build_context_and_sources. +""" + +import pytest +from unittest.mock import AsyncMock, patch + +from lamb.lamb_classes import Assistant + + +def _make_assistant(**overrides): + defaults = { + "id": 1, + "name": "Test", + "description": "", + "system_prompt": "", + "prompt_template": "", + "RAG_collections": "ks-1", + "RAG_Top_k": 3, + "owner": "user@test.com", + "api_callback": "", + "pre_retrieval_endpoint": "", + "post_retrieval_endpoint": "", + "RAG_endpoint": "", + } + defaults.update(overrides) + return Assistant(**defaults) + + +def _ks_success(chunks): + return {"status": "success", "data": {"results": chunks}} + + +@pytest.mark.asyncio +async def test_returns_numbered_context_and_aligned_sources(): + from lamb.completions.rag import knowledge_store_rag + + assistant = _make_assistant() + messages = [{"role": "user", "content": "What is photosynthesis?"}] + resp = _ks_success([ + {"text": "Photosynthesis converts light.", "score": 0.9, + "metadata": {"source_title": "Bio"}}, + {"text": "It happens in chloroplasts.", "score": 0.8, + "metadata": {"source_title": "Bio"}}, + ]) + + with patch.object(knowledge_store_rag, "query_one_ks", + new=AsyncMock(return_value=resp)) as mock_query: + result = await knowledge_store_rag.rag_processor(messages, assistant) + + mock_query.assert_called_once_with("ks-1", "What is photosynthesis?", 3, "user@test.com") + # Context chunks are numbered and sources carry the matching n. + assert result["context"] == ( + "[1] Photosynthesis converts light.\n\n[2] It happens in chloroplasts." + ) + assert [s["n"] for s in result["sources"]] == [1, 2] + + +@pytest.mark.asyncio +async def test_uses_last_user_message_as_query(): + from lamb.completions.rag import knowledge_store_rag + + assistant = _make_assistant() + messages = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "answer"}, + {"role": "user", "content": "the real question"}, + ] + with patch.object(knowledge_store_rag, "query_one_ks", + new=AsyncMock(return_value=_ks_success([]))) as mock_query: + await knowledge_store_rag.rag_processor(messages, assistant) + mock_query.assert_called_once_with("ks-1", "the real question", 3, "user@test.com") + + +@pytest.mark.asyncio +async def test_multiple_knowledge_stores_queried_and_numbered_continuously(): + from lamb.completions.rag import knowledge_store_rag + + assistant = _make_assistant(RAG_collections="ks-1,ks-2") + messages = [{"role": "user", "content": "q"}] + + async def fake_query(ks_id, *_a, **_k): + if ks_id == "ks-1": + return _ks_success([{"text": "from one", "score": 0.9, "metadata": {}}]) + return _ks_success([{"text": "from two", "score": 0.7, "metadata": {}}]) + + with patch.object(knowledge_store_rag, "query_one_ks", new=AsyncMock(side_effect=fake_query)) as mock_query: + result = await knowledge_store_rag.rag_processor(messages, assistant) + + assert mock_query.call_count == 2 + assert {c.args[0] for c in mock_query.call_args_list} == {"ks-1", "ks-2"} + # Citation numbers are continuous across stores. + assert result["context"] == "[1] from one\n\n[2] from two" + assert result["sources"][0]["knowledge_store_id"] == "ks-1" + assert result["sources"][1]["knowledge_store_id"] == "ks-2" + + +@pytest.mark.asyncio +async def test_failed_knowledge_store_contributes_nothing(): + from lamb.completions.rag import knowledge_store_rag + + assistant = _make_assistant(RAG_collections="ks-1,ks-2") + messages = [{"role": "user", "content": "q"}] + + async def fake_query(ks_id, *_a, **_k): + if ks_id == "ks-1": + return {"status": "error", "error": "boom"} + return _ks_success([{"text": "good", "score": 0.7, "metadata": {}}]) + + with patch.object(knowledge_store_rag, "query_one_ks", new=AsyncMock(side_effect=fake_query)): + result = await knowledge_store_rag.rag_processor(messages, assistant) + + # Only the successful KS contributes; numbering restarts cleanly at 1. + assert result["context"] == "[1] good" + assert len(result["sources"]) == 1 + # The raw responses still record both stores for diagnostics. + assert set(result["raw_responses"].keys()) == {"ks-1", "ks-2"} + + +@pytest.mark.asyncio +async def test_no_collections_returns_early(): + from lamb.completions.rag import knowledge_store_rag + + assistant = _make_assistant(RAG_collections="") + messages = [{"role": "user", "content": "hi"}] + result = await knowledge_store_rag.rag_processor(messages, assistant) + assert "No Knowledge Stores specified" in result["context"] + assert result["sources"] == [] + + +@pytest.mark.asyncio +async def test_whitespace_only_collections_returns_early(): + from lamb.completions.rag import knowledge_store_rag + + assistant = _make_assistant(RAG_collections=" , ,") + messages = [{"role": "user", "content": "hi"}] + result = await knowledge_store_rag.rag_processor(messages, assistant) + assert "RAG_collections is empty" in result["context"] + assert result["sources"] == [] + + +@pytest.mark.asyncio +async def test_no_user_message_returns_early(): + from lamb.completions.rag import knowledge_store_rag + + assistant = _make_assistant() + messages = [{"role": "assistant", "content": "only assistant"}] + result = await knowledge_store_rag.rag_processor(messages, assistant) + assert "No user message found" in result["context"] + assert result["sources"] == [] + + +@pytest.mark.asyncio +async def test_no_assistant_returns_early(): + from lamb.completions.rag import knowledge_store_rag + + result = await knowledge_store_rag.rag_processor([{"role": "user", "content": "x"}], None) + assert "No Knowledge Stores specified" in result["context"] + assert result["sources"] == [] diff --git a/backend/tests/test_knowledge_store_router_branches.py b/backend/tests/test_knowledge_store_router_branches.py new file mode 100644 index 000000000..dcd31af50 --- /dev/null +++ b/backend/tests/test_knowledge_store_router_branches.py @@ -0,0 +1,536 @@ +"""Branch/error-path tests for ``creator_interface.knowledge_store_router``. + +Complements the happy-path integration tests by driving the error and edge +branches: validation/conflict/rollback on create, the KB-Server-failure +fallbacks on get/update/delete, the per-item ingestion guards in add_content, +and the best-effort teardown in remove_content. Uses the conftest fixtures +(``client``, ``ks_db``, ``ks_client``, ``ks_library_client``, ``async_*``). +""" + +from __future__ import annotations + +import pytest +from fastapi import HTTPException + +import creator_interface.knowledge_store_router as ksr + + +# --------------------------------------------------------------------------- +# pure helpers +# --------------------------------------------------------------------------- + + +def test_build_permalinks_variants(): + # original_filename branch + pages branch. + p = ksr._build_permalinks(1, "lib", "item", pages_count=2, + original_filename="f.pdf") + assert p["full_markdown"] == "/docs/1/lib/item/content" + assert p["original"] == "/docs/1/lib/item/original/f.pdf" + assert p["pages"] == ["/docs/1/lib/item/content/pages/1", + "/docs/1/lib/item/content/pages/2"] + # source_url branch (no filename). + p2 = ksr._build_permalinks(1, "lib", "item", source_url="https://ext/x") + assert p2["original"] == "https://ext/x" + assert "pages" not in p2 + + +def test_flatten_pages(): + payload = {"pages": [ + {"page_number": 1, "text": "a"}, + {"number": 2, "markdown": "b"}, + "not-a-dict", # skipped + ]} + out = ksr._flatten_pages(payload) + assert out == [ + {"page_number": 1, "text": "a"}, + {"page_number": 2, "text": "b"}, + ] + assert ksr._flatten_pages({}) == [] + + +# --------------------------------------------------------------------------- +# /options + /llm-vendors +# --------------------------------------------------------------------------- + + +def test_options_success(client, ks_client, async_return): + ks_client.get_org_options = async_return({"vector_db_backends": []}) + resp = client.get("/creator/knowledge-stores/options") + assert resp.status_code == 200 + + +def test_options_unavailable_503(client, ks_client, async_raise): + from creator_interface.knowledge_store_client import KnowledgeStoreUnavailable + + ks_client.get_org_options = async_raise(KnowledgeStoreUnavailable("down")) + resp = client.get("/creator/knowledge-stores/options") + assert resp.status_code == 503 + assert resp.json()["error"] == "knowledge_store_unavailable" + + +def test_llm_vendors(client, ks_client, async_return): + ks_client.get_llm_vendors = async_return({"vendors": []}) + assert client.get("/creator/knowledge-stores/llm-vendors").status_code == 200 + + +# --------------------------------------------------------------------------- +# create +# --------------------------------------------------------------------------- + + +def _create_body(**over): + base = dict( + name="KS", chunking_strategy="simple", embedding_vendor="openai", + embedding_model="m", vector_db_backend="chromadb", + embedding_endpoint="http://e", + ) + base.update(over) + return base + + +def test_create_validation_error_400(client, ks_client): + ks_client.validate_against_allow_list.return_value = "vendor not allowed" + resp = client.post("/creator/knowledge-stores", json=_create_body()) + assert resp.status_code == 400 + assert "not allowed" in resp.json()["detail"] + + +def test_create_conflict_409(client, ks_db, ks_client): + ks_client.validate_against_allow_list.return_value = None + ks_db.create_knowledge_store.return_value = None # name taken + resp = client.post("/creator/knowledge-stores", json=_create_body()) + assert resp.status_code == 409 + + +def test_create_collection_failure_rolls_back_502(client, ks_db, ks_client, async_raise): + ks_client.validate_against_allow_list.return_value = None + ks_db.create_knowledge_store.return_value = True + ks_client.create_collection = async_raise(RuntimeError("kb boom")) + resp = client.post("/creator/knowledge-stores", json=_create_body()) + assert resp.status_code == 502 + ks_db.delete_knowledge_store.assert_called_once() + + +def test_create_success_resolves_endpoint(client, ks_db, ks_client, async_return, + monkeypatch): + import lamb.completions.org_config_resolver as ocr + + class _Resolver: + def __init__(self, email): + pass + + def get_provider_endpoint(self, vendor): + return "http://org-endpoint" + + monkeypatch.setattr(ocr, "OrganizationConfigResolver", _Resolver) + ks_client.validate_against_allow_list.return_value = None + ks_db.create_knowledge_store.return_value = True + ks_client.create_collection = async_return({}) + ks_db.get_knowledge_store.return_value = {"id": "ks1", "name": "KS"} + # No endpoint supplied -> resolver fallback path. + resp = client.post("/creator/knowledge-stores", json=_create_body(embedding_endpoint=None)) + assert resp.status_code == 200 + assert resp.json()["id"] == "ks1" + + +def test_create_endpoint_resolver_valueerror(client, ks_db, ks_client, async_return, + monkeypatch): + import lamb.completions.org_config_resolver as ocr + + class _Resolver: + def __init__(self, email): + pass + + def get_provider_endpoint(self, vendor): + raise ValueError("no endpoint") + + monkeypatch.setattr(ocr, "OrganizationConfigResolver", _Resolver) + ks_client.validate_against_allow_list.return_value = None + ks_db.create_knowledge_store.return_value = True + ks_client.create_collection = async_return({}) + ks_db.get_knowledge_store.return_value = {"id": "ks1"} + resp = client.post("/creator/knowledge-stores", json=_create_body(embedding_endpoint=None)) + assert resp.status_code == 200 + + +# --------------------------------------------------------------------------- +# get +# --------------------------------------------------------------------------- + + +def test_get_not_found_404(client, ks_db): + ks_db.get_knowledge_store.return_value = None + assert client.get("/creator/knowledge-stores/ks1").status_code == 404 + + +def test_get_success_with_server_data(client, ks_db, ks_client, async_return): + ks_db.get_knowledge_store.return_value = {"id": "ks1", "owner_user_id": 1} + ks_client.get_collection = async_return({ + "status": "ready", "document_count": 3, "chunk_count": 9, + "graph_enabled": True, "extraction": {"vendor": "openai"}, + }) + ks_db.get_kb_content_links_for_ks.return_value = [] + body = client.get("/creator/knowledge-stores/ks1").json() + assert body["document_count"] == 3 + assert body["graph_enabled"] is True + assert body["is_owner"] is True + + +def test_get_server_error_fallback(client, ks_db, ks_client, async_raise): + ks_db.get_knowledge_store.return_value = {"id": "ks1", "owner_user_id": 2} + ks_client.get_collection = async_raise(RuntimeError("kb down")) + ks_db.get_kb_content_links_for_ks.return_value = [] + body = client.get("/creator/knowledge-stores/ks1").json() + assert body["server_status"] is None + assert body["graph_enabled"] is False + + +# --------------------------------------------------------------------------- +# update +# --------------------------------------------------------------------------- + + +def test_update_nothing_to_update_400(client): + resp = client.put("/creator/knowledge-stores/ks1", json={}) + assert resp.status_code == 400 + + +def test_update_success(client, ks_db, ks_client, async_return): + ks_client.update_collection = async_return({}) + ks_db.update_knowledge_store.return_value = True + ks_db.get_knowledge_store.return_value = {"id": "ks1", "name": "New"} + resp = client.put("/creator/knowledge-stores/ks1", json={"name": "New"}) + assert resp.status_code == 200 + + +def test_update_collection_404_is_swallowed(client, ks_db, ks_client, async_raise): + ks_client.update_collection = async_raise(HTTPException(404, "missing")) + ks_db.update_knowledge_store.return_value = True + ks_db.get_knowledge_store.return_value = {"id": "ks1"} + resp = client.put("/creator/knowledge-stores/ks1", json={"description": "d"}) + assert resp.status_code == 200 + + +def test_update_collection_other_error_propagates(client, ks_client, async_raise): + ks_client.update_collection = async_raise(HTTPException(500, "boom")) + resp = client.put("/creator/knowledge-stores/ks1", json={"name": "x"}) + assert resp.status_code == 500 + + +def test_update_db_missing_404(client, ks_db, ks_client, async_return): + ks_client.update_collection = async_return({}) + ks_db.update_knowledge_store.return_value = False + resp = client.put("/creator/knowledge-stores/ks1", json={"name": "x"}) + assert resp.status_code == 404 + + +# --------------------------------------------------------------------------- +# delete +# --------------------------------------------------------------------------- + + +def test_delete_success(client, ks_db, ks_client, async_return): + ks_client.delete_collection = async_return({}) + resp = client.delete("/creator/knowledge-stores/ks1") + assert resp.status_code == 200 + ks_db.delete_knowledge_store.assert_called_once() + + +def test_delete_collection_404_proceeds(client, ks_db, ks_client, async_raise): + ks_client.delete_collection = async_raise(HTTPException(404, "gone")) + resp = client.delete("/creator/knowledge-stores/ks1") + assert resp.status_code == 200 + + +def test_delete_collection_5xx_502(client, ks_client, async_raise): + ks_client.delete_collection = async_raise(HTTPException(503, "down")) + resp = client.delete("/creator/knowledge-stores/ks1") + assert resp.status_code == 502 + + +def test_delete_collection_4xx_propagates(client, ks_client, async_raise): + ks_client.delete_collection = async_raise(HTTPException(403, "forbidden")) + resp = client.delete("/creator/knowledge-stores/ks1") + assert resp.status_code == 403 + + +# --------------------------------------------------------------------------- +# share + list content +# --------------------------------------------------------------------------- + + +def test_toggle_sharing_on_and_off(client, ks_db): + on = client.put("/creator/knowledge-stores/ks1/share", json={"is_shared": True}) + assert on.json()["is_shared"] is True + off = client.put("/creator/knowledge-stores/ks1/share", json={"is_shared": False}) + assert "private" in off.json()["message"] + + +def test_list_knowledge_stores(client, ks_db): + ks_db.get_accessible_knowledge_stores.return_value = [{"id": "ks1"}] + body = client.get("/creator/knowledge-stores").json() + assert body["knowledge_stores"][0]["id"] == "ks1" + + +def test_list_ks_content(client, ks_db): + ks_db.get_kb_content_links_for_ks.return_value = [ + {"library_item_id": "i1", "status": "ready"}, + ] + body = client.get("/creator/knowledge-stores/ks1/content").json() + assert body["items"][0]["library_item_id"] == "i1" + + +# --------------------------------------------------------------------------- +# add_content +# --------------------------------------------------------------------------- + + +def _add_body(item_ids=None): + return {"library_id": "lib1", "item_ids": item_ids or ["i1"]} + + +def test_add_content_ks_not_found(client, ks_db): + ks_db.get_knowledge_store.return_value = None + resp = client.post("/creator/knowledge-stores/ks1/content", json=_add_body()) + assert resp.status_code == 404 + + +def test_add_content_library_not_found(client, ks_db, ks_client): + ks_db.get_knowledge_store.return_value = {"organization_id": 1, "embedding_vendor": "openai"} + ks_db.get_library.return_value = None + ks_client.resolve_embedding_api_key.return_value = "sk" + resp = client.post("/creator/knowledge-stores/ks1/content", json=_add_body()) + assert resp.status_code == 404 + + +def test_add_content_org_mismatch_403(client, ks_db, ks_client): + ks_db.get_knowledge_store.return_value = {"organization_id": 1, "embedding_vendor": "openai"} + ks_db.get_library.return_value = {"organization_id": 2} + ks_client.resolve_embedding_api_key.return_value = "sk" + resp = client.post("/creator/knowledge-stores/ks1/content", json=_add_body()) + assert resp.status_code == 403 + + +def test_add_content_all_already_linked_noop(client, ks_db, ks_client): + ks_db.get_knowledge_store.return_value = {"organization_id": 1, "embedding_vendor": "openai"} + ks_db.get_library.return_value = {"organization_id": 1, "name": "Lib"} + ks_client.resolve_embedding_api_key.return_value = "sk" + ks_db.get_kb_content_link.return_value = {"status": "ready"} + resp = client.post("/creator/knowledge-stores/ks1/content", json=_add_body()) + assert resp.json()["status"] == "noop" + + +def test_add_content_item_fetch_error_400(client, ks_db, ks_client, ks_library_client, + async_raise): + ks_db.get_knowledge_store.return_value = {"organization_id": 1, "embedding_vendor": "openai"} + ks_db.get_library.return_value = {"organization_id": 1, "name": "Lib"} + ks_client.resolve_embedding_api_key.return_value = "sk" + ks_db.get_kb_content_link.return_value = None + ks_library_client.get_item = async_raise(HTTPException(404, "no item")) + resp = client.post("/creator/knowledge-stores/ks1/content", json=_add_body()) + assert resp.status_code == 400 + + +def test_add_content_item_not_ready_409(client, ks_db, ks_client, ks_library_client, + async_return): + ks_db.get_knowledge_store.return_value = {"organization_id": 1, "embedding_vendor": "openai"} + ks_db.get_library.return_value = {"organization_id": 1, "name": "Lib"} + ks_client.resolve_embedding_api_key.return_value = "sk" + ks_db.get_kb_content_link.return_value = None + ks_library_client.get_item = async_return({"status": "processing"}) + resp = client.post("/creator/knowledge-stores/ks1/content", json=_add_body()) + assert resp.status_code == 409 + + +def test_add_content_empty_text_400(client, ks_db, ks_client, ks_library_client, + async_return): + ks_db.get_knowledge_store.return_value = {"organization_id": 1, "embedding_vendor": "openai"} + ks_db.get_library.return_value = {"organization_id": 1, "name": "Lib"} + ks_client.resolve_embedding_api_key.return_value = "sk" + ks_db.get_kb_content_link.return_value = None + ks_library_client.get_item = async_return({"status": "ready", "title": "T"}) + + class _Resp: + text = "" + + ks_library_client.proxy_content = async_return(_Resp()) + resp = client.post("/creator/knowledge-stores/ks1/content", json=_add_body()) + assert resp.status_code == 400 + + +def test_add_content_proxy_content_error_400(client, ks_db, ks_client, ks_library_client, + async_return, async_raise): + ks_db.get_knowledge_store.return_value = {"organization_id": 1, "embedding_vendor": "openai"} + ks_db.get_library.return_value = {"organization_id": 1, "name": "Lib"} + ks_client.resolve_embedding_api_key.return_value = "sk" + ks_db.get_kb_content_link.return_value = None + ks_library_client.get_item = async_return({"status": "ready", "title": "T"}) + ks_library_client.proxy_content = async_raise(HTTPException(404, "no content")) + resp = client.post("/creator/knowledge-stores/ks1/content", json=_add_body()) + assert resp.status_code == 400 + + +def test_add_content_bad_page_count_defaults_zero(client, ks_db, ks_client, + ks_library_client, async_return): + ks_db.get_knowledge_store.return_value = {"organization_id": 1, "embedding_vendor": "openai", + "embedding_endpoint": ""} + ks_db.get_library.return_value = {"organization_id": 1, "name": "Lib"} + ks_client.resolve_embedding_api_key.return_value = "sk" + ks_db.get_kb_content_link.return_value = None + # page_count is non-numeric -> the int() guard falls back to 0. + ks_library_client.get_item = async_return( + {"status": "ready", "title": "Doc", "page_count": "not-a-number"}) + + class _Resp: + text = "content" + + ks_library_client.proxy_content = async_return(_Resp()) + ks_client.add_content = async_return({"job_id": "j", "status": "processing"}) + ks_db.register_kb_content_link.return_value = 1 + resp = client.post("/creator/knowledge-stores/ks1/content", json=_add_body()) + assert resp.status_code == 200 + + +def test_add_content_success_registers_links(client, ks_db, ks_client, ks_library_client, + async_return): + ks_db.get_knowledge_store.return_value = {"organization_id": 1, "embedding_vendor": "openai", + "embedding_endpoint": ""} + ks_db.get_library.return_value = {"organization_id": 1, "name": "Lib"} + ks_client.resolve_embedding_api_key.return_value = "sk" + ks_db.get_kb_content_link.return_value = None + ks_library_client.get_item = async_return( + {"status": "ready", "title": "Doc", "page_count": 2, + "original_filename": "f.pdf"}) + + class _Resp: + text = "real content" + + ks_library_client.proxy_content = async_return(_Resp()) + ks_client.add_content = async_return({"job_id": "job1", "status": "processing", + "documents_total": 1}) + ks_db.register_kb_content_link.return_value = 99 + resp = client.post("/creator/knowledge-stores/ks1/content", json=_add_body()) + body = resp.json() + assert body["job_id"] == "job1" + assert body["links"][0]["id"] == 99 + + +def test_add_content_ingestion_failure_502(client, ks_db, ks_client, ks_library_client, + async_return, async_raise): + ks_db.get_knowledge_store.return_value = {"organization_id": 1, "embedding_vendor": "openai", + "embedding_endpoint": ""} + ks_db.get_library.return_value = {"organization_id": 1, "name": "Lib"} + ks_client.resolve_embedding_api_key.return_value = "sk" + ks_db.get_kb_content_link.return_value = None + ks_library_client.get_item = async_return({"status": "ready", "title": "Doc"}) + + class _Resp: + text = "content" + + ks_library_client.proxy_content = async_return(_Resp()) + ks_client.add_content = async_raise(HTTPException(500, "ingest boom")) + resp = client.post("/creator/knowledge-stores/ks1/content", json=_add_body()) + assert resp.status_code == 502 + + +# --------------------------------------------------------------------------- +# get_content_link / remove_content +# --------------------------------------------------------------------------- + + +def test_get_content_link_404(client, ks_db): + ks_db.get_kb_content_link.return_value = None + resp = client.get("/creator/knowledge-stores/ks1/content/i1") + assert resp.status_code == 404 + + +def test_get_content_link_polls_job(client, ks_db, ks_client, async_return): + ks_db.get_kb_content_link.side_effect = [ + {"id": 1, "status": "processing", "kb_job_id": "job1"}, + {"id": 1, "status": "ready", "kb_job_id": "job1"}, + ] + ks_client.get_job_status = async_return({"status": "completed", "chunks_created": 5}) + body = client.get("/creator/knowledge-stores/ks1/content/i1").json() + assert body["status"] == "ready" + + +def test_get_content_link_poll_failure_warns(client, ks_db, ks_client, async_raise): + ks_db.get_kb_content_link.return_value = {"id": 1, "status": "processing", + "kb_job_id": "job1"} + ks_client.get_job_status = async_raise(RuntimeError("poll boom")) + # Returns the (unchanged) link despite the poll failure. + body = client.get("/creator/knowledge-stores/ks1/content/i1").json() + assert body["status"] == "processing" + + +def test_remove_content_404(client, ks_db): + ks_db.get_kb_content_link.return_value = None + resp = client.delete("/creator/knowledge-stores/ks1/content/i1") + assert resp.status_code == 404 + + +def test_remove_content_success_cancels_inflight(client, ks_db, ks_client, async_return): + ks_db.get_kb_content_link.return_value = {"id": 1, "status": "processing", + "kb_job_id": "job1"} + ks_client.cancel_job = async_return({}) + ks_client.delete_content_by_source = async_return({}) + resp = client.delete("/creator/knowledge-stores/ks1/content/i1") + assert resp.status_code == 200 + ks_db.delete_kb_content_link.assert_called_once() + + +def test_remove_content_cancel_job_failure_warns(client, ks_db, ks_client, + async_return, async_raise): + ks_db.get_kb_content_link.return_value = {"id": 1, "status": "pending", + "kb_job_id": "job1"} + # cancel_job raises -> warning, but teardown proceeds. + ks_client.cancel_job = async_raise(RuntimeError("cancel boom")) + ks_client.delete_content_by_source = async_return({}) + resp = client.delete("/creator/knowledge-stores/ks1/content/i1") + assert resp.status_code == 200 + ks_db.delete_kb_content_link.assert_called_once() + + +def test_remove_content_delete_best_effort_5xx(client, ks_db, ks_client, async_raise): + ks_db.get_kb_content_link.return_value = {"id": 1, "status": "ready", "kb_job_id": None} + ks_client.delete_content_by_source = async_raise(HTTPException(503, "down")) + # Best-effort: link still torn down, returns 200. + resp = client.delete("/creator/knowledge-stores/ks1/content/i1") + assert resp.status_code == 200 + + +def test_remove_content_delete_4xx_propagates(client, ks_db, ks_client, async_raise): + ks_db.get_kb_content_link.return_value = {"id": 1, "status": "ready", "kb_job_id": None} + ks_client.delete_content_by_source = async_raise(HTTPException(403, "forbidden")) + resp = client.delete("/creator/knowledge-stores/ks1/content/i1") + assert resp.status_code == 403 + + +# --------------------------------------------------------------------------- +# query + job status +# --------------------------------------------------------------------------- + + +def test_query_not_found_404(client, ks_db): + ks_db.get_knowledge_store.return_value = None + resp = client.post("/creator/knowledge-stores/ks1/query", json={"query_text": "q"}) + assert resp.status_code == 404 + + +def test_query_success(client, ks_db, ks_client, async_return): + ks_db.get_knowledge_store.return_value = {"embedding_vendor": "openai", "embedding_endpoint": ""} + ks_client.resolve_embedding_api_key.return_value = "sk" + ks_client.query = async_return({"results": [{"text": "hit"}]}) + body = client.post("/creator/knowledge-stores/ks1/query", + json={"query_text": "q", "top_k": 3}).json() + assert body["results"][0]["text"] == "hit" + + +def test_job_status_syncs_links(client, ks_db, ks_client, async_return): + ks_client.get_job_status = async_return({"status": "completed", "chunks_created": 7}) + ks_db.get_kb_content_links_for_ks.return_value = [ + {"id": 1, "kb_job_id": "job1", "status": "processing"}, + ] + body = client.get("/creator/knowledge-stores/ks1/jobs/job1").json() + assert body["status"] == "completed" + ks_db.update_kb_content_link_status.assert_called_once() diff --git a/backend/tests/test_ks_query_helpers.py b/backend/tests/test_ks_query_helpers.py new file mode 100644 index 000000000..5c539a337 --- /dev/null +++ b/backend/tests/test_ks_query_helpers.py @@ -0,0 +1,110 @@ +import pytest +from lamb.completions.rag._ks_query_helpers import ( + _extract_sources, + build_context_and_sources, +) + + +class TestExtractSources: + def test_empty_results_returns_empty_list(self): + assert _extract_sources("ks-1", []) == [] + + def test_extracts_title_from_source_title(self): + results = [{"text": "chunk", "score": 0.9, "metadata": {"source_title": "My Doc"}}] + sources = _extract_sources("ks-1", results) + assert len(sources) == 1 + assert sources[0]["title"] == "My Doc" + assert sources[0]["knowledge_store_id"] == "ks-1" + assert sources[0]["score"] == 0.9 + + def test_falls_back_to_library_name_for_title(self): + results = [{"text": "chunk", "score": 0.8, "metadata": {"library_name": "Lib A"}}] + sources = _extract_sources("ks-1", results) + assert sources[0]["title"] == "Lib A" + + def test_falls_back_to_source_for_title(self): + results = [{"text": "chunk", "score": 0.8, "metadata": {}}] + sources = _extract_sources("ks-1", results) + assert sources[0]["title"] == "Source" + + def test_permalink_page_is_primary_url(self): + results = [{"text": "c", "score": 0.5, "metadata": { + "permalink_original": "/docs/1/lib/item/original.pdf", + "permalink_markdown": "/docs/1/lib/item/full.md", + "permalink_page": "/docs/1/lib/item/pages/1.md", + }}] + sources = _extract_sources("ks-1", results) + assert sources[0]["url"] == "/docs/1/lib/item/pages/1.md" + assert sources[0]["permalink_original"] == "/docs/1/lib/item/original.pdf" + assert sources[0]["permalink_markdown"] == "/docs/1/lib/item/full.md" + assert sources[0]["permalink_page"] == "/docs/1/lib/item/pages/1.md" + + def test_permalink_markdown_is_primary_when_no_page(self): + results = [{"text": "c", "score": 0.5, "metadata": { + "permalink_original": "/docs/1/lib/item/original.pdf", + "permalink_markdown": "/docs/1/lib/item/full.md", + }}] + sources = _extract_sources("ks-1", results) + assert sources[0]["url"] == "/docs/1/lib/item/full.md" + + def test_includes_library_metadata(self): + results = [{"text": "c", "score": 0.5, "metadata": { + "library_id": "lib-1", + "library_name": "My Library", + "source_item_id": "item-1", + }}] + sources = _extract_sources("ks-1", results) + assert sources[0]["library_id"] == "lib-1" + assert sources[0]["library_name"] == "My Library" + assert sources[0]["source_item_id"] == "item-1" + + def test_handles_none_metadata(self): + results = [{"text": "c", "score": 0.5, "metadata": None}] + sources = _extract_sources("ks-1", results) + assert len(sources) == 1 + assert sources[0]["title"] == "Source" + + +class TestBuildContextAndSources: + def test_empty_input_returns_empty(self): + context, sources = build_context_and_sources([]) + assert context == "" + assert sources == [] + + def test_numbers_chunks_and_aligns_sources(self): + ks_results = [ + ("ks-1", [ + {"text": "first chunk", "score": 0.9, "metadata": {"source_title": "Doc A"}}, + {"text": "second chunk", "score": 0.8, "metadata": {"source_title": "Doc B"}}, + ]), + ] + context, sources = build_context_and_sources(ks_results) + # Context blocks are prefixed with their 1-based citation number. + assert context == "[1] first chunk\n\n[2] second chunk" + # Each source carries the matching n. + assert [s["n"] for s in sources] == [1, 2] + assert sources[0]["title"] == "Doc A" + assert sources[1]["title"] == "Doc B" + + def test_numbering_is_continuous_across_knowledge_stores(self): + ks_results = [ + ("ks-1", [{"text": "a", "score": 0.9, "metadata": {}}]), + ("ks-2", [{"text": "b", "score": 0.7, "metadata": {}}]), + ] + context, sources = build_context_and_sources(ks_results) + assert context == "[1] a\n\n[2] b" + assert sources[0]["n"] == 1 and sources[0]["knowledge_store_id"] == "ks-1" + assert sources[1]["n"] == 2 and sources[1]["knowledge_store_id"] == "ks-2" + + def test_empty_text_chunks_are_skipped(self): + ks_results = [ + ("ks-1", [ + {"text": "", "score": 0.9, "metadata": {}}, + {"text": " ", "score": 0.9, "metadata": {}}, + {"text": "real", "score": 0.9, "metadata": {}}, + ]), + ] + context, sources = build_context_and_sources(ks_results) + assert context == "[1] real" + assert len(sources) == 1 + assert sources[0]["n"] == 1 diff --git a/backend/tests/test_kvcache_augment.py b/backend/tests/test_kvcache_augment.py new file mode 100644 index 000000000..8f8461b00 --- /dev/null +++ b/backend/tests/test_kvcache_augment.py @@ -0,0 +1,114 @@ +"""Tests for the citation handling in the kvcache_augment PPS. + +Citations are opt-in per assistant via ``capabilities.expose_sources``. When +off (the default), no inline ``[N]`` markers or instruction reach the model and +the ``[N]`` context prefixes are stripped. When on, the numbered context and +the cite instruction are kept (the clickable source list is rendered +separately). +""" + +import json + +from lamb.completions.pps.kvcache_augment import ( + CITATION_INSTRUCTION, + _build_full_context, + _exposes_sources, + prompt_processor, +) +from lamb.lamb_classes import Assistant + + +def _make_assistant(expose_sources=False, **overrides): + defaults = { + "id": 1, + "name": "Test", + "description": "", + "system_prompt": "You are helpful.", + "prompt_template": "", + "RAG_collections": "ks-1", + "RAG_Top_k": 3, + "owner": "user@test.com", + "api_callback": json.dumps({"capabilities": {"expose_sources": expose_sources}}), + "pre_retrieval_endpoint": "", + "post_retrieval_endpoint": "", + "RAG_endpoint": "", + } + defaults.update(overrides) + return Assistant(**defaults) + + +class TestExposesSources: + def test_defaults_false(self): + assert _exposes_sources(_make_assistant()) is False + assert _exposes_sources(None) is False + + def test_true_when_enabled(self): + assert _exposes_sources(_make_assistant(expose_sources=True)) is True + + def test_false_on_missing_or_bad_metadata(self): + assert _exposes_sources(_make_assistant(api_callback="")) is False + assert _exposes_sources(_make_assistant(api_callback="not json")) is False + + +class TestBuildFullContext: + def test_cite_true_appends_instruction_and_keeps_markers(self): + rag_context = { + "context": "[1] some retrieved text", + "sources": [{"n": 1, "title": "Doc A", "score": 0.9}], + } + full = _build_full_context(rag_context, cite=True) + assert "[1] some retrieved text" in full + assert CITATION_INSTRUCTION in full + assert "## Available Sources" not in full + + def test_cite_false_strips_markers_and_omits_instruction(self): + rag_context = { + "context": "[1] first chunk\n\n[2] second chunk", + "sources": [{"n": 1}, {"n": 2}], + } + full = _build_full_context(rag_context, cite=False) + assert CITATION_INSTRUCTION not in full + # The [N] prefixes are removed so the model has no cue to cite. + assert "[1]" not in full and "[2]" not in full + assert "first chunk" in full and "second chunk" in full + + def test_no_sources_means_no_instruction(self): + rag_context = {"context": "plain context", "sources": []} + assert _build_full_context(rag_context, cite=True) == "plain context" + + def test_non_dict_context_is_stringified(self): + assert _build_full_context("raw") == "raw" + assert _build_full_context(None) == "" + + +class TestPromptProcessorGating: + _RAG = { + "context": "[1] X is a thing.", + "sources": [{"n": 1, "title": "Doc A", "score": 0.9}], + } + + def test_opted_in_assistant_gets_citation_instruction(self): + assistant = _make_assistant(expose_sources=True, prompt_template="") + request = {"messages": [{"role": "user", "content": "What is X?"}]} + out = prompt_processor(request, assistant=assistant, rag_context=self._RAG) + user_msg = out[-1]["content"] + assert CITATION_INSTRUCTION in user_msg + assert "[1] X is a thing." in user_msg + + def test_default_assistant_gets_no_citation_markers(self): + assistant = _make_assistant(expose_sources=False, prompt_template="") + request = {"messages": [{"role": "user", "content": "What is X?"}]} + out = prompt_processor(request, assistant=assistant, rag_context=self._RAG) + user_msg = out[-1]["content"] + assert CITATION_INSTRUCTION not in user_msg + # The [1] prefix is stripped; the underlying text remains. + assert "[1]" not in user_msg + assert "X is a thing." in user_msg + + def test_template_branch_respects_opt_in(self): + assistant = _make_assistant( + expose_sources=True, prompt_template="Context:\n{context}\n\nQ: {user_input}" + ) + request = {"messages": [{"role": "user", "content": "What is X?"}]} + out = prompt_processor(request, assistant=assistant, rag_context=self._RAG) + assert CITATION_INSTRUCTION in out[-1]["content"] diff --git a/backend/tests/test_library_manager_client_diffcov.py b/backend/tests/test_library_manager_client_diffcov.py new file mode 100644 index 000000000..ad85e122e --- /dev/null +++ b/backend/tests/test_library_manager_client_diffcov.py @@ -0,0 +1,181 @@ +"""Unit tests for ``creator_interface.library_manager_client.LibraryManagerClient``. + +Covers the thin async proxy methods that resolve org config and delegate to +``_request`` / ``_fetch_bytes``. Both internals are mocked (via AsyncMock and a +stubbed ``_get_library_config``) so nothing leaves the process. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock + +import pytest + +from creator_interface.library_manager_client import LibraryManagerClient + + +def run(coro): + return asyncio.run(coro) + + +def _cfg(**over): + base = {"url": "http://lib", "token": "tok", "allowed_plugins": [], + "external_keys": {}} + base.update(over) + return base + + +@pytest.fixture +def client(monkeypatch): + c = LibraryManagerClient() + monkeypatch.setattr(c, "_get_library_config", lambda creator_user=None: _cfg()) + c._request = AsyncMock(return_value={"ok": True}) + c._fetch_bytes = AsyncMock(return_value="raw-response") + return c + + +def _last_request(c): + args, kwargs = c._request.call_args + return args, kwargs + + +def _last_fetch(c): + args, kwargs = c._fetch_bytes.call_args + return args, kwargs + + +class _FakeUploadFile: + def __init__(self, filename="f.txt", content=b"data", content_type="text/plain"): + self.filename = filename + self.content_type = content_type + self._content = content + self.seek_calls = [] + + async def read(self): + return self._content + + async def seek(self, pos): + self.seek_calls.append(pos) + + +# --------------------------------------------------------------------------- +# import_file (lines 215-216 — folder_id present branch) +# --------------------------------------------------------------------------- + + +def test_import_file_with_folder_id(client): + f = _FakeUploadFile() + run(client.import_file("lib1", f, "simple", "Title", + plugin_params={"a": 1}, api_keys={"k": "v"}, + folder_id="fold1")) + args, kwargs = _last_request(client) + assert args[0] == "POST" + assert args[1] == "/libraries/lib1/import/file" + assert kwargs["data"]["folder_id"] == "fold1" + assert "file" in kwargs["files"] + assert f.seek_calls == [0] + + +def test_import_file_without_folder_id(client): + f = _FakeUploadFile() + run(client.import_file("lib1", f, "simple", "Title")) + _, kwargs = _last_request(client) + assert "folder_id" not in kwargs["data"] + + +# --------------------------------------------------------------------------- +# import_youtube (lines 251-252, 255) +# --------------------------------------------------------------------------- + + +def test_import_youtube_default_language(client): + run(client.import_youtube("lib1", "http://yt", "youtube", "Title")) + args, kwargs = _last_request(client) + assert args[1] == "/libraries/lib1/import/youtube" + assert kwargs["json"]["plugin_params"]["language"] == "en" + assert kwargs["json"]["language"] == "en" + + +def test_import_youtube_plugin_params_language_wins(client): + run(client.import_youtube("lib1", "http://yt", "youtube", "Title", + language="en", plugin_params={"language": "es"})) + _, kwargs = _last_request(client) + assert kwargs["json"]["plugin_params"]["language"] == "es" + assert kwargs["json"]["language"] == "es" + + +# --------------------------------------------------------------------------- +# get_tree (273-274) +# --------------------------------------------------------------------------- + + +def test_get_tree(client): + run(client.get_tree("lib1")) + args, _ = _last_request(client) + assert (args[0], args[1]) == ("GET", "/libraries/lib1/tree") + + +# --------------------------------------------------------------------------- +# folder ops (279-280, 287-288, 296-297, 304-305) + move_items (312-313) +# --------------------------------------------------------------------------- + + +def test_create_folder(client): + run(client.create_folder("lib1", "New", parent_folder_id="p1")) + args, kwargs = _last_request(client) + assert (args[0], args[1]) == ("POST", "/libraries/lib1/folders") + assert kwargs["json"] == {"name": "New", "parent_folder_id": "p1"} + + +def test_rename_folder(client): + run(client.rename_folder("lib1", "fold1", "Renamed")) + args, kwargs = _last_request(client) + assert (args[0], args[1]) == ("PUT", "/libraries/lib1/folders/fold1") + assert kwargs["json"] == {"name": "Renamed"} + + +def test_move_folder(client): + run(client.move_folder("lib1", "fold1", "parent2")) + args, kwargs = _last_request(client) + assert (args[0], args[1]) == ("PUT", "/libraries/lib1/folders/fold1/move") + assert kwargs["json"] == {"parent_folder_id": "parent2"} + + +def test_delete_folder(client): + run(client.delete_folder("lib1", "fold1")) + args, _ = _last_request(client) + assert (args[0], args[1]) == ("DELETE", "/libraries/lib1/folders/fold1") + + +def test_move_items(client): + run(client.move_items("lib1", ["i1", "i2"], "fold1")) + args, kwargs = _last_request(client) + assert (args[0], args[1]) == ("POST", "/libraries/lib1/items/move") + assert kwargs["json"] == {"item_ids": ["i1", "i2"], "folder_id": "fold1"} + + +# --------------------------------------------------------------------------- +# capabilities (367-368, 373-374) + item content (388-389) +# --------------------------------------------------------------------------- + + +def test_get_capabilities(client): + run(client.get_capabilities()) + args, _ = _last_request(client) + assert (args[0], args[1]) == ("GET", "/capabilities") + + +def test_get_item_capabilities(client): + run(client.get_item_capabilities("lib1", "item1")) + args, _ = _last_request(client) + assert args[0] == "GET" + assert args[1] == "/libraries/lib1/items/item1/capabilities" + + +def test_get_item_content(client): + out = run(client.get_item_content("lib1", "item1", "audio")) + args, _ = _last_fetch(client) + assert args[0] == "GET" + assert args[1] == "/libraries/lib1/items/item1/content/audio" + assert out == "raw-response" diff --git a/backend/tests/test_library_router_diffcov.py b/backend/tests/test_library_router_diffcov.py new file mode 100644 index 000000000..c2003b6e1 --- /dev/null +++ b/backend/tests/test_library_router_diffcov.py @@ -0,0 +1,298 @@ +"""Diff-coverage tests for ``creator_interface.library_router``. + +Targets the error/edge branches added by the feature branch: capabilities +proxy, FR-10 library-delete blocking, youtube audit-language resolution, +capability content dispatch + caps, image-file traversal guard, original-file +resolution, and the kb-links / knowledge-stores aggregation endpoints. + +These reuse the shared ``client``, ``lib_client``, ``lib_db``, ``auth_ctx``, +``async_return`` fixtures from ``conftest.py``. +""" + +from __future__ import annotations + +import httpx + +from creator_interface import library_router + + +def _resp(content: bytes, content_type: str = "application/octet-stream"): + """Build an async stub mimicking the LibraryManagerClient httpx.Response.""" + + async def _proxy(*args, **kwargs): # noqa: ARG001 + return httpx.Response( + status_code=200, + content=content, + headers={"content-type": content_type}, + ) + + return _proxy + + +# --------------------------------------------------------------------------- +# /capabilities (line 123) +# --------------------------------------------------------------------------- + + +def test_list_capabilities(client, lib_client, async_return): + lib_client.get_capabilities = async_return({"handlers": ["text", "pages"]}) + resp = client.get("/creator/libraries/capabilities") + assert resp.status_code == 200 + assert resp.json()["handlers"] == ["text", "pages"] + + +# --------------------------------------------------------------------------- +# delete_library FR-10 blocking (lines 269-279, 284) +# --------------------------------------------------------------------------- + + +def test_delete_library_blocked_by_active_links(client, lib_client, lib_db, async_return): + lib_db.get_kb_content_links_for_library.return_value = [ + { + "knowledge_store_id": "KS-1", + "knowledge_store_name": "Course", + "status": "ready", + "library_item_id": "I1", + "item_title": "Doc", + }, + # duplicate KS id -> exercises the "already in blocking_stores" skip + { + "knowledge_store_id": "KS-1", + "knowledge_store_name": "Course", + "status": "processing", + "library_item_id": "I2", + "item_title": "Doc2", + }, + # link without ks_id -> exercises the falsy-ks_id branch + { + "knowledge_store_id": None, + "knowledge_store_name": None, + "status": "ready", + "library_item_id": "I3", + "item_title": "Doc3", + }, + # failed link -> filtered out before blocking + { + "knowledge_store_id": "KS-9", + "knowledge_store_name": "Bad", + "status": "failed", + "library_item_id": "I4", + "item_title": "Doc4", + }, + ] + lib_client.delete_library = async_return({}) + + resp = client.delete("/creator/libraries/lib-1") + assert resp.status_code == 409 + detail = resp.json()["detail"] + ks_ids = {ks["id"] for ks in detail["knowledge_stores"]} + assert ks_ids == {"KS-1"} + # Failed link excluded from items list (active_links only). + assert len(detail["items"]) == 3 + + +# --------------------------------------------------------------------------- +# import_youtube audit_language from plugin_params (line 482) +# --------------------------------------------------------------------------- + + +def test_import_youtube_audit_language_from_plugin_params(client, lib_client, lib_db, async_return): + lib_db.get_library.return_value = { + "id": "lib-1", "organization_id": 1, "owner_user_id": 1, "is_shared": False, + "name": "lib", "description": "", "import_config": {}, "status": "active", + } + lib_db.register_library_item.return_value = None + lib_client.import_youtube = async_return( + {"item_id": "i-1", "job_id": "j-1", "status": "processing"} + ) + + resp = client.post( + "/creator/libraries/lib-1/import-youtube", + json={ + "video_url": "https://youtu.be/x", + "plugin_params": {"language": "es"}, + }, + ) + assert resp.status_code == 200 + # Audit was written with the plugin_params language. + call = lib_db.write_audit_log.call_args + assert call.kwargs["details"]["language"] == "es" + + +# --------------------------------------------------------------------------- +# get_item_capabilities (lines 720-721) +# --------------------------------------------------------------------------- + + +def test_get_item_capabilities(client, lib_client, async_return): + lib_client.get_item_capabilities = async_return({"capabilities": ["text", "images"]}) + resp = client.get("/creator/libraries/lib-1/items/i-1/capabilities") + assert resp.status_code == 200 + assert resp.json()["capabilities"] == ["text", "images"] + + +# --------------------------------------------------------------------------- +# get_item_capability_content (lines 739-740, 743-744, 750) + _content_disposition +# --------------------------------------------------------------------------- + + +def test_get_item_capability_content_happy(client, lib_client): + lib_client.get_item_content = _resp(b"hello", "text/plain") + resp = client.get("/creator/libraries/lib-1/items/i-1/content/text") + assert resp.status_code == 200 + assert resp.content == b"hello" + assert resp.headers["content-type"].startswith("text/plain") + + +def test_get_item_capability_content_too_large(client, lib_client): + huge = b"x" * (library_router.MAX_CONTENT_BYTES + 1) + lib_client.get_item_content = _resp(huge) + resp = client.get("/creator/libraries/lib-1/items/i-1/content/images") + assert resp.status_code == 413 + + +# --------------------------------------------------------------------------- +# get_item_image_file (lines 773, 777-778, 780, 786) + _content_disposition (703-706) +# --------------------------------------------------------------------------- + + +def test_get_item_image_file_happy(client, lib_client): + lib_client.proxy_content = _resp(b"\x89PNG", "image/png") + resp = client.get("/creator/libraries/lib-1/items/i-1/content/images/file/pic.png") + assert resp.status_code == 200 + assert resp.content == b"\x89PNG" + assert resp.headers["content-type"] == "image/png" + + +def test_get_item_image_file_traversal_rejected(client, lib_client): + lib_client.proxy_content = _resp(b"unreachable") + # Backslash traversal hits the validation branch (a forward slash would + # change the route and not reach this handler). + resp = client.get("/creator/libraries/lib-1/items/i-1/content/images/file/..%5Cetc") + assert resp.status_code == 400 + assert resp.json()["detail"] == "Invalid filename" + + +# --------------------------------------------------------------------------- +# get_item_original (lines 811, 813-815, 817-819, 821-822, 831, 837) +# --------------------------------------------------------------------------- + + +def test_get_item_original_binary(client, lib_client, async_return): + # Filename with a space + parens exercises the ascii_safe sanitizer (chars + # replaced with "_") and the RFC 5987 quote() path. Avoid raw non-ASCII + # bytes in the header itself (the TestClient decodes headers as latin-1). + lib_client.get_item = async_return({ + "original_filename": "my report (1).pdf", + "source_url": None, + "source_type": "file", + }) + lib_client.proxy_content = _resp(b"%PDF-1.4", "application/pdf") + + resp = client.get("/creator/libraries/lib-1/items/i-1/original") + assert resp.status_code == 200 + assert resp.content == b"%PDF-1.4" + cd = resp.headers["content-disposition"] + assert "filename*=UTF-8''" in cd + assert "filename=" in cd + # Parens were sanitized out of the ascii_safe variant but survive in quoted. + assert "%281%29" in cd + + +def test_get_item_original_not_found(client, lib_client, async_return): + lib_client.get_item = async_return(None) + resp = client.get("/creator/libraries/lib-1/items/i-1/original") + assert resp.status_code == 404 + assert resp.json()["detail"] == "Item not found" + + +def test_get_item_original_no_binary(client, lib_client, async_return): + lib_client.get_item = async_return({ + "original_filename": None, + "source_url": "https://youtu.be/x", + "source_type": "youtube", + }) + resp = client.get("/creator/libraries/lib-1/items/i-1/original") + assert resp.status_code == 404 + detail = resp.json()["detail"] + assert detail["source_url"] == "https://youtu.be/x" + assert detail["source_type"] == "youtube" + + +# --------------------------------------------------------------------------- +# get_library_kb_links (lines 860-868) +# --------------------------------------------------------------------------- + + +def test_get_library_kb_links(client, lib_db): + lib_db.get_kb_content_links_for_library.return_value = [ + { + "knowledge_store_id": "KS-1", + "knowledge_store_name": "Course", + "status": "ready", + "library_item_id": "I1", + "item_title": "Doc", + }, + # duplicate KS id -> dedup branch + { + "knowledge_store_id": "KS-1", + "knowledge_store_name": "Course", + "status": "processing", + "library_item_id": "I2", + "item_title": None, + }, + # failed -> filtered out + { + "knowledge_store_id": "KS-2", + "knowledge_store_name": "Bad", + "status": "failed", + "library_item_id": "I3", + "item_title": "Doc3", + }, + ] + resp = client.get("/creator/libraries/lib-1/kb-links") + assert resp.status_code == 200 + body = resp.json() + assert body["library_id"] == "lib-1" + assert len(body["items"]) == 2 + assert [ks["id"] for ks in body["knowledge_stores"]] == ["KS-1"] + # item_title falls back to library_item_id when None. + titles = {it["id"]: it["title"] for it in body["items"]} + assert titles["I2"] == "I2" + + +# --------------------------------------------------------------------------- +# get_library_knowledge_stores (lines 939, 941-947, 961) +# --------------------------------------------------------------------------- + + +def test_get_library_knowledge_stores_filters_inaccessible(client, lib_db, auth_ctx): + lib_db.get_knowledge_stores_for_library.return_value = [ + { + "id": "KS-vis", "name": "Visible", "description": "d", + "chunking_strategy": "simple", "embedding_vendor": "openai", + "embedding_model": "m", "vector_db_backend": "chroma", + "is_shared": True, "item_count": 3, "ready_count": 2, "failed_count": 1, + }, + { + "id": "KS-hidden", "name": "Hidden", "description": "d", + "chunking_strategy": "simple", "embedding_vendor": "openai", + "embedding_model": "m", "vector_db_backend": "chroma", + "is_shared": False, "item_count": 0, "ready_count": 0, "failed_count": 0, + }, + ] + + def _access(ks_id): + return "none" if ks_id == "KS-hidden" else "owner" + + auth_ctx.can_access_knowledge_store = _access # type: ignore[assignment] + + resp = client.get("/creator/libraries/lib-1/knowledge-stores") + assert resp.status_code == 200 + body = resp.json() + assert body["library_id"] == "lib-1" + assert len(body["knowledge_stores"]) == 1 + visible = body["knowledge_stores"][0] + assert visible["id"] == "KS-vis" + assert visible["access"] == "owner" + assert visible["item_count"] == 3 diff --git a/backend/tests/test_org_config_resolver_diffcov.py b/backend/tests/test_org_config_resolver_diffcov.py new file mode 100644 index 000000000..54ac3ac83 --- /dev/null +++ b/backend/tests/test_org_config_resolver_diffcov.py @@ -0,0 +1,189 @@ +"""Branch-coverage tests for org_config_resolver Knowledge Store / provider helpers. + +These exercise ``get_knowledge_store_config``, ``get_provider_api_key`` and +``get_provider_endpoint`` directly by constructing an +``OrganizationConfigResolver`` and pre-seeding its lazily loaded ``_org`` so +no real DB is touched (``LambDatabaseManager`` is also patched at construction). +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from lamb.completions.org_config_resolver import OrganizationConfigResolver + + +def make_resolver(org: dict, setup_name: str = "default") -> OrganizationConfigResolver: + """Build a resolver whose organization is pre-seeded (no DB access).""" + with patch( + "lamb.completions.org_config_resolver.LambDatabaseManager", + return_value=MagicMock(name="db_manager"), + ): + resolver = OrganizationConfigResolver("owner@example.com", setup_name) + resolver._org = org # skip lazy DB load + return resolver + + +# --------------------------------------------------------------------------- +# get_knowledge_store_config (lines 165-168, 170-171, 176-177, 190) +# --------------------------------------------------------------------------- + + +def test_ks_config_from_setup(): + """ks_config present in setup -> normalized dict returned (176-189).""" + org = { + "is_system": False, + "config": { + "setups": { + "default": { + "knowledge_store": { + "url": "http://ks:9092", + "api_key": "secret", + "allowed_vector_db_backends": ["chromadb"], + "allowed_chunking_strategies": ["simple"], + "allowed_embedding_vendors": ["openai"], + "allowed_embedding_models": {"openai": ["text-embedding-3-small"]}, + } + } + } + }, + } + result = make_resolver(org).get_knowledge_store_config() + assert result["server_url"] == "http://ks:9092" + assert result["api_token"] == "secret" + assert result["allowed_vector_db_backends"] == ["chromadb"] + assert result["allowed_chunking_strategies"] == ["simple"] + assert result["allowed_embedding_vendors"] == ["openai"] + assert result["allowed_embedding_models"] == {"openai": ["text-embedding-3-small"]} + + +def test_ks_config_token_precedence(): + """server_url/api_token keys win over url/api_key/token aliases.""" + org = { + "is_system": False, + "config": { + "setups": { + "default": { + "knowledge_store": { + "server_url": "http://primary:9092", + "api_token": "primary-token", + } + } + } + }, + } + result = make_resolver(org).get_knowledge_store_config() + assert result["server_url"] == "http://primary:9092" + assert result["api_token"] == "primary-token" + + +def test_ks_config_empty_returns_empty_non_system(): + """No ks_config and not system org -> empty dict (line 190).""" + org = {"is_system": False, "config": {"setups": {"default": {}}}} + assert make_resolver(org).get_knowledge_store_config() == {} + + +def test_ks_config_env_fallback_system_org(): + """No ks_config but system org -> env fallback (lines 170-174).""" + org = {"is_system": True, "config": {"setups": {"default": {}}}} + env = { + "LAMB_KB_SERVER_V2": "http://env-ks:9092", + "LAMB_KB_SERVER_V2_TOKEN": "env-token", + } + with patch.dict("os.environ", env, clear=False): + result = make_resolver(org).get_knowledge_store_config() + assert result["server_url"] == "http://env-ks:9092" + assert result["api_token"] == "env-token" + + +def test_ks_config_env_fallback_defaults(): + """System org, env vars unset -> defaulted url and empty token.""" + org = {"is_system": True, "config": {}} + with patch.dict( + "os.environ", + {}, + clear=True, + ): + result = make_resolver(org).get_knowledge_store_config() + assert result["server_url"] == "http://kb-server:9092" + # empty-string env default falls through the ``or`` chain to None + assert result["api_token"] is None + + +# --------------------------------------------------------------------------- +# get_provider_api_key (lines 205-206) +# --------------------------------------------------------------------------- + + +def test_provider_api_key_present(): + org = { + "is_system": False, + "config": { + "setups": { + "default": {"providers": {"openai": {"api_key": "sk-123"}}} + } + }, + } + assert make_resolver(org).get_provider_api_key("openai") == "sk-123" + + +def test_provider_api_key_missing_provider_returns_empty(): + """get_provider_config returns {} -> empty string (the falsy branch).""" + org = {"is_system": False, "config": {"setups": {"default": {"providers": {}}}}} + assert make_resolver(org).get_provider_api_key("openai") == "" + + +def test_provider_api_key_provider_without_key(): + """Provider config exists but no api_key key -> empty string.""" + org = { + "is_system": False, + "config": { + "setups": {"default": {"providers": {"openai": {"models": ["gpt-4"]}}}} + }, + } + assert make_resolver(org).get_provider_api_key("openai") == "" + + +# --------------------------------------------------------------------------- +# get_provider_endpoint (lines 219-226) +# --------------------------------------------------------------------------- + + +def test_provider_endpoint_missing_provider_returns_empty(): + """No provider config -> early empty return (lines 220-221).""" + org = {"is_system": False, "config": {"setups": {"default": {"providers": {}}}}} + assert make_resolver(org).get_provider_endpoint("openai") == "" + + +@pytest.mark.parametrize( + "provider_cfg,expected", + [ + ({"endpoint": "http://e1"}, "http://e1"), + ({"base_url": "http://e2"}, "http://e2"), + ({"api_endpoint": "http://e3"}, "http://e3"), + # precedence: endpoint wins over later keys + ({"endpoint": "http://e1", "base_url": "http://e2"}, "http://e1"), + ], +) +def test_provider_endpoint_resolves_keys(provider_cfg, expected): + """First of endpoint/base_url/api_endpoint that is set wins (222-225).""" + org = { + "is_system": False, + "config": { + "setups": {"default": {"providers": {"openai": provider_cfg}}} + }, + } + assert make_resolver(org).get_provider_endpoint("openai") == expected + + +def test_provider_endpoint_no_matching_key_returns_empty(): + """Provider config present but none of the endpoint keys set -> '' (line 226).""" + org = { + "is_system": False, + "config": { + "setups": {"default": {"providers": {"openai": {"api_key": "sk"}}}} + }, + } + assert make_resolver(org).get_provider_endpoint("openai") == "" diff --git a/backend/tests/test_permalink_signing.py b/backend/tests/test_permalink_signing.py new file mode 100644 index 000000000..1042348c0 --- /dev/null +++ b/backend/tests/test_permalink_signing.py @@ -0,0 +1,39 @@ +"""Tests for the per-org HMAC permalink signing used by public citation links.""" + +from creator_interface.permalink_signing import sign, verify + + +class TestSignVerify: + def test_roundtrip(self): + sig = sign("3", "3/lib-1/item-1/view") + assert verify("3", "3/lib-1/item-1/view", sig) is True + + def test_rejects_tampered_resource(self): + sig = sign("3", "3/lib-1/item-1/view") + assert verify("3", "3/lib-1/item-2/view", sig) is False + + def test_rejects_empty_signature(self): + assert verify("3", "3/lib-1/item-1/view", "") is False + assert verify("3", "3/lib-1/item-1/view", None) is False + + def test_rejects_garbage_signature(self): + assert verify("3", "3/lib-1/item-1/view", "deadbeef") is False + + def test_per_org_isolation(self): + # A signature minted for org 3 must not validate for org 7, even for an + # identical resource path — the signing key is org-derived. + sig_org3 = sign("3", "3/lib-1/item-1/view") + assert verify("7", "3/lib-1/item-1/view", sig_org3) is False + # And the two orgs produce different signatures for the same resource. + assert sign("3", "x/y/z/view") != sign("7", "x/y/z/view") + + def test_org_id_int_or_str_consistent(self): + # Callers may pass org_id as int or str; signing normalizes to str. + assert sign(3, "3/lib/item/view") == sign("3", "3/lib/item/view") + + def test_original_download_resource_distinct_from_view(self): + view = sign("3", "3/lib/item/view") + original = sign("3", "3/lib/item/original/cv.pdf") + assert view != original + assert verify("3", "3/lib/item/original/cv.pdf", original) is True + assert verify("3", "3/lib/item/original/cv.pdf", view) is False diff --git a/backend/tests/test_query_rewriting_helper.py b/backend/tests/test_query_rewriting_helper.py new file mode 100644 index 000000000..33082dbdf --- /dev/null +++ b/backend/tests/test_query_rewriting_helper.py @@ -0,0 +1,47 @@ +import pytest +from lamb.completions.rag._query_rewriting_helper import get_last_user_message_text + + +class TestGetLastUserMessageText: + def test_empty_messages_returns_empty_string(self): + assert get_last_user_message_text([]) == "" + + def test_no_user_messages_returns_empty_string(self): + messages = [ + {"role": "assistant", "content": "Hello"}, + {"role": "system", "content": "System prompt"}, + ] + assert get_last_user_message_text(messages) == "" + + def test_extracts_last_user_message(self): + messages = [ + {"role": "user", "content": "First question"}, + {"role": "assistant", "content": "First answer"}, + {"role": "user", "content": "Second question"}, + ] + assert get_last_user_message_text(messages) == "Second question" + + def test_handles_multimodal_content(self): + messages = [ + {"role": "user", "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": "http://example.com/img.png"}}, + ]}, + ] + assert get_last_user_message_text(messages) == "Describe this" + + def test_handles_multimodal_with_multiple_text_parts(self): + messages = [ + {"role": "user", "content": [ + {"type": "text", "text": "Part one"}, + {"type": "text", "text": "Part two"}, + ]}, + ] + assert get_last_user_message_text(messages) == "Part one Part two" + + def test_handles_empty_content(self): + messages = [{"role": "user", "content": ""}] + assert get_last_user_message_text(messages) == "" + + def test_handles_none_messages(self): + assert get_last_user_message_text(None) == "" diff --git a/backend/tests/test_query_rewriting_ks_rag.py b/backend/tests/test_query_rewriting_ks_rag.py new file mode 100644 index 000000000..e1fde951f --- /dev/null +++ b/backend/tests/test_query_rewriting_ks_rag.py @@ -0,0 +1,161 @@ +import pytest +from unittest.mock import patch, AsyncMock, MagicMock +from lamb.lamb_classes import Assistant + + +def _make_assistant(**overrides): + defaults = { + "id": 1, + "name": "Test", + "description": "", + "system_prompt": "", + "prompt_template": "", + "RAG_collections": "ks-1", + "RAG_Top_k": 3, + "owner": "user@test.com", + "api_callback": "", + "pre_retrieval_endpoint": "", + "post_retrieval_endpoint": "", + "RAG_endpoint": "", + } + defaults.update(overrides) + return Assistant(**defaults) + + +@pytest.mark.asyncio +async def test_fallback_to_last_user_message_when_no_small_fast_model(): + """When small-fast-model is not configured, use last user message as query.""" + from lamb.completions.rag import query_rewriting_ks_rag + + assistant = _make_assistant() + messages = [ + {"role": "user", "content": "What is photosynthesis?"}, + {"role": "assistant", "content": "It is a process..."}, + {"role": "user", "content": "How does it work in detail?"}, + ] + + mock_ks_response = { + "status": "success", + "data": {"results": [{"text": "chunk text", "score": 0.9, "metadata": {"source_title": "Bio"}}]}, + } + + with patch.object(query_rewriting_ks_rag, "generate_optimal_query", new=AsyncMock(return_value="How does it work in detail?")), \ + patch.object(query_rewriting_ks_rag, "query_one_ks", new=AsyncMock(return_value=mock_ks_response)) as mock_query: + result = await query_rewriting_ks_rag.rag_processor(messages, assistant) + + mock_query.assert_called_once_with("ks-1", "How does it work in detail?", 3, "user@test.com") + assert "chunk text" in result["context"] + assert len(result["sources"]) == 1 + + +@pytest.mark.asyncio +async def test_uses_small_fast_model_query_when_configured(): + """When small-fast-model is configured, use its output as the query.""" + from lamb.completions.rag import query_rewriting_ks_rag + + assistant = _make_assistant() + messages = [ + {"role": "user", "content": "What is ML?"}, + {"role": "assistant", "content": "Machine learning is..."}, + {"role": "user", "content": "How does it differ from DL?"}, + ] + + mock_ks_response = { + "status": "success", + "data": {"results": [{"text": "ML vs DL chunk", "score": 0.85, "metadata": {}}]}, + } + + with patch.object(query_rewriting_ks_rag, "generate_optimal_query", new=AsyncMock(return_value="machine learning vs deep learning comparison")), \ + patch.object(query_rewriting_ks_rag, "query_one_ks", new=AsyncMock(return_value=mock_ks_response)) as mock_query: + result = await query_rewriting_ks_rag.rag_processor(messages, assistant) + + mock_query.assert_called_once_with("ks-1", "machine learning vs deep learning comparison", 3, "user@test.com") + assert "ML vs DL chunk" in result["context"] + + +@pytest.mark.asyncio +async def test_no_collections_returns_early(): + """When RAG_collections is empty, return early with info message.""" + from lamb.completions.rag import query_rewriting_ks_rag + + assistant = _make_assistant(RAG_collections="") + messages = [{"role": "user", "content": "Hello"}] + + result = await query_rewriting_ks_rag.rag_processor(messages, assistant) + assert "No Knowledge Stores specified" in result["context"] + assert result["sources"] == [] + + +@pytest.mark.asyncio +async def test_no_user_message_returns_early(): + """When there is no user message, return early.""" + from lamb.completions.rag import query_rewriting_ks_rag + + assistant = _make_assistant() + messages = [{"role": "assistant", "content": "Hello"}] + + with patch.object(query_rewriting_ks_rag, "generate_optimal_query", new=AsyncMock(return_value="")): + result = await query_rewriting_ks_rag.rag_processor(messages, assistant) + assert "No user message found" in result["context"] + + +@pytest.mark.asyncio +async def test_multiple_ks_queried_concurrently(): + """When multiple KS IDs are specified, all are queried.""" + from lamb.completions.rag import query_rewriting_ks_rag + + assistant = _make_assistant(RAG_collections="ks-1,ks-2") + messages = [{"role": "user", "content": "Test query"}] + + mock_response = {"status": "success", "data": {"results": []}} + + with patch.object(query_rewriting_ks_rag, "generate_optimal_query", new=AsyncMock(return_value="Test query")), \ + patch.object(query_rewriting_ks_rag, "query_one_ks", new=AsyncMock(return_value=mock_response)) as mock_query: + result = await query_rewriting_ks_rag.rag_processor(messages, assistant) + + assert mock_query.call_count == 2 + called_ids = {call.args[0] for call in mock_query.call_args_list} + assert called_ids == {"ks-1", "ks-2"} + + +@pytest.mark.asyncio +async def test_sfm_error_falls_back_to_last_user_message(): + """When small-fast-model raises an exception, fall back to last user message.""" + from lamb.completions.rag import query_rewriting_ks_rag + + assistant = _make_assistant() + messages = [ + {"role": "user", "content": "What is AI?"}, + {"role": "assistant", "content": "AI is..."}, + {"role": "user", "content": "Give me examples"}, + ] + + mock_ks_response = {"status": "success", "data": {"results": []}} + + with patch.object(query_rewriting_ks_rag, "generate_optimal_query", new=AsyncMock(return_value="Give me examples")), \ + patch.object(query_rewriting_ks_rag, "query_one_ks", new=AsyncMock(return_value=mock_ks_response)) as mock_query: + result = await query_rewriting_ks_rag.rag_processor(messages, assistant) + + mock_query.assert_called_once_with("ks-1", "Give me examples", 3, "user@test.com") + + +@pytest.mark.asyncio +async def test_handles_multimodal_content_in_fallback(): + """When user message is multimodal (list), extract text parts.""" + from lamb.completions.rag import query_rewriting_ks_rag + + assistant = _make_assistant() + messages = [ + {"role": "user", "content": [ + {"type": "text", "text": "Describe this image"}, + {"type": "image_url", "image_url": {"url": "http://example.com/img.png"}}, + ]}, + ] + + mock_ks_response = {"status": "success", "data": {"results": []}} + + with patch.object(query_rewriting_ks_rag, "generate_optimal_query", new=AsyncMock(return_value="Describe this image")), \ + patch.object(query_rewriting_ks_rag, "query_one_ks", new=AsyncMock(return_value=mock_ks_response)) as mock_query: + result = await query_rewriting_ks_rag.rag_processor(messages, assistant) + + mock_query.assert_called_once_with("ks-1", "Describe this image", 3, "user@test.com") diff --git a/backend/tests/test_simple_augment.py b/backend/tests/test_simple_augment.py index d60a62f3d..1ad9c8bbd 100644 --- a/backend/tests/test_simple_augment.py +++ b/backend/tests/test_simple_augment.py @@ -1,17 +1,9 @@ """ Tests for the simple_augment prompt processor. - -Regression test for defect D3 (lifecycle 2026-05-03): when an assistant -has an empty ``prompt_template`` but a non-empty ``rag_context`` is -supplied, simple_augment must apply a default template so the retrieved -chunks actually reach the LLM. Previously it silently dropped them. """ from lamb.lamb_classes import Assistant -from lamb.completions.pps.simple_augment import ( - DEFAULT_RAG_PROMPT_TEMPLATE, - prompt_processor, -) +from lamb.completions.pps.simple_augment import prompt_processor def _make_assistant(prompt_template: str = "", system_prompt: str = "You are a helper.") -> Assistant: @@ -33,35 +25,6 @@ def _make_assistant(prompt_template: str = "", system_prompt: str = "You are a h ) -def test_empty_template_with_rag_context_uses_default_template(): - """D3 regression: empty prompt_template + non-empty rag_context must - fall back to the module's DEFAULT_RAG_PROMPT_TEMPLATE so the retrieved - chunks reach the LLM. The augmented user message must contain both the - user's question and the retrieved context text.""" - assistant = _make_assistant(prompt_template="") - request = {"messages": [{"role": "user", "content": "What is mitochondria?"}]} - rag_context = {"context": "Mitochondria are the powerhouse of the cell."} - - result = prompt_processor(request, assistant=assistant, rag_context=rag_context) - - # System prompt preserved. - assert result[0]["role"] == "system" - assert result[0]["content"] == "You are a helper." - - # Last message is the augmented user message. - last = result[-1] - assert last["role"] == "user" - augmented = last["content"] - assert "Mitochondria are the powerhouse of the cell." in augmented, ( - "RAG context was dropped — defect D3 has regressed." - ) - assert "What is mitochondria?" in augmented - # And the default template was used (verifies the fallback path, not - # accidental string match). - assert "{context}" not in augmented # placeholder must be substituted - assert "{user_input}" not in augmented - - def test_empty_template_without_rag_context_passthrough(): """If the template is empty AND there's no rag_context, the user message is passed through unchanged — preserves the prior no-RAG behaviour.""" @@ -75,7 +38,7 @@ def test_empty_template_without_rag_context_passthrough(): def test_explicit_template_with_context_placeholder_unchanged(): """When the assistant defines its own template containing {context}, - the default fallback must NOT kick in — the user-provided template wins. + the template is applied correctly. """ custom = "Refer to context: {context}\n\nQ: {user_input}" assistant = _make_assistant(prompt_template=custom) @@ -87,15 +50,3 @@ def test_explicit_template_with_context_placeholder_unchanged(): augmented = result[-1]["content"] assert "Refer to context:" in augmented # custom template preserved assert "Some chunk." in augmented - - -def test_default_template_constant_matches_cli(): - """Sanity check: the module-level constant matches the CLI default - (kept in lamb-cli/src/lamb_cli/commands/assistant.py). If either side - changes, the other should be updated to match.""" - expected = ( - "Use the following context to answer the question. " - "If the context does not contain the answer, say you do not know.\n\n" - "Context:\n{context}\n\nQuestion: {user_input}" - ) - assert DEFAULT_RAG_PROMPT_TEMPLATE == expected diff --git a/backend/tests/test_smallfry_diffcov.py b/backend/tests/test_smallfry_diffcov.py new file mode 100644 index 000000000..d0735661b --- /dev/null +++ b/backend/tests/test_smallfry_diffcov.py @@ -0,0 +1,205 @@ +"""Diff-coverage tests for three small feature-added gaps (label: smallfry). + +Covers, in one file: + +1. ``creator_interface.main`` lines 3-6 — the feature-added + ``knowledge_store_router`` / ``knowledge_store_graph_router`` import + lines. See ``test_creator_interface_main_imports`` for why these are + import-time-only and the fallback handling. +2. ``creator_interface.knowledges_router`` line 237 — + ``load_dotenv(override=True)`` (module top-level). Reached by importing + the module. +3. ``lamb.completions.pps.simple_augment`` line 144 — the vision / + multimodal augmentation branch (``effective_template.replace`` of + ``{user_input}`` for a list-type message on a vision-capable assistant). +""" + +from __future__ import annotations + +import importlib +import json +import os +import sqlite3 +import sys +from pathlib import Path + +import pytest + +# Make ``backend/`` importable (mirrors conftest) so ``config`` etc. resolve. +_BACKEND_ROOT = Path(__file__).parent.parent +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + + +# --------------------------------------------------------------------------- +# Target 3: simple_augment line 144 (vision multimodal {user_input} replace) +# --------------------------------------------------------------------------- + + +def _make_vision_assistant(prompt_template: str): + from lamb.lamb_classes import Assistant + + return Assistant( + id=1, + organization_id=1, + name="Vision Assistant", + description="", + owner="user@example.com", + # metadata maps to api_callback; capabilities.vision=True turns on the + # multimodal branch in simple_augment. + api_callback=json.dumps({"capabilities": {"vision": True}}), + system_prompt="You are a helper.", + prompt_template=prompt_template, + pre_retrieval_endpoint="", + post_retrieval_endpoint="", + RAG_endpoint="", + RAG_Top_k=3, + RAG_collections="", + ) + + +def test_simple_augment_vision_multimodal_user_input_substitution(): + """Line 144: a vision-enabled assistant whose last message is a list of + content parts must substitute ``{user_input}`` with the joined text parts + while preserving non-text (image) parts.""" + from lamb.completions.pps.simple_augment import prompt_processor + + assistant = _make_vision_assistant( + prompt_template="Answer: {user_input}\nContext: {context}" + ) + + request = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What"}, + {"type": "text", "text": "is this?"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,AAAA"}, + }, + ], + } + ] + } + + result = prompt_processor( + request, + assistant=assistant, + rag_context={"context": "Some retrieved context."}, + ) + + # Last processed message is the augmented multimodal user message. + augmented = result[-1] + assert augmented["role"] == "user" + content = augmented["content"] + assert isinstance(content, list) + + # First element is the augmented text with {user_input} replaced (line 144). + text_part = content[0] + assert text_part["type"] == "text" + assert "What is this?" in text_part["text"] + assert "Some retrieved context." in text_part["text"] + + # Original image part is preserved. + assert any(part.get("type") == "image_url" for part in content) + + +def test_simple_augment_vision_multimodal_without_rag_context(): + """Same multimodal branch but with no rag_context, so {context} is + stripped — still exercises the {user_input} replace on line 144.""" + from lamb.completions.pps.simple_augment import prompt_processor + + assistant = _make_vision_assistant( + prompt_template="Q: {user_input} / {context}" + ) + + request = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe the picture"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,BBBB"}, + }, + ], + } + ] + } + + result = prompt_processor(request, assistant=assistant, rag_context=None) + + text_part = result[-1]["content"][0] + assert "Describe the picture" in text_part["text"] + # No context placeholder left dangling. + assert "{context}" not in text_part["text"] + + +# --------------------------------------------------------------------------- +# Target 2: knowledges_router line 237 (load_dotenv(override=True)) +# --------------------------------------------------------------------------- + + +def _ensure_owi_db_exists() -> None: + """``OwiDatabaseManager.__init__`` blocks in a poll loop until the OWI + ``webui.db`` file exists. Create an (empty) sqlite file at the configured + path so importing knowledges_router does not hang. The managers tolerate a + table-less DB at import time (they only log errors).""" + import config + + owi_path = config.OWI_PATH + assert owi_path, "OWI_PATH must be configured for this test" + os.makedirs(owi_path, exist_ok=True) + db_path = os.path.join(owi_path, "webui.db") + if not os.path.exists(db_path): + sqlite3.connect(db_path).close() + + +def test_knowledges_router_module_imports(): + """Line 237: importing the module executes the top-level + ``load_dotenv(override=True)`` and the rest of the module body.""" + os.environ.setdefault("LAMB_KB_SERVER_TOKEN", "test-token") + _ensure_owi_db_exists() + + mod = importlib.import_module("creator_interface.knowledges_router") + # Re-import to be explicit even if another test already loaded it. + importlib.reload(mod) if mod.__name__ in sys.modules else None + + assert mod.router is not None + assert hasattr(mod, "KB_SERVER_CONFIGURED") + + +# --------------------------------------------------------------------------- +# Target 1: creator_interface.main lines 3-6 (knowledge_store router imports) +# --------------------------------------------------------------------------- + + +def test_creator_interface_main_imports(): + """Lines 3-6: feature-added ``knowledge_store_router`` / + ``knowledge_store_graph_router`` import statements at the top of + ``creator_interface/main.py``. + + Importing the module executes those lines, but the import chain reaches + ``creator_interface.evaluaitor_router`` -> ``lamb.evaluaitor.ai_generator`` + -> ``from openai import OpenAI``, and the ``openai`` package is not + installed in the test environment. If it ever becomes importable here this + test covers the lines; otherwise it is skipped and the lines carry a + ``# pragma: no cover`` justified by this exact constraint.""" + os.environ.setdefault("LAMB_KB_SERVER_TOKEN", "test-token") + _ensure_owi_db_exists() + + try: + import openai # noqa: F401 + except Exception: + pytest.skip( + "openai package is not installed in the test environment; " + "creator_interface.main cannot be imported, so lines 3-6 carry " + "# pragma: no cover." + ) + + mod = importlib.import_module("creator_interface.main") + assert hasattr(mod, "knowledge_store_router") + assert hasattr(mod, "knowledge_store_graph_router") diff --git a/backend/utils/pipelines/auth.py b/backend/utils/pipelines/auth.py index 1ec718f10..8ef54970f 100644 --- a/backend/utils/pipelines/auth.py +++ b/backend/utils/pipelines/auth.py @@ -19,7 +19,8 @@ import requests import uuid -SESSION_SECRET = os.getenv("SESSION_SECRET", " ") +# No insecure fallback: signing without a configured secret must fail closed. +SESSION_SECRET = os.getenv("SESSION_SECRET") ALGORITHM = "HS256" ############## @@ -41,6 +42,8 @@ def get_password_hash(password): def create_token(data: dict, expires_delta: Union[timedelta, None] = None) -> str: + if not SESSION_SECRET: + raise RuntimeError("SESSION_SECRET environment variable is required to sign tokens") payload = data.copy() if expires_delta: @@ -52,6 +55,8 @@ def create_token(data: dict, expires_delta: Union[timedelta, None] = None) -> st def decode_token(token: str) -> Optional[dict]: + if not SESSION_SECRET: + return None try: decoded = jwt.decode(token, SESSION_SECRET, algorithms=[ALGORITHM]) return decoded @@ -65,6 +70,21 @@ def extract_token_from_auth_header(auth_header: str): def get_current_user( credentials: HTTPAuthorizationCredentials = Depends(bearer_security), -) -> Optional[dict]: +) -> str: + """Validate the bearer token against the system API key. + + Previously this returned the token unconditionally, so any + ``Authorization: Bearer `` header was accepted. The + endpoints that depend on this (legacy LTI-users routes, the + pipelines reload endpoint) are system/admin surfaces and require + the LAMB system token. (#410) + """ + import config + token = credentials.credentials + if not token or token != config.API_KEY: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid or missing authentication token", + ) return token diff --git a/docker-compose-example.yaml b/docker-compose-example.yaml index 175d32ca4..f3557db33 100644 --- a/docker-compose-example.yaml +++ b/docker-compose-example.yaml @@ -63,6 +63,7 @@ services: - PYTHONUNBUFFERED=1 - GLOBAL_LOG_LEVEL=WARNING - KB_LOG_LEVEL=WARNING + - LAMB_API_KEY=0p3n-w3bu! volumes: - ${LAMB_PROJECT_PATH}:${LAMB_PROJECT_PATH} - pip-cache:/root/.cache/pip @@ -129,6 +130,10 @@ services: # per-request by LAMB on add-content/query and override this fallback. - OPENAI_API_KEY=placeholder-set-per-request - EMBEDDINGS_APIKEY=placeholder-set-per-request + - HF_TOKEN=none + - HF_HOME=/tmp/huggingface + - TRANSFORMERS_CACHE=/tmp/huggingface + - HF_HUB_CACHE=/tmp/huggingface/hub volumes: - ${LAMB_PROJECT_PATH}:${LAMB_PROJECT_PATH} - pip-cache:/root/.cache/pip @@ -152,9 +157,11 @@ services: chown -R ${HOST_UID:-1000}:${HOST_GID:-1001} ${LAMB_PROJECT_PATH}/lamb-kb-server/data 2>/dev/null || true && \ python -m pip install --upgrade pip && \ pip install -e '${LAMB_PROJECT_PATH}/lamb-kb-server[all]' && \ + mkdir -p /root/.cache/huggingface && \ + touch /root/.cache/huggingface/token && \ + chown -R ${HOST_UID:-1000}:${HOST_GID:-1001} /root/.cache/huggingface 2>/dev/null || true && \ LOG_LEVEL=$$(echo \"$${KB_LOG_LEVEL:-$${GLOBAL_LOG_LEVEL:-WARNING}}\" | tr '[:upper:]' '[:lower:]') && \ - exec setpriv --reuid=${HOST_UID:-1000} --regid=${HOST_GID:-1001} --clear-groups \ - uvicorn main:app --host 0.0.0.0 --port 9092 --reload --log-level $$LOG_LEVEL --no-access-log" + exec setpriv --reuid=${HOST_UID:-1000} --regid=${HOST_GID:-1001} --clear-groups uvicorn main:app --host 0.0.0.0 --port 9092 --reload --log-level $$LOG_LEVEL --no-access-log" backend: image: python:3.11-slim extra_hosts: diff --git a/docker-compose-workers.yaml b/docker-compose-workers.yaml index 2d8b44270..665048c23 100644 --- a/docker-compose-workers.yaml +++ b/docker-compose-workers.yaml @@ -51,6 +51,8 @@ services: kb: image: python:3.11-slim working_dir: ${LAMB_PROJECT_PATH}/lamb-kb-server-stable/backend + extra_hosts: + - "host.docker.internal:host-gateway" environment: - PIP_DISABLE_PIP_VERSION_CHECK=1 - PIP_CACHE_DIR=/root/.cache/pip diff --git a/docker-compose.firecrawl.yaml b/docker-compose.firecrawl.yaml index 026f08b00..355434bee 100644 --- a/docker-compose.firecrawl.yaml +++ b/docker-compose.firecrawl.yaml @@ -1,6 +1,7 @@ services: firecrawl: - build: https://github.com/firecrawl/firecrawl.git#main:apps/api + # Use the prebuilt image (no BuildKit/buildx needed). Was: build from git apps/api. + image: ghcr.io/firecrawl/firecrawl:latest environment: HOST: 0.0.0.0 PORT: ${FIRECRAWL_INTERNAL_PORT:-3002} @@ -144,7 +145,8 @@ services: firecrawl-postgres: # Firecrawl's repo builds a custom image with pg_cron and init SQL. - build: https://github.com/firecrawl/firecrawl.git#main:apps/nuq-postgres + # Use the prebuilt image (no BuildKit/buildx needed). Was: build from git apps/nuq-postgres. + image: ghcr.io/firecrawl/nuq-postgres:latest command: postgres -c log_min_messages=WARNING -c log_statement=none -c cron.log_run=off -c cron.log_statement=off environment: POSTGRES_USER: ${FIRECRAWL_POSTGRES_USER:-postgres} diff --git a/docker-compose.next.yaml b/docker-compose.next.yaml index 3c94102c2..1162deab8 100644 --- a/docker-compose.next.yaml +++ b/docker-compose.next.yaml @@ -121,6 +121,32 @@ services: volumes: - ollama-data:/root/.ollama + # Optional Neo4j for KG-RAG / semantic graph. Only starts with the + # ``kg-rag`` profile so existing deployments keep their current behavior. + # + # docker compose -f docker-compose.next.yaml --profile kg-rag up -d + # + # Set KG_RAG_ENABLED=true in the kb service env, plus KG_RAG_NEO4J_URI=bolt://neo4j:7687. + neo4j: + image: neo4j:5-community + profiles: ["kg-rag"] + restart: unless-stopped + environment: + NEO4J_AUTH: ${KG_RAG_NEO4J_USER:-neo4j}/${KG_RAG_NEO4J_PASSWORD:-change-this-password} + NEO4J_PLUGINS: '[]' + ports: + - "7474:7474" + - "7687:7687" + volumes: + - neo4j-data:/data + - neo4j-logs:/logs + healthcheck: + test: ["CMD", "wget", "-q", "--spider", "http://localhost:7474/"] + interval: 10s + timeout: 5s + retries: 12 + start_period: 30s + volumes: lamb-data: kb-data: @@ -128,3 +154,5 @@ volumes: openwebui-data: library-manager-data: ollama-data: + neo4j-data: + neo4j-logs: diff --git a/frontend/svelte-app/package.json b/frontend/svelte-app/package.json index 90b28133b..d49630070 100644 --- a/frontend/svelte-app/package.json +++ b/frontend/svelte-app/package.json @@ -49,9 +49,20 @@ "ai": "^4.3.7", "axios": "^1.8.1", "dompurify": "^3.4.1", + "graphology": "^0.26.0", + "graphology-layout-forceatlas2": "^0.10.1", + "katex": "^0.17.0", "lucide-svelte": "^0.460.0", "marked": "^15.0.12", "openai": "^4.53.3", - "svelte-i18n": "^4.0.1" + "rehype-katex": "^7.0.1", + "rehype-stringify": "^10.0.1", + "remark-gfm": "^4.0.1", + "remark-math": "^6.0.0", + "remark-parse": "^11.0.0", + "remark-rehype": "^11.1.2", + "sigma": "^3.0.3", + "svelte-i18n": "^4.0.1", + "unified": "^11.0.5" } } diff --git a/frontend/svelte-app/src/lib/components/ChatInterface.svelte b/frontend/svelte-app/src/lib/components/ChatInterface.svelte index 37e829c49..224688ab5 100644 --- a/frontend/svelte-app/src/lib/components/ChatInterface.svelte +++ b/frontend/svelte-app/src/lib/components/ChatInterface.svelte @@ -1,7 +1,8 @@ -
- - {#if loading} -
-
- {$_('knowledgeBases.detail.loading', { default: 'Loading knowledge base details...' })} -
-
- - - {:else if error && !kb} -
-
- {#if serverOffline} -
-

- {$_('knowledgeBases.detail.serverOffline', { - default: 'Knowledge Base Server Offline' - })} -

-

- {$_('knowledgeBases.detail.tryAgainLater', { - default: 'Please try again later or contact an administrator.' - })} -

-
- {:else} - {error} - {/if} -
- -
- - - {:else if kb} -
- -
-

- {kb.name} -

-

- {kb.description || - $_('knowledgeBases.detail.noDescription', { default: 'No description provided.' })} -

-
- - -
-
-
-
- {$_('knowledgeBases.detail.idLabel', { default: 'ID' })} -
-
- {kb.id} -
-
- - {#if kb.owner} -
-
- {$_('knowledgeBases.detail.ownerLabel', { default: 'Owner' })} -
-
- {kb.owner} -
-
- {/if} - - {#if kb.created_at} -
-
- {$_('knowledgeBases.detail.createdLabel', { default: 'Created' })} -
-
- {new Date(kb.created_at * 1000).toLocaleString()} -
-
- {/if} - - {#if kb.metadata?.access_control} -
-
- {$_('knowledgeBases.detail.accessLabel', { default: 'Access Control' })} -
-
- - {kb.metadata.access_control} - -
-
- {/if} -
-
- - -
- -
- -
- - -
- - {#if activeTab === 'files'} -
- -
-
- {#if loadingJobs} - - - - - - Loading status... - - {:else if jobsError} - ⚠ {jobsError} - {/if} -
- -
- - {#if kb.files && kb.files.length > 0} -
- - - - - - - - - - - - {#each kb.files as file (file.id)} - {@const job = getJobForFile(file.filename)} - {@const statusColors = job - ? getStatusColors(job.status) - : getStatusColors('unknown')} - - - - - - - - {/each} - -
- {$_('knowledgeBases.detail.fileNameColumn', { default: 'File Name' })} - - {$_('knowledgeBases.detail.fileSizeColumn', { default: 'Size' })} - - {$_('knowledgeBases.detail.fileTypeColumn', { default: 'Type' })} - - {$_('knowledgeBases.detail.fileStatusColumn', { default: 'Status' })} - - {$_('knowledgeBases.detail.fileActionsColumn', { default: 'Actions' })} -
-
-
- {#if file.file_url} - - {file.filename} - - {:else} - {file.filename} - {/if} -
-
- 		-
- {file.size ? formatFileSize(file.size) : 'N/A'} -
- 		-
- {file.content_type || 'Unknown'} -
- 		- {#if job} - - {:else if loadingJobs} - Loading... - {:else} - - Unknown - - {/if} - - {#if kb.can_modify === true} - - {/if} -
-
- {:else} -
- {$_('knowledgeBases.detail.noFiles', { - default: 'No files have been uploaded to this knowledge base.' - })} -
- {/if} -
- {/if} - - - {#if activeTab === 'ingest'} - {#key selectedPlugin?.kind} -
-

- {selectedPlugin?.kind === 'file-ingest' - ? $_('knowledgeBases.fileUpload.sectionTitle', { - default: 'Upload and Ingest New File' - }) - : $_('knowledgeBases.fileUpload.sectionTitleBase', { - default: 'Configure and Run Ingestion' - })} -

- - {#if loadingPlugins} - - {:else if pluginsError} - - {:else if plugins.length === 0} - - {:else} - -
{ - e.preventDefault(); - handleSubmitIngestion(); - }} - class="space-y-6" - > - - {#if pluginCanUseFileUpload(selectedPlugin)} -
- -
- -
- {#if selectedPlugin?.kind !== 'file-ingest'} -

- Optional: If a file is selected, ingestion runs through the file-upload - path. Leave empty to run direct ingestion. -

- {/if} - {#if selectedFile} -

- {selectedFile.name} ({formatFileSize(selectedFile.size)}) -

- {/if} -
- {/if} - - -
- -
- -
-
- - - {#if selectedPlugin?.parameters} - {#each Object.entries(selectedPlugin.parameters).filter( ([name]) => name.startsWith('_') ) as [paramName, paramDef] (paramName)} -
-
- - - -
-

- {paramName.replace(/^_/, '').replace(/_/g, ' ')} -

-

- {paramDef.description || paramDef.default || ''} -

-
-
-
- {/each} - {/if} - - - {#if selectedPlugin && selectedPlugin.parameters && visibleParamKeys.length > 0} -
-
-
- {$_('knowledgeBases.fileUpload.parametersLabel', { - default: 'Plugin Parameters' - })} -
- {#if hasAdvancedParams} - - {/if} -
- - {#each displayedParamKeys as paramName (paramName)} - {@const paramDef = selectedPlugin.parameters[paramName]} - -
- - - -
- {#if paramDef.enum && Array.isArray(paramDef.enum)} - - - {:else if paramDef.type === 'boolean'} - -
- updateParamValue(paramName, e)} - /> - -
- {:else if paramDef.type === 'integer' || paramDef.type === 'number'} - - updateParamValue(paramName, e)} - /> - {#if paramDef.min !== undefined || paramDef.max !== undefined} -

- {#if paramDef.min !== undefined && paramDef.max !== undefined} - Range: {paramDef.min} - {paramDef.max} - {:else if paramDef.min !== undefined} - Minimum: {paramDef.min} - {:else} - Maximum: {paramDef.max} - {/if} -

- {/if} - {:else if paramDef.type === 'array'} - - -

- Enter values separated by new lines. -

- {:else if paramDef.type === 'long-string'} - - - {:else} - - updateParamValue(paramName, e)} - /> - {/if} -
- - - {#if paramDef.help_text} -

{paramDef.help_text}

- {/if} -
- {/each} -
- {/if} - - -
- - {#if uploadSuccess} -
-
- - - - {$_('knowledgeBases.fileUpload.success', { - default: 'File uploaded and ingestion started successfully!' - })} -
-
- {/if} - {#if uploadError} -
-
- - - - {uploadError} -
-
- {/if} - -
-
- {/if} -
- {/key} - {/if} - - - {#if activeTab === 'query'} -
-

- {$_('knowledgeBases.detail.query.title', { default: 'Query Knowledge Base' })} -

- -
{ - e.preventDefault(); - handleQuerySubmit(); - }} - class="space-y-4" - > -
- - -
- -
- -
-
- - - {#if queryLoading} -
- {$_('knowledgeBases.detail.query.loadingResults', { - default: 'Fetching results...' - })} -
- {/if} - - {#if queryError} -
-

- {$_('knowledgeBases.detail.query.errorTitle', { default: 'Error:' })} -

- 
{queryError}
-
- {/if} - - {#if queryResult} -
-

- {$_('knowledgeBases.detail.query.resultsTitle', { default: 'Query Results:' })} -

- - {#if queryResult.results && queryResult.results.length > 0} -
- {#each queryResult.results as result, i} -
- - {#if result.metadata?.file_url && result.metadata?.filename} -
- - - - - {result.metadata.filename} - -
- {/if} - -
-
- {$_('knowledgeBases.detail.query.resultItemTitle', { - default: 'Result' - })} - {i + 1} -
- - {$_('knowledgeBases.detail.query.similarityLabel', { - default: 'Similarity:' - })} - {result.similarity.toFixed(4)} - -
- -
-

- {$_('knowledgeBases.detail.query.contentLabel', { - default: 'Content:' - })} -

-
- {result.data} -
-
- -
-

- {$_('knowledgeBases.detail.query.metadataLabel', { - default: 'Metadata:' - })} -

- 
{JSON.stringify(
-																result.metadata,
-																null,
-																2
-															)}
-
-
- {/each} -
- {:else} -

- {$_('knowledgeBases.detail.query.noResults', { - default: 'No relevant results found for your query.' - })} -

- {/if} - - -
- {$_('knowledgeBases.detail.query.showRawJson', { - default: 'Show Raw JSON Response' - })} - 
{JSON.stringify(
-												queryResult,
-												null,
-												2
-											)}
-
-
- {/if} -
- {/if} -
-
-
- {:else} -
- {$_('knowledgeBases.detail.noData', { default: 'No knowledge base data available.' })} -
- {/if} +
+ + {#if loading} +
+
+ {$_('knowledgeBases.detail.loading', { default: 'Loading knowledge base details...' })} +
+
+ + + {:else if error && !kb} +
+
+ {#if serverOffline} +
+

+ {$_('knowledgeBases.detail.serverOffline', { default: 'Knowledge Base Server Offline' })} +

+

+ {$_('knowledgeBases.detail.tryAgainLater', { default: 'Please try again later or contact an administrator.' })} +

+
+ {:else} + {error} + {/if} +
+ +
+ + + {:else if kb} +
+ +
+

+ {kb.name} +

+

+ {kb.description || $_('knowledgeBases.detail.noDescription', { default: 'No description provided.' })} +

+
+ + +
+
+
+
+ {$_('knowledgeBases.detail.idLabel', { default: 'ID' })} +
+
+ {kb.id} +
+
+ + {#if kb.owner} +
+
+ {$_('knowledgeBases.detail.ownerLabel', { default: 'Owner' })} +
+
+ {kb.owner} +
+
+ {/if} + + {#if kb.created_at} +
+
+ {$_('knowledgeBases.detail.createdLabel', { default: 'Created' })} +
+
+ {new Date(kb.created_at * 1000).toLocaleString()} +
+
+ {/if} + + {#if kb.metadata?.access_control} +
+
+ {$_('knowledgeBases.detail.accessLabel', { default: 'Access Control' })} +
+
+ + {kb.metadata.access_control} + +
+
+ {/if} +
+
+ + +
+ +
+ +
+ + +
+ + {#if activeTab === 'files'} +
+ +
+
+ {#if loadingJobs} + + + + + + Loading status... + + {:else if jobsError} + ⚠ {jobsError} + {/if} +
+ +
+ + {#if kb.files && kb.files.length > 0} +
+ + + + + + + + + + + + {#each kb.files as file (file.id)} + {@const job = getJobForFile(file.filename)} + {@const statusColors = job ? getStatusColors(job.status) : getStatusColors('unknown')} + + + + + + + + {/each} + +
+ {$_('knowledgeBases.detail.fileNameColumn', { default: 'File Name' })} + + {$_('knowledgeBases.detail.fileSizeColumn', { default: 'Size' })} + + {$_('knowledgeBases.detail.fileTypeColumn', { default: 'Type' })} + + {$_('knowledgeBases.detail.fileStatusColumn', { default: 'Status' })} + + {$_('knowledgeBases.detail.fileActionsColumn', { default: 'Actions' })} +
+
+
+ {#if file.file_url} + + {file.filename} + + {:else} + {file.filename} + {/if} +
+
+ 		+
+ {file.size ? formatFileSize(file.size) : 'N/A'} +
+ 		+
+ {file.content_type || 'Unknown'} +
+ 		+ {#if job} + + {:else if loadingJobs} + Loading... + {:else} + + Unknown + + {/if} + + {#if kb.can_modify === true} + + {/if} +
+
+ {:else} +
+ {$_('knowledgeBases.detail.noFiles', { default: 'No files have been uploaded to this knowledge base.' })} +
+ {/if} +
+ {/if} + + + {#if activeTab === 'ingest'} + {#key selectedPlugin?.kind} +
+

+ {selectedPlugin?.kind === 'file-ingest' + ? $_('knowledgeBases.fileUpload.sectionTitle', { default: 'Upload and Ingest New File' }) + : $_('knowledgeBases.fileUpload.sectionTitleBase', { default: 'Configure and Run Ingestion' })} +

+ + {#if loadingPlugins} + + {:else if pluginsError} + + {:else if plugins.length === 0} + + {:else} + +
{ e.preventDefault(); handleSubmitIngestion(); }} class="space-y-6"> + + {#if pluginCanUseFileUpload(selectedPlugin)} +
+ +
+ +
+ {#if selectedPlugin?.kind !== 'file-ingest'} +

+ Optional: If a file is selected, ingestion runs through the file-upload path. Leave empty to run direct ingestion. +

+ {/if} + {#if selectedFile} +

+ {selectedFile.name} ({formatFileSize(selectedFile.size)}) +

+ {/if} +
+ {/if} + + +
+ +
+ +
+
+ + + {#if selectedPlugin?.parameters} + {#each Object.entries(selectedPlugin.parameters).filter(([name]) => name.startsWith('_')) as [paramName, paramDef] (paramName)} +
+
+ + + +
+

+ {paramName.replace(/^_/, '').replace(/_/g, ' ')} +

+

+ {paramDef.description || paramDef.default || ''} +

+
+
+
+ {/each} + {/if} + + + {#if selectedPlugin && selectedPlugin.parameters && visibleParamKeys.length > 0} +
+
+
+ {$_('knowledgeBases.fileUpload.parametersLabel', { default: 'Plugin Parameters' })} +
+ {#if hasAdvancedParams} + + {/if} +
+ + {#each displayedParamKeys as paramName (paramName)} + {@const paramDef = selectedPlugin.parameters[paramName]} + +
+ + + +
+ {#if paramDef.enum && Array.isArray(paramDef.enum)} + + + {:else if paramDef.type === 'boolean'} + +
+ updateParamValue(paramName, e)} + /> + +
+ {:else if paramDef.type === 'integer' || paramDef.type === 'number'} + + updateParamValue(paramName, e)} + /> + {#if paramDef.min !== undefined || paramDef.max !== undefined} +

+ {#if paramDef.min !== undefined && paramDef.max !== undefined} + Range: {paramDef.min} - {paramDef.max} + {:else if paramDef.min !== undefined} + Minimum: {paramDef.min} + {:else} + Maximum: {paramDef.max} + {/if} +

+ {/if} + {:else if paramDef.type === 'array'} + + +

Enter values separated by new lines.

+ {:else if paramDef.type === 'long-string'} + + + {:else} + + updateParamValue(paramName, e)} + /> + {/if} +
+ + + {#if paramDef.help_text} +

{paramDef.help_text}

+ {/if} +
+ {/each} +
+ {/if} + + +
+ + {#if uploadSuccess} +
+
+ + + + {$_('knowledgeBases.fileUpload.success', { default: 'File uploaded and ingestion started successfully!' })} +
+
+ {/if} + {#if uploadError} +
+
+ + + + {uploadError} +
+
+ {/if} + +
+
+ {/if} +
+ {/key} + {/if} + + + {#if activeTab === 'query'} +
+

+ {$_('knowledgeBases.detail.query.title', { default: 'Query Knowledge Base' })} +

+ +
{ e.preventDefault(); handleQuerySubmit(); }} class="space-y-4"> +
+ + +
+ +
+ +
+
+ + + {#if queryLoading} +
+ {$_('knowledgeBases.detail.query.loadingResults', { default: 'Fetching results...' })} +
+ {/if} + + {#if queryError} +
+

{$_('knowledgeBases.detail.query.errorTitle', { default: 'Error:' })}

+ 
{queryError}
+
+ {/if} + + {#if queryResult} +
+

+ {$_('knowledgeBases.detail.query.resultsTitle', { default: 'Query Results:' })} +

+ + {#if queryResult.results && queryResult.results.length > 0} +
+ {#each queryResult.results as result, i} +
+ + + {#if result.metadata?.file_url && result.metadata?.filename} +
+ + + + + + {result.metadata.filename} + +
+ {/if} + +
+
+ {$_('knowledgeBases.detail.query.resultItemTitle', { default: 'Result' })} {i + 1} +
+ + {$_('knowledgeBases.detail.query.similarityLabel', { default: 'Similarity:' })} {result.similarity.toFixed(4)} + +
+ +
+

{$_('knowledgeBases.detail.query.contentLabel', { default: 'Content:' })}

+
+ {result.data} +
+
+ +
+

{$_('knowledgeBases.detail.query.metadataLabel', { default: 'Metadata:' })}

+ 
{JSON.stringify(result.metadata, null, 2)}
+
+
+ {/each} +
+ {:else} +

{$_('knowledgeBases.detail.query.noResults', { default: 'No relevant results found for your query.' })}

+ {/if} + + +
+ {$_('knowledgeBases.detail.query.showRawJson', { default: 'Show Raw JSON Response' })} + 
{JSON.stringify(queryResult, null, 2)}
+
+
+ {/if} + +
+ {/if} +
+
+ +
+ {:else} +
+ {$_('knowledgeBases.detail.noData', { default: 'No knowledge base data available.' })} +
+ {/if}
{#if showJobModal && selectedJob} - {@const colors = getStatusColors(selectedJob.status)} - - + {@const colors = getStatusColors(selectedJob.status)} + + {/if} + + + { notification.isOpen = false; }} /> diff --git a/frontend/svelte-app/src/lib/components/aac/AacTerminal.svelte b/frontend/svelte-app/src/lib/components/aac/AacTerminal.svelte index eb37772b5..1fcfe847e 100644 --- a/frontend/svelte-app/src/lib/components/aac/AacTerminal.svelte +++ b/frontend/svelte-app/src/lib/components/aac/AacTerminal.svelte @@ -2,7 +2,7 @@ import { onMount, onDestroy, tick } from 'svelte'; import { sendMessageStream, getSession, sendMessage } from '$lib/services/aacService'; import { recordTabActivity } from '$lib/stores/aacStore.svelte'; - import { marked } from 'marked'; + import { renderMarkdownWithMath } from '$lib/utils/renderMarkdown.js'; // Abort any in-flight stream when the component unmounts so the fetch // and getReader() loop stop running in the background (#352, H3). @@ -280,17 +280,15 @@ darkMode = !darkMode; } - // Configure marked for terminal-style rendering - marked.setOptions({ breaks: true, gfm: true }); - /** - * Render markdown to HTML using marked. + * Render markdown to sanitized HTML (DOMPurify-backed; #417). * @param {string} text * @returns {string} */ function renderMarkdown(text) { if (!text) return ''; - return marked.parse(text); + // Sanitized render — agent/LLM output is untrusted (#417). + return renderMarkdownWithMath(text); } diff --git a/frontend/svelte-app/src/lib/components/analytics/ChatAnalytics.svelte b/frontend/svelte-app/src/lib/components/analytics/ChatAnalytics.svelte index b13c64bed..f28a831f0 100644 --- a/frontend/svelte-app/src/lib/components/analytics/ChatAnalytics.svelte +++ b/frontend/svelte-app/src/lib/components/analytics/ChatAnalytics.svelte @@ -15,6 +15,7 @@ import { onMount } from 'svelte'; import { _ } from '$lib/i18n'; + import { renderMarkdownWithMath } from '$lib/utils/renderMarkdown.js'; import Pagination from '$lib/components/common/Pagination.svelte'; import { getAssistantChats, @@ -654,7 +655,7 @@
- {message.content} + {@html renderMarkdownWithMath(message.content)}
diff --git a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte index 72a33e37b..3ce38ab37 100644 --- a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte @@ -1,291 +1,109 @@ + -
- -
-
-

- {#if formState === 'create'} - {$_('assistants.form.titleCreate', { default: 'Create New Assistant' })} - {:else} - {$_('assistants.form.titleViewEdit', { default: 'Assistant Details' })} - {#if initialAssistantData?.id} - (ID: {initialAssistantData.id}){/if} - {/if} -

- - - {#if formState === 'create'} -
- -
- {/if} -
- - - {#if importError} -
-
-
- - - -
-
-

- {$_('assistants.form.import.error', { default: 'Import Error' })} -

-

{importError}

-
-
-
- {/if} -
+
- {#if $assistantConfigStore.loading && !configInitialized} -

- {$_('assistants.loadingConfig', { default: 'Loading configuration...' })} -

+ + + {#if $assistantConfigStore.loading && !form.configInitialized} +

{$_('assistants.loadingConfig', { default: 'Loading configuration...' })}

{:else if $assistantConfigStore.error} -

- {$_('assistants.errorConfig', { default: 'Error loading configuration:' })} - {$assistantConfigStore.error} -

- {:else if !configInitialized} -

- {$_('assistants.initializingForm', { default: 'Initializing form...' })} -

+

{$_('assistants.errorConfig', { default: 'Error loading configuration:' })} {$assistantConfigStore.error}

+ {:else if !form.configInitialized} +

{$_('assistants.initializingForm', { default: 'Initializing form...' })}

{:else} -
+
-
- -
- - {#if formState === 'edit'} - - {:else} - - {#if showSanitizationPreview} -
-

- {$_('assistants.form.name.willBeSaved', { - default: 'Will be saved as:' - })} - {sanitizedNameInfo.sanitized} -

-
- {:else if !name.trim()} -

- {$_('assistants.form.name.hint', { - default: 'Special characters and spaces will be converted to underscores' - })} -

- {/if} - {/if} -
- - -
- -
- - - -
-

- {$_('assistants.form.description.help', { - default: 'Click Generate after filling in name and prompts.' - })} -

-
- - -
-
- - {#if formState === 'create'} - - {/if} -
- -
- - -
- -
- {$_('insert_placeholder') || 'Insert placeholder:'}: - {#each ragPlaceholders as placeholder} - - {/each} -
- - - - {#if prompt_template} -
-
- {$_('preview') || 'Preview with highlighted placeholders:'} -
-
- {@html highlightPlaceholders(prompt_template)} -
-
- {/if} - - {#if selectedPromptProcessor === 'template_validator_processor'} -

- {$_('assistants.form.promptTemplate.help', { - default: 'This processor requires a valid prompt template.' - })} -

- {/if} -
- - - {#if showRubricSelector} -
-

- {$_('assistants.form.rubric.label', { default: 'Select Rubric' })} -

- - {#if loadingRubrics} -

- {$_('assistants.form.rubric.loading', { default: 'Loading rubrics...' })} -

- {:else if rubricError} -

- {$_('assistants.form.rubric.error', { default: 'Error loading rubrics:' })} - {rubricError} -

- {:else if accessibleRubrics.length === 0} -

- {$_('assistants.form.rubric.noneFound', { default: 'No rubrics available.' })} - - {$_('assistants.form.rubric.createLink', { default: 'Create a rubric' })} → - -

- {:else} - -
- - -
- - - {#if selectedRubricId} - {@const selectedRubric = accessibleRubrics.find( - (r) => r.rubric_id === selectedRubricId - )} - {#if selectedRubric} -
-
-
-
- {$_('assistants.form.rubric.selected', { - default: 'Selected:' - })} - {selectedRubric.title} - {#if selectedRubric.is_showcase} - 🌟 - {/if} -
- {#if selectedRubric.description} -

{selectedRubric.description}

- {/if} -
- - {$_('assistants.form.rubric.view', { default: 'View' })} → - -
-
- {/if} - {/if} - - -
-
- - - - - - - - - - - - {#if filteredRubrics().length === 0} - - - - {:else} - {#each filteredRubrics() as rubric (rubric.rubric_id)} - - - - - - - - {/each} - {/if} - -
- {$_('assistants.form.rubric.table.select', { default: 'Select' })} - - {$_('assistants.form.rubric.table.title', { default: 'Title' })} - - {$_('assistants.form.rubric.table.description', { - default: 'Description' - })} - - {$_('assistants.form.rubric.table.type', { default: 'Type' })} - - {$_('assistants.form.rubric.table.actions', { default: 'Actions' })} -
- {$_('assistants.form.rubric.noMatches', { - default: 'No rubrics match your search.' - })} -
- - -
- {rubric.title} - {#if rubric.is_showcase} - 🌟 - {/if} -
- 		-

- {rubric.description || ''} -

- 		- {#if rubric.is_mine} - - {$_('assistants.form.rubric.mine', { default: 'Mine' })} - - {:else} - - {$_('assistants.form.rubric.public', { default: 'Public' })} - - {/if} - - - {$_('assistants.form.rubric.view', { default: 'View' })} → - -
-
-
- - {#if !selectedRubricId} -

- {$_('assistants.form.rubric.required', { default: 'Please select a rubric' })} -

- {/if} - - -
-
- {$_('assistants.form.rubric.format.label', { - default: 'Rubric Format for LLM' - })} -
-
- - -
-

- {$_('assistants.form.rubric.format.help', { - default: - 'Choose the format that works best with your selected LLM. You can test both to see which produces better results.' - })} -

-
- {/if} -
- {/if} -
- - -
- - {#if formState === 'create'} -
- -
- {/if} - - -
- {$_('assistants.form.configSection.title', { default: 'Configuration' })} - - - {#if isAdvancedMode || formState === 'edit'} -
- - -
- {/if} - - - {#if isAdvancedMode || formState === 'edit'} -
- - - - {#if currentConnectorMetadata?.description} -

- {currentConnectorMetadata.description} -

- {/if} -
- {/if} - - -
- - -
- - - {#if selectedConnector === 'openai' || selectedConnector === 'banana_img' || visionEnabled} -
- -
- {/if} - - - {#if selectedConnector === 'banana_img' || imageGenerationEnabled || currentConnectorMetadata?.capabilities?.image_generation} -
- -
- {/if} - - -
- - -
- - - {#if showRagOptions} -
-

- {$_('assistants.form.ragOptions.title', { default: 'RAG Options' })} -

- - - {#if selectedRagProcessor === 'rubric_rag'} -
-

- 📋 {$_('assistants.form.rubric.configLocation', { - default: 'See rubric options below the Prompt Template section' - })} -

-
- {/if} - - - {#if selectedRagProcessor === 'simple_rag' || selectedRagProcessor === 'context_aware_rag' || selectedRagProcessor === 'hierarchical_rag'} -
- - -

- {$_('assistants.form.ragTopK.help', { - default: 'Number of relevant documents to retrieve (1-10).' - })} -

-
- {/if} - - - {#if showKnowledgeBaseSelector} -
-

- {$_('assistants.form.knowledgeBases.label', { default: 'Knowledge Bases' })} -

- {#if loadingKnowledgeBases} -

- {$_('assistants.form.knowledgeBases.loading', { - default: 'Loading knowledge bases...' - })} -

- {:else if knowledgeBaseError} -

- {$_('assistants.form.knowledgeBases.error', { - default: 'Error loading knowledge bases:' - })} - {knowledgeBaseError} -

- {:else if accessibleKnowledgeBases.length === 0} -

- {$_('assistants.form.knowledgeBases.noneFound', { - default: 'No accessible knowledge bases found.' - })} -

- {:else} -
- - {#if ownedKnowledgeBases.length > 0} -
-
- {$_('assistants.form.knowledgeBases.myKB', { - default: 'My Knowledge Bases' - })} -
-
- {$_('assistants.form.knowledgeBases.myKB', { - default: 'My Knowledge Bases' - })} - {#each ownedKnowledgeBases as kb (kb.id)} - - {/each} -
-
- {/if} - - - {#if sharedKnowledgeBases.length > 0} -
-
- {$_('assistants.form.knowledgeBases.sharedKB', { - default: 'Shared Knowledge Bases' - })} -
-
- {$_('assistants.form.knowledgeBases.sharedKB', { - default: 'Shared Knowledge Bases' - })} - {#each sharedKnowledgeBases as kb (kb.id)} - - {/each} -
-
- {/if} -
- {/if} -
- {/if} - - - {#if showSingleFileSelector} -
-

- {$_('assistants.form.singleFile.label', { default: 'Select File' })} -

- - -
- -
- - {#if fileUploadLoading} - {$_('assistants.form.singleFile.uploading', { - default: 'Uploading...' - })} - {/if} -
- {#if fileUploadError} -

{fileUploadError}

- {/if} -
- - {#if loadingFiles} -

- {$_('assistants.form.singleFile.loading', { default: 'Loading files...' })} -

- {:else if fileError} -

- {$_('assistants.form.singleFile.error', { - default: 'Error loading files:' - })} - {fileError} -

- {:else if userFiles.length === 0} -

- {$_('assistants.form.singleFile.noneFound', { - default: 'No files found. Please upload a file.' - })} -

- {:else} -
- {#each userFiles as file (file.path)} - - {/each} -
- {#if !selectedFilePath && formState === 'edit'} -

- {$_('assistants.form.singleFile.required', { - default: 'Please select a file' - })} -

- {/if} - {/if} -
- {/if} -
- {/if} -
-
+
+ handleFieldChange(form)} /> + + handleFieldChange(form)} + /> + + handleFieldChange(form)} + onTemplateApplied={() => { + form.formDirty = true; + }} + /> + + {#if showRubricSelector} + + {/if}
- - {#if formError} -

- Error: {formError} -

- {/if} - {#if successMessage && formState !== 'edit'} -

- {successMessage} -

- {/if} - - -
-
- {#if formState === 'edit'} - - - {/if} - - - {#if showSanitizationPreview} -
-
- - - -
-

- {$_('assistants.form.name.willBeSaved', { default: 'Will be saved as:' })} -

- {sanitizedNameInfo.sanitized} -
-
-
- {/if} - - - -
+ +
+ handleFieldChange(form)} + /> +
+ + + {/if} - - - - -
diff --git a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte.test.js b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte.test.js index 508efc70e..59dd9873f 100644 --- a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte.test.js +++ b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte.test.js @@ -80,6 +80,11 @@ vi.mock('$lib/services/apiClient', () => ({ apiJson: vi.fn().mockResolvedValue([]) })); +vi.mock('$lib/services/libraryService', () => ({ + getLibraries: vi.fn().mockResolvedValue([]), + getItems: vi.fn().mockResolvedValue({ items: [], total: 0 }) +})); + vi.mock('$lib/stores/templateStore', async () => { const { writable } = await import('svelte/store'); return { diff --git a/frontend/svelte-app/src/lib/components/assistants/AssistantSharing.svelte b/frontend/svelte-app/src/lib/components/assistants/AssistantSharing.svelte index 7295b28a6..d45f7cf15 100644 --- a/frontend/svelte-app/src/lib/components/assistants/AssistantSharing.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/AssistantSharing.svelte @@ -329,7 +329,7 @@ - {#each sharedUsers as share} + {#each sharedUsers as share (share.user_id)} {share.user_name} {share.user_email} @@ -368,7 +368,7 @@

{:else}
- {#each organizationUsers as user} + {#each organizationUsers as user (user.id)} {@const isAlreadyShared = sharedUsers.some((s) => s.user_id === user.id)}
@@ -610,10 +701,111 @@ /> {/if} + + {#if graphEnabled && vectorDb === 'chromadb' && llmVendors.length > 0} +
+
+ {$_('knowledgeStores.createModal.extractionTitle', { + default: 'Graph RAG: concept extraction LLM' + })} +
+

+ {$_('knowledgeStores.createModal.extractionHint', { + default: + 'The model used at ingest time to extract entities and typed relationships from each chunk. Locked at creation, like embedding.' + })} +

+
+ + +
+
+ + {#if extractionModelChoices.length > 0} + + {:else} + + {/if} +
+
+ + +
+
+ {/if} {/if}
+ +
+ + {$_('knowledgeStores.createModal.otherIndexingLabel', { + default: 'Other indexing strategy (coming soon)' + })} + +
+ {$_('knowledgeStores.createModal.otherIndexingHint', { + default: + 'Additional indexing strategies (beyond vector indexing) will be selectable here in a future version.' + })} +
+
+
- -
  • - -
    • -
  • - -
    • -
  • - -
    • -
  • - -
    • - -
    - - - {#if currentView === 'dashboard'} - - {:else if currentView === 'users'} - -
    -

    - {localeLoaded ? $_('admin.users.title', { default: 'User Management' }) : 'User Management'} -

    - -
    - - {#if isLoadingUsers} -

    - {localeLoaded - ? $_('admin.users.loading', { default: 'Loading users...' }) - : 'Loading users...'} -

    - {:else if usersError} - - -
    - -
    - {:else} - - 0 - ? [ - { - key: 'organization', - label: localeLoaded - ? $_('admin.users.table.organization', { default: 'Organization' }) - : 'Organization', - options: organizationsForUsers.map((org) => ({ - value: String(org.id), - label: org.name - })) - } - ] - : []) - ]} - filterValues={{ - user_type: usersFilterType, - enabled: usersFilterEnabled, - organization: usersFilterOrg - }} - showSort={false} - on:searchChange={handleUsersSearchChange} - on:filterChange={handleUsersFilterChange} - on:clearFilters={handleUsersClearFilters} - /> - - -
    -
    - {#if usersSearch || usersFilterType || usersFilterEnabled || usersFilterOrg} - {localeLoaded - ? $_('admin.users.resultsCount.showing', { - default: 'Showing {filtered} of {total} users', - values: { filtered: usersTotalItems, total: allUsers.length } - }) - : `Showing ${usersTotalItems} of ${allUsers.length} users`} - {:else} - {localeLoaded - ? $_('admin.users.resultsCount.total', { - default: '{count} users', - values: { count: usersTotalItems } - }) - : `${usersTotalItems} users`} - {/if} -
    -
    - - {#if displayUsers.length === 0} - {#if allUsers.length === 0} - -
    -

    - {localeLoaded - ? $_('admin.users.noUsers', { default: 'No users found.' }) - : 'No users found.'} -

    -
    - {:else} - -
    -

    - {localeLoaded - ? $_('admin.users.resultsCount.noMatch', { default: 'No users match your filters' }) - : 'No users match your filters'} -

    - -
    - {/if} - {:else} - - {#if selectedUsers.length > 0} -
    -
    - - {selectedUsers.length === 1 - ? localeLoaded - ? $_('admin.users.bulkActions.selected', { - default: '{count} user selected', - values: { count: selectedUsers.length } - }) - : `${selectedUsers.length} user selected` - : localeLoaded - ? $_('admin.users.bulkActions.selectedPlural', { - default: '{count} users selected', - values: { count: selectedUsers.length } - }) - : `${selectedUsers.length} users selected`} - -
    - - - -
    -
    -
    - {/if} - - -
    - - - - - - - - - - - - {#each users as user (user.id)} - - - - - - - - {/each} - -
    - 0 && - selectedUsers.length === - users.filter((u) => !(currentUserData && currentUserData.email === u.email)) - .length} - onchange={handleSelectAll} - class="h-5 w-5 cursor-pointer rounded border-2 border-gray-400 bg-white text-blue-600 checked:border-blue-600 checked:bg-blue-600 focus:ring-2 focus:ring-blue-500" - style="accent-color: #2563eb;" - aria-label="Select all users" - /> - -
    - - -
    -     		- {localeLoaded - ? $_('admin.users.table.actions', { default: 'Actions' }) - : 'Actions'} -
    - - -
    - -
    -
    {user.email}
    -     		-
    - - {#if user.auth_provider !== 'lti_creator'} - - {/if} - - {#if !(currentUserData && currentUserData.email === user.email)} - {#if user.enabled} - - {:else} - - {/if} - {/if} - - {#if !user.enabled && !(currentUserData && currentUserData.email === user.email)} - - {/if} -
    -
    -
    - - - - {/if} - {/if} - {:else if currentView === 'organizations'} - -
    -

    - {localeLoaded - ? $_('admin.organizations.title', { default: 'Organization Management' }) - : 'Organization Management'} -

    -
    - - -
    -
    - - {#if isLoadingOrganizations} -

    - {localeLoaded - ? $_('admin.organizations.loading', { default: 'Loading organizations...' }) - : 'Loading organizations...'} -

    - {:else if organizationsError} - - -
    - -
    - {:else if organizations.length === 0} -

    - {localeLoaded - ? $_('admin.organizations.noOrganizations', { default: 'No organizations found.' }) - : 'No organizations found.'} -

    - {:else} - -
    - - - - - - - - - - - - {#each organizations as org (org.id)} - - - - - - - - {/each} - -
    - {localeLoaded ? $_('admin.organizations.table.name', { default: 'Name' }) : 'Name'} - - {localeLoaded ? $_('admin.organizations.table.slug', { default: 'Slug' }) : 'Slug'} - - {localeLoaded - ? $_('admin.organizations.table.actions', { default: 'Actions' }) - : 'Actions'} -
    - - -
    {org.slug}
    -     		-
    - - {#if !org.is_system} - - {/if} - - {#if !org.is_system} - - - {/if} -
    -
    -
    - {/if} - {/if} + +
    +
      +
    • + +
      • +
    • + +
      • +
    • + +
      • +
    • + +
      • +
    • + +
      • +
    +
    + + + {#if currentView === 'dashboard'} + + {:else if currentView === 'users'} + +
    +

    {localeLoaded ? $_('admin.users.title', { default: 'User Management' }) : 'User Management'}

    + +
    + + {#if isLoadingUsers} +

    {localeLoaded ? $_('admin.users.loading', { default: 'Loading users...' }) : 'Loading users...'}

    + {:else if usersError} + + +
    + +
    + {:else} + + 0 ? [{ + key: 'organization', + label: localeLoaded ? $_('admin.users.table.organization', { default: 'Organization' }) : 'Organization', + options: organizationsForUsers.map(org => ({ + value: String(org.id), + label: org.name + })) + }] : []) + ]} + filterValues={{ + user_type: usersFilterType, + enabled: usersFilterEnabled, + organization: usersFilterOrg + }} + showSort={false} + on:searchChange={handleUsersSearchChange} + on:filterChange={handleUsersFilterChange} + on:clearFilters={handleUsersClearFilters} + /> + + +
    +
    + {#if usersSearch || usersFilterType || usersFilterEnabled || usersFilterOrg} + {localeLoaded ? $_('admin.users.resultsCount.showing', { default: 'Showing {filtered} of {total} users', values: { filtered: usersTotalItems, total: allUsers.length } }) : `Showing ${usersTotalItems} of ${allUsers.length} users`} + {:else} + {localeLoaded ? $_('admin.users.resultsCount.total', { default: '{count} users', values: { count: usersTotalItems } }) : `${usersTotalItems} users`} + {/if} +
    +
    + + {#if displayUsers.length === 0} + {#if allUsers.length === 0} + +
    +

    {localeLoaded ? $_('admin.users.noUsers', { default: 'No users found.' }) : 'No users found.'}

    +
    + {:else} + +
    +

    {localeLoaded ? $_('admin.users.resultsCount.noMatch', { default: 'No users match your filters' }) : 'No users match your filters'}

    + +
    + {/if} + {:else} + + {#if selectedUsers.length > 0} +
    +
    + + {selectedUsers.length === 1 + ? (localeLoaded ? $_('admin.users.bulkActions.selected', { default: '{count} user selected', values: { count: selectedUsers.length } }) : `${selectedUsers.length} user selected`) + : (localeLoaded ? $_('admin.users.bulkActions.selectedPlural', { default: '{count} users selected', values: { count: selectedUsers.length } }) : `${selectedUsers.length} users selected`) + } + +
    + + + +
    +
    +
    + {/if} + + +
    + + + + + + + + + + + + {#each users as user (user.id)} + + + + + + + + {/each} + +
    + 0 && selectedUsers.length === users.filter(u => !(currentUserData && currentUserData.email === u.email)).length} + onchange={handleSelectAll} + class="w-5 h-5 text-blue-600 bg-white border-2 border-gray-400 rounded cursor-pointer focus:ring-2 focus:ring-blue-500 checked:bg-blue-600 checked:border-blue-600" + style="accent-color: #2563eb;" + aria-label="Select all users" + /> + +
    + + +
    +     		+ {localeLoaded ? $_('admin.users.table.actions', { default: 'Actions' }) : 'Actions'} +
    + + +
    + +
    +
    {user.email}
    +     		+
    + + {#if user.auth_provider !== 'lti_creator'} + + {/if} + + {#if !(currentUserData && currentUserData.email === user.email)} + {#if user.enabled} + + {:else} + + {/if} + {/if} + + {#if !user.enabled && !(currentUserData && currentUserData.email === user.email)} + + {/if} +
    +
    +
    + + + + {/if} + {/if} + {:else if currentView === 'organizations'} + +
    +

    {localeLoaded ? $_('admin.organizations.title', { default: 'Organization Management' }) : 'Organization Management'}

    +
    + + +
    +
    + + {#if isLoadingOrganizations} +

    {localeLoaded ? $_('admin.organizations.loading', { default: 'Loading organizations...' }) : 'Loading organizations...'}

    + {:else if organizationsError} + + +
    + +
    + {:else if organizations.length === 0} +

    {localeLoaded ? $_('admin.organizations.noOrganizations', { default: 'No organizations found.' }) : 'No organizations found.'}

    + {:else} + +
    + + + + + + + + + + + + {#each organizations as org (org.id)} + + + + + + + + {/each} + +
    + {localeLoaded ? $_('admin.organizations.table.name', { default: 'Name' }) : 'Name'} + + {localeLoaded ? $_('admin.organizations.table.slug', { default: 'Slug' }) : 'Slug'} + + {localeLoaded ? $_('admin.organizations.table.actions', { default: 'Actions' }) : 'Actions'} +
    + + +
    {org.slug}
    +     		+
    + + {#if !org.is_system} + + {/if} + + {#if !org.is_system} + + + {/if} +
    +
    +
    + {/if} + {/if}
    { - passwordChangeData.new_password = pwd; - }} + isOpen={isChangePasswordModalOpen} + userName={selectedUserName} + userEmail={passwordChangeData.email} + newPassword={passwordChangeData.new_password} + isChanging={isChangingPassword} + error={changePasswordError} + success={changePasswordSuccess} + {localeLoaded} + onSubmit={handleChangePassword} + onClose={closeChangePasswordModal} + onPasswordChange={(pwd) => { passwordChangeData.new_password = pwd; }} /> { - newUser = user; - }} + isOpen={isCreateUserModalOpen} + isSuperAdmin={true} + {newUser} + organizations={organizationsForUsers} + isLoadingOrganizations={isLoadingOrganizationsForUsers} + organizationsError={organizationsForUsersError} + isCreating={isCreatingUser} + error={createUserError} + success={createUserSuccess} + {localeLoaded} + onSubmit={handleCreateUser} + onClose={closeCreateUserModal} + onUserChange={(user) => { newUser = user; }} /> fetchOrganizations()} - onClose={closeCreateOrgModal} + bind:isOpen={isCreateOrgModalOpen} + {localeLoaded} + {getAuthToken} + onSuccess={() => fetchOrganizations()} + onClose={closeCreateOrgModal} /> {#if isViewConfigModalOpen && selectedOrg} -
    -
    -
    -
    -

    - Configuration: {selectedOrg.name} ({selectedOrg.slug}) -

    - -
    - - {#if isLoadingOrgConfig} -
    -

    Loading configuration...

    -
    - {:else if configError} - - {:else if selectedOrgConfig} -
    -

    Configuration JSON:

    - 
    {JSON.stringify(selectedOrgConfig, null, 2)}
    - - {#if selectedOrgConfig.setups} -
    -

    Setups Summary:

    -
    - {#each Object.entries(selectedOrgConfig.setups) as [setupName, setup]} -
    -
    {setup.name || setupName}
    -

    - {setup.is_default ? '(Default)' : ''} -

    - {#if setup.providers} -
    - Providers: - {Object.keys(setup.providers).join(', ') || 'None'} -
    - {/if} -
    - {/each} -
    -
    - {/if} - - {#if selectedOrgConfig.features} -
    -

    Features:

    -
    - {#each Object.entries(selectedOrgConfig.features) as [feature, enabled]} -
    - - {feature.replace('_', ' ')} -
    - {/each} -
    -
    - {/if} - - {#if selectedOrgConfig.limits && selectedOrgConfig.limits.usage} -
    -

    Usage Limits:

    -
    -
    - Tokens/Month: - {selectedOrgConfig.limits.usage.tokens_per_month?.toLocaleString() || - 'Unlimited'} -
    -
    - Max Assistants: - {selectedOrgConfig.limits.usage.max_assistants || 'Unlimited'} -
    -
    - Storage: - {selectedOrgConfig.limits.usage.storage_gb || 'Unlimited'} GB -
    -
    -
    - {/if} -
    - {/if} - -
    - -
    -
    -
    -
    +
    +
    +
    +
    +

    + Configuration: {selectedOrg.name} ({selectedOrg.slug}) +

    + +
    + + {#if isLoadingOrgConfig} +
    +

    Loading configuration...

    +
    + {:else if configError} + + {:else if selectedOrgConfig} +
    +

    Configuration JSON:

    + 
    {JSON.stringify(selectedOrgConfig, null, 2)}
    + + {#if selectedOrgConfig.setups} +
    +

    Setups Summary:

    +
    + {#each Object.entries(selectedOrgConfig.setups) as [setupName, setup]} +
    +
    {setup.name || setupName}
    +

    + {setup.is_default ? '(Default)' : ''} +

    + {#if setup.providers} +
    + Providers: {Object.keys(setup.providers).join(', ') || 'None'} +
    + {/if} +
    + {/each} +
    +
    + {/if} + + {#if selectedOrgConfig.features} +
    +

    Features:

    +
    + {#each Object.entries(selectedOrgConfig.features) as [feature, enabled]} +
    + + {feature.replace('_', ' ')} +
    + {/each} +
    +
    + {/if} + + {#if selectedOrgConfig.limits && selectedOrgConfig.limits.usage} +
    +

    Usage Limits:

    +
    +
    + Tokens/Month: {selectedOrgConfig.limits.usage.tokens_per_month?.toLocaleString() || 'Unlimited'} +
    +
    + Max Assistants: {selectedOrgConfig.limits.usage.max_assistants || 'Unlimited'} +
    +
    + Storage: {selectedOrgConfig.limits.usage.storage_gb || 'Unlimited'} GB +
    +
    +
    + {/if} +
    + {/if} + +
    + +
    +
    +
    +
    {/if} { - showDisableConfirm = false; - }} + isOpen={showDisableConfirm} + action="disable" + isBulk={actionType === 'bulk'} + targetUser={actionType === 'single' ? targetUser : null} + selectedCount={selectedUsers.length} + {localeLoaded} + onConfirm={confirmDisable} + onClose={() => { showDisableConfirm = false; }} /> { - showEnableConfirm = false; - }} + isOpen={showEnableConfirm} + action="enable" + isBulk={actionType === 'bulk'} + targetUser={actionType === 'single' ? targetUser : null} + selectedCount={selectedUsers.length} + {localeLoaded} + onConfirm={confirmEnable} + onClose={() => { showEnableConfirm = false; }} /> {#if showDeleteConfirm} -
    -
    -
    -
    - - - -

    Delete User

    -
    -
    - {#if isCheckingDependencies} -

    Checking user dependencies...

    - {:else if userDependencies} -

    - Are you sure you want to permanently delete {deleteTargetUser?.name} - ({deleteTargetUser?.email})? -

    - - {#if userDependencies.has_dependencies} -
    -

    - ⚠️ Cannot delete user - has dependencies: -

    - {#if userDependencies.assistant_count > 0} -
    -

    - {userDependencies.assistant_count} Assistant(s): -

    -
      - {#each userDependencies.assistants.slice(0, 5) as assistant} -
    • {assistant.name}
      • - {/each} - {#if userDependencies.assistants.length > 5} -
    • ... and {userDependencies.assistants.length - 5} more
      • - {/if} -
    -
    - {/if} - {#if userDependencies.kb_count > 0} -
    -

    - {userDependencies.kb_count} Knowledge Base(s): -

    -
      - {#each userDependencies.kbs.slice(0, 5) as kb} -
    • {kb.name}
      • - {/each} - {#if userDependencies.kbs.length > 5} -
    • ... and {userDependencies.kbs.length - 5} more
      • - {/if} -
    -
    - {/if} -

    - Please delete or reassign these resources before deleting the user. -

    -
    - {:else} -
    -

    - ✓ User has no dependencies and can be safely deleted. -

    -
    - {/if} - -

    - Note: This action cannot be undone. -

    - {/if} -
    -
    - - -
    -
    -
    -
    +
    +
    +
    +
    + + + +

    Delete User

    +
    +
    + {#if isCheckingDependencies} +

    Checking user dependencies...

    + {:else if userDependencies} +

    + Are you sure you want to permanently delete {deleteTargetUser?.name} ({deleteTargetUser?.email})? +

    + + {#if userDependencies.has_dependencies} +
    +

    + ⚠️ Cannot delete user - has dependencies: +

    + {#if userDependencies.assistant_count > 0} +
    +

    + {userDependencies.assistant_count} Assistant(s): +

    +
      + {#each userDependencies.assistants.slice(0, 5) as assistant} +
    • {assistant.name}
      • + {/each} + {#if userDependencies.assistants.length > 5} +
    • ... and {userDependencies.assistants.length - 5} more
      • + {/if} +
    +
    + {/if} + {#if userDependencies.kb_count > 0} +
    +

    + {userDependencies.kb_count} Knowledge Base(s): +

    +
      + {#each userDependencies.kbs.slice(0, 5) as kb} +
    • {kb.name}
      • + {/each} + {#if userDependencies.kbs.length > 5} +
    • ... and {userDependencies.kbs.length - 5} more
      • + {/if} +
    +
    + {/if} +

    + Please delete or reassign these resources before deleting the user. +

    +
    + {:else} +
    +

    + ✓ User has no dependencies and can be safely deleted. +

    +
    + {/if} + +

    + Note: This action cannot be undone. +

    + {/if} +
    +
    + + +
    +
    +
    +
    {/if} {#if isMigrationModalOpen && migrationSourceOrg} -
    -
    -
    -
    -

    - {localeLoaded - ? $_('admin.organizations.migration.title', { default: 'Migrate Organization' }) - : 'Migrate Organization'} -

    - -
    - -

    - {localeLoaded - ? $_('admin.organizations.migration.description', { - default: 'Migrate all resources from' - }) - : 'Migrate all resources from'} {migrationSourceOrg.name} to another organization. -

    - - {#if migrationSuccess && migrationReport} - - - {:else} - -
    - -
    - - - -
    - - - {#if migrationValidationError} - - {/if} - - {#if migrationValidationResult && migrationValidationResult.can_migrate} - -
    -

    - {localeLoaded - ? $_('admin.organizations.migration.resources', { - default: 'Resources to Migrate' - }) - : 'Resources to Migrate'} -

    -
      -
    • • {migrationValidationResult.resources?.users || 0} users
      • -
    • • {migrationValidationResult.resources?.assistants || 0} assistants
      • -
    • • {migrationValidationResult.resources?.templates || 0} templates
      • -
    • • {migrationValidationResult.resources?.kbs || 0} knowledge bases
      • -
    • • {migrationValidationResult.resources?.usage_logs || 0} usage logs
      • -
    - - {#if migrationValidationResult.conflicts.assistants.length > 0 || migrationValidationResult.conflicts.templates.length > 0} -
    -
    - {localeLoaded - ? $_('admin.organizations.migration.conflicts', { - default: 'Conflicts Detected' - }) - : 'Conflicts Detected'} -
    - {#if migrationValidationResult.conflicts.assistants.length > 0} -

    - {migrationValidationResult.conflicts.assistants.length} assistant(s) will be - renamed: -

    -
      - {#each migrationValidationResult.conflicts.assistants.slice(0, 3) as conflict} -
    • {conflict.name} (by {conflict.owner})
      • - {/each} - {#if migrationValidationResult.conflicts.assistants.length > 3} -
    • - ... and {migrationValidationResult.conflicts.assistants.length - 3} more -
      • - {/if} -
    - {/if} - {#if migrationValidationResult.conflicts.templates.length > 0} -

    - {migrationValidationResult.conflicts.templates.length} template(s) will be renamed: -

    -
      - {#each migrationValidationResult.conflicts.templates.slice(0, 3) as conflict} -
    • {conflict.name} (by {conflict.owner_email})
      • - {/each} - {#if migrationValidationResult.conflicts.templates.length > 3} -
    • - ... and {migrationValidationResult.conflicts.templates.length - 3} more -
      • - {/if} -
    - {/if} -
    - {/if} -
    - - -
    -

    - {localeLoaded - ? $_('admin.organizations.migration.options', { default: 'Migration Options' }) - : 'Migration Options'} -

    - - -
    - - -
    - - -
    - -
    - -

    - {localeLoaded - ? $_('admin.organizations.migration.preserveAdminRolesHint', { - default: - 'If unchecked, admins from the source organization will be migrated as regular members. If checked, they will keep their admin privileges in the target organization.' - }) - : 'If unchecked, admins from the source organization will be migrated as regular members. If checked, they will keep their admin privileges in the target organization.'} -

    -
    -
    - - -
    - -
    - -

    - {localeLoaded - ? $_('admin.organizations.migration.deleteSourceHint', { - default: - 'If checked, the source organization will be permanently deleted after successful migration. This action cannot be undone.' - }) - : 'If checked, the source organization will be permanently deleted after successful migration. This action cannot be undone.'} -

    -
    -
    -
    - - - {#if migrationError} - - {/if} - - -
    - - -
    - {/if} -
    - {/if} -
    -
    -
    -{:else if currentView === 'lti-settings'} - -
    -

    LTI Tool Configuration

    -

    - Use these values when creating an External Tool (LTI 1.1) in your LMS (Moodle, Canvas, etc.). -

    -
    - - {#if isLoadingLtiGlobal} -
    -
    -
    - Loading LTI configuration... -
    -
    - {:else} - - {#if ltiGlobalError} - - {/if} - - - {#if ltiGlobalSuccess} - - {/if} - - -
    -
    -

    LMS Setup Information

    -

    - Copy these three values into your LMS External Tool configuration. -

    - - -
    - -
    - - -
    -

    - This is the URL the LMS sends the LTI launch POST to. -

    -
    - - -
    - -
    - - -
    -
    - - -
    -

    Shared Secret

    -
    - {#if ltiHasSecret} - - - Configured ({ltiGlobalConfig.oauth_consumer_secret_masked}) - - {:else} - - - Not set — configure below - - {/if} - - Source: {ltiGlobalConfig.source === 'database' ? 'Database' : '.env file'} - {#if ltiGlobalConfig.updated_at} - · Updated {new Date(ltiGlobalConfig.updated_at * 1000).toLocaleDateString()} - {/if} - -
    -

    - The secret is never displayed in full. Enter it in your LMS exactly as configured below. -

    -
    -
    -
    - - -
    -
    -

    Edit LTI Credentials

    -

    - Change the consumer key or secret. Saving stores them in the database and overrides any .env values. -

    - -
    - -
    - - -

    - The LMS sends this as oauth_consumer_key. Must match on both sides. -

    -
    - - -
    - -
    - - {#if ltiGlobalForm.consumer_secret} - - {/if} -
    -

    - Used to sign and verify LTI launch requests (HMAC-SHA1). Must match in the LMS. -

    -
    - - -
    - -
    -
    -
    -
    - - -
    -

    How the Unified LTI Endpoint Works

    -
      -
    1. In your LMS, create a new External Tool (LTI 1.1).
      2. -
    2. - Paste the Launch URL, Consumer Key, and - Secret from above. -
      3. -
    3. - When an instructor launches the tool for the first time, they choose which - assistants to assign. -
      4. -
    4. - Subsequent student launches go directly to Open WebUI with those assistants - available. -
      5. -
    -
    - {/if} -{:else if currentView === 'cost-management'} - -
    -
    -

    - {localeLoaded - ? $_('admin.costManagement.title', { default: 'Cost Management' }) - : 'Cost Management'} -

    -

    - {localeLoaded - ? $_('admin.costManagement.subtitle', { - default: 'Token usage and estimated cost per assistant across the platform.' - }) - : 'Token usage and estimated cost per assistant across the platform.'} -

    -
    - -
    - - {#if isLoadingCostData} -
    -
    -
    - {localeLoaded - ? $_('admin.costManagement.loading', { default: 'Loading usage data...' }) - : 'Loading usage data...'} -
    -
    - {:else if costDataError} - - {:else} - -
    -
    -

    Total Estimated Cost

    -

    ${costTotals.total_cost.toFixed(4)}

    -
    -
    -

    Total Tokens

    -

    {costTotals.total_tokens.toLocaleString()}

    -

    - Prompt: {costTotals.prompt_tokens.toLocaleString()} · Completion: {costTotals.completion_tokens.toLocaleString()} -

    -
    -
    -

    Assistants

    -

    {costData.length}

    -

    - {costData.filter((a) => a.quota_exceeded).length} quota exceeded -

    -
    -
    - - -
    - -
    - - {#if filteredCostData.length === 0} -
    - {costData.length === 0 - ? localeLoaded - ? $_('admin.costManagement.noData', { default: 'No assistants found.' }) - : 'No assistants found.' - : 'No assistants match your search.'} -
    - {:else} -
    -
    - - - - - - - - - - - - - - - {#each filteredCostData as assistant (assistant.id)} - openQuotaEditModal(assistant)} - title="Click to edit quota" - > - - - - - - - - - - {/each} - -
    - {localeLoaded - ? $_('admin.costManagement.table.assistant', { default: 'Assistant' }) - : 'Assistant'} - - {localeLoaded - ? $_('admin.costManagement.table.organization', { default: 'Organization' }) - : 'Organization'} - - {localeLoaded - ? $_('admin.costManagement.table.model', { default: 'Model' }) - : 'Model'} - - {localeLoaded - ? $_('admin.costManagement.table.promptTokens', { default: 'Prompt Tokens' }) - : 'Prompt Tokens'} - - {localeLoaded - ? $_('admin.costManagement.table.completionTokens', { - default: 'Completion Tokens' - }) - : 'Completion Tokens'} - - {localeLoaded - ? $_('admin.costManagement.table.cost', { default: 'Estimated Cost' }) - : 'Estimated Cost'} - - {localeLoaded - ? $_('admin.costManagement.table.quota', { default: 'Quota' }) - : 'Quota'} - - {localeLoaded - ? $_('admin.costManagement.table.status', { default: 'Status' }) - : 'Status'} -
    -
    {assistant.name}
    -
    {assistant.owner}
    -     		{assistant.organization_name || '—'} - {#if assistant.model_name} - {assistant.model_name} - {:else} - — - {/if} - {assistant.prompt_tokens.toLocaleString()}{assistant.completion_tokens.toLocaleString()} - ${assistant.cost_usd.toFixed(4)} - - {#if !assistant.quota_enabled} - {localeLoaded - ? $_('admin.costManagement.quota.noQuota', { default: 'No quota' }) - : 'No quota'} - {:else if assistant.cost_limit_usd != null} - {@const tablePct = (assistant.cost_usd / assistant.cost_limit_usd) * 100} - {@const hasAlert = - assistant.alert_thresholds && - assistant.alert_thresholds.some((t) => t <= tablePct)} - ${assistant.cost_limit_usd.toFixed(2)} -
    -
    -
    -
    - {((assistant.cost_usd / assistant.cost_limit_usd) * 100).toFixed(1)}% used -
    - {/if} -     		- {#if assistant.quota_exceeded} - - {localeLoaded - ? $_('admin.costManagement.quota.exceeded', { default: 'Exceeded' }) - : 'Exceeded'} - - {:else if assistant.quota_enabled} - - {localeLoaded - ? $_('admin.costManagement.quota.active', { default: 'Active' }) - : 'Active'} - - {:else} - - {localeLoaded - ? $_('admin.costManagement.quota.disabled', { default: 'Disabled' }) - : 'No quota'} - - {/if} -
    -
    -
    - {/if} - {/if} -{:else if currentView === 'user-detail'} - -
    - -
    - { - if (userDetailId) fetchUserDetail(userDetailId); - }} - /> -{/if} +
    +
    +
    +
    +

    + {localeLoaded ? $_('admin.organizations.migration.title', { default: 'Migrate Organization' }) : 'Migrate Organization'} +

    + +
    + +

    + {localeLoaded + ? $_('admin.organizations.migration.description', { default: 'Migrate all resources from' }) + : 'Migrate all resources from'} {migrationSourceOrg.name} to another organization. +

    + + {#if migrationSuccess && migrationReport} + + + {:else} + +
    + +
    + + + +
    + + + {#if migrationValidationError} + + {/if} + + {#if migrationValidationResult && migrationValidationResult.can_migrate} + +
    +

    + {localeLoaded ? $_('admin.organizations.migration.resources', { default: 'Resources to Migrate' }) : 'Resources to Migrate'} +

    +
      +
    • • {migrationValidationResult.resources?.users || 0} users
      • +
    • • {migrationValidationResult.resources?.assistants || 0} assistants
      • +
    • • {migrationValidationResult.resources?.templates || 0} templates
      • +
    • • {migrationValidationResult.resources?.kbs || 0} knowledge bases
      • +
    • • {migrationValidationResult.resources?.usage_logs || 0} usage logs
      • +
    + + {#if migrationValidationResult.conflicts.assistants.length > 0 || migrationValidationResult.conflicts.templates.length > 0} +
    +
    + {localeLoaded ? $_('admin.organizations.migration.conflicts', { default: 'Conflicts Detected' }) : 'Conflicts Detected'} +
    + {#if migrationValidationResult.conflicts.assistants.length > 0} +

    + {migrationValidationResult.conflicts.assistants.length} assistant(s) will be renamed: +

    +
      + {#each migrationValidationResult.conflicts.assistants.slice(0, 3) as conflict} +
    • {conflict.name} (by {conflict.owner})
      • + {/each} + {#if migrationValidationResult.conflicts.assistants.length > 3} +
    • ... and {migrationValidationResult.conflicts.assistants.length - 3} more
      • + {/if} +
    + {/if} + {#if migrationValidationResult.conflicts.templates.length > 0} +

    + {migrationValidationResult.conflicts.templates.length} template(s) will be renamed: +

    +
      + {#each migrationValidationResult.conflicts.templates.slice(0, 3) as conflict} +
    • {conflict.name} (by {conflict.owner_email})
      • + {/each} + {#if migrationValidationResult.conflicts.templates.length > 3} +
    • ... and {migrationValidationResult.conflicts.templates.length - 3} more
      • + {/if} +
    + {/if} +
    + {/if} +
    + + +
    +

    + {localeLoaded ? $_('admin.organizations.migration.options', { default: 'Migration Options' }) : 'Migration Options'} +

    + + +
    + + +
    + + +
    + +
    + +

    + {localeLoaded + ? $_('admin.organizations.migration.preserveAdminRolesHint', { default: 'If unchecked, admins from the source organization will be migrated as regular members. If checked, they will keep their admin privileges in the target organization.' }) + : 'If unchecked, admins from the source organization will be migrated as regular members. If checked, they will keep their admin privileges in the target organization.'} +

    +
    +
    + + +
    + +
    + +

    + {localeLoaded + ? $_('admin.organizations.migration.deleteSourceHint', { default: 'If checked, the source organization will be permanently deleted after successful migration. This action cannot be undone.' }) + : 'If checked, the source organization will be permanently deleted after successful migration. This action cannot be undone.'} +

    +
    +
    +
    + + + {#if migrationError} + + {/if} + + +
    + + +
    + {/if} +
    + {/if} +
    +
    +
    + + {:else if currentView === 'lti-settings'} + +
    +

    LTI Tool Configuration

    +

    + Use these values when creating an External Tool (LTI 1.1) in your LMS (Moodle, Canvas, etc.). +

    +
    + + {#if isLoadingLtiGlobal} +
    +
    +
    + Loading LTI configuration... +
    +
    + {:else} + + {#if ltiGlobalError} + + {/if} + + + {#if ltiGlobalSuccess} + + {/if} + + +
    +
    +

    LMS Setup Information

    +

    Copy these three values into your LMS External Tool configuration.

    + + +
    + +
    + + +
    +

    This is the URL the LMS sends the LTI launch POST to.

    +
    + + +
    + +
    + + +
    +
    + + +
    +

    Shared Secret

    +
    + {#if ltiHasSecret} + + + Configured ({ltiGlobalConfig.oauth_consumer_secret_masked}) + + {:else} + + + Not set — configure below + + {/if} + + Source: {ltiGlobalConfig.source === 'database' ? 'Database' : '.env file'} + {#if ltiGlobalConfig.updated_at} + · Updated {new Date(ltiGlobalConfig.updated_at * 1000).toLocaleDateString()} + {/if} + +
    +

    + The secret is never displayed in full. Enter it in your LMS exactly as configured below. +

    +
    +
    +
    + + +
    +
    +

    Edit LTI Credentials

    +

    + Change the consumer key or secret. Saving stores them in the database and overrides any .env values. +

    + +
    + +
    + + +

    The LMS sends this as oauth_consumer_key. Must match on both sides.

    +
    + + +
    + +
    + + {#if ltiGlobalForm.consumer_secret} + + {/if} +
    +

    Used to sign and verify LTI launch requests (HMAC-SHA1). Must match in the LMS.

    +
    + + +
    + +
    +
    +
    +
    + + +
    +

    How the Unified LTI Endpoint Works

    +
      +
    1. In your LMS, create a new External Tool (LTI 1.1).
      2. +
    2. Paste the Launch URL, Consumer Key, and Secret from above.
      3. +
    3. When an instructor launches the tool for the first time, they choose which assistants to assign.
      4. +
    4. Subsequent student launches go directly to Open WebUI with those assistants available.
      5. +
    +
    + {/if} + + {:else if currentView === 'cost-management'} + +
    +
    +

    {localeLoaded ? $_('admin.costManagement.title', { default: 'Cost Management' }) : 'Cost Management'}

    +

    {localeLoaded ? $_('admin.costManagement.subtitle', { default: 'Token usage and estimated cost per assistant across the platform.' }) : 'Token usage and estimated cost per assistant across the platform.'}

    +
    + +
    + + {#if isLoadingCostData} +
    +
    +
    + {localeLoaded ? $_('admin.costManagement.loading', { default: 'Loading usage data...' }) : 'Loading usage data...'} +
    +
    + {:else if costDataError} + + {:else} + +
    +
    +

    Total Estimated Cost

    +

    ${costTotals.total_cost.toFixed(4)}

    +
    +
    +

    Total Tokens

    +

    {costTotals.total_tokens.toLocaleString()}

    +

    Prompt: {costTotals.prompt_tokens.toLocaleString()} · Completion: {costTotals.completion_tokens.toLocaleString()}

    +
    +
    +

    Assistants

    +

    {costData.length}

    +

    {costData.filter(a => a.quota_exceeded).length} quota exceeded

    +
    +
    + + +
    + +
    + + {#if filteredCostData.length === 0} +
    + {costData.length === 0 + ? (localeLoaded ? $_('admin.costManagement.noData', { default: 'No assistants found.' }) : 'No assistants found.') + : 'No assistants match your search.'} +
    + {:else} +
    +
    + + + + + + + + + + + + + + + {#each filteredCostData as assistant (assistant.id)} + openQuotaEditModal(assistant)} + title="Click to edit quota" + > + + + + + + + + + + {/each} + +
    + {localeLoaded ? $_('admin.costManagement.table.assistant', { default: 'Assistant' }) : 'Assistant'} + + {localeLoaded ? $_('admin.costManagement.table.organization', { default: 'Organization' }) : 'Organization'} + + {localeLoaded ? $_('admin.costManagement.table.model', { default: 'Model' }) : 'Model'} + + {localeLoaded ? $_('admin.costManagement.table.promptTokens', { default: 'Prompt Tokens' }) : 'Prompt Tokens'} + + {localeLoaded ? $_('admin.costManagement.table.completionTokens', { default: 'Completion Tokens' }) : 'Completion Tokens'} + + {localeLoaded ? $_('admin.costManagement.table.cost', { default: 'Estimated Cost' }) : 'Estimated Cost'} + + {localeLoaded ? $_('admin.costManagement.table.quota', { default: 'Quota' }) : 'Quota'} + + {localeLoaded ? $_('admin.costManagement.table.status', { default: 'Status' }) : 'Status'} +
    +
    {assistant.name}
    +
    {assistant.owner}
    +     		{assistant.organization_name || '—'} + {#if assistant.model_name} + {assistant.model_name} + {:else} + — + {/if} + {assistant.prompt_tokens.toLocaleString()}{assistant.completion_tokens.toLocaleString()} + ${assistant.cost_usd.toFixed(4)} + + {#if !assistant.quota_enabled} + {localeLoaded ? $_('admin.costManagement.quota.noQuota', { default: 'No quota' }) : 'No quota'} + {:else if assistant.cost_limit_usd != null} + {@const tablePct = (assistant.cost_usd / assistant.cost_limit_usd) * 100} + {@const hasAlert = assistant.alert_thresholds && assistant.alert_thresholds.some(t => t <= tablePct)} + ${assistant.cost_limit_usd.toFixed(2)} +
    +
    +
    +
    {((assistant.cost_usd / assistant.cost_limit_usd) * 100).toFixed(1)}% used
    + {/if} +     		+ {#if assistant.quota_exceeded} + + {localeLoaded ? $_('admin.costManagement.quota.exceeded', { default: 'Exceeded' }) : 'Exceeded'} + + {:else if assistant.quota_enabled} + + {localeLoaded ? $_('admin.costManagement.quota.active', { default: 'Active' }) : 'Active'} + + {:else} + + {localeLoaded ? $_('admin.costManagement.quota.disabled', { default: 'Disabled' }) : 'No quota'} + + {/if} +
    +
    +
    + {/if} + {/if} + + {:else if currentView === 'user-detail'} + +
    + +
    + { if (userDetailId) fetchUserDetail(userDetailId); }} + /> + {/if} {#if isMembersModalOpen && membersModalOrg} -
    -
    -
    -

    - Members: {membersModalOrg.name} - ({membersModalOrg.slug}) -

    - -
    - - {#if roleUpdateSuccess} -
    - {roleUpdateSuccess} -
    - {/if} - - {#if membersError} -
    - {membersError} -
    - {/if} - - {#if isLoadingMembers} -
    Loading members...
    - {:else if orgMembers.length === 0} -
    -

    No members in this organization.

    -

    - Users can join via signup or be assigned by a system admin. -

    -
    - {:else} -
    - - - - - - - - - - - - {#each orgMembers as member (member.id)} - - - - - - - - {/each} - -
    UserTypeRoleStatusActions
    -
    {member.name || '-'}
    -
    {member.email}
    -     		- {#if member.auth_provider === 'lti_creator'} - LTI Creator - {:else} - Creator - {/if} - - {#if member.role === 'admin'} - Admin - {:else} - Member - {/if} - - {#if member.enabled} - Active - {:else} - Disabled - {/if} - - {#if member.role === 'admin'} - - {:else} - - {/if} -
    -
    -

    - Any user, including LTI Creator users, can be promoted to organization admin. -

    - {/if} - -
    - -
    -
    -
    +
    +
    +
    +

    + Members: {membersModalOrg.name} + ({membersModalOrg.slug}) +

    + +
    + + {#if roleUpdateSuccess} +
    + {roleUpdateSuccess} +
    + {/if} + + {#if membersError} +
    + {membersError} +
    + {/if} + + {#if isLoadingMembers} +
    Loading members...
    + {:else if orgMembers.length === 0} +
    +

    No members in this organization.

    +

    Users can join via signup or be assigned by a system admin.

    +
    + {:else} +
    + + + + + + + + + + + + {#each orgMembers as member (member.id)} + + + + + + + + {/each} + +
    UserTypeRoleStatusActions
    +
    {member.name || '-'}
    +
    {member.email}
    +     		+ {#if member.auth_provider === 'lti_creator'} + LTI Creator + {:else} + Creator + {/if} + + {#if member.role === 'admin'} + Admin + {:else} + Member + {/if} + + {#if member.enabled} + Active + {:else} + Disabled + {/if} + + {#if member.role === 'admin'} + + {:else} + + {/if} +
    +
    +

    + Any user, including LTI Creator users, can be promoted to organization admin. +

    + {/if} + +
    + +
    +
    +
    {/if} + + + { notification.isOpen = false; }} /> {#if quotaEditAssistant} - + {/if} + /* Add specific styles if needed, though Tailwind should cover most */ + \ No newline at end of file diff --git a/frontend/svelte-app/src/routes/agent/+page.svelte b/frontend/svelte-app/src/routes/agent/+page.svelte index 48f02f30a..645e8241e 100644 --- a/frontend/svelte-app/src/routes/agent/+page.svelte +++ b/frontend/svelte-app/src/routes/agent/+page.svelte @@ -1,206 +1,195 @@ -
    -
    -

    - 🤖 {$_('home.dashboard.agent.title', { default: 'LAMB Agent' })} -

    -
    - - - - - {$_('agent.history.title', { default: 'History' })} - - {#if sessionId && !loading} - - {/if} -
    -
    +
    +
    +

    + 🤖 {$_('home.dashboard.agent.title', { default: 'LAMB Agent' })} +

    +
    + + + + + {$_('agent.history.title', { default: 'History' })} + + {#if sessionId && !loading} + + {/if} +
    +
    - {#if loading} -
    -
    ⏳
    -

    {$_('home.dashboard.agent.starting', { default: 'Starting...' })}

    -
    - {:else if error} -
    -

    {error}

    -
    - {:else if sessionId} -
    - {#key sessionId} -
    - -
    - {/key} -
    - {/if} + {#if loading} +
    +
    ⏳
    +

    {$_('home.dashboard.agent.starting', { default: 'Starting...' })}

    +
    + {:else if error} +
    +

    {error}

    +
    + {:else if sessionId} +
    + {#key sessionId} +
    + +
    + {/key} +
    + {/if}
    + + + { showNewConversationModal = false; }} +/> diff --git a/frontend/svelte-app/src/routes/agent/history/+page.svelte b/frontend/svelte-app/src/routes/agent/history/+page.svelte index 07821a847..cebd20caf 100644 --- a/frontend/svelte-app/src/routes/agent/history/+page.svelte +++ b/frontend/svelte-app/src/routes/agent/history/+page.svelte @@ -1,204 +1,211 @@ -
    -
    -

    - {$_('agent.history.title', { default: 'Agent Sessions' })} -

    - - ← {$_('agent.history.backToAgent', { default: 'Back to Agent' })} - -
    +
    +
    +

    + {$_('agent.history.title', { default: 'Agent Sessions' })} +

    + + ← {$_('agent.history.backToAgent', { default: 'Back to Agent' })} + +
    - -
    -
    - {#each [['all', 'All'], ['today', 'Today'], ['active', 'Active'], ['errors', 'Errors']] as [val, label]} - - {/each} -
    - - -
    + +
    +
    + {#each [['all', 'All'], ['today', 'Today'], ['active', 'Active'], ['errors', 'Errors']] as [val, label]} + + {/each} +
    + + +
    - - {#if loading} -
    Loading sessions...
    - {:else} - {@const filtered = filteredSessions()} -
    {filtered.length} sessions
    -
    - - - - - - - - - - - - - - - {#each filtered as s} - - - - - - - - - - - {/each} - {#if filtered.length === 0} - - {/if} - -
    TitleSkillTurnsToolsErrorsUpdatedActions
    {skillIcon(s.skill_id)} - - {s.title || s.id.slice(0, 8)} - - {s.skill_id || '-'}{s.turn_count || 0}{s.tool_calls || 0}{s.tool_errors || 0}{formatDate(s.updated_at)} -
    - {#if s.status === 'active'} - - {/if} - Review - -
    -
    No sessions match your filters.
    -
    - {/if} + + {#if loading} +
    Loading sessions...
    + {:else} + {@const filtered = filteredSessions()} +
    {filtered.length} sessions
    +
    + + + + + + + + + + + + + + + {#each filtered as s} + + + + + + + + + + + {/each} + {#if filtered.length === 0} + + {/if} + +
    TitleSkillTurnsToolsErrorsUpdatedActions
    {skillIcon(s.skill_id)} + + {s.title || s.id.slice(0, 8)} + + {s.skill_id || '-'}{s.turn_count || 0}{s.tool_calls || 0}{s.tool_errors || 0}{formatDate(s.updated_at)} +
    + {#if s.status === 'active'} + + {/if} + Review + +
    +
    No sessions match your filters.
    +
    + {/if}
    + + + { showDeleteModal = false; sessionToDelete = null; }} +/> diff --git a/frontend/svelte-app/src/routes/agent/history/[id]/+page.svelte b/frontend/svelte-app/src/routes/agent/history/[id]/+page.svelte index d2becab49..d8d78241b 100644 --- a/frontend/svelte-app/src/routes/agent/history/[id]/+page.svelte +++ b/frontend/svelte-app/src/routes/agent/history/[id]/+page.svelte @@ -1,184 +1,163 @@ -
    -
    - - ← {$_('agent.history.backToHistory', { default: 'Back to sessions' })} - -
    +
    +
    + + ← {$_('agent.history.backToHistory', { default: 'Back to sessions' })} + +
    - {#if loading} -
    Loading...
    - {:else if error} -
    -

    {error}

    -
    - {:else if session} - -
    -
    -
    -

    - {session.title || session.id.slice(0, 12)} -

    -
    - Status: {session.status} - Created: {formatDate(session.created_at)} - Updated: {formatDate(session.updated_at)} - {#if session.assistant_id} - Assistant: #{session.assistant_id} - {/if} -
    -
    - {#if session.status === 'active'} - - {/if} -
    -
    + {#if loading} +
    Loading...
    + {:else if error} +
    +

    {error}

    +
    + {:else if session} + +
    +
    +
    +

    {session.title || session.id.slice(0, 12)}

    +
    + Status: {session.status} + Created: {formatDate(session.created_at)} + Updated: {formatDate(session.updated_at)} + {#if session.assistant_id} + Assistant: #{session.assistant_id} + {/if} +
    +
    + {#if session.status === 'active'} + + {/if} +
    +
    - - {#if session.tool_audit?.length} -
    - - {#if showAudit} -
    - - - - - - - - - - - {#each session.tool_audit as e} - - - - - - - {/each} - -
    TimeStatusDurationCommand
    {(e.ts || '').slice(11, 19)}{e.success ? 'ok' : 'FAIL'}{Math.round(e.elapsed_ms || 0)}ms{e.command || '?'}
    -
    - {/if} -
    - {/if} + + {#if session.tool_audit?.length} +
    + + {#if showAudit} +
    + + + + + + + + + + + {#each session.tool_audit as e} + + + + + + + {/each} + +
    TimeStatusDurationCommand
    {(e.ts || '').slice(11, 19)}{e.success ? 'ok' : 'FAIL'}{Math.round(e.elapsed_ms || 0)}ms{e.command || '?'}
    +
    + {/if} +
    + {/if} - -
    -
    -

    Conversation

    -
    -
    - {#each getVisibleMessages() as msg} - {#if msg.role === 'user'} -
    -
    -
    - $ - {msg.content} -
    -
    -
    - {:else if msg.role === 'assistant'} -
    - {@html renderMarkdown(msg.content)} -
    - {/if} - {/each} - {#if getVisibleMessages().length === 0} -

    No messages in this session.

    - {/if} -
    -
    - {/if} + +
    +
    +

    Conversation

    +
    +
    + {#each getVisibleMessages() as msg} + {#if msg.role === 'user'} +
    +
    +
    + $ + {msg.content} +
    +
    +
    + {:else if msg.role === 'assistant'} +
    + {@html renderMarkdown(msg.content)} +
    + {/if} + {/each} + {#if getVisibleMessages().length === 0} +

    No messages in this session.

    + {/if} +
    +
    + {/if}
    diff --git a/frontend/svelte-app/src/routes/assistants/+page.svelte b/frontend/svelte-app/src/routes/assistants/+page.svelte index 880be3292..f1e4cf8a3 100644 --- a/frontend/svelte-app/src/routes/assistants/+page.svelte +++ b/frontend/svelte-app/src/routes/assistants/+page.svelte @@ -1,727 +1,710 @@ -

    - {currentLocale ? $_('assistants.title') : 'Learning Assistants'} -

    +

    {currentLocale ? $_('assistants.title') : 'Learning Assistants'}

    - +
    {#if currentView === 'list'} -
    -
    - -
    -
    +
    +
    + +
    +
    {:else if currentView === 'create'} - - + + {:else if currentView === 'detail'} - -
    - - {#if isOwner} - - {/if} - {#if canShare && (isOwner || canManageSharing)} - - {/if} - - {#if canViewAnalytics} - - {/if} - {#if isOwner} - - {/if} - - - {#each aacTabs.filter((t) => t.assistantId === selectedAssistantData?.id) as tab} - - {/each} - - - {#if isOwner} - - - - {/if} -
    - - -
    - {#if loadingDetail} -

    {currentLocale ? $_('assistants.loadingDetail') : 'Loading assistant details...'}

    - {:else if detailError} - - {:else if selectedAssistantData} - - {#if detailSubView === 'properties'} - - - {#key selectedAssistantData?.id} - - {#if !isOwner && accessLevel === 'read_only'} -
    - - - - - {currentLocale - ? $_('assistants.detail.readOnlyBanner', { - default: 'Read-only view — This assistant belongs to {owner}', - values: { owner: selectedAssistantData?.owner || '' } - }) - : `Read-only view — This assistant belongs to ${selectedAssistantData?.owner || ''}`} - -
    - {/if} - -
    -

    - {currentLocale - ? $_('assistants.detail.propertiesTitle', { default: 'Assistant Properties' }) - : 'Assistant Properties'} -

    -
    - {#if isOwner} - - - {/if} - - - - - {#if isOwner} - - - {/if} - - {#if canDelete} - - - {/if} -
    -
    - - -
    - -
    - -
    -
    - {$_('assistants.form.name.label')} * -
    -
    - {selectedAssistantData.name?.replace(/^\d+_/, '') || - (currentLocale ? $_('common.notAvailable', { default: 'N/A' }) : 'N/A')} -
    -
    - -
    -
    - {$_('assistants.form.description.label', { default: 'Description' })} -
    -
    - {selectedAssistantData.description || - (currentLocale - ? $_('common.notSpecified', { default: 'Not specified' }) - : 'Not specified')} -
    -
    - -
    -
    - {$_('assistants.form.systemPrompt.label', { default: 'System Prompt' })} -
    -
    - {selectedAssistantData.system_prompt || - (currentLocale - ? $_('common.notSpecified', { default: 'Not specified' }) - : 'Not specified')} -
    -
    - -
    -
    - {$_('assistants.form.promptTemplate.label', { default: 'Prompt Template' })} -
    -
    - {selectedAssistantData.prompt_template || - (currentLocale - ? $_('common.notSpecified', { default: 'Not specified' }) - : 'Not specified')} -
    -
    - - - {#if selectedAssistantData} - {@const apiCallback = getAssistantMetadataObject(selectedAssistantData)} - {#if apiCallback.rag_processor === 'rubric_rag'} -
    -
    - {$_('assistants.form.rubric.selectedLabel', { default: 'Selected Rubric' })} -
    - {#if loadingRubric} -
    - {$_('assistants.form.rubric.loading', { default: 'Loading rubric...' })} -
    - {:else if rubricError} -
    - {$_('assistants.form.rubric.error', { - default: 'Error loading rubric:' - })} - {rubricError} -
    - {:else if rubricMarkdown} -
    - -
    -
    -
    - {$_('assistants.form.rubric.id', { default: 'Rubric ID:' })} - {apiCallback.rubric_id || 'N/A'} -
    -
    - {$_('assistants.form.rubric.format.label', { - default: 'Format:' - })} - - {apiCallback.rubric_format || 'markdown'} - -
    -
    -
    - -
    - - {$_('assistants.form.rubric.viewRubric', { - default: 'View Rubric Content' - })} - -
    - 
    {rubricMarkdown}
    -
    -
    -
    - {:else} -
    - {apiCallback.rubric_id || - (currentLocale - ? $_('common.notSpecified', { default: 'Not specified' }) - : 'Not specified')} -
    - {/if} -
    - {/if} - {/if} -
    - - - {#if selectedAssistantData} -
    - - {#if selectedAssistantData.published} -
    -

    - {currentLocale - ? $_('assistants.detail.ltiTitle', { default: 'LTI Publish Details' }) - : 'LTI Publish Details'} -

    - {console.log( - 'Checking selectedAssistantData for LTI box:', - selectedAssistantData - )} -
    -
    - {currentLocale - ? $_('assistants.detail.ltiAssistantName', { - default: 'Assistant Name' - }) - : 'Assistant Name'}: - {selectedAssistantData.name} -
    -
    - {currentLocale - ? $_('assistants.detail.ltiModelId', { default: 'Model ID' }) - : 'Model ID'}: - {selectedAssistantData.id} -
    -
    - {currentLocale - ? $_('assistants.detail.ltiToolUrl', { default: 'Tool URL' }) - : 'Tool URL'}: - {lambServerUrl}/lamb/v1/lti_users/lti -
    -
    - {currentLocale - ? $_('assistants.detail.ltiConsumerKey', { default: 'Consumer Key' }) - : 'Consumer Key'}: - {selectedAssistantData.oauth_consumer_name || - (currentLocale - ? $_('common.notAvailable', { default: 'N/A' }) - : 'N/A')} -
    -
    - {currentLocale - ? $_('assistants.detail.ltiSecret', { default: 'Secret' }) - : 'Secret'}: - ASK ADMIN FOR THE SECRET -
    -
    -
    - {/if} - - -
    -

    - {currentLocale - ? $_('assistants.detail.configuration', { default: 'Configuration' }) - : 'Configuration'} -

    - - - {#if selectedAssistantData} - {@const apiCallback = getAssistantMetadataObject(selectedAssistantData)} - -
    - -
    -
    - {currentLocale - ? $_('assistants.form.promptProcessor.label', { - default: 'Prompt Processor' - }) - : 'Prompt Processor'} -
    -
    - {apiCallback.prompt_processor || - (currentLocale - ? $_('common.notSpecified', { default: 'Not specified' }) - : 'Not specified')} -
    -
    - - -
    -
    - {currentLocale - ? $_('assistants.form.connector.label', { default: 'Connector' }) - : 'Connector'} -
    -
    - {apiCallback.connector || - (currentLocale - ? $_('common.notSpecified', { default: 'Not specified' }) - : 'Not specified')} -
    -
    - - -
    -
    - {currentLocale - ? $_('assistants.form.llm.label', { default: 'Language Model (LLM)' }) - : 'Language Model (LLM)'} -
    -
    - {apiCallback.llm || - (currentLocale - ? $_('common.notSpecified', { default: 'Not specified' }) - : 'Not specified')} -
    -
    - - -
    -
    - {currentLocale - ? $_('assistants.form.vision.label', { default: 'Vision Capability' }) - : 'Vision Capability'} -
    -
    - {#if apiCallback.capabilities?.vision} - - - - - {currentLocale - ? $_('assistants.form.vision.enabled', { default: 'Enabled' }) - : 'Enabled'} - - {:else} - - {currentLocale - ? $_('assistants.form.vision.disabled', { default: 'Disabled' }) - : 'Disabled'} - - {/if} -
    -
    - - -
    -
    Image Generation
    -
    - {#if apiCallback.capabilities?.image_generation} - - - - - Enabled - - {:else} - Disabled - {/if} -
    -
    - - -
    -
    - {currentLocale - ? $_('assistants.form.ragProcessor.label', { default: 'RAG Processor' }) - : 'RAG Processor'} -
    -
    - {apiCallback.rag_processor - ?.replace(/_/g, ' ') - .replace(/\b\w/g, (l) => l.toUpperCase()) || - (currentLocale - ? $_('common.notSpecified', { default: 'Not specified' }) - : 'Not specified')} -
    -
    - - - {#if apiCallback.rag_processor && apiCallback.rag_processor !== 'no_rag'} -
    -

    - {$_('assistants.form.ragOptions.title', { default: 'RAG Options' })} -

    - - - {#if apiCallback.rag_processor !== 'single_file_rag' && apiCallback.rag_processor !== 'rubric_rag'} -
    -
    - {$_('assistants.form.ragTopK.label', { default: 'RAG Top K' })} -
    -
    - {selectedAssistantData.RAG_Top_k ?? 3} -
    -
    - {/if} - - - {#if apiCallback.rag_processor === 'simple_rag'} -
    -
    - {$_('assistants.form.knowledgeBases.label', { - default: 'Knowledge Bases' - })} -
    - {#if loadingKnowledgeBases} -
    - {$_('assistants.form.knowledgeBases.loading', { - default: 'Loading...' - })} -
    - {:else if knowledgeBaseError} -
    - {$_('assistants.form.knowledgeBases.error', { - default: 'Error loading KBs' - })} -
    - {:else if selectedAssistantData.RAG_collections} -
    - {#each selectedAssistantData.RAG_collections.split(',') as kbId} - {@const kb = accessibleKnowledgeBases.find( - (k) => k.id === kbId - )} - - {kb ? kb.name : `${kbId} (Not Found)`} - - {/each} -
    - {:else} -
    - {$_('common.none', { default: 'None' })} -
    - {/if} -
    - {/if} - - - {#if apiCallback.rag_processor === 'single_file_rag'} -
    -
    - {$_('assistants.form.singleFile.selectedLabel', { - default: 'Selected File' - })} -
    -
    - {apiCallback.file_path || - (currentLocale - ? $_('common.notSpecified', { default: 'Not specified' }) - : 'Not specified')} -
    -
    - {/if} -
    - {/if} -
    - {/if} -
    -
    - {/if} -
    - - - {#if selectedAssistantData} -
    -
    -
    - {currentLocale - ? $_('assistants.detail.created', { default: 'Created:' }) - : 'Created:'} - {formatDateForTable(selectedAssistantData.created_at)} -
    -
    - {currentLocale - ? $_('assistants.detail.updated', { default: 'Updated:' }) - : 'Updated:'} - {formatDateForTable(selectedAssistantData.updated_at)} -
    -
    -
    - {/if} - {/key} - {:else if detailSubView === 'edit'} - -
    -
    -

    - {currentLocale - ? $_('assistants.detail.editTitle', { default: 'Edit Assistant' }) - : 'Edit Assistant'} -

    -
    - - -
    - { - // Refresh the assistant data after successful update with forceRefresh - fetchAssistantDetail(selectedAssistantData.id, true); - detailSubView = 'properties'; // Switch back to properties view - }} - on:cancel={() => { - // Switch back to properties view when user cancels - detailSubView = 'properties'; - }} - id="assistant-edit-form" - /> -
    -
    - {:else if detailSubView === 'share'} - -
    -
    -

    - {currentLocale - ? $_('assistants.sharing.title', { default: 'Shared Users' }) - : 'Shared Users'} -

    -

    - {currentLocale - ? $_('assistants.sharing.description', { - default: 'Manage who has access to this assistant' - }) - : 'Manage who has access to this assistant'} -

    -
    - - {#if loadingShares} -
    -
    Loading shared users...
    -
    - {:else if sharedUsers.length === 0} -
    -

    - {currentLocale - ? $_('assistants.sharing.noShares', { - default: 'This assistant is not shared with anyone yet.' - }) - : 'This assistant is not shared with anyone yet.'} -

    -
    - {:else} -
    -
    - {#each sharedUsers as share (share.user_id)} -
    -
    -
    {share.user_name}
    -
    {share.user_email}
    -
    -
    - Shared by {share.shared_by_name} -
    -
    - {/each} -
    -
    - {/if} - -
    - -
    -
    - - - {#if showSharingModal} - - {/if} - {:else if detailSubView === 'chat'} - {#if configError} - - {:else if lambServerUrl && userToken} - -
    -

    - {currentLocale ? $_('assistants.detail.chatTitle', { default: 'Chat' }) : 'Chat'} -

    -
    - - {:else} - -

    - {currentLocale ? $_('assistants.chatLoadingConfig') : 'Loading chat configuration...'} -

    - {/if} - {:else if detailSubView === 'analytics'} - -
    - -
    - {:else if detailSubView === 'tests'} - - launchAacSkill(skill)} - /> - {:else if detailSubView === 'aac' && activeAacSessionId} - - {#key `${activeAacSessionId}-${aacSkillStartup}`} -
    - -
    - {/key} - {/if} - {/if} -
    - - + +
    + + {#if isOwner} + + {/if} + {#if canShare && (isOwner || canManageSharing)} + + {/if} + + {#if canViewAnalytics} + + {/if} + {#if isOwner} + + {/if} + + + {#each aacTabs.filter(t => t.assistantId === selectedAssistantData?.id) as tab} + + {/each} + + + {#if isOwner} + + + + {/if} +
    + + +
    + {#if loadingDetail} +

    {currentLocale ? $_('assistants.loadingDetail') : 'Loading assistant details...'}

    + {:else if detailError} + + {:else if selectedAssistantData} + + {#if detailSubView === 'properties'} + + + {#key selectedAssistantData?.id} + + + {#if !isOwner && accessLevel === 'read_only'} +
    + + + + + {currentLocale ? $_('assistants.detail.readOnlyBanner', { default: 'Read-only view — This assistant belongs to {owner}', values: { owner: selectedAssistantData?.owner || '' } }) : `Read-only view — This assistant belongs to ${selectedAssistantData?.owner || ''}`} + +
    + {/if} + +
    +

    + {currentLocale ? $_('assistants.detail.propertiesTitle', { default: 'Assistant Properties' }) : 'Assistant Properties'} +

    +
    + {#if isOwner} + + + {/if} + + + + + + {#if isOwner} + + + {/if} + + {#if canDelete} + + + {/if} +
    +
    + + +
    + +
    + +
    +
    {$_('assistants.form.name.label')} *
    +
    + {selectedAssistantData.name?.replace(/^\d+_/, '') || (currentLocale ? $_('common.notAvailable', { default: 'N/A' }) : 'N/A')} +
    +
    + +
    +
    {$_('assistants.form.description.label', { default: 'Description' })}
    +
    + {selectedAssistantData.description || (currentLocale ? $_('common.notSpecified', { default: 'Not specified' }) : 'Not specified')} +
    +
    + +
    +
    {$_('assistants.form.systemPrompt.label', { default: 'System Prompt' })}
    +
    + {selectedAssistantData.system_prompt || (currentLocale ? $_('common.notSpecified', { default: 'Not specified' }) : 'Not specified')} +
    +
    + +
    +
    {$_('assistants.form.promptTemplate.label', { default: 'Prompt Template' })}
    +
    + {selectedAssistantData.prompt_template || (currentLocale ? $_('common.notSpecified', { default: 'Not specified' }) : 'Not specified')} +
    +
    + + + {#if selectedAssistantData} + {@const apiCallback = getAssistantMetadataObject(selectedAssistantData)} + {#if apiCallback.rag_processor === 'rubric_rag'} +
    +
    {$_('assistants.form.rubric.selectedLabel', { default: 'Selected Rubric' })}
    + {#if loadingRubric} +
    + {$_('assistants.form.rubric.loading', { default: 'Loading rubric...' })} +
    + {:else if rubricError} +
    + {$_('assistants.form.rubric.error', { default: 'Error loading rubric:' })} {rubricError} +
    + {:else if rubricMarkdown} +
    + +
    +
    +
    + {$_('assistants.form.rubric.id', { default: 'Rubric ID:' })} + {apiCallback.rubric_id || 'N/A'} +
    +
    + {$_('assistants.form.rubric.format.label', { default: 'Format:' })} + + {apiCallback.rubric_format || 'markdown'} + +
    +
    +
    + +
    + + {$_('assistants.form.rubric.viewRubric', { default: 'View Rubric Content' })} + +
    + 
    {rubricMarkdown}
    +
    +
    +
    + {:else} +
    + {apiCallback.rubric_id || (currentLocale ? $_('common.notSpecified', { default: 'Not specified' }) : 'Not specified')} +
    + {/if} +
    + {/if} + {/if} +
    + + + {#if selectedAssistantData} +
    + + {#if selectedAssistantData.published} +
    +

    + {currentLocale ? $_('assistants.detail.ltiTitle', { default: 'LTI Publish Details' }) : 'LTI Publish Details'} +

    + {console.log('Checking selectedAssistantData for LTI box:', selectedAssistantData)} +
    +
    + {currentLocale ? $_('assistants.detail.ltiAssistantName', { default: 'Assistant Name' }) : 'Assistant Name'}: + {selectedAssistantData.name} +
    +
    + {currentLocale ? $_('assistants.detail.ltiModelId', { default: 'Model ID' }) : 'Model ID'}: + {selectedAssistantData.id} +
    +
    + {currentLocale ? $_('assistants.detail.ltiToolUrl', { default: 'Tool URL' }) : 'Tool URL'}: + {lambServerUrl}/lamb/v1/lti_users/lti +
    +
    + {currentLocale ? $_('assistants.detail.ltiConsumerKey', { default: 'Consumer Key' }) : 'Consumer Key'}: + {selectedAssistantData.oauth_consumer_name || (currentLocale ? $_('common.notAvailable', { default: 'N/A' }) : 'N/A')} +
    +
    + {currentLocale ? $_('assistants.detail.ltiSecret', { default: 'Secret' }) : 'Secret'}: + ASK ADMIN FOR THE SECRET +
    +
    +
    + {/if} + + +
    +

    + {currentLocale ? $_('assistants.detail.configuration', { default: 'Configuration' }) : 'Configuration'} +

    + + + {#if selectedAssistantData} + {@const apiCallback = getAssistantMetadataObject(selectedAssistantData)} + +
    + +
    +
    + {currentLocale ? $_('assistants.form.promptProcessor.label', { default: 'Prompt Processor' }) : 'Prompt Processor'} +
    +
    + {apiCallback.prompt_processor || (currentLocale ? $_('common.notSpecified', { default: 'Not specified' }) : 'Not specified')} +
    +
    + + +
    +
    + {currentLocale ? $_('assistants.form.connector.label', { default: 'Connector' }) : 'Connector'} +
    +
    + {apiCallback.connector || (currentLocale ? $_('common.notSpecified', { default: 'Not specified' }) : 'Not specified')} +
    +
    + + +
    +
    + {currentLocale ? $_('assistants.form.llm.label', { default: 'Language Model (LLM)' }) : 'Language Model (LLM)'} +
    +
    + {apiCallback.llm || (currentLocale ? $_('common.notSpecified', { default: 'Not specified' }) : 'Not specified')} +
    +
    + + +
    +
    + {currentLocale ? $_('assistants.form.vision.label', { default: 'Vision Capability' }) : 'Vision Capability'} +
    +
    + {#if apiCallback.capabilities?.vision} + + + + + {currentLocale ? $_('assistants.form.vision.enabled', { default: 'Enabled' }) : 'Enabled'} + + {:else} + + {currentLocale ? $_('assistants.form.vision.disabled', { default: 'Disabled' }) : 'Disabled'} + + {/if} +
    +
    + + +
    +
    + Image Generation +
    +
    + {#if apiCallback.capabilities?.image_generation} + + + + + Enabled + + {:else} + + Disabled + + {/if} +
    +
    + + +
    +
    + {currentLocale ? $_('assistants.form.ragProcessor.label', { default: 'RAG Processor' }) : 'RAG Processor'} +
    +
    + {apiCallback.rag_processor?.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase()) || (currentLocale ? $_('common.notSpecified', { default: 'Not specified' }) : 'Not specified')} +
    +
    + + + {#if apiCallback.rag_processor && apiCallback.rag_processor !== 'no_rag'} +
    +

    {$_('assistants.form.ragOptions.title', { default: 'RAG Options' })}

    + + + {#if apiCallback.rag_processor !== 'single_file_rag' && apiCallback.rag_processor !== 'rubric_rag'} +
    +
    {$_('assistants.form.ragTopK.label', { default: 'RAG Top K' })}
    +
    + {selectedAssistantData.RAG_Top_k ?? 3} +
    +
    + {/if} + + + {#if apiCallback.rag_processor === 'simple_rag'} +
    +
    {$_('assistants.form.knowledgeBases.label', { default: 'Knowledge Bases' })}
    + {#if loadingKnowledgeBases} +
    + {$_('assistants.form.knowledgeBases.loading', { default: 'Loading...' })} +
    + {:else if knowledgeBaseError} +
    + {$_('assistants.form.knowledgeBases.error', { default: 'Error loading KBs' })} +
    + {:else if selectedAssistantData.RAG_collections} +
    + {#each selectedAssistantData.RAG_collections.split(',') as kbId} + {@const kb = accessibleKnowledgeBases.find(k => k.id === kbId)} + + {kb ? kb.name : `${kbId} (Not Found)`} + + {/each} +
    + {:else} +
    + {$_('common.none', { default: 'None' })} +
    + {/if} +
    + {/if} + + + {#if apiCallback.rag_processor === 'single_file_rag'} +
    +
    {$_('assistants.form.singleFile.selectedLabel', { default: 'Selected File' })}
    +
    + {apiCallback.file_path || (currentLocale ? $_('common.notSpecified', { default: 'Not specified' }) : 'Not specified')} +
    +
    + {/if} +
    + {/if} +
    + {/if} +
    +
    + {/if} +
    + + + {#if selectedAssistantData} +
    +
    +
    + {currentLocale ? $_('assistants.detail.created', { default: 'Created:' }) : 'Created:'} + {formatDateForTable(selectedAssistantData.created_at)} +
    +
    + {currentLocale ? $_('assistants.detail.updated', { default: 'Updated:' }) : 'Updated:'} + {formatDateForTable(selectedAssistantData.updated_at)} +
    +
    +
    + {/if} + {/key} + + {:else if detailSubView === 'edit'} + +
    +
    +

    + {currentLocale ? $_('assistants.detail.editTitle', { default: 'Edit Assistant' }) : 'Edit Assistant'} +

    +
    + + +
    + { + // Refresh the assistant data after successful update with forceRefresh + fetchAssistantDetail(selectedAssistantData.id, true); + detailSubView = 'properties'; // Switch back to properties view + }} + onCancel={() => { + // Switch back to properties view when user cancels + detailSubView = 'properties'; + }} + /> +
    +
    + {:else if detailSubView === 'share'} + +
    +
    +

    + {currentLocale ? $_('assistants.sharing.title', { default: 'Shared Users' }) : 'Shared Users'} +

    +

    + {currentLocale ? $_('assistants.sharing.description', { default: 'Manage who has access to this assistant' }) : 'Manage who has access to this assistant'} +

    +
    + + {#if loadingShares} +
    +
    Loading shared users...
    +
    + {:else if sharedUsers.length === 0} +
    +

    + {currentLocale ? $_('assistants.sharing.noShares', { default: 'This assistant is not shared with anyone yet.' }) : 'This assistant is not shared with anyone yet.'} +

    +
    + {:else} +
    +
    + {#each sharedUsers as share (share.user_id)} +
    +
    +
    {share.user_name}
    +
    {share.user_email}
    +
    +
    + Shared by {share.shared_by_name} +
    +
    + {/each} +
    +
    + {/if} + +
    + +
    +
    + + + {#if showSharingModal} + + {/if} + {:else if detailSubView === 'chat'} + {#if configError} + + {:else if lambServerUrl && userToken} + +
    +

    + {currentLocale ? $_('assistants.detail.chatTitle', { default: 'Chat' }) : 'Chat'} +

    +
    + + {:else} + +

    {currentLocale ? $_('assistants.chatLoadingConfig') : 'Loading chat configuration...'}

    + {/if} + {:else if detailSubView === 'analytics'} + +
    + +
    + {:else if detailSubView === 'tests'} + + launchAacSkill(skill)} + /> + {:else if detailSubView === 'aac' && activeAacSessionId} + + {#key `${activeAacSessionId}-${aacSkillStartup}`} +
    + +
    + {/key} + {/if} + {/if} +
    + {:else if currentView === 'shared'} - -
    -
    - -
    -
    + +
    +
    + +
    +
    {:else if currentView === 'templates'} - -
    - -
    + +
    + +
    {:else} - -

    {currentLocale ? $_('assistants.noAssistantData') : 'Assistant data not available.'}

    -{/if} + +

    {currentLocale ? $_('assistants.noAssistantData') : 'Assistant data not available.'}

    +{/if} + { - isDeleteModalOpen = false; - assistantToDeleteId = null; - assistantToDeleteName = null; - deleteError = ''; // Clear errors on close - }} + bind:isOpen={isDeleteModalOpen} + bind:isLoading={isDeletingAssistant} + title={currentLocale ? $_('assistants.deleteModal.title', { default: 'Delete Assistant' }) : 'Delete Assistant'} + message={currentLocale ? $_('assistants.deleteModal.confirmation', { values: { name: assistantToDeleteName || '' }, default: `Are you sure you want to delete the assistant "${assistantToDeleteName}"? This action cannot be undone.` }) : `Are you sure you want to delete the assistant "${assistantToDeleteName}"? This action cannot be undone.`} + confirmText={currentLocale ? $_('assistants.deleteModal.confirmButton', { default: 'Delete' }) : 'Delete'} + variant="danger" + onconfirm={handleDeleteConfirm} + oncancel={() => { + isDeleteModalOpen = false; + assistantToDeleteId = null; + assistantToDeleteName = null; + deleteError = ''; // Clear errors on close + }} +/> + + + { notification.isOpen = false; }} /> {#if loadingDetail} -

    {currentLocale ? $_('assistants.loadingDetail') : 'Loading assistant details...'}

    -{/if} +

    {currentLocale ? $_('assistants.loadingDetail') : 'Loading assistant details...'}

    +{/if} \ No newline at end of file diff --git a/frontend/svelte-app/src/routes/org-admin/+page.svelte b/frontend/svelte-app/src/routes/org-admin/+page.svelte index 6d8bef9d3..312d81fcc 100644 --- a/frontend/svelte-app/src/routes/org-admin/+page.svelte +++ b/frontend/svelte-app/src/routes/org-admin/+page.svelte @@ -1,5579 +1,4835 @@ - Organization Admin - LAMB + Organization Admin - LAMB
    - - - - -
    -
    - - {#if error} - - {/if} - - - {#if currentView === 'dashboard'} -
    - - {#if dashboardData} -
    -
    -
    -
    -

    - {dashboardData.organization.name} -

    -

    Organization Administration

    -
    -
    - {#if targetOrgSlug} -
    - - - - System Admin View -
    - {/if} -
    Organization ID
    -
    {dashboardData.organization.slug}
    -
    -
    - - - {#if dashboardData.settings.signup_key_set} -
    -
    -
    -

    - Organization Signup Key -

    -
    - {showSignupKey ? signupKey : '••••••••••••••••'} -
    -
    - -
    -

    - Share this key with users to allow them to sign up for your organization -

    -
    - {:else} -
    -

    - - - - No signup key configured. Users cannot self-register for this organization. -

    -
    - {/if} -
    - {#if dashboardData.settings.lti_creator_enabled} -

    - - - - LTI Creator access is enabled. Users can join this organization via LTI from their - LMS. -

    - {:else} -

    - - - - LTI Creator access is not configured. Configure it in Settings to allow LMS-based - signup. -

    - {/if} -
    -
    -
    - {/if} - - {#if isLoadingDashboard} -
    -
    -

    Loading dashboard...

    -
    - {:else if dashboardData} - -
    - - - - - - - - - - -
    -
    -
    -
    - - - -
    -
    -
    API Status
    -
    - {#if dashboardData.api_status} - {#if dashboardData.api_status.overall_status === 'working'} - Working - {:else if dashboardData.api_status.overall_status === 'partial'} - Partial - {:else if dashboardData.api_status.overall_status === 'error'} - Error - {:else} - Not Configured - {/if} - {:else} - Unknown - {/if} -
    -
    -
    - {#if dashboardData.api_status} -
    - {dashboardData.api_status.summary.total_models} models available -
    - {/if} -
    -
    -
    - - - {#if dashboardData.default_models} -
    -
    -

    Default Models

    -
    -
    -
    - Global Default Model -
    - {#if dashboardData.default_models.global_default?.model} -
    - {dashboardData.default_models.global_default.model} -
    -
    - Provider: {dashboardData.default_models.global_default.provider} -
    - {:else} -
    Not configured
    - {/if} -
    -
    -
    - Small / Fast Model -
    - {#if dashboardData.default_models.small_fast?.model} -
    - {dashboardData.default_models.small_fast.model} -
    -
    - Provider: {dashboardData.default_models.small_fast.provider} -
    - {:else} -
    Not configured
    - {/if} -
    -
    -
    -
    - {/if} - - - {#if dashboardData.api_status && Object.keys(dashboardData.api_status.providers).length > 0} -
    -
    -

    - Enabled Models by Connector -

    -
    - {#each Object.entries(dashboardData.api_status.providers) as [providerName, providerStatus]} -
    -
    -
    {providerName}
    - - {providerStatus.status} - -
    - {#if providerStatus.enabled_models && providerStatus.enabled_models.length > 0} -
    -
    - {providerStatus.enabled_models.length} models enabled - {#if providerStatus.default_model} - • Default: {providerStatus.default_model} - {/if} -
    -
    - {#each providerStatus.enabled_models.slice(0, 8) as model} - - {model} - - {/each} - {#if providerStatus.enabled_models.length > 8} - - +{providerStatus.enabled_models.length - 8} more - - {/if} -
    -
    - {:else} -
    No models enabled
    - {/if} -
    - {/each} -
    -
    -
    - {/if} - - -
    -
    -

    Quick Settings

    -
    -
    - Signup Enabled - - {dashboardData.settings.signup_enabled ? 'Yes' : 'No'} - -
    -
    - Signup Key Set - - {dashboardData.settings.signup_key_set ? 'Yes' : 'No'} - -
    -
    - API Configured - - {dashboardData.settings.api_configured ? 'Yes' : 'No'} - -
    -
    -
    -
    - {/if} -
    - {/if} - - - {#if currentView === 'users'} -
    - - {#if dashboardData} -
    -
    -
    -
    -

    - {dashboardData.organization.name} -

    -

    User Management

    -
    -
    - {#if targetOrgSlug} -
    - - - - System Admin View -
    - {/if} - {usersTotalItems} users total -
    -
    -
    -
    - {/if} - - -
    -

    Manage Users

    - -
    - - {#if usersError} - - {/if} - - -
    -
    - -
    - - -
    - - -
    - - -
    - - -
    - - -
    -
    -
    - - {#if isLoadingUsers} -
    -
    -

    Loading users...

    -
    - {:else} - - {#if selectedUsers.length > 0} -
    -
    -
    - - - - - {selectedUsers.length} user{selectedUsers.length > 1 ? 's' : ''} selected - -
    -
    - - - -
    -
    -
    - {/if} - - -
    - - - - - - - - - - - - - {#if displayUsers.length === 0} - - - - {:else} - {#each displayUsers as user (user.id)} - - - - - - - - - - - - - {/each} - {/if} - -
    - 0 && - selectedUsers.length === - displayUsers.filter((u) => !(userData && userData.email === u.email)) - .length} - onchange={(e) => { - const checked = e.target?.checked; - displayUsers = displayUsers.map((u) => { - // Don't allow selection of current user - if (userData && userData.email === u.email) { - return u; - } - return { ...u, selected: checked }; - }); - }} - class="text-brand focus:ring-brand h-4 w-4 rounded border-gray-300" - /> - -
    - Name - Email -
    -     		-
    - User Type - Status -
    -     		- Actions - - Can Share -
    - {#if usersSearchQuery || usersFilterUserType !== 'all' || usersFilterStatus !== 'all'} - No users match your filters. Try adjusting the search or filters. - {:else} - No users found in your organization. - {/if} -
    - - -
    {user.name || '-'}
    -
    - {user.email} - {#if userData && userData.email === user.email} - (You) - {/if} -
    -     		-
    - {#if user.auth_provider === 'lti_creator'} - - LTI Creator - - {:else if user.user_type === 'end_user'} - - End User - - {:else} - - Creator - - {/if} -
    -
    - - {user.enabled ? 'Enabled' : 'Disabled'} - -
    -     		-
    - - {#if user.auth_provider !== 'lti_creator'} - - {/if} - - - {#if !(userData && userData.email === user.email)} - {#if user.enabled} - - {:else} - - {/if} - {/if} -
    -     		- -
    -
    - - - {#if usersTotalPages > 1} - - {/if} - {/if} -
    - {/if} - - - {#if currentView === 'assistants'} -
    - - {#if dashboardData} -
    -
    -
    -
    -

    - {dashboardData.organization.name} -

    -

    Assistants Access

    -
    -
    - {#if targetOrgSlug} - System Admin View - {/if} -
    ID: {dashboardData.organization.slug}
    -
    -
    -
    -
    - {/if} - - -
    -
    -
    -

    Organization Assistants

    -

    - View and manage user access to all assistants in your organization -

    -
    -
    -
    - - -
    -
    -
    - -
    -
    - - -
    - -
    -
    - - - {#if assistantsError} -
    - {assistantsError} -
    - {/if} - - - {#if isLoadingAssistants} -
    -
    Loading assistants...
    -
    - {:else if orgAssistants.filter((asst) => { - const matchesSearch = !assistantsSearchQuery || asst.name - .toLowerCase() - .includes(assistantsSearchQuery.toLowerCase()) || asst.owner - .toLowerCase() - .includes(assistantsSearchQuery.toLowerCase()) || (asst.description && asst.description - .toLowerCase() - .includes(assistantsSearchQuery.toLowerCase())); - const matchesPublished = assistantsFilterPublished === 'all' || (assistantsFilterPublished === 'published' && asst.published) || (assistantsFilterPublished === 'unpublished' && !asst.published); - return matchesSearch && matchesPublished; - }).length === 0} - -
    -
    - {#if assistantsSearchQuery || assistantsFilterPublished !== 'all'} -

    No assistants match your search criteria.

    - - {:else} -

    No assistants found in your organization.

    - {/if} -
    -
    - {:else} - -
    - - - - - - - - - - - - - - {#each orgAssistants.filter((asst) => { - const matchesSearch = !assistantsSearchQuery || asst.name - .toLowerCase() - .includes(assistantsSearchQuery.toLowerCase()) || asst.owner - .toLowerCase() - .includes(assistantsSearchQuery.toLowerCase()) || (asst.description && asst.description - .toLowerCase() - .includes(assistantsSearchQuery.toLowerCase())); - const matchesPublished = assistantsFilterPublished === 'all' || (assistantsFilterPublished === 'published' && asst.published) || (assistantsFilterPublished === 'unpublished' && !asst.published); - return matchesSearch && matchesPublished; - }) as assistant} - - - - - - - - - - - {/each} - -
    NameDescriptionOwnerStatusSharedCreatedActions
    - {assistant.name} - - {assistant.description || '—'} - - {assistant.owner} - - {#if assistant.published} - - Published - - {:else} - - Unpublished - - {/if} - - {#if assistantShareCounts[assistant.id] !== undefined} - {#if assistantShareCounts[assistant.id] === 0} - Not shared - {:else} - - Shared with {assistantShareCounts[assistant.id]} - {assistantShareCounts[assistant.id] === 1 ? 'user' : 'users'} - - {/if} - {:else} - — - {/if} - - {formatDate(assistant.created_at)} - - -
    -
    - {/if} -
    - {/if} - - - {#if currentView === 'lti-activities'} -
    - - {#if dashboardData} -
    -
    -
    -
    -

    - {dashboardData.organization.name} -

    -

    LTI Activities

    -
    -
    - {#if targetOrgSlug} - System Admin View - {/if} -
    ID: {dashboardData.organization.slug}
    -
    -
    -
    -
    - {/if} - -
    -
    -
    -

    LTI Activities

    - -
    -

    - Manage LTI activities configured through the unified LTI endpoint. Each activity - represents an LTI placement in an LMS course with its own set of assistants. -

    - - {#if isLoadingLtiActivities} -
    -
    - Loading activities... -
    - {:else} - {#if ltiActivitiesError} - - {/if} - {#if ltiActivitiesSuccess} - - {/if} - - {#if ltiActivities.length === 0} -
    - - - -

    No LTI activities yet

    -

    - Activities are created when an instructor launches the unified LTI tool for - the first time. -

    -
    - {:else} -
    - - - - - - - - - - - - - - - {#each ltiActivities as activity} - - - - - - - - - - - {/each} - -
    ActivityCourseAssistantsOwnerOptionsStatusCreatedActions
    -
    - {activity.activity_name || '(unnamed)'} -
    -
    - {activity.resource_link_id} -
    -     		- {activity.context_title || activity.context_id || '-'} - - - {activity.assistant_count || 0} assistant{(activity.assistant_count || - 0) !== 1 - ? 's' - : ''} - - - {activity.owner_name || - activity.owner_email || - activity.configured_by_name || - activity.configured_by_email || - '-'} - - {#if activity.chat_visibility_enabled} - - Chat visible - - {:else} - — - {/if} - - - {activity.status} - - - {activity.created_at - ? new Date(activity.created_at * 1000).toLocaleDateString() - : '-'} - - -
    -
    - {/if} - {/if} -
    -
    -
    - {/if} - - - {#if currentView === 'settings'} -
    - - {#if dashboardData} -
    -
    -
    -
    -

    - {dashboardData.organization.name} -

    -

    Organization Settings

    -
    -
    - {#if targetOrgSlug} -
    - - - - System Admin View -
    - {/if} - ID: {dashboardData.organization.slug} -
    -
    -
    -
    - {/if} - -

    Configuration

    - - {#if settingsError} - - {/if} - - {#if isLoadingSettings} -
    -
    -

    Loading settings...

    -
    - {:else} - -
    -
    - -
    -
    - - - {#if settingsSubView === 'general'} - -
    -
    -

    Signup Settings

    - - - {#if signupSettingsError} - - {/if} - - - {#if signupSettingsSuccess} - - {/if} - -
    -
    - - -
    - - {#if newSignupSettings.signup_enabled} -
    - - -

    - Unique key for users to signup to this organization (8-64 characters, - letters, numbers, hyphens, and underscores only) -

    -
    - {/if} - - -
    -
    -
    - {/if} - - - {#if settingsSubView === 'api'} - - {#if isLoadingDashboard} -
    -
    -

    - API Status Overview -

    -
    -
    - Checking API connections... -
    -
    -
    - {:else if dashboardData && dashboardData.api_status} -
    -
    -

    - API Status Overview -

    - -
    - {#each Object.entries(dashboardData.api_status.providers) as [providerName, providerStatus]} -
    -
    -

    {providerName}

    - - {providerStatus.status} - -
    - - {#if providerStatus.status === 'working'} -
    - {providerStatus.model_count} models available - {#if providerStatus.enabled_models && providerStatus.enabled_models.length > 0} - • {providerStatus.enabled_models.length} selected - {#if providerStatus.default_model} - • Default: {providerStatus.default_model} - {/if} - {:else} - • No models selected - {/if} -
    - - - {#if providerStatus.warning} -
    -
    - ⚠️ - {providerStatus.warning} -
    - {#if providerStatus.needs_model_selection} -
    - -
    - {/if} -
    - {/if} - - {#if providerStatus.models && providerStatus.models.length > 0} -
    -
    - {#each providerStatus.models.slice(0, 10) as model} -
    {model}
    - {/each} - {#if providerStatus.models.length > 10} -
    - ...and {providerStatus.models.length - 10} more -
    - {/if} -
    -
    - {/if} - {:else if providerStatus.error} -
    - Error: - {providerStatus.error} -
    - {#if providerStatus.error_code} -
    - Error code: {providerStatus.error_code} -
    - {/if} - - {#if providerStatus.models && providerStatus.models.length > 0} -
    -
    - {providerStatus.model_count} models available (connection - works, but test failed) -
    - -
    - {/if} - {/if} - - {#if providerStatus.api_base} -
    - {providerStatus.api_base} -
    - {/if} -
    - {/each} -
    - - {#if Object.keys(dashboardData.api_status.providers).length === 0} -
    -

    No API providers configured

    -

    Configure OpenAI or Ollama below to get started

    -
    - {/if} -
    -
    - {:else} -
    -
    -

    - API Status Overview -

    -
    -

    API status not available

    -

    - Configure your API settings below, then save to check connection status -

    -
    -
    -
    - {/if} - - -
    -
    -

    - API Configuration -

    - -
    -
    - - - newApiSettings.openai_api_key && - addPendingChange('OpenAI API key updated')} - /> -

    - {#if apiSettings.openai_api_key_set} - API key is currently set. Enter a new key to replace it. - {:else} - Enter your OpenAI API key to enable AI features. - {/if} -

    -
    - -
    - - -

    - Custom OpenAI API endpoint. Leave empty to use default - (https://api.openai.com/v1). -

    -
    - -
    - - -

    - URL of your Ollama server. This can be a local or remote Ollama - installation. -

    -
    - - -
    -

    Model Selection

    -

    - Configure which models are available to users in your organization -

    - - - {#if isLoadingDashboard} -
    -
    -
    - Loading available models... -
    -
    - - {:else if dashboardData?.api_status?.providers && Object.keys(dashboardData.api_status.providers).length > 0} - {#each Object.entries(dashboardData.api_status.providers) as [providerName, providerStatus]} - {@const freshModels = providerStatus.models || []} -
    -
    -

    {providerName}

    - -
    - -
    - {newApiSettings.selected_models?.[providerName]?.length || - 0} - of {freshModels.length} models enabled -
    - -
    - {#if newApiSettings.selected_models?.[providerName]?.length > 0} -
    - {#each newApiSettings.selected_models[providerName] as model} - - {model} - - {/each} -
    - {:else} -

    No models enabled

    - {/if} -
    - - - {#if newApiSettings.selected_models?.[providerName]?.length > 0} -
    - - -

    - This model will be used as the default when creating new - assistants for this provider. -

    -
    - {/if} -
    - {/each} - {:else if apiSettings.available_models && Object.keys(apiSettings.available_models || {}).length > 0} - - {#each Object.entries(apiSettings.available_models) as [providerName, models]} -
    -
    -

    {providerName}

    - -
    - -
    - {newApiSettings.selected_models?.[providerName]?.length || - 0} - of {models.length} models enabled - (cached - refresh to get latest) -
    - -
    - {#if newApiSettings.selected_models?.[providerName]?.length > 0} -
    - {#each newApiSettings.selected_models[providerName] as model} - - {model} - - {/each} -
    - {:else} -

    No models enabled

    - {/if} -
    - - - {#if newApiSettings.selected_models?.[providerName]?.length > 0} -
    - - -

    - This model will be used as the default when creating new - assistants for this provider. -

    -
    - {/if} -
    - {/each} - {:else} -
    -
    - - - -

    No API providers configured

    -

    - Enter your API credentials above and save to see available models -

    -
    -
    - {/if} -
    - - -
    -

    - Global Model Configuration -

    -

    - Configure organization-wide default models that apply across all assistants - and operations. -

    - - -
    -

    - Global Default Model -

    -

    - The primary model for this organization. Used for assistants and - completions when no specific model is configured. This overrides - per-provider defaults. -

    - -
    - -
    - - -
    - - -
    - - -
    -
    - - {#if newApiSettings.global_default_model.provider && newApiSettings.global_default_model.model} -
    - ✓ Global default model configured: - {newApiSettings.global_default_model.provider}/{newApiSettings - .global_default_model.model} -
    - {:else} -
    - ℹ️ No global default model configured. Per-provider defaults will be - used. -
    - {/if} -
    - - -
    -

    - Small Fast Model (Optional) -

    -

    - Configure a lightweight model for auxiliary plugin operations like query - rewriting, classification, and data extraction. This can reduce costs and - improve performance for internal processing tasks. -

    - -
    - -
    - - -
    - - -
    - - -
    -
    - - {#if newApiSettings.small_fast_model.provider && newApiSettings.small_fast_model.model} -
    - ✓ Small fast model configured: - {newApiSettings.small_fast_model.provider}/{newApiSettings - .small_fast_model.model} -
    - {:else} -
    - ℹ️ No small fast model configured. Plugins will use default models for - auxiliary operations. -
    - {/if} -
    -
    - - - {#if hasUnsavedChanges} -
    -
    -
    - - - -
    -
    -

    Unsaved Changes

    -
    - {#each pendingChanges as change} -
    - - {change} -
    - {/each} -
    -

    - Click "Commit Changes" below to save all modifications to the - database. -

    -
    -
    -
    - {/if} - - -
    -
    -
    - {/if} - - - {#if settingsSubView === 'kb'} -
    -
    -

    - Knowledge Base Server -

    -

    - Configure the Knowledge Base server connection for RAG (Retrieval-Augmented - Generation) functionality. Connection will be tested before saving. -

    - - - {#if kbSettingsError} - - {/if} - - - {#if kbSettingsSuccess} - - {/if} - - - {#if kbTestResult} - - {/if} - -
    - -
    - - (kbTestResult = null)} - /> -

    - URL of the Knowledge Base server (e.g., http://kb:9090 or - http://192.168.1.100:9090) -

    -
    - - -
    -
    - - {#if kbSettings.api_key_set || newKbSettings.api_key} - - {/if} -
    - { - kbTestResult = null; - // Auto-show API key when user starts typing - if (newKbSettings.api_key && !showKbApiKey) { - showKbApiKey = true; - } - }} - /> -

    - {#if kbSettings.api_key_set} - API key is currently set. {showKbApiKey - ? 'You can view it above.' - : 'Click "Show API Key" to view or leave empty to keep existing.'} - {:else} - Enter the API key for KB server authentication. - {/if} -

    -
    - - -
    - - -
    - - {#if showKbAdvanced} - -
    -
    - - {#if kbEmbeddingsConfig.apikey_configured || newKbSettings.embedding_api_key} - - {/if} -
    - - {#if !kbSettings.url} - -
    -
    - ℹ️ -
    -

    - KB Server Connection Required -

    -

    - Before setting an embedding API key, you must first configure and - save the KB Server URL and API Key above, then click "Update KB - Settings". -

    -
    -
    -
    - {/if} - - { - kbTestResult = null; - // Auto-show API key when user starts typing - if (newKbSettings.embedding_api_key && !showEmbeddingApiKey) { - showEmbeddingApiKey = true; - } - // Track if field is dirty (different from original) - embeddingApiKeyDirty = - newKbSettings.embedding_api_key !== embeddingApiKeyOriginal; - }} - /> -

    - {#if kbEmbeddingsConfig.apikey_configured} - API key is currently set on KB server ({kbEmbeddingsConfig.config_source === - 'file' - ? 'persisted config' - : 'from env vars'}). {showEmbeddingApiKey - ? 'You can view it above.' - : 'Click "Show API Key" to view or leave empty to keep existing.'} - {:else} - Enter the API key for embeddings service (e.g., OpenAI). This will be - set on the KB server. - {/if} -

    - - - {#if embeddingApiKeyDirty} -
    -
    - -
    - -

    - ⚠️ - This will update the embedding API key for all knowledge base - collections. Use this when rotating API keys. -

    -
    -
    -
    - {/if} - -
    - - {#if kbEmbeddingsConfig.config_source === 'file'} - - {/if} -
    - {#if kbEmbeddingsConfigError} -

    {kbEmbeddingsConfigError}

    - {/if} - - {#if applyToAllKbResult} -

    {applyToAllKbResult}

    - {/if} -
    - - -
    - - -

    - Default embedding model for new collections (leave empty for KB server - default) -

    -
    - - -
    - Collection Defaults (Optional) -
    -
    - - -
    -
    - - -
    -
    -

    - Default chunking parameters for new collections -

    -
    - {/if} - - -
    - - -
    - -

    - 💡 Tip: You can test the current configuration by clicking "Test - Connection" without entering a new API key. -

    -

    - ⚠️ Connection will be tested automatically before saving. Changes will not be - saved if the test fails. -

    -
    -
    -
    - {/if} - - - {#if settingsSubView === 'defaults'} - -
    -
    -

    - Assistant Defaults -

    -

    - These values seed the Create/Edit Assistant form for users in this organization. - Edit as raw JSON to add or change fields dynamically. -

    - - {#if assistantDefaultsError} - - {/if} - - {#if assistantDefaultsSuccess} - - {/if} - -
    - - - -
    - - -
    -

    - Tip: Fields are dynamic. Unknown keys will be preserved. Ensure the - `connector` and `llm` are enabled in this organization. -

    -
    -
    -
    - {/if} - - - {#if settingsSubView === 'lti-creator'} -
    -
    -

    - LTI Creator Access -

    -

    - Configure LTI-based login for the Creator Interface. This allows educators to - access LAMB directly from their LMS (e.g., Moodle) without needing separate - credentials. -

    - - - {#if isLoadingLtiCreatorSettings} -
    -
    - Loading LTI settings... -
    - {:else} - - {#if ltiCreatorSettingsError} - - {/if} - - - {#if ltiCreatorSettingsSuccess} - - {/if} - - -
    -

    Current Status

    - {#if ltiCreatorSettings.has_key} -
    - - {ltiCreatorSettings.enabled ? 'Enabled' : 'Disabled'} - -
    -

    - Consumer Key: - {ltiCreatorSettings.oauth_consumer_key} -

    - {#if ltiCreatorSettings.launch_url} -

    - Launch URL: - {ltiCreatorSettings.launch_url} -

    - {/if} - {:else} -

    - No LTI creator key configured yet. -

    - {/if} -
    - - -
    - - {#if ltiCreatorSettings.has_key} -
    - - -
    - {/if} - - -
    - - -

    - Must be unique across all organizations. Recommended format: {dashboardData?.organization?.slug || 'orgslug'}_creator -

    -
    - - -
    - -
    - - {#if newLtiCreatorSettings.oauth_consumer_secret} - - {/if} -
    -

    - Use a secure, random string. This will be used to sign LTI requests. -

    -
    - - -
    -

    - Setting up LTI in your LMS -

    -
      -
    • 1. In your LMS, create a new External Tool (LTI 1.1)
      • -
    • - 2. Set the Launch URL to: - {ltiCreatorSettings.launch_url || - '/lamb/v1/lti_creator/launch'} -
      • -
    • - 3. Enter the Consumer Key and - Consumer Secret from above -
      • -
    • - 4. Users who access via LTI will automatically get creator accounts -
      • -
    -
    - - -
    - - {#if ltiCreatorSettings.has_key} - - {/if} -
    - - -
    -

    - Note: LTI creator users cannot change their passwords (they - authenticate via LTI). They are identified by their LMS user ID and can access - LAMB from any LTI instance on their LMS. LTI creator users can be promoted - to organization admin by a system administrator. -

    -
    -
    - {/if} -
    -
    - {/if} - {/if} -
    - {/if} -
    -
    + + + + +
    +
    + + + {#if error} + + {/if} + + + {#if currentView === 'dashboard'} +
    + + {#if dashboardData} +
    +
    +
    +
    +

    + {dashboardData.organization.name} +

    +

    Organization Administration

    +
    +
    + {#if targetOrgSlug} +
    + + + + System Admin View +
    + {/if} +
    Organization ID
    +
    {dashboardData.organization.slug}
    +
    +
    + + + {#if dashboardData.settings.signup_key_set} +
    +
    +
    +

    Organization Signup Key

    +
    + {showSignupKey ? signupKey : '••••••••••••••••'} +
    +
    + +
    +

    + Share this key with users to allow them to sign up for your organization +

    +
    + {:else} +
    +

    + + + + No signup key configured. Users cannot self-register for this organization. +

    +
    + {/if} +
    + {#if dashboardData.settings.lti_creator_enabled} +

    + + + + LTI Creator access is enabled. Users can join this organization via LTI from their LMS. +

    + {:else} +

    + + + + LTI Creator access is not configured. Configure it in Settings to allow LMS-based signup. +

    + {/if} +
    +
    +
    + {/if} + + {#if isLoadingDashboard} +
    +
    +

    Loading dashboard...

    +
    + {:else if dashboardData} + +
    + + + + + + + + + + +
    +
    +
    +
    + + + +
    +
    +
    API Status
    +
    + {#if dashboardData.api_status} + {#if dashboardData.api_status.overall_status === 'working'} + Working + {:else if dashboardData.api_status.overall_status === 'partial'} + Partial + {:else if dashboardData.api_status.overall_status === 'error'} + Error + {:else} + Not Configured + {/if} + {:else} + Unknown + {/if} +
    +
    +
    + {#if dashboardData.api_status} +
    + {dashboardData.api_status.summary.total_models} models available +
    + {/if} +
    +
    +
    + + + {#if dashboardData.default_models} +
    +
    +

    Default Models

    +
    +
    +
    Global Default Model
    + {#if dashboardData.default_models.global_default?.model} +
    {dashboardData.default_models.global_default.model}
    +
    Provider: {dashboardData.default_models.global_default.provider}
    + {:else} +
    Not configured
    + {/if} +
    +
    +
    Small / Fast Model
    + {#if dashboardData.default_models.small_fast?.model} +
    {dashboardData.default_models.small_fast.model}
    +
    Provider: {dashboardData.default_models.small_fast.provider}
    + {:else} +
    Not configured
    + {/if} +
    +
    +
    +
    + {/if} + + + {#if dashboardData.api_status && Object.keys(dashboardData.api_status.providers).length > 0} +
    +
    +

    Enabled Models by Connector

    +
    + {#each Object.entries(dashboardData.api_status.providers) as [providerName, providerStatus]} +
    +
    +
    {providerName}
    + + {providerStatus.status} + +
    + {#if providerStatus.enabled_models && providerStatus.enabled_models.length > 0} +
    +
    + {providerStatus.enabled_models.length} models enabled + {#if providerStatus.default_model} + • Default: {providerStatus.default_model} + {/if} +
    +
    + {#each providerStatus.enabled_models.slice(0, 8) as model} + + {model} + + {/each} + {#if providerStatus.enabled_models.length > 8} + + +{providerStatus.enabled_models.length - 8} more + + {/if} +
    +
    + {:else} +
    No models enabled
    + {/if} +
    + {/each} +
    +
    +
    + {/if} + + +
    +
    +

    Quick Settings

    +
    +
    + Signup Enabled + + {dashboardData.settings.signup_enabled ? 'Yes' : 'No'} + +
    +
    + Signup Key Set + + {dashboardData.settings.signup_key_set ? 'Yes' : 'No'} + +
    +
    + API Configured + + {dashboardData.settings.api_configured ? 'Yes' : 'No'} + +
    +
    +
    +
    + {/if} +
    + {/if} + + + {#if currentView === 'users'} +
    + + {#if dashboardData} +
    +
    +
    +
    +

    {dashboardData.organization.name}

    +

    User Management

    +
    +
    + {#if targetOrgSlug} +
    + + + + System Admin View +
    + {/if} + {usersTotalItems} users total +
    +
    +
    +
    + {/if} + + +
    +

    Manage Users

    + +
    + + {#if usersError} + + {/if} + + +
    +
    + +
    + + +
    + + +
    + + +
    + + +
    + + +
    +
    +
    + + {#if isLoadingUsers} +
    +
    +

    Loading users...

    +
    + {:else} + + {#if selectedUsers.length > 0} +
    +
    +
    + + + + + {selectedUsers.length} user{selectedUsers.length > 1 ? 's' : ''} selected + +
    +
    + + + +
    +
    +
    + {/if} + + +
    + + + + + + + + + + + + + {#if displayUsers.length === 0} + + + + {:else} + {#each displayUsers as user (user.id)} + + + + + + + + + + + + + {/each} + {/if} + +
    + 0 && selectedUsers.length === displayUsers.filter(u => !(userData && userData.email === u.email)).length} + onchange={(e) => { + const checked = e.target?.checked; + displayUsers = displayUsers.map(u => { + // Don't allow selection of current user + if (userData && userData.email === u.email) { + return u; + } + return {...u, selected: checked}; + }); + }} + class="h-4 w-4 text-brand focus:ring-brand border-gray-300 rounded" + /> + +
    + Name + Email +
    +     		+
    + User Type + Status +
    +     		+ Actions + + Can Share +
    + {#if usersSearchQuery || usersFilterUserType !== 'all' || usersFilterStatus !== 'all'} + No users match your filters. Try adjusting the search or filters. + {:else} + No users found in your organization. + {/if} +
    + + +
    {user.name || '-'}
    +
    + {user.email} + {#if userData && userData.email === user.email} + (You) + {/if} +
    +     		+
    + {#if user.auth_provider === 'lti_creator'} + + LTI Creator + + {:else if user.user_type === 'end_user'} + + End User + + {:else} + + Creator + + {/if} +
    +
    + + {user.enabled ? 'Enabled' : 'Disabled'} + +
    +     		+
    + + {#if user.auth_provider !== 'lti_creator'} + + {/if} + + + {#if !(userData && userData.email === user.email)} + {#if user.enabled} + + {:else} + + {/if} + {/if} +
    +     		+ +
    +
    + + + {#if usersTotalPages > 1} + + {/if} + {/if} +
    + {/if} + + + {#if currentView === 'assistants'} +
    + + {#if dashboardData} +
    +
    +
    +
    +

    {dashboardData.organization.name}

    +

    Assistants Access

    +
    +
    + {#if targetOrgSlug} + System Admin View + {/if} +
    ID: {dashboardData.organization.slug}
    +
    +
    +
    +
    + {/if} + + +
    +
    +
    +

    Organization Assistants

    +

    View and manage user access to all assistants in your organization

    +
    +
    +
    + + +
    +
    +
    + +
    +
    + + +
    + +
    +
    + + + {#if assistantsError} +
    + {assistantsError} +
    + {/if} + + + {#if isLoadingAssistants} +
    +
    Loading assistants...
    +
    + {:else} + {#if orgAssistants.filter(asst => { + const matchesSearch = !assistantsSearchQuery || + asst.name.toLowerCase().includes(assistantsSearchQuery.toLowerCase()) || + asst.owner.toLowerCase().includes(assistantsSearchQuery.toLowerCase()) || + (asst.description && asst.description.toLowerCase().includes(assistantsSearchQuery.toLowerCase())); + const matchesPublished = + assistantsFilterPublished === 'all' || + (assistantsFilterPublished === 'published' && asst.published) || + (assistantsFilterPublished === 'unpublished' && !asst.published); + return matchesSearch && matchesPublished; + }).length === 0} + +
    +
    + {#if assistantsSearchQuery || assistantsFilterPublished !== 'all'} +

    No assistants match your search criteria.

    + + {:else} +

    No assistants found in your organization.

    + {/if} +
    +
    + {:else} + +
    + + + + + + + + + + + + + + {#each orgAssistants.filter(asst => { + const matchesSearch = !assistantsSearchQuery || + asst.name.toLowerCase().includes(assistantsSearchQuery.toLowerCase()) || + asst.owner.toLowerCase().includes(assistantsSearchQuery.toLowerCase()) || + (asst.description && asst.description.toLowerCase().includes(assistantsSearchQuery.toLowerCase())); + const matchesPublished = + assistantsFilterPublished === 'all' || + (assistantsFilterPublished === 'published' && asst.published) || + (assistantsFilterPublished === 'unpublished' && !asst.published); + return matchesSearch && matchesPublished; + }) as assistant} + + + + + + + + + + + {/each} + +
    NameDescriptionOwnerStatusSharedCreatedActions
    + {assistant.name} + + {assistant.description || '—'} + + {assistant.owner} + + {#if assistant.published} + + Published + + {:else} + + Unpublished + + {/if} + + {#if assistantShareCounts[assistant.id] !== undefined} + {#if assistantShareCounts[assistant.id] === 0} + Not shared + {:else} + + Shared with {assistantShareCounts[assistant.id]} {assistantShareCounts[assistant.id] === 1 ? 'user' : 'users'} + + {/if} + {:else} + — + {/if} + + {formatDate(assistant.created_at)} + + +
    +
    + {/if} + {/if} +
    + {/if} + + + {#if currentView === 'lti-activities'} +
    + + {#if dashboardData} +
    +
    +
    +
    +

    {dashboardData.organization.name}

    +

    LTI Activities

    +
    +
    + {#if targetOrgSlug} + System Admin View + {/if} +
    ID: {dashboardData.organization.slug}
    +
    +
    +
    +
    + {/if} + +
    +
    +
    +

    LTI Activities

    + +
    +

    + Manage LTI activities configured through the unified LTI endpoint. Each activity represents + an LTI placement in an LMS course with its own set of assistants. +

    + + {#if isLoadingLtiActivities} +
    +
    + Loading activities... +
    + {:else} + {#if ltiActivitiesError} + + {/if} + {#if ltiActivitiesSuccess} + + {/if} + + {#if ltiActivities.length === 0} +
    + + + +

    No LTI activities yet

    +

    + Activities are created when an instructor launches the unified LTI tool for the first time. +

    +
    + {:else} +
    + + + + + + + + + + + + + + + {#each ltiActivities as activity} + + + + + + + + + + + {/each} + +
    ActivityCourseAssistantsOwnerOptionsStatusCreatedActions
    +
    + {activity.activity_name || '(unnamed)'} +
    +
    + {activity.resource_link_id} +
    +     		+ {activity.context_title || activity.context_id || '-'} + + + {activity.assistant_count || 0} assistant{(activity.assistant_count || 0) !== 1 ? 's' : ''} + + + {activity.owner_name || activity.owner_email || activity.configured_by_name || activity.configured_by_email || '-'} + + {#if activity.chat_visibility_enabled} + + Chat visible + + {:else} + — + {/if} + + + {activity.status} + + + {activity.created_at ? new Date(activity.created_at * 1000).toLocaleDateString() : '-'} + + +
    +
    + {/if} + {/if} +
    +
    +
    + {/if} + + + {#if currentView === 'settings'} +
    + + {#if dashboardData} +
    +
    +
    +
    +

    {dashboardData.organization.name}

    +

    Organization Settings

    +
    +
    + {#if targetOrgSlug} +
    + + + + System Admin View +
    + {/if} + ID: {dashboardData.organization.slug} +
    +
    +
    +
    + {/if} + +

    Configuration

    + + {#if settingsError} + + {/if} + + {#if isLoadingSettings} +
    +
    +

    Loading settings...

    +
    + {:else} + +
    +
    + +
    +
    + + + {#if settingsSubView === 'general'} + +
    +
    +

    Signup Settings

    + + + {#if signupSettingsError} + + {/if} + + + {#if signupSettingsSuccess} + + {/if} + +
    +
    + + +
    + + {#if newSignupSettings.signup_enabled} +
    + + +

    + Unique key for users to signup to this organization (8-64 characters, letters, numbers, hyphens, and underscores only) +

    +
    + {/if} + + +
    +
    +
    + {/if} + + + {#if settingsSubView === 'api'} + + {#if isLoadingDashboard} +
    +
    +

    API Status Overview

    +
    +
    + Checking API connections... +
    +
    +
    + {:else if dashboardData && dashboardData.api_status} +
    +
    +

    API Status Overview

    + +
    + {#each Object.entries(dashboardData.api_status.providers) as [providerName, providerStatus]} +
    +
    +

    {providerName}

    + + {providerStatus.status} + +
    + + {#if providerStatus.status === 'working'} +
    + {providerStatus.model_count} models available + {#if providerStatus.enabled_models && providerStatus.enabled_models.length > 0} + • {providerStatus.enabled_models.length} selected + {#if providerStatus.default_model} + • Default: {providerStatus.default_model} + {/if} + {:else} + • No models selected + {/if} +
    + + + {#if providerStatus.warning} +
    +
    + ⚠️ + {providerStatus.warning} +
    + {#if providerStatus.needs_model_selection} +
    + +
    + {/if} +
    + {/if} + + {#if providerStatus.models && providerStatus.models.length > 0} +
    +
    + {#each providerStatus.models.slice(0, 10) as model} +
    {model}
    + {/each} + {#if providerStatus.models.length > 10} +
    ...and {providerStatus.models.length - 10} more
    + {/if} +
    +
    + {/if} + {:else if providerStatus.error} +
    + Error: {providerStatus.error} +
    + {#if providerStatus.error_code} +
    + Error code: {providerStatus.error_code} +
    + {/if} + + {#if providerStatus.models && providerStatus.models.length > 0} +
    +
    + {providerStatus.model_count} models available (connection works, but test failed) +
    + +
    + {/if} + {/if} + + {#if providerStatus.api_base} +
    + {providerStatus.api_base} +
    + {/if} +
    + {/each} +
    + + {#if Object.keys(dashboardData.api_status.providers).length === 0} +
    +

    No API providers configured

    +

    Configure OpenAI or Ollama below to get started

    +
    + {/if} +
    +
    + {:else} +
    +
    +

    API Status Overview

    +
    +

    API status not available

    +

    Configure your API settings below, then save to check connection status

    +
    +
    +
    + {/if} + + +
    +
    +

    API Configuration

    + +
    +
    + + newApiSettings.openai_api_key && addPendingChange('OpenAI API key updated')} + > +

    + {#if apiSettings.openai_api_key_set} + API key is currently set. Enter a new key to replace it. + {:else} + Enter your OpenAI API key to enable AI features. + {/if} +

    +
    + +
    + + +

    + Custom OpenAI API endpoint. Leave empty to use default (https://api.openai.com/v1). +

    +
    + +
    + + +

    + URL of your Ollama server. This can be a local or remote Ollama installation. +

    +
    + + +
    +

    Model Selection

    +

    + Configure which models are available to users in your organization +

    + + + {#if isLoadingDashboard} +
    +
    +
    + Loading available models... +
    +
    + + {:else if dashboardData?.api_status?.providers && Object.keys(dashboardData.api_status.providers).length > 0} + {#each Object.entries(dashboardData.api_status.providers) as [providerName, providerStatus]} + {@const freshModels = providerStatus.models || []} +
    +
    +

    {providerName}

    + +
    + +
    + {newApiSettings.selected_models?.[providerName]?.length || 0} of {freshModels.length} models enabled +
    + +
    + {#if newApiSettings.selected_models?.[providerName]?.length > 0} +
    + {#each newApiSettings.selected_models[providerName] as model} + + {model} + + {/each} +
    + {:else} +

    No models enabled

    + {/if} +
    + + + {#if newApiSettings.selected_models?.[providerName]?.length > 0} +
    + + +

    + This model will be used as the default when creating new assistants for this provider. +

    +
    + {/if} +
    + {/each} + {:else if apiSettings.available_models && Object.keys(apiSettings.available_models || {}).length > 0} + + {#each Object.entries(apiSettings.available_models) as [providerName, models]} +
    +
    +

    {providerName}

    + +
    + +
    + {newApiSettings.selected_models?.[providerName]?.length || 0} of {models.length} models enabled + (cached - refresh to get latest) +
    + +
    + {#if newApiSettings.selected_models?.[providerName]?.length > 0} +
    + {#each newApiSettings.selected_models[providerName] as model} + + {model} + + {/each} +
    + {:else} +

    No models enabled

    + {/if} +
    + + + {#if newApiSettings.selected_models?.[providerName]?.length > 0} +
    + + +

    + This model will be used as the default when creating new assistants for this provider. +

    +
    + {/if} +
    + {/each} + {:else} +
    +
    + + + +

    No API providers configured

    +

    Enter your API credentials above and save to see available models

    +
    +
    + {/if} +
    + + +
    +

    + Global Model Configuration +

    +

    + Configure organization-wide default models that apply across all assistants and operations. +

    + + +
    +

    + Global Default Model +

    +

    + The primary model for this organization. Used for assistants and completions when + no specific model is configured. This overrides per-provider defaults. +

    + +
    + +
    + + +
    + + +
    + + +
    +
    + + {#if newApiSettings.global_default_model.provider && newApiSettings.global_default_model.model} +
    + ✓ Global default model configured: + {newApiSettings.global_default_model.provider}/{newApiSettings.global_default_model.model} +
    + {:else} +
    + ℹ️ No global default model configured. Per-provider defaults will be used. +
    + {/if} +
    + + +
    +

    + Small Fast Model (Optional) +

    +

    + Configure a lightweight model for auxiliary plugin operations like query rewriting, + classification, and data extraction. This can reduce costs and improve performance + for internal processing tasks. +

    + +
    + +
    + + +
    + + +
    + + +
    +
    + + {#if newApiSettings.small_fast_model.provider && newApiSettings.small_fast_model.model} +
    + ✓ Small fast model configured: + {newApiSettings.small_fast_model.provider}/{newApiSettings.small_fast_model.model} +
    + {:else} +
    + ℹ️ No small fast model configured. Plugins will use default models for auxiliary operations. +
    + {/if} +
    +
    + + + {#if hasUnsavedChanges} +
    +
    +
    + + + +
    +
    +

    + Unsaved Changes +

    +
    + {#each pendingChanges as change} +
    + + {change} +
    + {/each} +
    +

    + Click "Commit Changes" below to save all modifications to the database. +

    +
    +
    +
    + {/if} + + +
    +
    +
    + {/if} + + + {#if settingsSubView === 'kb'} +
    +
    +

    Knowledge Base Server

    +

    + Configure the Knowledge Base server connection for RAG (Retrieval-Augmented Generation) functionality. + Connection will be tested before saving. +

    + + + {#if kbSettingsError} + + {/if} + + + {#if kbSettingsSuccess} + + {/if} + + + {#if kbTestResult} + + {/if} + +
    + +
    + + kbTestResult = null} + > +

    + URL of the Knowledge Base server (e.g., http://kb:9090 or http://192.168.1.100:9090) +

    +
    + + +
    +
    + + {#if kbSettings.api_key_set || newKbSettings.api_key} + + {/if} +
    + { + kbTestResult = null; + // Auto-show API key when user starts typing + if (newKbSettings.api_key && !showKbApiKey) { + showKbApiKey = true; + } + }} + > +

    + {#if kbSettings.api_key_set} + API key is currently set. {showKbApiKey ? 'You can view it above.' : 'Click "Show API Key" to view or leave empty to keep existing.'} + {:else} + Enter the API key for KB server authentication. + {/if} +

    +
    + + +
    + + +
    + + {#if showKbAdvanced} + +
    +
    + + {#if kbEmbeddingsConfig.apikey_configured || newKbSettings.embedding_api_key} + + {/if} +
    + + {#if !kbSettings.url} + +
    +
    + ℹ️ +
    +

    KB Server Connection Required

    +

    + Before setting an embedding API key, you must first configure and save the KB Server URL and API Key above, then click "Update KB Settings". +

    +
    +
    +
    + {/if} + + { + kbTestResult = null; + // Auto-show API key when user starts typing + if (newKbSettings.embedding_api_key && !showEmbeddingApiKey) { + showEmbeddingApiKey = true; + } + // Track if field is dirty (different from original) + embeddingApiKeyDirty = newKbSettings.embedding_api_key !== embeddingApiKeyOriginal; + }} + > +

    + {#if kbEmbeddingsConfig.apikey_configured} + API key is currently set on KB server ({kbEmbeddingsConfig.config_source === 'file' ? 'persisted config' : 'from env vars'}). {showEmbeddingApiKey ? 'You can view it above.' : 'Click "Show API Key" to view or leave empty to keep existing.'} + {:else} + Enter the API key for embeddings service (e.g., OpenAI). This will be set on the KB server. + {/if} +

    + + + {#if embeddingApiKeyDirty} +
    +
    + +
    + +

    + ⚠️ + This will update the embedding API key for all knowledge base collections. Use this when rotating API keys. +

    +
    +
    +
    + {/if} + +
    + + {#if kbEmbeddingsConfig.config_source === 'file'} + + {/if} +
    + {#if kbEmbeddingsConfigError} +

    {kbEmbeddingsConfigError}

    + {/if} + + {#if applyToAllKbResult} +

    {applyToAllKbResult}

    + {/if} +
    + + +
    + + +

    + Default embedding model for new collections (leave empty for KB server default) +

    +
    + + +
    + Collection Defaults (Optional) +
    +
    + + +
    +
    + + +
    +
    +

    + Default chunking parameters for new collections +

    +
    + {/if} + + +
    + + +
    + +

    + 💡 Tip: You can test the current configuration by clicking "Test Connection" without entering a new API key. +

    +

    + ⚠️ Connection will be tested automatically before saving. Changes will not be saved if the test fails. +

    +
    +
    +
    + {/if} + + + {#if settingsSubView === 'defaults'} + +
    +
    +

    Assistant Defaults

    +

    These values seed the Create/Edit Assistant form for users in this organization. Edit as raw JSON to add or change fields dynamically.

    + + {#if assistantDefaultsError} + + {/if} + + {#if assistantDefaultsSuccess} + + {/if} + +
    + + + +
    + + +
    +

    Tip: Fields are dynamic. Unknown keys will be preserved. Ensure the `connector` and `llm` are enabled in this organization.

    +
    +
    +
    + {/if} + + + {#if settingsSubView === 'lti-creator'} +
    +
    +

    LTI Creator Access

    +

    + Configure LTI-based login for the Creator Interface. This allows educators to access LAMB + directly from their LMS (e.g., Moodle) without needing separate credentials. +

    + + + {#if isLoadingLtiCreatorSettings} +
    +
    + Loading LTI settings... +
    + {:else} + + {#if ltiCreatorSettingsError} + + {/if} + + + {#if ltiCreatorSettingsSuccess} + + {/if} + + +
    +

    Current Status

    + {#if ltiCreatorSettings.has_key} +
    + + {ltiCreatorSettings.enabled ? 'Enabled' : 'Disabled'} + +
    +

    + Consumer Key: {ltiCreatorSettings.oauth_consumer_key} +

    + {#if ltiCreatorSettings.launch_url} +

    + Launch URL: {ltiCreatorSettings.launch_url} +

    + {/if} + {:else} +

    No LTI creator key configured yet.

    + {/if} +
    + + +
    + + {#if ltiCreatorSettings.has_key} +
    + + +
    + {/if} + + +
    + + +

    + Must be unique across all organizations. Recommended format: {dashboardData?.organization?.slug || 'orgslug'}_creator +

    +
    + + +
    + +
    + + {#if newLtiCreatorSettings.oauth_consumer_secret} + + {/if} +
    +

    + Use a secure, random string. This will be used to sign LTI requests. +

    +
    + + +
    +

    Setting up LTI in your LMS

    +
      +
    • 1. In your LMS, create a new External Tool (LTI 1.1)
      • +
    • 2. Set the Launch URL to: {ltiCreatorSettings.launch_url || '/lamb/v1/lti_creator/launch'}
      • +
    • 3. Enter the Consumer Key and Consumer Secret from above
      • +
    • 4. Users who access via LTI will automatically get creator accounts
      • +
    +
    + + +
    + + {#if ltiCreatorSettings.has_key} + + {/if} +
    + + +
    +

    + Note: LTI creator users cannot change their passwords (they authenticate via LTI). + They are identified by their LMS user ID and can access LAMB from any LTI instance on their LMS. + LTI creator users can be promoted to organization admin by a system administrator. +

    +
    +
    + {/if} +
    +
    + {/if} + {/if} +
    + {/if} + +
    +
    {#if showApplyToAllKbConfirmation} -
    -
    -
    -
    - - - -
    - -

    - Apply embedding API key to all knowledge bases? -

    - -
    -

    - This will update the embedding API key for all existing KB collections in - this organization. -

    -
    -

    - Use this when rotating keys. Model/vendor/endpoint are unchanged — only the key is - updated. -

    -
    -
    - -
    - - -
    -
    -
    -
    +
    +
    +
    +
    + + + +
    + +

    + Apply embedding API key to all knowledge bases? +

    + +
    +

    + This will update the embedding API key for all existing KB collections in this organization. +

    +
    +

    + Use this when rotating keys. Model/vendor/endpoint are unchanged — only the key is updated. +

    +
    +
    + +
    + + +
    +
    +
    +
    {/if} { - newUser = user; - }} + isOpen={isCreateUserModalOpen} + isSuperAdmin={false} + {newUser} + isCreating={isCreatingUser} + error={createUserError} + success={createUserSuccess} + onSubmit={handleCreateUser} + onClose={closeCreateUserModal} + onUserChange={(user) => { newUser = user; }} /> { - passwordChangeData.new_password = pwd; - }} + isOpen={isChangePasswordModalOpen} + userName={passwordChangeData.user_name} + userEmail={passwordChangeData.user_email} + newPassword={passwordChangeData.new_password} + isChanging={isChangingPassword} + error={changePasswordError} + success={changePasswordSuccess} + onSubmit={handleChangePassword} + onClose={closeChangePasswordModal} + onPasswordChange={(pwd) => { passwordChangeData.new_password = pwd; }} /> {#if isModelModalOpen} -
    -
    -
    - -
    -

    - Manage {modalProviderName} Models -

    - -
    - - -
    - -
    -
    -

    Enabled Models

    - ({modalEnabledModels.length}) -
    - - - - - -
    - {#each modalEnabledModels.filter((model) => model - .toLowerCase() - .includes(enabledSearchTerm.toLowerCase())) as model} - - {/each} - {#if modalEnabledModels.length === 0} -
    No models enabled
    - {/if} -
    -
    - - -
    - - - - -
    - - -
    -
    -

    Available Models

    - ({modalDisabledModels.length}) -
    - - - - - -
    - {#each modalDisabledModels.filter((model) => model - .toLowerCase() - .includes(disabledSearchTerm.toLowerCase())) as model} - - {/each} - {#if modalDisabledModels.length === 0} -
    All models enabled
    - {/if} -
    -
    -
    - - -
    - - -
    -
    -
    -
    +
    +
    +
    + +
    +

    + Manage {modalProviderName} Models +

    + +
    + + +
    + +
    +
    +

    Enabled Models

    + ({modalEnabledModels.length}) +
    + + + + + +
    + {#each modalEnabledModels.filter(model => model.toLowerCase().includes(enabledSearchTerm.toLowerCase())) as model} + + {/each} + {#if modalEnabledModels.length === 0} +
    + No models enabled +
    + {/if} +
    +
    + + +
    + + + + +
    + + +
    +
    +

    Available Models

    + ({modalDisabledModels.length}) +
    + + + + + +
    + {#each modalDisabledModels.filter(model => model.toLowerCase().includes(disabledSearchTerm.toLowerCase())) as model} + + {/each} + {#if modalDisabledModels.length === 0} +
    + All models enabled +
    + {/if} +
    +
    +
    + + +
    + + +
    +
    +
    +
    {/if} + {#if showSharingModal && modalAssistant} - + {/if} { - await updateKbEmbeddingsConfig(); - showResetKbConfigModal = false; - }} - oncancel={() => { - showResetKbConfigModal = false; - }} + bind:isOpen={showResetKbConfigModal} + title="Reset KB Configuration" + message="Reset to environment variables? This will remove the persisted configuration and use the defaults from your environment settings." + confirmText="Reset" + variant="warning" + onconfirm={async () => { + await updateKbEmbeddingsConfig(); + showResetKbConfigModal = false; + }} + oncancel={() => { showResetKbConfigModal = false; }} /> + /* Custom scrollbar styles */ + .overflow-x-auto::-webkit-scrollbar { + height: 8px; + } + + .overflow-x-auto::-webkit-scrollbar-track { + background: #f1f1f1; + border-radius: 4px; + } + + .overflow-x-auto::-webkit-scrollbar-thumb { + background: #c1c1c1; + border-radius: 4px; + } + + .overflow-x-auto::-webkit-scrollbar-thumb:hover { + background: #a8a8a8; + } + \ No newline at end of file diff --git a/lamb-kb-server-stable/backend/dependencies.py b/lamb-kb-server-stable/backend/dependencies.py index e9858c66b..fd9a5d17f 100644 --- a/lamb-kb-server-stable/backend/dependencies.py +++ b/lamb-kb-server-stable/backend/dependencies.py @@ -2,8 +2,18 @@ from fastapi import Depends, HTTPException, status from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials -# Get API key from environment variable or use default -API_KEY = os.getenv("LAMB_API_KEY", "0p3n-w3bu!") +# API key from env. Backwards-compatible: fall back to the historical default +# so existing installs keep working, but warn loudly — it must be overridden in +# production (#413). +import sys +API_KEY = os.getenv("LAMB_API_KEY") +if not API_KEY: + API_KEY = "0p3n-w3bu!" + print( + "WARNING [kb-server]: LAMB_API_KEY not set; using the insecure default token. " + "Set LAMB_API_KEY to a strong value in production.", + file=sys.stderr, + ) # Security scheme security = HTTPBearer( diff --git a/lamb-kb-server-stable/backend/main.py b/lamb-kb-server-stable/backend/main.py index 1f401b63d..fb1cdcfc1 100644 --- a/lamb-kb-server-stable/backend/main.py +++ b/lamb-kb-server-stable/backend/main.py @@ -63,8 +63,18 @@ QueryPluginInfo ) -# Get API key from environment variable or use default -API_KEY = os.getenv("LAMB_API_KEY", "0p3n-w3bu!") +# API key from env. Backwards-compatible: fall back to the historical default +# so existing installs keep working, but warn loudly — it must be overridden in +# production (#413). +import sys as _sys +API_KEY = os.getenv("LAMB_API_KEY") +if not API_KEY: + API_KEY = "0p3n-w3bu!" + print( + "WARNING [kb-server]: LAMB_API_KEY not set; using the insecure default token. " + "Set LAMB_API_KEY to a strong value in production.", + file=_sys.stderr, + ) # Get default embeddings model configuration from environment variables # Default to using Ollama with nomic-embed-text model @@ -83,12 +93,12 @@ ## Authentication - All API endpoints are secured with Bearer token authentication. The token must match - the `LAMB_API_KEY` environment variable (default: `0p3n-w3bu!`). - + All API endpoints are secured with Bearer token authentication. The token must match + the `LAMB_API_KEY` environment variable, which must be set to a strong, unique value. + Example: ``` - curl -H 'Authorization: Bearer 0p3n-w3bu!' http://localhost:9090/ + curl -H "Authorization: Bearer $LAMB_API_KEY" http://localhost:9090/ ``` ## Features @@ -210,7 +220,7 @@ async def get_ingestion_config(): Example: ```bash curl -X GET 'http://localhost:9090/ingestion/plugins' \ - -H 'Authorization: Bearer 0p3n-w3bu!' + -H "Authorization: Bearer $LAMB_API_KEY" ``` """, tags=["Ingestion"], diff --git a/lamb-kb-server/README.md b/lamb-kb-server/README.md index 4101d6ad0..de109d5d2 100644 --- a/lamb-kb-server/README.md +++ b/lamb-kb-server/README.md @@ -107,7 +107,7 @@ ruff check backend/ tests/ The e2e tier requires Docker (Qdrant + Ollama containers brought up automatically by the session fixture). If `QDRANT_TEST_PORT` and `OLLAMA_TEST_PORT` env vars are set, the fixture uses those pre-started containers instead of spinning up its own. If Docker is unavailable, the entire e2e tier is skipped with a clear message. -Combined coverage hits **99%** line + branch on `backend/` (572 tests). The remaining 1% is structurally unreachable: a `sys.exit(1)` startup guard, an `ImportError` re-raise inside an `except ImportError` block, and a defensive branch the splitter algorithm never enters. +Combined coverage is **~94%** line + branch on `backend/` (594 tests). The uncovered remainder is dominated by code paths that the e2e tier exercises only over real HTTP — that work runs in a uvicorn subprocess the parent `pytest-cov` process cannot instrument — together with a few structurally unreachable guards: a `sys.exit(1)` startup guard, an `ImportError` re-raise inside an `except ImportError` block, and a defensive branch the splitter algorithm never enters. ### Docker @@ -156,7 +156,15 @@ SQLite at `$DATA_DIR/kb-server.db`, WAL mode enabled at connection time. Two tab - `collections` — one row per KB (immutable store setup). - `ingestion_jobs` — persistent queue (document text lives here until processed; credentials do not). -Schema is managed with `Base.metadata.create_all` (no Alembic migrations). +Schema is managed with Alembic. Migrations live in `backend/migrations/` and are +run automatically up to `head` at startup (`init_db` → `_run_migrations`). Databases +created before Alembic adoption are detected (tables present, no `alembic_version`) +and stamped to the baseline revision before upgrading, so existing data is preserved. +To create a new migration after changing the models: + +```bash +cd backend && alembic -c alembic.ini revision --autogenerate -m "describe change" +``` Per-org vector storage lives under `$DATA_DIR/storage/{organization_id}/{collection_id}/`. diff --git a/lamb-kb-server/backend/.env.example b/lamb-kb-server/backend/.env.example index 942bcdbca..3a9999484 100644 --- a/lamb-kb-server/backend/.env.example +++ b/lamb-kb-server/backend/.env.example @@ -49,3 +49,43 @@ INGESTION_TASK_TIMEOUT_SECONDS=1800 # --- Qdrant backend (only used if VECTOR_DB_QDRANT is enabled) --- # QDRANT_URL=http://localhost:6333 # QDRANT_API_KEY= + +# --------------------------------------------------------------------------- +# KG-RAG / Semantic Graph (optional) +# --------------------------------------------------------------------------- +# When KG_RAG_ENABLED=true the server mounts the /graph and /benchmarks +# routers and can extract concepts into a Neo4j graph for graph-augmented +# retrieval. When false (default) those routers are never registered and the +# server behaves as a plain vector KB — existing flows are unaffected. +# +# Requires the optional dependency group: pip install -e '.[kg-rag]' +# (the docker-compose kb-v2 service installs '.[kg-rag]' automatically). +# Needs a reachable Neo4j 5 instance — the root docker-compose `kg-rag` +# profile starts one. See ../../Documentation/kg-rag-deployment.md. + +# Master switch. +# KG_RAG_ENABLED=false + +# Run concept extraction + graph write on every ingestion job (1 OpenAI call +# per parent text). Set false to keep ingestion fast and back-populate later +# with POST /graph/collections/{id}/migrate. +# KG_RAG_INDEX_ON_INGEST=true + +# OpenAI credentials for extraction. Per-request X-OpenAI-Api-Key headers are +# preferred; this is the ingestion-time fallback. Falls back to OPENAI_API_KEY. +# KG_RAG_OPENAI_API_KEY= + +# Models. KG_RAG_EXTRACTION_MODEL falls back to KG_RAG_CHAT_MODEL / gpt-4o-mini. +# KG_RAG_CHAT_MODEL=gpt-4o-mini +# KG_RAG_EXTRACTION_MODEL= + +# Neo4j connection (also accepts the bare NEO4J_URI/USER/PASSWORD fallbacks). +# KG_RAG_NEO4J_URI=bolt://localhost:7687 +# KG_RAG_NEO4J_USER=neo4j +# KG_RAG_NEO4J_PASSWORD= + +# Graph traversal tuning (clamped: depth 1-4, limit_factor 1-20, workers 1-16). +# KG_RAG_GRAPH_DEPTH=2 +# KG_RAG_LIMIT_FACTOR=4 +# KG_RAG_EXTRACTION_MAX_WORKERS=4 +# KG_RAG_OPENAI_TIMEOUT_SECONDS=60 diff --git a/lamb-kb-server/backend/alembic.ini b/lamb-kb-server/backend/alembic.ini new file mode 100644 index 000000000..241c046a2 --- /dev/null +++ b/lamb-kb-server/backend/alembic.ini @@ -0,0 +1,48 @@ +# Alembic configuration for the LAMB KB Server. +# +# The database URL is NOT set here — env.py derives it from config.DB_PATH so +# the DATA_DIR env override (used by Docker and tests) stays authoritative. + +[alembic] +# Migration scripts live in backend/migrations (deliberately NOT named +# "alembic" so the directory cannot shadow the installed alembic package on +# sys.path). +script_location = migrations +prepend_sys_path = . + +# Use a stable, sortable file naming scheme. +file_template = %%(year)d%%(month).2d%%(day).2d_%%(rev)s_%%(slug)s + +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARNING +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARNING +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/lamb-kb-server/backend/config.py b/lamb-kb-server/backend/config.py index 5174b65c3..4156b05bd 100644 --- a/lamb-kb-server/backend/config.py +++ b/lamb-kb-server/backend/config.py @@ -9,8 +9,12 @@ plugin registration is skipped gracefully for anything that fails to import. """ +import logging import os from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) # --- Server --- HOST: str = os.getenv("HOST", "0.0.0.0") @@ -38,6 +42,14 @@ RESPLIT_CHUNK_SIZE: int = int(os.getenv("RESPLIT_CHUNK_SIZE", "4000")) RESPLIT_OVERLAP: int = 200 MAX_JOB_ATTEMPTS: int = int(os.getenv("KB_MAX_JOB_ATTEMPTS", "3")) +# How long (minutes) to retain a failed job's embedding credentials in memory +# so it can be retried without the client re-sending them. Sized to span the +# initial attempt plus the remaining retries (≈ MAX_JOB_ATTEMPTS × the task +# timeout). After this window the credentials are purged and a retry requires +# a fresh add-content request (ADR-4: credentials are never persisted). +RETRY_CACHE_TTL_MINUTES: int = int( + os.getenv("KB_RETRY_CACHE_TTL_MINUTES", "90") +) # --- Payload limits --- # Hard cap on add-content request bodies. Default 200 MB. @@ -56,6 +68,87 @@ def ensure_directories() -> None: STORAGE_DIR.mkdir(parents=True, exist_ok=True) +# --- KG-RAG (optional graph augmentation) --- +# When ``KG_RAG_ENABLED=true`` the server exposes ``/graph`` and ``/benchmarks`` +# routers and the ``kg_rag_query`` query plugin path. Per-request OpenAI +# credentials are preferred over a permanent ``KG_RAG_OPENAI_API_KEY``; the +# env var remains supported as a fallback for ingestion-time extraction when +# the API caller did not supply one. + + +def _env_bool(name: str, default: bool = False) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on", "enable", "enabled"} + + +def _env_int(name: str, default: int, minimum: int, maximum: int) -> int: + value = os.getenv(name) + if value is None or value.strip() == "": + return default + try: + parsed = int(value) + except ValueError: + logger.warning( + "Invalid integer for %s=%s. Using default %s.", name, value, default + ) + return default + return max(minimum, min(maximum, parsed)) + + +def get_kg_rag_config() -> dict[str, Any]: + """Read the KG-RAG configuration from the environment. + + KG-RAG is disabled by default. The graph pipeline is intentionally + environment-driven so existing deployments keep their current vector + ingestion/query behavior unless they opt in explicitly. + + The ``openai_api_key`` field is the *fallback* extraction key. Where the + API caller can pass a per-request credential (graph migration, benchmark, + KG-RAG query) the per-request value takes precedence (ADR-4). + """ + enabled = _env_bool("KG_RAG_ENABLED", False) + chat_model = os.getenv("KG_RAG_CHAT_MODEL") or os.getenv( + "OPENAI_CHAT_MODEL", "gpt-4o-mini" + ) + + return { + "enabled": enabled, + "index_on_ingest": _env_bool("KG_RAG_INDEX_ON_INGEST", True), + "openai_api_key": ( + os.getenv("KG_RAG_OPENAI_API_KEY") or os.getenv("OPENAI_API_KEY", "") + ), + "chat_model": chat_model, + "extraction_model": ( + os.getenv("KG_RAG_EXTRACTION_MODEL") + or os.getenv("OPENAI_EXTRACTION_MODEL") + or chat_model + ), + "neo4j_uri": os.getenv("KG_RAG_NEO4J_URI") or os.getenv("NEO4J_URI", ""), + "neo4j_user": ( + os.getenv("KG_RAG_NEO4J_USER") or os.getenv("NEO4J_USER", "neo4j") + ), + "neo4j_password": ( + os.getenv("KG_RAG_NEO4J_PASSWORD") or os.getenv("NEO4J_PASSWORD", "") + ), + "graph_depth": _env_int("KG_RAG_GRAPH_DEPTH", 2, 1, 4), + "limit_factor": _env_int("KG_RAG_LIMIT_FACTOR", 4, 1, 20), + "extraction_max_workers": _env_int( + "KG_RAG_EXTRACTION_MAX_WORKERS", 4, 1, 16 + ), + # Per-request OpenAI timeout. Default 60s is generous for the + # extraction prompt; raise it for very large chunks, lower it if + # you want ingestion to fail fast. + "openai_timeout_seconds": _env_int( + "KG_RAG_OPENAI_TIMEOUT_SECONDS", 60, 5, 600 + ), + } + + +KG_RAG_ENABLED: bool = _env_bool("KG_RAG_ENABLED", False) + + def plugin_mode(category: str, name: str) -> str: """Read the ENABLE/DISABLE mode for a plugin from environment. diff --git a/lamb-kb-server/backend/database/connection.py b/lamb-kb-server/backend/database/connection.py index 078b93946..61da96b2b 100644 --- a/lamb-kb-server/backend/database/connection.py +++ b/lamb-kb-server/backend/database/connection.py @@ -1,19 +1,19 @@ """Database connection management. Provides a single engine and session factory for the KB Server's SQLite -metadata database. All tables are created on first call to ``init_db``. +metadata database. The schema is brought to ``head`` with Alembic on the +first call to ``init_db`` (see ``_run_migrations``). """ import logging from collections.abc import Generator +from pathlib import Path from config import DB_PATH -from sqlalchemy import create_engine, event +from sqlalchemy import create_engine, event, inspect from sqlalchemy.engine import Engine from sqlalchemy.orm import Session, sessionmaker -from database.models import Base - logger = logging.getLogger(__name__) _engine: Engine | None = None @@ -31,14 +31,42 @@ def _enable_sqlite_wal(dbapi_conn, _connection_record) -> None: # noqa: ANN001 _lock_file = None +def _run_migrations() -> None: + """Bring the database schema to ``head`` using Alembic. + + Databases created by the historical ``create_all`` path already have the + tables but no ``alembic_version`` row. Those are stamped to the baseline + revision first so the baseline migration is not re-applied on top of the + existing schema; any later revisions then run normally. + """ + from alembic import command # noqa: PLC0415 + from alembic.config import Config # noqa: PLC0415 + from alembic.script import ScriptDirectory # noqa: PLC0415 + + backend_dir = Path(__file__).resolve().parents[1] + cfg = Config(str(backend_dir / "alembic.ini")) + cfg.set_main_option("script_location", str(backend_dir / "migrations")) + + inspector = inspect(_engine) + has_schema = inspector.has_table("collections") + has_version = inspector.has_table("alembic_version") + if has_schema and not has_version: + base_rev = ScriptDirectory.from_config(cfg).get_base() + command.stamp(cfg, base_rev) + logger.info("Stamped pre-Alembic database to baseline revision %s", base_rev) + + command.upgrade(cfg, "head") + logger.info("Database schema migrated to head") + + def init_db() -> None: - """Create the engine, enable SQLite optimizations, and create all tables. + """Create the engine, enable SQLite optimizations, and migrate the schema. Acquires an exclusive file lock on the data directory to prevent two - instances from running against the same storage simultaneously. + instances from running against the same storage simultaneously, then runs + Alembic migrations up to ``head``. - Safe to call multiple times — tables are created only if they do not - already exist. + Safe to call multiple times. Raises: RuntimeError: If another instance holds the lock. @@ -71,13 +99,85 @@ def init_db() -> None: event.listen(_engine, "connect", _enable_sqlite_wal) - Base.metadata.create_all(bind=_engine) + # Bring the schema to head with Alembic, then apply the KG-RAG columns + # via the idempotent lightweight path. Those columns are declared on the + # model but are not yet part of the Alembic baseline, so we add them here + # for fresh DBs; the migration skips any column that already exists, so it + # is also safe on installations created by the historical create_all path. + _run_migrations() + _run_lightweight_migrations(_engine) _SessionLocal = sessionmaker(bind=_engine, expire_on_commit=False) logger.info("Database initialized at %s", DB_PATH) +def _run_lightweight_migrations(engine: Engine) -> None: + """Apply forward-compatible ALTER TABLEs that ``create_all`` cannot do. + + ``Base.metadata.create_all`` only creates missing *tables*; it never + adds missing *columns* to existing tables. When this repo evolves the + schema with a new column on an already-populated DB, we apply the + change here so existing deployments don't crash on the next query. + + Add new entries below as additive, idempotent ``ALTER TABLE ADD + COLUMN`` statements; never delete data or change types here — those + need a real migration tool. + """ + additions = [ + # (table, column, ddl-snippet) + ( + "collections", + "graph_enabled", + "INTEGER NOT NULL DEFAULT 0", + ), + ( + "collections", + "extraction_vendor", + "TEXT", + ), + ( + "collections", + "extraction_model", + "TEXT", + ), + ( + "collections", + "extraction_endpoint", + "TEXT", + ), + ] + # Use a direct sqlite3 connection rather than the SQLAlchemy engine. + # The engine's pool keeps the underlying connection around between + # calls (with WAL state intact), and that lingering state has been + # observed to interfere with fork-based tests that re-acquire the + # data-directory lock. A short-lived ``sqlite3.connect`` opened and + # closed entirely inside this function sidesteps that. + import sqlite3 # noqa: PLC0415 + + db_url = str(engine.url) + if not db_url.startswith("sqlite:///"): + return # Only SQLite needs this hand-rolled path right now. + db_file = db_url[len("sqlite:///") :] + conn = sqlite3.connect(db_file) + try: + for table, column, ddl in additions: + cur = conn.execute(f"PRAGMA table_info({table})") + existing = {row[1] for row in cur.fetchall()} + if column in existing: + continue + conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {ddl}") + conn.commit() + logger.info( + "Schema migration applied: ALTER TABLE %s ADD COLUMN %s %s", + table, + column, + ddl, + ) + finally: + conn.close() + + def get_session() -> Generator[Session, None, None]: """Yield a SQLAlchemy session and ensure it is closed afterward. diff --git a/lamb-kb-server/backend/database/models.py b/lamb-kb-server/backend/database/models.py index fa29ed1e7..6076a4f9e 100644 --- a/lamb-kb-server/backend/database/models.py +++ b/lamb-kb-server/backend/database/models.py @@ -16,6 +16,7 @@ from datetime import UTC, datetime from sqlalchemy import ( + Boolean, Column, DateTime, Index, @@ -70,6 +71,23 @@ class Collection(Base): # Relative path (under STORAGE_DIR) for this collection's persistent data. storage_path = Column(String, nullable=False) + # --- Optional KG-RAG graph augmentation --- + # Per-collection opt-in: when True and ``KG_RAG_ENABLED`` is set at the + # server level, ingestion also indexes extracted concepts/relations into + # Neo4j (services/graph_store.py). Default False keeps existing flows + # untouched. + graph_enabled = Column(Boolean, nullable=False, default=False) + + # Locked LLM extraction config (only meaningful when graph_enabled=true). + # Defaults to NULL on collections that pre-date the field; the extractor + # falls back to the server-level ``KG_RAG_EXTRACTION_MODEL`` env in that + # case. Like chunking/embedding/vector-DB these are immutable after + # creation: changing them mid-flight would produce a graph indexed with + # one vendor's notion of an entity and queried against another's. + extraction_vendor = Column(String, nullable=True) + extraction_model = Column(String, nullable=True) + extraction_endpoint = Column(String, nullable=True) + # --- Status tracking --- status = Column(String, nullable=False, default="ready") # ready / error diff --git a/lamb-kb-server/backend/main.py b/lamb-kb-server/backend/main.py index 8965f5a24..3303ee0ce 100644 --- a/lamb-kb-server/backend/main.py +++ b/lamb-kb-server/backend/main.py @@ -33,7 +33,7 @@ # --- Startup checks --- -if not config.LAMB_API_TOKEN: +if not config.LAMB_API_TOKEN: # pragma: no cover - import-time startup guard logger.critical("LAMB_API_TOKEN is not set. Refusing to start.") sys.exit(1) @@ -101,6 +101,17 @@ async def log_requests(request: Request, call_next: Callable): app.include_router(query.router) app.include_router(jobs.router) +# KG-RAG graph + benchmark routers are mounted only when the feature flag +# is set. Keeping them off the OpenAPI surface by default means callers +# see a clean 404 rather than a 503 on every endpoint. +if config.KG_RAG_ENABLED: # pragma: no cover - import-time, env-gated router wiring + from routers import benchmarks as _bench_router # noqa: PLC0415 + from routers import graph as _graph_router # noqa: PLC0415 + + app.include_router(_graph_router.router) + app.include_router(_bench_router.router) + logger.info("KG-RAG enabled: mounted /graph and /benchmarks routers") + # --- Plugin discovery --- def _discover_plugins(package: str = "plugins") -> None: diff --git a/lamb-kb-server/backend/migrations/env.py b/lamb-kb-server/backend/migrations/env.py new file mode 100644 index 000000000..43eb23c86 --- /dev/null +++ b/lamb-kb-server/backend/migrations/env.py @@ -0,0 +1,83 @@ +"""Alembic migration environment for the LAMB KB Server. + +Derives the database URL from ``config.DB_PATH`` (so the ``DATA_DIR`` env +override stays authoritative), targets ``database.models.Base.metadata`` for +autogenerate, and enables SQLite batch mode + WAL/foreign-key pragmas so +migrations behave identically to the running application. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +from alembic import context +from sqlalchemy import create_engine, event, pool + +# Make the backend package importable when Alembic runs from the CLI. +_BACKEND_DIR = Path(__file__).resolve().parents[1] +if str(_BACKEND_DIR) not in sys.path: + sys.path.insert(0, str(_BACKEND_DIR)) + +from config import DB_PATH # noqa: E402 +from database.models import Base # noqa: E402 + +config = context.config + +# NB: deliberately not calling logging.config.fileConfig here — migrations run +# in-process at startup (see database.connection._run_migrations) and we must +# not reconfigure the application's logging. + +target_metadata = Base.metadata + +_DB_URL = f"sqlite:///{DB_PATH}" + + +def _enable_sqlite_pragmas(dbapi_conn, _record) -> None: # noqa: ANN001 + """Match the application's per-connection SQLite pragmas.""" + cursor = dbapi_conn.cursor() + cursor.execute("PRAGMA journal_mode=WAL") + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode (emit SQL without a DBAPI connection).""" + context.configure( + url=_DB_URL, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + render_as_batch=True, + ) + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations in 'online' mode against a live SQLite connection.""" + engine = create_engine( + _DB_URL, + poolclass=pool.NullPool, + connect_args={"check_same_thread": False}, + ) + event.listen(engine, "connect", _enable_sqlite_pragmas) + + with engine.connect() as connection: + context.configure( + connection=connection, + target_metadata=target_metadata, + # SQLite cannot ALTER most things in place; batch mode rebuilds + # tables via copy-and-swap so future migrations work. + render_as_batch=True, + ) + with context.begin_transaction(): + context.run_migrations() + + engine.dispose() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/lamb-kb-server/backend/migrations/script.py.mako b/lamb-kb-server/backend/migrations/script.py.mako new file mode 100644 index 000000000..5716ac760 --- /dev/null +++ b/lamb-kb-server/backend/migrations/script.py.mako @@ -0,0 +1,25 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} +""" +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: str | None = ${repr(down_revision)} +branch_labels: str | Sequence[str] | None = ${repr(branch_labels)} +depends_on: str | Sequence[str] | None = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/lamb-kb-server/backend/migrations/versions/20260624_4d0c6668e3b9_baseline_schema.py b/lamb-kb-server/backend/migrations/versions/20260624_4d0c6668e3b9_baseline_schema.py new file mode 100644 index 000000000..85ea0eb1b --- /dev/null +++ b/lamb-kb-server/backend/migrations/versions/20260624_4d0c6668e3b9_baseline_schema.py @@ -0,0 +1,86 @@ +"""baseline schema + +Revision ID: 4d0c6668e3b9 +Revises: +Create Date: 2026-06-24 12:32:39.522458 +""" +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '4d0c6668e3b9' +down_revision: str | None = None +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('collections', + sa.Column('id', sa.String(), nullable=False), + sa.Column('organization_id', sa.String(), nullable=False), + sa.Column('name', sa.String(), nullable=False), + sa.Column('description', sa.Text(), nullable=True), + sa.Column('chunking_strategy', sa.String(), nullable=False), + sa.Column('chunking_params', sa.Text(), nullable=True), + sa.Column('embedding_vendor', sa.String(), nullable=False), + sa.Column('embedding_model', sa.String(), nullable=False), + sa.Column('embedding_endpoint', sa.String(), nullable=True), + sa.Column('vector_db_backend', sa.String(), nullable=False), + sa.Column('backend_collection_id', sa.String(), nullable=True), + sa.Column('storage_path', sa.String(), nullable=False), + sa.Column('status', sa.String(), nullable=False), + sa.Column('error_message', sa.Text(), nullable=True), + sa.Column('document_count', sa.Integer(), nullable=False), + sa.Column('chunk_count', sa.Integer(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.Column('updated_at', sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('organization_id', 'name', name='uq_collection_org_name') + ) + with op.batch_alter_table('collections', schema=None) as batch_op: + batch_op.create_index('idx_collections_org', ['organization_id'], unique=False) + batch_op.create_index('idx_collections_status', ['status'], unique=False) + + op.create_table('ingestion_jobs', + sa.Column('id', sa.String(), nullable=False), + sa.Column('collection_id', sa.String(), nullable=False), + sa.Column('organization_id', sa.String(), nullable=False), + sa.Column('documents_json', sa.Text(), nullable=False), + sa.Column('documents_total', sa.Integer(), nullable=False), + sa.Column('documents_processed', sa.Integer(), nullable=False), + sa.Column('chunks_created', sa.Integer(), nullable=False), + sa.Column('status', sa.String(), nullable=False), + sa.Column('error_message', sa.Text(), nullable=True), + sa.Column('attempts', sa.Integer(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.Column('updated_at', sa.DateTime(), nullable=False), + sa.Column('started_at', sa.DateTime(), nullable=True), + sa.Column('completed_at', sa.DateTime(), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + with op.batch_alter_table('ingestion_jobs', schema=None) as batch_op: + batch_op.create_index('idx_ingestion_jobs_collection', ['collection_id'], unique=False) + batch_op.create_index('idx_ingestion_jobs_status', ['status'], unique=False) + batch_op.create_index('idx_ingestion_jobs_status_created', ['status', 'created_at'], unique=False) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('ingestion_jobs', schema=None) as batch_op: + batch_op.drop_index('idx_ingestion_jobs_status_created') + batch_op.drop_index('idx_ingestion_jobs_status') + batch_op.drop_index('idx_ingestion_jobs_collection') + + op.drop_table('ingestion_jobs') + with op.batch_alter_table('collections', schema=None) as batch_op: + batch_op.drop_index('idx_collections_status') + batch_op.drop_index('idx_collections_org') + + op.drop_table('collections') + # ### end Alembic commands ### diff --git a/lamb-kb-server/backend/plugins/base.py b/lamb-kb-server/backend/plugins/base.py index 2a1ddc937..bf52a9dc4 100644 --- a/lamb-kb-server/backend/plugins/base.py +++ b/lamb-kb-server/backend/plugins/base.py @@ -178,6 +178,57 @@ def query( ) -> list[QueryResult]: """Embed ``query_text`` and return the top ``top_k`` similar chunks.""" + def get_chunks_by_id( + self, + *, + collection_id: str, + storage_path: str, + chunk_ids: list[str], + embedding_function: EmbeddingFunction, + ) -> list[QueryResult]: + """Fetch chunks by their backend-side ID without similarity scoring. + + Used by KG-RAG to materialize chunks discovered through graph + traversal. The default implementation returns an empty list — a + backend that supports ID lookup (ChromaDB) overrides this. + Callers must treat an empty return as "lookup not supported" and + degrade gracefully. + """ + return [] + + def get_chunks_by_source( + self, + *, + collection_id: str, + storage_path: str, + source_item_id: str, + embedding_function: EmbeddingFunction, + ) -> list[QueryResult]: + """Fetch every chunk produced from one source item. + + Used by the ingestion-time graph indexer to grab the chunks it + just inserted (so it can run extraction with stable IDs). Default + returns an empty list — a backend that supports metadata-where + filters (ChromaDB) overrides this. + """ + return [] + + def iter_all_chunks( + self, + *, + collection_id: str, + storage_path: str, + embedding_function: EmbeddingFunction, + batch_size: int = 500, + ): + """Yield every stored chunk in batches, for migration / re-indexing. + + Default raises ``NotImplementedError`` — a backend that supports + cursor-style scrolling (ChromaDB's ``get`` with offset) overrides + this. The graph-migration route checks for support before calling. + """ + raise NotImplementedError + def get_parameters(self) -> list[PluginParameter]: """Return the backend-specific configuration schema (usually empty).""" return [] @@ -248,6 +299,58 @@ def get_parameters(self) -> list[PluginParameter]: return [] +class LLMExtractionFunction(abc.ABC): + """Abstract base for chat-completion vendors used by KG-RAG extraction. + + Extractors are constructed fresh per ingestion job because credentials + are request-scoped (same pattern as embeddings — ADR-4). Each + implementation wraps a vendor's chat-completion call and returns a + parsed JSON object that the concept extractor consumes. + """ + + name: str = "base" + description: str = "Base LLM extraction backend" + + def __init__( + self, + *, + model: str, + api_key: str = "", + api_endpoint: str = "", + timeout_seconds: float = 60.0, + ) -> None: + self.model = model + self.api_key = api_key + self.api_endpoint = api_endpoint + self.timeout_seconds = timeout_seconds + + @abc.abstractmethod + def chat_json( + self, + *, + system: str, + user: str, + fallback_model: str | None = None, + ) -> dict[str, Any]: + """Run a chat-completion with JSON output. + + Args: + system: System prompt. + user: User message (already JSON-encoded by the caller). + fallback_model: Optional model to retry with if the primary + model rejects the JSON-mode request (some smaller / older + models don't support ``response_format=json_object``). + + Returns: + Parsed dict from the model's JSON output, or an empty dict + shape on parse / API failure. + """ + + def get_parameters(self) -> list[PluginParameter]: + """Return the parameter schema for this vendor (model, endpoint).""" + return [] + + # --------------------------------------------------------------------------- # Registry infrastructure # --------------------------------------------------------------------------- @@ -392,3 +495,33 @@ def build( if plugin_class is None: raise ValueError(f"Embedding vendor '{name}' is not registered.") return plugin_class(model=model, api_key=api_key, api_endpoint=api_endpoint) + + +class LLMExtractionRegistry(_BaseRegistry): + category = "LLM_EXTRACTION" + _plugins: dict[str, type[LLMExtractionFunction]] = {} + + @classmethod + def build( + cls, + name: str, + *, + model: str, + api_key: str = "", + api_endpoint: str = "", + timeout_seconds: float = 60.0, + ) -> LLMExtractionFunction: + """Construct an LLM extraction backend for ``name`` with credentials. + + Raises: + ValueError: If the vendor is not registered. + """ + plugin_class = cls._plugins.get(name) + if plugin_class is None: + raise ValueError(f"LLM extraction vendor '{name}' is not registered.") + return plugin_class( + model=model, + api_key=api_key, + api_endpoint=api_endpoint, + timeout_seconds=timeout_seconds, + ) diff --git a/lamb-kb-server/backend/plugins/embedding/local.py b/lamb-kb-server/backend/plugins/embedding/local.py index ec948fb45..646d1e616 100644 --- a/lamb-kb-server/backend/plugins/embedding/local.py +++ b/lamb-kb-server/backend/plugins/embedding/local.py @@ -13,9 +13,14 @@ from __future__ import annotations import logging +import os + +# Prevent huggingface_hub from trying to read /root/.cache/huggingface/token +# when running as a non-root user inside the container. +os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1") # Guard import: if sentence_transformers is missing the module will raise -# ImportError and the plugin will be skipped by the discovery mechanism. +# ImportError at import time so the plugin will be skipped by the discovery mechanism. try: from sentence_transformers import SentenceTransformer except ImportError: diff --git a/lamb-kb-server/backend/plugins/kg_rag_query.py b/lamb-kb-server/backend/plugins/kg_rag_query.py new file mode 100644 index 000000000..0bac1297f --- /dev/null +++ b/lamb-kb-server/backend/plugins/kg_rag_query.py @@ -0,0 +1,428 @@ +"""KG-RAG query augmentation: vector seed retrieval + Neo4j graph expansion. + +In the legacy KB server this was a top-level ``QueryPlugin`` registered on +the monolithic ``PluginRegistry``. The new KB server's plugin registries +(``VectorDBRegistry``, ``ChunkingRegistry``, ``EmbeddingRegistry``) own +ingestion-time plugins only; query-time augmentation here is invoked +directly from :mod:`services.query_service` when the caller selects +``plugin_name='kg_rag_query'`` (see ``query_with_plugin``). + +The augmentation is a no-op fallback when: + +* ``KG_RAG_ENABLED`` is false at the server level, or +* the collection's ``graph_enabled`` flag is false, or +* Neo4j is not configured / reachable, or +* the baseline vector search returned no seeds. + +In every failure path the baseline vector results are returned unchanged +with a ``warnings`` entry attached to the trace. +""" + +from __future__ import annotations + +import logging +import time +from typing import Any + +from database.models import Collection +from plugins.base import EmbeddingFunction, VectorDBBackend + +logger = logging.getLogger(__name__) + + +class KGRAGQueryPlugin: + """Augment vector retrieval with Neo4j graph expansion.""" + + name = "kg_rag_query" + description = "KG-RAG query with vector seed retrieval and graph expansion" + + def augment( + self, + *, + db: Any, + collection: Collection, + backend: VectorDBBackend, + embedding_function: EmbeddingFunction, + query_text: str, + baseline_results: list[dict[str, Any]], + params: dict[str, Any], + ) -> list[dict[str, Any]]: + import config as config_module + + top_k = int(params.get("top_k", 5) or 5) + return_parent_context = self._as_bool(params.get("return_parent_context", True)) + include_trace = self._as_bool(params.get("include_trace", True)) + + kg_config = config_module.get_kg_rag_config() + graph_depth = int( + params.get("graph_depth") or kg_config.get("graph_depth") or 2 + ) + graph_depth = max(1, min(graph_depth, 4)) + + seed_chunk_ids = [ + cid + for cid in (self._result_chunk_id(result) for result in baseline_results) + if cid + ] + + trace: dict[str, Any] = { + "mode": "kg_rag", + "enabled": bool(kg_config.get("enabled")), + "graph_expanded": False, + "seed_chunk_ids": seed_chunk_ids, + "entry_concepts": [], + "expanded_chunk_ids": [], + "vector_latency_ms": 0.0, + "graph_latency_ms": 0.0, + "warnings": [], + } + + if not kg_config.get("enabled"): + trace["warnings"].append("KG-RAG is disabled; returning vector baseline") + return self._attach_trace(baseline_results, trace, include_trace) + if not getattr(collection, "graph_enabled", False): + trace["warnings"].append( + "Collection has graph_enabled=false; returning vector baseline" + ) + return self._attach_trace(baseline_results, trace, include_trace) + if not seed_chunk_ids: + trace["warnings"].append( + "No vector seed chunks found; graph expansion skipped" + ) + return self._attach_trace(baseline_results, trace, include_trace) + + # Defer the graph import until we actually need it so the plugin + # module stays loadable when the optional ``neo4j`` package is + # missing. + from services.graph_store import get_graph_store + + graph_store = get_graph_store() + if not graph_store.is_configured(): + trace["warnings"].append( + "Neo4j is not configured; returning vector baseline" + ) + return self._attach_trace(baseline_results, trace, include_trace) + + # Question-entity seeding: extract named-entity-like tokens from + # the question itself and let the graph contribute chunks that + # MENTION them, even if vector retrieval missed those chunks. + # This is the ``local search'' pattern from Microsoft GraphRAG: + # we start from the concepts that match the question's entities + # and follow RELATES_TO edges to a bounded depth. + graph_start = time.perf_counter() + question_entities = self._extract_question_entities(query_text, collection) + question_expansion: dict[str, Any] = {} + if question_entities: + try: + question_expansion = graph_store.expand_from_concept_names( + collection_id=str(collection.id), + org_id=str(collection.organization_id), + concept_names=question_entities, + depth=graph_depth, + limit=max(top_k, 2 * top_k), + ) + except Exception as exc: # noqa: BLE001 + logger.warning("Question-entity expansion failed: %s", exc) + trace["warnings"].append(f"Graph expansion failed: {exc}") + + merged_entry_concepts = list( + dict.fromkeys(question_expansion.get("entry_concepts", [])) + ) + merged_expanded_ids = list( + dict.fromkeys(question_expansion.get("expanded_chunk_ids", [])) + ) + + trace.update( + { + "graph_expanded": bool(merged_expanded_ids), + "entry_concepts": merged_entry_concepts, + "expanded_chunk_ids": merged_expanded_ids, + "question_entities": question_entities, + "graph_latency_ms": question_expansion.get( + "graph_latency_ms", + (time.perf_counter() - graph_start) * 1000, + ), + } + ) + + expanded_ids = [ + cid for cid in merged_expanded_ids if cid not in seed_chunk_ids + ] + expanded_results = self._fetch_expanded_results( + backend=backend, + collection=collection, + embedding_function=embedding_function, + expanded_ids=expanded_ids, + return_parent_context=return_parent_context, + ) + + merged = self._rrf_merge( + baseline_results=baseline_results, + expanded_results=expanded_results, + top_k=top_k, + rrf_k=int(params.get("rrf_k", 20) or 20), + graph_weight=float(params.get("graph_weight", 0.85) or 0.85), + ) + if not expanded_results and not trace["graph_expanded"]: + trace["warnings"].append("Graph returned no additional chunks") + return self._attach_trace(merged, trace, include_trace) + + # ------------------------------------------------------------------ + # Question-entity extraction (LLM-based) + # ------------------------------------------------------------------ + # + # Pulling named entities from a free-text question is the exact kind + # of short, structured task small LLMs handle reliably. The previous + # regex-based heuristic missed lowercased multi-word entities + # (``unsupervised learning``, ``parent-child chunking``) and accepted + # any capitalized token as if it were a name. A single small-model + # call with a JSON-schema prompt is more accurate, costs cents per + # thousand queries on ``gpt-4o-mini``, and finishes in well under + # 500~ms, comparable to the embedding round-trip. + # + # The call is cached per-question text inside this process so + # benchmark re-runs and identical user queries don't pay the round + # trip twice. The cache is bounded by ``_QUESTION_CACHE_SIZE`` and + # uses simple FIFO eviction. + + _QUESTION_CACHE_SIZE = 512 + _question_cache: dict[str, list[str]] = {} + + @classmethod + def _extract_question_entities( + cls, question: str, collection: Collection + ) -> list[str]: + """Use an LLM to extract named entities from the question. + + Uses the same vendor / model / endpoint configured for the + collection's build-time extraction (``extraction_vendor`` / + ``extraction_model`` / ``extraction_endpoint``). Keeping the + question extractor in lockstep with the build extractor means + the entity names produced at query time match the canonical + form stored in the graph — different models can normalize + entity names differently, and a mismatch silently zeroes the + question-entity seeding path. + + Falls back to the server-level KG-RAG defaults (then OpenAI + gpt-4o-mini) when the collection has no per-collection + extraction config — preserves back-compat for collections + created before the picker existed. + + Returns an empty list if the backend isn't available or the + call fails — the graph expansion silently degrades to the + baseline-seeded path in that case. + """ + question = (question or "").strip() + if not question: + return [] + + import config as config_module # noqa: PLC0415 + + kg = config_module.get_kg_rag_config() + coll_vendor = getattr(collection, "extraction_vendor", None) + coll_model = getattr(collection, "extraction_model", None) + coll_endpoint = getattr(collection, "extraction_endpoint", None) + vendor = coll_vendor or "openai" + model = coll_model or kg.get("chat_model") or "gpt-4o-mini" + endpoint = coll_endpoint or "" + + # Cache keyed by (vendor, model, question) so different + # collections sharing the same question text don't poison each + # other when they use different models. + cache_key = f"{vendor}::{model}::{question}" + if cache_key in cls._question_cache: + return cls._question_cache[cache_key] + + api_key = "" + if vendor == "openai": + api_key = (kg.get("openai_api_key") or "").strip() + if not api_key: + return [] + + from plugins.base import LLMExtractionRegistry # noqa: PLC0415 + + try: + backend = LLMExtractionRegistry.build( + vendor, + model=model, + api_key=api_key, + api_endpoint=endpoint, + timeout_seconds=15.0, + ) + except ValueError as exc: + logger.debug("Question-entity backend %s unavailable: %s", vendor, exc) + return [] + + system = ( + "Extract every named entity from the user question. " + "Return STRICT JSON of the form {\"entities\": [\"...\", \"...\"]}. " + "Keep multi-word names whole. Lowercase common nouns are fine if " + "they would plausibly identify a knowledge-graph concept (e.g. " + "'parent-child chunking', 'reciprocal rank fusion'). " + "Do not output entity types, descriptions or explanations — only " + "the surface strings." + ) + + parsed = backend.chat_json(system=system, user=question) + entities_raw = parsed.get("entities") if isinstance(parsed, dict) else None + if not isinstance(entities_raw, list): + entities_raw = [] + entities = [str(e).strip() for e in entities_raw if str(e).strip()] + + # FIFO cache eviction + if len(cls._question_cache) >= cls._QUESTION_CACHE_SIZE: + try: + cls._question_cache.pop(next(iter(cls._question_cache))) + except StopIteration: # pragma: no cover - cache is non-empty here + pass + cls._question_cache[cache_key] = entities + return entities + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + @staticmethod + def _as_bool(value: Any) -> bool: + if isinstance(value, bool): + return value + if value is None: + return False + return str(value).strip().lower() in { + "1", + "true", + "yes", + "on", + "enable", + "enabled", + } + + @staticmethod + def _result_chunk_id(result: dict[str, Any]) -> str | None: + metadata = result.get("metadata") or {} + return ( + metadata.get("document_id") + or metadata.get("child_chunk_id") + or metadata.get("chunk_id") + ) + + def _fetch_expanded_results( + self, + *, + backend: VectorDBBackend, + collection: Collection, + embedding_function: EmbeddingFunction, + expanded_ids: list[str], + return_parent_context: bool, + ) -> list[dict[str, Any]]: + """Look up chunks discovered by graph expansion via the public backend API. + + Uses :meth:`VectorDBBackend.get_chunks_by_id` so each backend can + implement ID lookup in its own way. A backend that doesn't + support it returns ``[]``, and KG-RAG degrades to "graph found + related chunks but we can't pull their text" — the trace surfaces + this as a warning. + """ + if not expanded_ids: + return [] + + try: + fetched = backend.get_chunks_by_id( + collection_id=collection.backend_collection_id or str(collection.id), + storage_path=collection.storage_path, + chunk_ids=expanded_ids, + embedding_function=embedding_function, + ) + except Exception as exc: # noqa: BLE001 + logger.debug("Could not fetch expanded chunks: %s", exc) + return [] + + results: list[dict[str, Any]] = [] + for item in fetched: + metadata = dict(item.metadata or {}) + metadata["kg_rag_origin"] = "graph_expansion" + data = metadata.pop("parent_text", None) if return_parent_context else None + results.append( + { + "similarity": item.score or 0.72, + "data": data or item.text, + "metadata": metadata, + } + ) + return results + + def _rrf_merge( + self, + *, + baseline_results: list[dict[str, Any]], + expanded_results: list[dict[str, Any]], + top_k: int, + rrf_k: int = 20, + graph_weight: float = 0.85, + ) -> list[dict[str, Any]]: + """Reciprocal Rank Fusion of the vector baseline and the graph expansion. + + Standard RRF assigns each item a score of ``1/(rrf_k + rank)`` per + ranked list it appears in, and sums across lists. Items appearing + in *both* the vector list and the graph-expanded list get boosted + (the graph confirms the vector signal), which is exactly the + regime where KG-RAG should beat the baseline. + + The graph list contributes with a smaller weight (``graph_weight``) + because chunks in the expanded list have no semantic similarity + score against the query --- they were reached via concept + traversal. The weight prevents pure-graph hits from displacing + the highest-rank vector hits when the vector signal is already + unambiguous. + + Each result's ``similarity`` field is replaced with the fused + score so downstream merging / sorting remains consistent. + """ + rankings: dict[str, dict[str, Any]] = {} + + def _accumulate(items, weight, list_name): + for rank, item in enumerate(items, start=1): + chunk_id = self._result_chunk_id(item) or str(item.get("data", ""))[:120] + entry = rankings.setdefault( + chunk_id, + { + "item": item, + "score": 0.0, + "lists": set(), + }, + ) + entry["score"] += weight * (1.0 / (rrf_k + rank)) + entry["lists"].add(list_name) + if list_name == "baseline" or "item" not in entry: + entry["item"] = item + + _accumulate(baseline_results, weight=1.0, list_name="baseline") + _accumulate(expanded_results, weight=graph_weight, list_name="graph") + + ordered = sorted(rankings.values(), key=lambda e: e["score"], reverse=True) + merged: list[dict[str, Any]] = [] + for entry in ordered[: max(top_k, 1) + 3]: + item = dict(entry["item"]) + # Preserve the original similarity for trace transparency but + # surface the fused RRF score in a separate field. + item["rrf_score"] = entry["score"] + item["fusion_lists"] = sorted(entry["lists"]) + merged.append(item) + return merged + + @staticmethod + def _attach_trace( + results: list[dict[str, Any]], + trace: dict[str, Any], + include_trace: bool, + ) -> list[dict[str, Any]]: + if not include_trace: + return results + traced_results: list[dict[str, Any]] = [] + for result in results: + item = dict(result) + metadata = dict(item.get("metadata") or {}) + metadata["kg_rag"] = trace + item["metadata"] = metadata + traced_results.append(item) + return traced_results diff --git a/lamb-kb-server/backend/plugins/llm_extraction/__init__.py b/lamb-kb-server/backend/plugins/llm_extraction/__init__.py new file mode 100644 index 000000000..ad1517e84 --- /dev/null +++ b/lamb-kb-server/backend/plugins/llm_extraction/__init__.py @@ -0,0 +1,6 @@ +"""LLM extraction backends for KG-RAG concept/relation extraction. + +Each backend wraps a vendor's chat-completion endpoint and returns a +parsed JSON object. Backends are enabled/disabled via tri-state env +vars (``LLM_EXTRACTION_OPENAI=DISABLE`` / ``LLM_EXTRACTION_OLLAMA=DISABLE``). +""" diff --git a/lamb-kb-server/backend/plugins/llm_extraction/ollama.py b/lamb-kb-server/backend/plugins/llm_extraction/ollama.py new file mode 100644 index 000000000..2c9f1bca6 --- /dev/null +++ b/lamb-kb-server/backend/plugins/llm_extraction/ollama.py @@ -0,0 +1,144 @@ +"""Ollama chat-completion backend for KG-RAG concept extraction. + +Talks to a local Ollama daemon (default http://localhost:11434) via the +official ``ollama`` Python SDK. Sends ``format='json'`` so the model is +required to emit a JSON object. + +No API key is needed for a default Ollama install; if a deployment puts +Ollama behind an auth proxy, the per-request token can be passed via +``api_key`` and is forwarded as a Bearer header. +""" + +from __future__ import annotations + +import json +import logging +import os +from typing import Any + +from plugins.base import ( + LLMExtractionFunction, + LLMExtractionRegistry, + PluginParameter, +) + +logger = logging.getLogger(__name__) + + +@LLMExtractionRegistry.register +class OllamaExtraction(LLMExtractionFunction): + """Ollama chat-completion extraction backend (local models).""" + + name = "ollama" + description = "Ollama local chat-completion extraction (llama3, qwen2, ...)" + + def __init__( + self, + *, + model: str = "llama3.1:8b", + api_key: str = "", + api_endpoint: str = "", + timeout_seconds: float = 60.0, + ) -> None: + super().__init__( + model=model, + api_key=api_key, + api_endpoint=api_endpoint, + timeout_seconds=timeout_seconds, + ) + self._model = model or "llama3.1:8b" + + def chat_json( + self, + *, + system: str, + user: str, + fallback_model: str | None = None, + ) -> dict[str, Any]: + try: + from ollama import Client # noqa: PLC0415 + except ImportError: + logger.warning("Ollama SDK is not installed") + return {} + + host = ( + self.api_endpoint + or os.getenv("OLLAMA_HOST") + or "http://localhost:11434" + ) + headers: dict[str, str] = {} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + + client = Client(host=host, timeout=self.timeout_seconds, headers=headers) + + messages = [ + {"role": "system", "content": system}, + {"role": "user", "content": user}, + ] + + try: + response = client.chat( + model=self._model, + messages=messages, + format="json", + options={"temperature": 0}, + ) + except Exception as exc: # noqa: BLE001 + if fallback_model and fallback_model != self._model: + logger.warning( + "Ollama extraction with %s failed (%s); retrying with %s", + self._model, + exc, + fallback_model, + ) + try: + response = client.chat( + model=fallback_model, + messages=messages, + format="json", + options={"temperature": 0}, + ) + except Exception as exc2: # noqa: BLE001 + logger.warning("Ollama extraction fallback failed: %s", exc2) + return {} + else: + logger.warning("Ollama extraction failed: %s", exc) + return {} + + content = response.get("message", {}).get("content", "{}") + try: + parsed = json.loads(content) + except json.JSONDecodeError as exc: + logger.warning("Ollama extraction returned invalid JSON: %s", exc) + return {} + return parsed if isinstance(parsed, dict) else {} + + @classmethod + def class_parameters(cls) -> list[PluginParameter]: + return [ + PluginParameter( + "model", + "string", + "Ollama model tag (must be pulled on the Ollama host)", + "llama3.1:8b", + choices=[ + "llama3.1:8b", + "llama3.1:70b", + "llama3.2:3b", + "llama3.3:70b", + "qwen2.5:7b", + "qwen2.5:14b", + "mistral:7b", + "mixtral:8x7b", + "phi3:medium", + "gemma2:9b", + ], + ), + PluginParameter( + "api_endpoint", + "string", + "Ollama host URL (leave empty for http://localhost:11434)", + "", + ), + ] diff --git a/lamb-kb-server/backend/plugins/llm_extraction/openai.py b/lamb-kb-server/backend/plugins/llm_extraction/openai.py new file mode 100644 index 000000000..83402832d --- /dev/null +++ b/lamb-kb-server/backend/plugins/llm_extraction/openai.py @@ -0,0 +1,148 @@ +"""OpenAI chat-completion backend for KG-RAG concept extraction. + +Uses the openai v1 SDK directly with ``response_format=json_object`` to +extract structured entity / relationship lists from text chunks. Falls +back to a secondary model when the primary refuses JSON mode (older +models or third-party OpenAI-compatible endpoints). + +API keys are passed per-request (ADR-4). Falls back to +``KG_RAG_OPENAI_API_KEY`` / ``OPENAI_API_KEY`` env var if none is provided +at call time. +""" + +from __future__ import annotations + +import json +import logging +import os +from typing import Any + +from plugins.base import ( + LLMExtractionFunction, + LLMExtractionRegistry, + PluginParameter, +) + +logger = logging.getLogger(__name__) + + +@LLMExtractionRegistry.register +class OpenAIExtraction(LLMExtractionFunction): + """OpenAI (or OpenAI-compatible) chat-completion extraction backend.""" + + name = "openai" + description = "OpenAI chat-completion extraction (gpt-4o-mini, gpt-4o, ...)" + + def __init__( + self, + *, + model: str = "gpt-4o-mini", + api_key: str = "", + api_endpoint: str = "", + timeout_seconds: float = 60.0, + ) -> None: + super().__init__( + model=model, + api_key=api_key, + api_endpoint=api_endpoint, + timeout_seconds=timeout_seconds, + ) + self._model = model or "gpt-4o-mini" + + def chat_json( + self, + *, + system: str, + user: str, + fallback_model: str | None = None, + ) -> dict[str, Any]: + try: + from openai import OpenAI # noqa: PLC0415 + except ImportError: + logger.warning( + "OpenAI SDK is not installed; install kb-server with [kg-rag]" + ) + return {} + + resolved_key = ( + self.api_key + or os.getenv("KG_RAG_OPENAI_API_KEY") + or os.getenv("OPENAI_API_KEY", "") + ) + if not resolved_key: + logger.warning("OpenAI extraction: no API key configured") + return {} + + kwargs: dict[str, Any] = { + "api_key": resolved_key, + "timeout": self.timeout_seconds, + } + if self.api_endpoint: + kwargs["base_url"] = self.api_endpoint.rstrip("/") + + client = OpenAI(**kwargs) + messages = [ + {"role": "system", "content": system}, + {"role": "user", "content": user}, + ] + + try: + response = client.chat.completions.create( + model=self._model, + messages=messages, + response_format={"type": "json_object"}, + temperature=0, + ) + except Exception as exc: # noqa: BLE001 + if fallback_model and fallback_model != self._model: + logger.warning( + "OpenAI extraction with %s failed (%s); retrying with %s", + self._model, + exc, + fallback_model, + ) + try: + response = client.chat.completions.create( + model=fallback_model, + messages=messages, + response_format={"type": "json_object"}, + temperature=0, + ) + except Exception as exc2: # noqa: BLE001 + logger.warning("OpenAI extraction fallback failed: %s", exc2) + return {} + else: + logger.warning("OpenAI extraction failed: %s", exc) + return {} + + content = response.choices[0].message.content or "{}" + try: + parsed = json.loads(content) + except json.JSONDecodeError as exc: + logger.warning("OpenAI extraction returned invalid JSON: %s", exc) + return {} + return parsed if isinstance(parsed, dict) else {} + + @classmethod + def class_parameters(cls) -> list[PluginParameter]: + return [ + PluginParameter( + "model", + "string", + "Chat-completion model name", + "gpt-4o-mini", + choices=[ + "gpt-4o-mini", + "gpt-4o", + "gpt-4-turbo", + "gpt-5-nano", + "gpt-3.5-turbo", + ], + ), + PluginParameter( + "api_endpoint", + "string", + "Custom OpenAI-compatible base URL (leave empty for api.openai.com)", + "", + ), + ] diff --git a/lamb-kb-server/backend/plugins/vector_db/chromadb_backend.py b/lamb-kb-server/backend/plugins/vector_db/chromadb_backend.py index 48085c803..54495c5b9 100644 --- a/lamb-kb-server/backend/plugins/vector_db/chromadb_backend.py +++ b/lamb-kb-server/backend/plugins/vector_db/chromadb_backend.py @@ -13,6 +13,7 @@ from __future__ import annotations import logging +import os import re import shutil from typing import Any @@ -130,6 +131,19 @@ def delete_collection(self, *, collection_id: str, storage_path: str) -> None: collection is already gone (idempotent). Then removes the storage directory entirely. """ + # Idempotent fast path: if there's no cached client AND the storage + # directory is already gone, the collection has been fully deleted. + # Returning here avoids re-creating a PersistentClient against a path + # that was just ``rmtree``-d — ChromaDB's process-global system cache + # would otherwise hand back a stale system pointing at the deleted + # sqlite file, surfacing as a spurious "disk I/O error". + if storage_path not in _clients and not os.path.isdir(storage_path): + logger.debug( + "ChromaDB delete_collection: '%s' storage already absent, skipping", + collection_id, + ) + return + client = _get_client(storage_path) try: client.delete_collection(name=collection_id) @@ -228,6 +242,127 @@ def delete_by_source( return count + def get_chunks_by_id( + self, + *, + collection_id: str, + storage_path: str, + chunk_ids: list[str], + embedding_function: EmbeddingFunction, + ) -> list[QueryResult]: + """Return chunks for ``chunk_ids`` in the order they were requested. + + Used by KG-RAG to materialize chunks discovered through graph + traversal. Score is set to a constant 0.72 sentinel so callers can + tell graph-sourced results apart from real similarity hits without + having to inspect metadata. + """ + if not chunk_ids: + return [] + client = _get_client(storage_path) + try: + collection = client.get_collection( + name=collection_id, + embedding_function=_to_chroma_ef(embedding_function), # type: ignore[arg-type] + ) + rows = collection.get( + ids=chunk_ids, + include=["documents", "metadatas"], + ) + except Exception as exc: # noqa: BLE001 + logger.debug("ChromaDB get_chunks_by_id failed: %s", exc) + return [] + + ids = rows.get("ids") or chunk_ids + documents = rows.get("documents") or [] + metadatas = rows.get("metadatas") or [] + results: list[QueryResult] = [] + for index, chunk_id in enumerate(ids): + document = documents[index] if index < len(documents) else "" + metadata: dict[str, Any] = ( + dict(metadatas[index] or {}) if index < len(metadatas) else {} + ) + metadata.setdefault("document_id", chunk_id) + text = metadata.pop("parent_text", None) or document + results.append(QueryResult(text=text, score=0.72, metadata=metadata)) + return results + + def get_chunks_by_source( + self, + *, + collection_id: str, + storage_path: str, + source_item_id: str, + embedding_function: EmbeddingFunction, + ) -> list[QueryResult]: + """Return every chunk whose ``source_item_id`` matches. + + Score is the 0.72 sentinel (these are exact lookups, not + similarity results). Empty list when no chunks are found. + """ + client = _get_client(storage_path) + try: + collection = client.get_collection( + name=collection_id, + embedding_function=_to_chroma_ef(embedding_function), # type: ignore[arg-type] + ) + rows = collection.get( + where={"source_item_id": source_item_id}, + include=["documents", "metadatas"], + ) + except Exception as exc: # noqa: BLE001 + logger.debug("ChromaDB get_chunks_by_source failed: %s", exc) + return [] + + ids = rows.get("ids") or [] + documents = rows.get("documents") or [] + metadatas = rows.get("metadatas") or [] + results: list[QueryResult] = [] + for index, chunk_id in enumerate(ids): + document = documents[index] if index < len(documents) else "" + metadata: dict[str, Any] = ( + dict(metadatas[index] or {}) if index < len(metadatas) else {} + ) + metadata.setdefault("chunk_id", chunk_id) + results.append(QueryResult(text=document, score=0.72, metadata=metadata)) + return results + + def iter_all_chunks( + self, + *, + collection_id: str, + storage_path: str, + embedding_function: EmbeddingFunction, + batch_size: int = 500, + ): + """Yield ``(ids, texts, metadatas)`` tuples until the collection is drained. + + Wraps ChromaDB's ``get(limit, offset)`` paginator so the graph + migration route can iterate without each caller re-implementing + the offset bookkeeping. + """ + client = _get_client(storage_path) + chroma_collection = client.get_collection( + name=collection_id, + embedding_function=_to_chroma_ef(embedding_function), # type: ignore[arg-type] + ) + offset = 0 + while True: + result = chroma_collection.get( + include=["documents", "metadatas"], + limit=batch_size, + offset=offset, + ) + ids = result.get("ids") or [] + if not ids: + return + yield ( + list(ids), + list(result.get("documents") or []), + [dict(m or {}) for m in (result.get("metadatas") or [])], + ) + offset += len(ids) + def query( self, *, @@ -261,12 +396,20 @@ def query( ) results: list[QueryResult] = [] + ids = (raw.get("ids") or [[]])[0] documents = (raw.get("documents") or [[]])[0] metadatas = (raw.get("metadatas") or [[]])[0] distances = (raw.get("distances") or [[]])[0] - for doc, meta, dist in zip(documents, metadatas, distances): + for idx, (doc, meta, dist) in enumerate(zip(documents, metadatas, distances)): meta_dict: dict[str, Any] = dict(meta) if meta else {} + # Propagate the backend's internal chunk ID so downstream + # consumers (KG-RAG plugin in particular) can look the chunk + # back up by ID. ChromaDB always returns ``ids`` alongside + # documents — we surface it as ``chunk_id`` for the plugin's + # seed-extraction step. + if idx < len(ids) and ids[idx]: # pragma: no branch - chromadb always aligns ids + meta_dict.setdefault("chunk_id", ids[idx]) # For hierarchical retrieval: return parent context if available text = meta_dict.pop("parent_text", None) or doc score = max(0.0, min(1.0, 1.0 - float(dist))) diff --git a/lamb-kb-server/backend/plugins/vector_db/qdrant_backend.py b/lamb-kb-server/backend/plugins/vector_db/qdrant_backend.py index fb6e55f4a..441850ebf 100644 --- a/lamb-kb-server/backend/plugins/vector_db/qdrant_backend.py +++ b/lamb-kb-server/backend/plugins/vector_db/qdrant_backend.py @@ -27,7 +27,7 @@ try: from qdrant_client import QdrantClient, models -except ImportError: +except ImportError: # pragma: no cover - optional-dependency import guard raise ImportError( "qdrant-client is not installed; " "the 'qdrant' vector DB plugin will not be available." diff --git a/lamb-kb-server/backend/routers/benchmarks.py b/lamb-kb-server/backend/routers/benchmarks.py new file mode 100644 index 000000000..ce341cfa1 --- /dev/null +++ b/lamb-kb-server/backend/routers/benchmarks.py @@ -0,0 +1,92 @@ +"""Benchmark endpoints for vector RAG versus KG-RAG evaluation.""" + +from __future__ import annotations + +from database.connection import get_session +from dependencies import verify_token +from fastapi import APIRouter, Depends +from schemas.benchmark import ( + BenchmarkDataset, + BenchmarkDatasetSummary, + BenchmarkRunAllRequest, + BenchmarkRunAllResponse, + BenchmarkRunRequest, + BenchmarkRunResponse, +) +from services.benchmark import BenchmarkService +from sqlalchemy.orm import Session + +router = APIRouter( + prefix="/benchmarks", + tags=["Benchmarks"], + dependencies=[Depends(verify_token)], +) + + +@router.get( + "/datasets", + response_model=list[BenchmarkDatasetSummary], + summary="List built-in benchmark datasets", +) +async def list_benchmark_datasets() -> list[BenchmarkDatasetSummary]: + return BenchmarkService.list_datasets() + + +@router.get( + "/datasets/{dataset_id}", + response_model=BenchmarkDataset, + summary="Get a built-in benchmark dataset", +) +async def get_benchmark_dataset(dataset_id: str) -> BenchmarkDataset: + return BenchmarkService.get_dataset(dataset_id) + + +@router.post( + "/collections/{collection_id}/run", + response_model=BenchmarkRunResponse, + summary="Run a benchmark dataset against a collection", +) +async def run_collection_benchmark( + collection_id: str, + body: BenchmarkRunRequest, + db: Session = Depends(get_session), +) -> BenchmarkRunResponse: + """Run one benchmark dataset. + + Single body parameter so FastAPI doesn't force the LAMB proxy to wrap + fields under ``request``. Embedded ``embedding_credentials`` carry the + per-request key for the baseline vector pass. + """ + creds = body.embedding_credentials + return BenchmarkService.run( + db=db, + collection_id=collection_id, + request=body, + embedding_credentials={ + "api_key": creds.api_key, + "api_endpoint": creds.api_endpoint, + }, + ) + + +@router.post( + "/collections/{collection_id}/run-all", + response_model=BenchmarkRunAllResponse, + summary="Run all selected benchmark datasets against a collection", +) +async def run_all_collection_benchmarks( + collection_id: str, + body: BenchmarkRunAllRequest, + db: Session = Depends(get_session), +) -> BenchmarkRunAllResponse: + creds = body.embedding_credentials + return BenchmarkService.run_all( + db=db, + collection_id=collection_id, + dataset_ids=body.dataset_ids, + threshold=body.threshold, + embedding_credentials={ + "api_key": creds.api_key, + "api_endpoint": creds.api_endpoint, + }, + ) diff --git a/lamb-kb-server/backend/routers/graph.py b/lamb-kb-server/backend/routers/graph.py new file mode 100644 index 000000000..da73bfae6 --- /dev/null +++ b/lamb-kb-server/backend/routers/graph.py @@ -0,0 +1,507 @@ +"""KG-RAG graph traceability endpoints (new KB server architecture). + +These endpoints are only mounted when ``KG_RAG_ENABLED=true`` (see +``main.py``). Per-request OpenAI credentials are accepted via the +``X-OpenAI-Api-Key`` header for migration jobs that need to extract +concepts at ingest time. Falls back to the ``KG_RAG_OPENAI_API_KEY`` +env var when the header is absent (ADR-4 spirit: prefer per-request +credentials but allow a server-level fallback for migration). +""" + +from __future__ import annotations + +from typing import Any + +from database.connection import get_session +from database.models import Collection +from dependencies import verify_token +from fastapi import APIRouter, Depends, Header, HTTPException, Query +from schemas.graph import ( + GraphChangeDetail, + GraphChangeEvent, + GraphConceptCurationRequest, + GraphConceptMergeRequest, + GraphConceptRenameRequest, + GraphManualOperationResponse, + GraphRelationshipCurationRequest, + GraphRelationshipEditRequest, + GraphRevertRequest, + GraphRevertResponse, + GraphSnapshotResponse, +) +from services.graph_store import get_graph_store +from sqlalchemy.orm import Session + +router = APIRouter(prefix="/graph", tags=["Graph Traceability"]) + + +def _get_collection_or_404(db: Session, collection_id: str) -> Collection: + collection = ( + db.query(Collection).filter(Collection.id == collection_id).first() + ) + if not collection: + raise HTTPException( + status_code=404, detail=f"Collection {collection_id} not found" + ) + return collection + + +def _graph_store_or_503(): + graph_store = get_graph_store() + if not graph_store.is_configured(): + raise HTTPException(status_code=503, detail="KG-RAG Neo4j is not configured") + if not graph_store.is_available(): + raise HTTPException(status_code=503, detail="KG-RAG Neo4j is not available") + return graph_store + + +def _graph_status_payload() -> dict[str, Any]: + import config as config_module + + kg_config = config_module.get_kg_rag_config() + graph_store = get_graph_store() + neo4j_configured = graph_store.is_configured() + neo4j_available = False + if neo4j_configured: + try: + neo4j_available = graph_store.is_available() + except Exception: + neo4j_available = False + + return { + "enabled": bool(kg_config.get("enabled")), + "index_on_ingest": bool(kg_config.get("index_on_ingest", True)), + "neo4j_configured": neo4j_configured, + "neo4j_available": neo4j_available, + } + + +@router.get("/status", summary="Get Graph RAG feature availability") +async def get_graph_status(token: str = Depends(verify_token)): + return _graph_status_payload() + + +@router.post( + "/collections/{collection_id}/migrate", + summary="Migrate existing collection chunks into Graph RAG", +) +async def migrate_collection_to_graph( + collection_id: str, + token: str = Depends(verify_token), + db: Session = Depends(get_session), + x_openai_api_key: str | None = Header(default=None), +): + """Re-index a collection's existing vectors into Neo4j. + + Reads chunks from the vector backend, runs LLM concept extraction, and + writes the resulting graph into Neo4j. Sets ``graph_enabled=True`` on + success. + """ + graph_status = _graph_status_payload() + if not graph_status["enabled"]: + raise HTTPException(status_code=503, detail="KG-RAG is disabled") + + collection = _get_collection_or_404(db, collection_id) + _graph_store_or_503() + + from plugins.base import EmbeddingRegistry, VectorDBRegistry # noqa: PLC0415 + + backend = VectorDBRegistry.get(collection.vector_db_backend) + if backend is None: + raise HTTPException( + status_code=503, + detail=( + f"Vector DB backend '{collection.vector_db_backend}' is not " + "available." + ), + ) + + # We need an embedding function to open the collection on backends that + # require it; no credentials are needed for a read-only scan. + embedding_function = EmbeddingRegistry.build( + collection.embedding_vendor, + model=collection.embedding_model, + api_key="", + api_endpoint=collection.embedding_endpoint or "", + ) + + ids: list[str] = [] + texts: list[str] = [] + metadatas: list[dict[str, Any]] = [] + + try: + for batch_ids, batch_texts, batch_metas in backend.iter_all_chunks( + collection_id=collection.backend_collection_id or collection.id, + storage_path=collection.storage_path, + embedding_function=embedding_function, + ): + for i, chunk_id in enumerate(batch_ids): + text = batch_texts[i] if i < len(batch_texts) else None + if not text: + continue + metadata = ( + batch_metas[i] + if i < len(batch_metas) and isinstance(batch_metas[i], dict) + else {} + ) + ids.append(chunk_id) + texts.append(text) + metadatas.append(metadata) + except NotImplementedError as exc: + raise HTTPException( + status_code=400, + detail=( + f"Graph migration is not supported for the " + f"'{collection.vector_db_backend}' backend yet. Migration " + "requires a backend that exposes iter_all_chunks() — " + "currently chromadb." + ), + ) from exc + except Exception as exc: # noqa: BLE001 + raise HTTPException( + status_code=404, + detail=f"Vector backend collection for {collection.id} not found", + ) from exc + + if not ids: + collection.graph_enabled = True + db.commit() + return { + "status": "success", + "collection_id": collection_id, + "graph_enabled": True, + "indexed": False, + "chunks": 0, + "reason": "no_chunks_found", + } + + # Run LLM extraction + Neo4j indexing. + from services.graph_indexing import index_chunks_for_collection + + indexed = index_chunks_for_collection( + collection=collection, + ids=ids, + texts=texts, + metadatas=metadatas, + openai_api_key=x_openai_api_key, + ) + + collection.graph_enabled = True + db.commit() + return { + "status": "success", + "collection_id": collection_id, + "graph_enabled": True, + "chunks_seen": len(ids), + **indexed, + } + + +@router.get( + "/collections/{collection_id}/snapshot", + response_model=GraphSnapshotResponse, + summary="Get graph snapshot for visualization", +) +async def get_graph_snapshot( + collection_id: str, + concept: str | None = Query(None), + document_id: str | None = Query(None), + chunk_id: str | None = Query(None), + filename: str | None = Query(None), + include_chunks: bool = Query(True), + limit: int = Query(60, ge=1, le=10000), + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + return graph_store.get_collection_graph( + collection_id=collection_id, + org_id=str(collection.organization_id), + concept=concept, + document_id=document_id, + chunk_id=chunk_id, + filename=filename, + include_chunks=include_chunks, + limit=limit, + ) + + +@router.get( + "/collections/{collection_id}/changes", + response_model=list[GraphChangeEvent], + summary="List graph change history", +) +async def list_graph_changes( + collection_id: str, + concept: str | None = Query(None), + relationship_source: str | None = Query(None), + relationship_target: str | None = Query(None), + relationship_relation: str | None = Query(None), + document_id: str | None = Query(None), + filename: str | None = Query(None), + operation: str | None = Query(None), + limit: int = Query(25, ge=1, le=200), + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + return graph_store.list_changes( + collection_id=collection_id, + org_id=str(collection.organization_id), + concept=concept, + relationship_source=relationship_source, + relationship_target=relationship_target, + relationship_relation=relationship_relation, + document_id=document_id, + filename=filename, + operation=operation, + limit=limit, + ) + + +@router.get( + "/collections/{collection_id}/changes/{event_id}", + response_model=GraphChangeDetail, + summary="Inspect a graph change event", +) +async def get_graph_change( + collection_id: str, + event_id: str, + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + change = graph_store.get_change( + collection_id=collection_id, + org_id=str(collection.organization_id), + event_id=event_id, + ) + if not change: + raise HTTPException( + status_code=404, detail=f"Graph change {event_id} not found" + ) + return change + + +@router.get( + "/collections/{collection_id}/concepts/{concept}/changes", + response_model=list[GraphChangeEvent], + summary="Inspect graph changes for a concept", +) +async def list_concept_changes( + collection_id: str, + concept: str, + limit: int = Query(25, ge=1, le=200), + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + return graph_store.list_changes( + collection_id=collection_id, + org_id=str(collection.organization_id), + concept=concept, + limit=limit, + ) + + +@router.get( + "/collections/{collection_id}/documents/{document_id}/changes", + response_model=list[GraphChangeEvent], + summary="Inspect graph changes for a document", +) +async def list_document_changes( + collection_id: str, + document_id: str, + limit: int = Query(25, ge=1, le=200), + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + return graph_store.list_changes( + collection_id=collection_id, + org_id=str(collection.organization_id), + document_id=document_id, + limit=limit, + ) + + +@router.post( + "/collections/{collection_id}/changes/{event_id}/revert", + response_model=GraphRevertResponse, + summary="Revert a supported graph change", +) +async def revert_graph_change( + collection_id: str, + event_id: str, + request: GraphRevertRequest, + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + result = graph_store.revert_change( + collection_id=collection_id, + org_id=str(collection.organization_id), + event_id=event_id, + actor=request.actor, + reason=request.reason, + ) + if not result.get("reverted") and result.get("reason") == "change_not_found": + raise HTTPException( + status_code=404, detail=f"Graph change {event_id} not found" + ) + if not result.get("reverted"): + raise HTTPException(status_code=400, detail=result) + return result + + +@router.patch( + "/collections/{collection_id}/concepts/{concept}/rename", + response_model=GraphManualOperationResponse, + summary="Rename a graph concept in a collection", +) +async def rename_graph_concept( + collection_id: str, + concept: str, + request: GraphConceptRenameRequest, + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + result = graph_store.rename_concept( + collection_id=collection_id, + org_id=str(collection.organization_id), + old_name=concept, + new_name=request.new_name, + actor=request.actor, + reason=request.reason, + ) + if not result.get("ok"): + raise HTTPException(status_code=400, detail=result) + return result + + +@router.post( + "/collections/{collection_id}/concepts/merge", + response_model=GraphManualOperationResponse, + summary="Merge graph concepts in a collection", +) +async def merge_graph_concepts( + collection_id: str, + request: GraphConceptMergeRequest, + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + result = graph_store.merge_concepts( + collection_id=collection_id, + org_id=str(collection.organization_id), + source_names=request.source_names, + target_name=request.target_name, + actor=request.actor, + reason=request.reason, + ) + if not result.get("ok"): + raise HTTPException(status_code=400, detail=result) + return result + + +@router.patch( + "/collections/{collection_id}/concepts/{concept}/curation", + response_model=GraphManualOperationResponse, + summary="Update graph concept curation metadata", +) +async def curate_graph_concept( + collection_id: str, + concept: str, + request: GraphConceptCurationRequest, + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + result = graph_store.update_concept_curation( + collection_id=collection_id, + org_id=str(collection.organization_id), + concept_name=concept, + notes=request.notes, + tags=request.tags, + verification_state=request.verification_state, + actor=request.actor, + reason=request.reason, + ) + if not result.get("ok"): + raise HTTPException(status_code=400, detail=result) + return result + + +@router.patch( + "/collections/{collection_id}/relationships", + response_model=GraphManualOperationResponse, + summary="Edit graph relationship type or weight", +) +async def edit_graph_relationship( + collection_id: str, + request: GraphRelationshipEditRequest, + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + result = graph_store.edit_relationship( + collection_id=collection_id, + org_id=str(collection.organization_id), + source_name=request.source_concept, + target_name=request.target_concept, + relation=request.relation, + new_relation=request.new_relation, + weight=request.weight, + description=request.description, + evidence=request.evidence, + notes=request.notes, + tags=request.tags, + verification_state=request.verification_state, + actor=request.actor, + reason=request.reason, + ) + if not result.get("ok"): + raise HTTPException(status_code=400, detail=result) + return result + + +@router.patch( + "/collections/{collection_id}/relationships/curation", + response_model=GraphManualOperationResponse, + summary="Update graph relationship curation metadata", +) +async def curate_graph_relationship( + collection_id: str, + request: GraphRelationshipCurationRequest, + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + result = graph_store.edit_relationship( + collection_id=collection_id, + org_id=str(collection.organization_id), + source_name=request.source_concept, + target_name=request.target_concept, + relation=request.relation, + notes=request.notes, + tags=request.tags, + verification_state=request.verification_state, + actor=request.actor, + reason=request.reason, + operation="manual_curate_relationship", + ) + if not result.get("ok"): + raise HTTPException(status_code=400, detail=result) + return result diff --git a/lamb-kb-server/backend/routers/jobs.py b/lamb-kb-server/backend/routers/jobs.py index 9e9e3586d..5b0498b84 100644 --- a/lamb-kb-server/backend/routers/jobs.py +++ b/lamb-kb-server/backend/routers/jobs.py @@ -2,13 +2,15 @@ import logging +from config import MAX_JOB_ATTEMPTS from database.connection import get_session from database.models import IngestionJob from dependencies import verify_token from fastapi import APIRouter, Depends, HTTPException, status -from schemas.jobs import JobStatusResponse +from schemas.jobs import JobStatusResponse, RetryAvailability from services.ingestion_service import cancel_job as cancel_job_service from sqlalchemy.orm import Session +from tasks.worker import retry_available logger = logging.getLogger(__name__) @@ -63,3 +65,86 @@ async def cancel_job( """ job = cancel_job_service(db, job_id) return JobStatusResponse.model_validate(job) + + +def _load_job(db: Session, job_id: str) -> IngestionJob: + job = db.query(IngestionJob).filter(IngestionJob.id == job_id).first() + if job is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Job '{job_id}' not found.", + ) + return job + + +@router.get("/{job_id}/retry-available", response_model=RetryAvailability) +async def get_retry_available( + job_id: str, + db: Session = Depends(get_session), +) -> RetryAvailability: + """Report whether a failed job can still be retried in place. + + A retry is possible only while the job is ``failed``, has attempts left, + and its embedding credentials are still cached in memory (within the + retention window and since the last restart). The UI uses this to decide + whether to show a Retry affordance. + """ + job = _load_job(db, job_id) + available = ( + job.status == "failed" + and job.attempts < MAX_JOB_ATTEMPTS + and retry_available(job_id) + ) + return RetryAvailability( + available=available, + attempts=job.attempts, + max_attempts=MAX_JOB_ATTEMPTS, + ) + + +@router.post("/{job_id}/retry", response_model=JobStatusResponse) +async def retry_job( + job_id: str, + db: Session = Depends(get_session), +) -> JobStatusResponse: + """Retry a failed ingestion job, reusing its cached credentials. + + The document payload is already persisted in the job row, and the + embedding credentials are kept in memory across attempts, so a retry does + not require the client to re-send anything. The job is flipped back to + ``pending`` and the worker picks it up on the next poll. + + Errors: + 404 — job not found. + 409 — job is not ``failed``, or has exhausted ``MAX_JOB_ATTEMPTS``. + 410 — credentials are no longer cached (retention window elapsed or + the service restarted); the client must re-submit add-content. + """ + job = _load_job(db, job_id) + + if job.status != "failed": + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Only failed jobs can be retried (status='{job.status}').", + ) + if job.attempts >= MAX_JOB_ATTEMPTS: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Job has exhausted its retries ({MAX_JOB_ATTEMPTS} attempts).", + ) + if not retry_available(job_id): + raise HTTPException( + status_code=status.HTTP_410_GONE, + detail=( + "Retry window expired — credentials are no longer held. " + "Re-submit the content to index it again." + ), + ) + + job.status = "pending" + job.error_message = None + job.completed_at = None + db.commit() + db.refresh(job) + logger.info("Job %s queued for retry (attempt %d)", job_id, job.attempts + 1) + return JobStatusResponse.model_validate(job) diff --git a/lamb-kb-server/backend/routers/query.py b/lamb-kb-server/backend/routers/query.py index 420713f6e..eaf729aa2 100644 --- a/lamb-kb-server/backend/routers/query.py +++ b/lamb-kb-server/backend/routers/query.py @@ -41,6 +41,17 @@ async def query_collection( """ results = query_service.query_collection(db, collection_id, body) + # When the collection has graph_enabled, query_service.query_collection + # routes through the KG-RAG plugin which attaches an identical + # ``kg_rag`` trace to every chunk's metadata. Surface its + # ``question_entities`` at the top level so callers don't need to dig + # through per-chunk metadata. ``None`` when the plugin didn't run. + entities: list[str] | None = None + if results: + trace = (results[0].metadata or {}).get("kg_rag") or {} + if "question_entities" in trace: + entities = list(trace.get("question_entities") or []) + return QueryResponse( results=[ QueryResultItem( @@ -52,4 +63,5 @@ async def query_collection( ], query=body.query_text, top_k=body.top_k, + entities=entities, ) diff --git a/lamb-kb-server/backend/routers/system.py b/lamb-kb-server/backend/routers/system.py index d81f606bc..9883fd8bf 100644 --- a/lamb-kb-server/backend/routers/system.py +++ b/lamb-kb-server/backend/routers/system.py @@ -5,7 +5,12 @@ from database.connection import get_session_direct from dependencies import verify_token from fastapi import APIRouter, Depends -from plugins.base import ChunkingRegistry, EmbeddingRegistry, VectorDBRegistry +from plugins.base import ( + ChunkingRegistry, + EmbeddingRegistry, + LLMExtractionRegistry, + VectorDBRegistry, +) from sqlalchemy import text from tasks.worker import is_worker_running @@ -81,3 +86,14 @@ async def list_embedding_vendors() -> dict: Dict with ``vendors`` list. """ return {"vendors": EmbeddingRegistry.list_plugins()} + + +@router.get("/llm-vendors", dependencies=[Depends(verify_token)]) +async def list_llm_vendors() -> dict: + """List all registered LLM extraction vendors (KG-RAG concept extractor). + + Returns: + Dict with ``vendors`` list. Each vendor entry includes ``name``, + ``description``, and a ``parameters`` schema (model, api_endpoint). + """ + return {"vendors": LLMExtractionRegistry.list_plugins()} diff --git a/lamb-kb-server/backend/schemas/benchmark.py b/lamb-kb-server/backend/schemas/benchmark.py new file mode 100644 index 000000000..bef93ce43 --- /dev/null +++ b/lamb-kb-server/backend/schemas/benchmark.py @@ -0,0 +1,147 @@ +"""Schemas for KG-RAG benchmark evaluation.""" + +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, ConfigDict, Field + + +class BenchmarkQuestion(BaseModel): + id: str = Field(..., description="Stable question ID") + question: str = Field(..., description="Question text to retrieve against") + expected_answer: Optional[str] = Field("", description="Reference answer") + relevant_files: List[str] = Field( + default_factory=list, description="Relevant source filenames" + ) + expected_concepts: List[str] = Field( + default_factory=list, description="Expected graph concepts" + ) + kind: str = Field("custom", description="Question category") + answerable: bool = Field(True, description="Whether the answer is present") + + +class BenchmarkDatasetSummary(BaseModel): + id: str + name: str + description: str + expected_behavior: str = Field( + ..., description="Expected KG-RAG behavior for this dataset" + ) + recommended_top_k: int + recommended_graph_depth: int + question_count: int + + +class BenchmarkDataset(BenchmarkDatasetSummary): + questions: List[BenchmarkQuestion] + + +class BenchmarkMetrics(BaseModel): + precision_at_k: float = 0.0 + recall_at_k: float = 0.0 + mrr: float = 0.0 + avg_vector_ms: float = 0.0 + avg_graph_ms: float = 0.0 + avg_total_ms: float = 0.0 + + +class BenchmarkQueryScore(BaseModel): + precision_at_k: float = 0.0 + recall_at_k: float = 0.0 + mrr: float = 0.0 + vector_ms: float = 0.0 + graph_ms: float = 0.0 + total_ms: float = 0.0 + retrieved_files: List[str] = Field(default_factory=list) + + +class BenchmarkQueryResult(BaseModel): + question_id: str + question: str + kind: str + relevant_files: List[str] + expected_concepts: List[str] = Field(default_factory=list) + baseline: BenchmarkQueryScore + kg_rag: BenchmarkQueryScore + + +class BenchmarkComparison(BaseModel): + delta_precision_at_k: float = 0.0 + delta_recall_at_k: float = 0.0 + delta_mrr: float = 0.0 + graph_overhead_ms: float = 0.0 + expected_behavior: str + + +class EmbeddingCredentialsBody(BaseModel): + """Embedded request-scoped embedding credentials. + + Mirrors :class:`schemas.content.EmbeddingCredentials` but lives on the + benchmark request models so the LAMB proxy can send a single flat body + instead of having to wrap fields under ``request`` (which FastAPI + requires when multiple ``Body`` parameters coexist on one route). + """ + + api_key: str = Field(default="", description="Vendor API key.") + api_endpoint: str = Field( + default="", description="Optional API base URL override." + ) + + +class BenchmarkRunRequest(BaseModel): + """Body for a single benchmark run. + + Extra fields are allowed so callers can pass plugin-specific + tuning knobs (``rrf_k``, ``graph_weight``, ``graph_limit_factor``) + without requiring a schema change. ``BenchmarkService.run`` + forwards any recognised extra field to the KG-RAG plugin via + ``plugin_params``. + """ + + model_config = ConfigDict(extra="allow") + + dataset_id: Optional[str] = Field( + "educational", description="Built-in dataset ID to use when questions are omitted" + ) + questions: Optional[List[BenchmarkQuestion]] = Field( + None, description="Custom benchmark questions" + ) + top_k: Optional[int] = Field(None, ge=1, le=50) + graph_depth: Optional[int] = Field(None, ge=1, le=4) + threshold: float = Field(0.0, ge=0.0, le=1.0) + embedding_credentials: EmbeddingCredentialsBody = Field( + default_factory=EmbeddingCredentialsBody, + description=( + "Per-request embedding credentials. LAMB resolves these from " + "``setups.default.providers.{vendor}.api_key``; the field is " + "optional so direct callers can omit it." + ), + ) + + +class BenchmarkRunResponse(BaseModel): + collection_id: str + dataset_id: str + dataset_name: str + top_k: int + graph_depth: int + baseline: BenchmarkMetrics + kg_rag: BenchmarkMetrics + comparison: BenchmarkComparison + results: List[BenchmarkQueryResult] + + +class BenchmarkRunAllRequest(BaseModel): + dataset_ids: List[str] = Field( + default_factory=lambda: ["educational", "control", "paper", "extreme"] + ) + threshold: float = Field(0.0, ge=0.0, le=1.0) + embedding_credentials: EmbeddingCredentialsBody = Field( + default_factory=EmbeddingCredentialsBody, + description="Per-request embedding credentials (same semantics as BenchmarkRunRequest).", + ) + + +class BenchmarkRunAllResponse(BaseModel): + collection_id: str + runs: List[BenchmarkRunResponse] + summary: Dict[str, Any] = Field(default_factory=dict) diff --git a/lamb-kb-server/backend/schemas/collection.py b/lamb-kb-server/backend/schemas/collection.py index 984fc83ec..7bf2e8744 100644 --- a/lamb-kb-server/backend/schemas/collection.py +++ b/lamb-kb-server/backend/schemas/collection.py @@ -23,6 +23,30 @@ class EmbeddingConfig(BaseModel): ) +class ExtractionConfig(BaseModel): + """Describes the collection-level KG-RAG extraction setup. + + Locked at creation, mirrors the embedding pattern: credentials are + request-scoped, the vendor + model + endpoint are persisted on the + collection. ``None`` everywhere means "fall back to server env + defaults" (backwards-compat for graph_enabled stores created before + this surface existed). + """ + + vendor: str | None = Field( + default=None, + description="LLM extraction vendor name (e.g. 'openai', 'ollama').", + ) + model: str | None = Field( + default=None, + description="Model identifier (e.g. 'gpt-4o-mini', 'llama3.1:8b').", + ) + api_endpoint: str | None = Field( + default=None, + description="Optional override for the vendor's API base URL.", + ) + + # --- Requests --- @@ -63,6 +87,22 @@ class CreateCollectionRequest(BaseModel): default_factory=dict, description="Schema-declared vector-DB backend extras (backend-specific knobs).", ) + graph_enabled: bool = Field( + default=False, + description=( + "Opt this collection into KG-RAG: extracted concepts/relations " + "are indexed into Neo4j alongside vector storage at ingestion " + "time. Requires ``KG_RAG_ENABLED=true`` at the server level." + ), + ) + extraction: ExtractionConfig | None = Field( + default=None, + description=( + "KG-RAG extraction vendor/model/endpoint. Only meaningful when " + "``graph_enabled=true``. If omitted, the server falls back to " + "the ``KG_RAG_EXTRACTION_MODEL`` env var (OpenAI by default)." + ), + ) class UpdateCollectionRequest(BaseModel): @@ -103,6 +143,8 @@ class CollectionResponse(BaseModel): chunking_params: dict[str, Any] embedding: EmbeddingConfig vector_db_backend: str + graph_enabled: bool = False + extraction: ExtractionConfig | None = None status: str document_count: int chunk_count: int @@ -132,6 +174,17 @@ def from_orm_row(cls, row: Any) -> "CollectionResponse": api_endpoint=row.embedding_endpoint or "", ), vector_db_backend=row.vector_db_backend, + graph_enabled=bool(getattr(row, "graph_enabled", False)), + extraction=( + ExtractionConfig( + vendor=getattr(row, "extraction_vendor", None), + model=getattr(row, "extraction_model", None), + api_endpoint=getattr(row, "extraction_endpoint", None), + ) + if getattr(row, "extraction_vendor", None) + or getattr(row, "extraction_model", None) + else None + ), status=row.status, document_count=row.document_count, chunk_count=row.chunk_count, diff --git a/lamb-kb-server/backend/schemas/graph.py b/lamb-kb-server/backend/schemas/graph.py new file mode 100644 index 000000000..37f55d73b --- /dev/null +++ b/lamb-kb-server/backend/schemas/graph.py @@ -0,0 +1,142 @@ +"""Schemas for KG-RAG graph traceability endpoints.""" + +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, Field + + +class GraphChangeEvent(BaseModel): + event_id: str = Field(..., description="Unique graph change event ID") + collection_id: str = Field(..., description="Collection ID") + org_id: str = Field(..., description="Organization or owner scope") + operation: str = Field(..., description="Graph operation name") + actor: Optional[str] = Field(None, description="Actor that produced the change") + timestamp: Optional[str] = Field(None, description="Change timestamp") + filename: Optional[str] = Field(None, description="Source filename") + concepts: List[str] = Field( + default_factory=list, description="Concepts touched by the change" + ) + payload_json: Optional[str] = Field( + None, description="Raw JSON payload stored in Neo4j" + ) + document_id: Optional[str] = Field(None, description="Related graph document ID") + file_id: Optional[int] = Field(None, description="Related file registry ID") + + +class GraphChangeDetail(GraphChangeEvent): + chunk_ids: List[str] = Field( + default_factory=list, description="Related graph chunk IDs" + ) + + +class GraphRevertRequest(BaseModel): + actor: str = Field( + "graph-traceability-api", description="Actor requesting the revert" + ) + reason: str = Field("", description="Human-readable reason for the revert") + + +class GraphRevertResponse(BaseModel): + reverted: bool = Field(..., description="Whether the revert was applied") + reason: Optional[str] = Field(None, description="Reason when no revert was applied") + event_id: Optional[str] = Field(None, description="Requested event ID") + revert_event_id: Optional[str] = Field( + None, description="Audit event created for the revert" + ) + operation: Optional[str] = Field(None, description="Original operation") + document_id: Optional[str] = Field(None, description="Reverted graph document ID") + chunk_ids: List[str] = Field( + default_factory=list, description="Reverted graph chunk IDs" + ) + + +class GraphNode(BaseModel): + id: str = Field(..., description="Stable frontend node ID") + type: str = Field(..., description="Node type, such as concept or chunk") + label: str = Field(..., description="Human-readable node label") + data: Dict[str, Any] = Field(default_factory=dict, description="Node metadata") + + +class GraphEdge(BaseModel): + id: str = Field(..., description="Stable frontend edge ID") + type: str = Field(..., description="Graph edge type") + source: str = Field(..., description="Source node ID") + target: str = Field(..., description="Target node ID") + label: Optional[str] = Field(None, description="Human-readable edge label") + weight: Optional[float] = Field(None, description="Edge weight") + data: Dict[str, Any] = Field(default_factory=dict, description="Edge metadata") + + +class GraphSnapshotResponse(BaseModel): + collection_id: str = Field(..., description="Collection ID") + nodes: List[GraphNode] = Field(default_factory=list, description="Graph nodes") + edges: List[GraphEdge] = Field(default_factory=list, description="Graph edges") + filters: Dict[str, Any] = Field(default_factory=dict, description="Applied filters") + counts: Dict[str, int] = Field( + default_factory=dict, description="Graph summary counts" + ) + + +class GraphManualOperationResponse(BaseModel): + ok: bool = Field(..., description="Whether the manual graph operation was applied") + operation: Optional[str] = Field(None, description="Recorded manual operation name") + event_id: Optional[str] = Field( + None, description="ChangeEvent ID recorded for the operation" + ) + reason: Optional[str] = Field( + None, description="Reason when the operation was not applied" + ) + details: Dict[str, Any] = Field( + default_factory=dict, description="Operation-specific details" + ) + + +class GraphConceptRenameRequest(BaseModel): + new_name: str = Field(..., description="New concept name") + actor: str = Field("graph-curation-api", description="Actor requesting the rename") + reason: str = Field("", description="Human-readable reason for the rename") + + +class GraphConceptMergeRequest(BaseModel): + source_names: List[str] = Field( + ..., description="Concept names to merge into the target" + ) + target_name: str = Field(..., description="Target concept name") + actor: str = Field("graph-curation-api", description="Actor requesting the merge") + reason: str = Field("", description="Human-readable reason for the merge") + + +class GraphConceptCurationRequest(BaseModel): + notes: Optional[str] = Field(None, description="Manual curation notes") + tags: Optional[List[str]] = Field(None, description="Manual curation tags") + verification_state: Optional[str] = Field(None, description="Verification state") + actor: str = Field("graph-curation-api", description="Actor requesting the update") + reason: str = Field("", description="Human-readable reason for the update") + + +class GraphRelationshipEditRequest(BaseModel): + source_concept: str = Field(..., description="Source concept name") + target_concept: str = Field(..., description="Target concept name") + relation: str = Field(..., description="Current relationship relation value") + new_relation: Optional[str] = Field( + None, description="New relationship relation value" + ) + weight: Optional[float] = Field(None, description="New relationship weight") + description: Optional[str] = Field(None, description="Relationship description") + evidence: Optional[str] = Field(None, description="Relationship evidence") + notes: Optional[str] = Field(None, description="Manual curation notes") + tags: Optional[List[str]] = Field(None, description="Manual curation tags") + verification_state: Optional[str] = Field(None, description="Verification state") + actor: str = Field("graph-curation-api", description="Actor requesting the update") + reason: str = Field("", description="Human-readable reason for the update") + + +class GraphRelationshipCurationRequest(BaseModel): + source_concept: str = Field(..., description="Source concept name") + target_concept: str = Field(..., description="Target concept name") + relation: str = Field(..., description="Current relationship relation value") + notes: Optional[str] = Field(None, description="Manual curation notes") + tags: Optional[List[str]] = Field(None, description="Manual curation tags") + verification_state: Optional[str] = Field(None, description="Verification state") + actor: str = Field("graph-curation-api", description="Actor requesting the update") + reason: str = Field("", description="Human-readable reason for the update") diff --git a/lamb-kb-server/backend/schemas/jobs.py b/lamb-kb-server/backend/schemas/jobs.py index c5bd32b66..feb2d5d96 100644 --- a/lamb-kb-server/backend/schemas/jobs.py +++ b/lamb-kb-server/backend/schemas/jobs.py @@ -22,3 +22,11 @@ class JobStatusResponse(BaseModel): updated_at: datetime started_at: datetime | None = None completed_at: datetime | None = None + + +class RetryAvailability(BaseModel): + """Whether a failed job can be retried in place, with attempt accounting.""" + + available: bool + attempts: int + max_attempts: int diff --git a/lamb-kb-server/backend/schemas/query.py b/lamb-kb-server/backend/schemas/query.py index 605d8b82e..d20a09081 100644 --- a/lamb-kb-server/backend/schemas/query.py +++ b/lamb-kb-server/backend/schemas/query.py @@ -40,3 +40,11 @@ class QueryResponse(BaseModel): results: list[QueryResultItem] query: str top_k: int + entities: list[str] | None = Field( + default=None, + description=( + "Named entities extracted from the question when the query was " + "routed through KG-RAG (collections with ``graph_enabled=true``). " + "``null`` for plain vector queries." + ), + ) diff --git a/lamb-kb-server/backend/services/benchmark.py b/lamb-kb-server/backend/services/benchmark.py new file mode 100644 index 000000000..206a4bb1d --- /dev/null +++ b/lamb-kb-server/backend/services/benchmark.py @@ -0,0 +1,714 @@ +"""Benchmark service for comparing vector RAG and KG-RAG retrieval.""" + +from __future__ import annotations + +from pathlib import PurePosixPath +from statistics import mean +from typing import Any, Dict, List, Optional, Sequence + +from fastapi import HTTPException + +from schemas.benchmark import ( + BenchmarkComparison, + BenchmarkDataset, + BenchmarkDatasetSummary, + BenchmarkMetrics, + BenchmarkQuestion, + BenchmarkQueryResult, + BenchmarkQueryScore, + BenchmarkRunAllResponse, + BenchmarkRunRequest, + BenchmarkRunResponse, +) + +_DATASETS: Dict[str, Dict[str, Any]] = { + "educational": { + "name": "LAMB KG-RAG educational QA starter set", + "description": "Starter benchmark for comparing classic vector RAG and KG-RAG on LAMB-style educational retrieval questions.", + "expected_behavior": "mixed: single-hop parity, multi-hop KG-RAG may improve recall", + "recommended_top_k": 5, + "recommended_graph_depth": 2, + "questions": [ + { + "id": "q1", + "question": "What does the current LAMB-style baseline use ChromaDB for?", + "expected_answer": "It stores semantic vectors for chunks and retrieves similar fragments.", + "relevant_files": ["rag_basics.md", "lamb_kb_server.md"], + "expected_concepts": ["chromadb", "vector rag", "embeddings"], + "kind": "single-hop", + }, + { + "id": "q2", + "question": "Why can a knowledge graph improve questions that require information from multiple course documents?", + "expected_answer": "It traverses related concepts and can recover evidence that is conceptually connected even when it is not the closest vector match.", + "relevant_files": ["rag_basics.md", "knowledge_graphs.md"], + "expected_concepts": [ + "multi-hop retrieval", + "knowledge graph", + "vector search", + ], + "kind": "multi-hop", + }, + { + "id": "q3", + "question": "Why does traceability matter when educators rename or merge concepts?", + "expected_answer": "Traceability records ChangeEvent nodes with operation, actor, timestamp, affected concepts, and payload so graph edits can be audited.", + "relevant_files": [ + "traceability_and_curation.md", + "knowledge_graphs.md", + ], + "expected_concepts": ["changeevent", "manual curation", "traceability"], + "kind": "multi-hop", + }, + { + "id": "q4", + "question": "Which technologies does the prototype replicate from the LAMB kb-server architecture?", + "expected_answer": "FastAPI, SQLite, ChromaDB, Markdown ingestion, parent-child chunking, and chat completion.", + "relevant_files": ["lamb_kb_server.md", "rag_basics.md"], + "expected_concepts": [ + "fastapi", + "sqlite", + "chromadb", + "parent-child chunking", + ], + "kind": "single-hop", + }, + { + "id": "q5", + "question": "How should graph expansion latency be evaluated in the project?", + "expected_answer": "Benchmark vector retrieval latency, graph expansion latency, and total retrieval latency, with a target of keeping graph expansion below 50 ms when possible.", + "relevant_files": ["evaluation_metrics.md", "knowledge_graphs.md"], + "expected_concepts": ["latency", "cypher", "graph expansion"], + "kind": "multi-hop", + }, + { + "id": "q6", + "question": "What is the purpose of parent-child chunking in a LAMB-style retrieval pipeline?", + "expected_answer": "Small child chunks improve semantic search while larger parent chunks give the language model enough context.", + "relevant_files": ["rag_basics.md"], + "expected_concepts": [ + "parent-child chunking", + "child chunks", + "parent chunks", + ], + "kind": "single-hop", + }, + { + "id": "q7", + "question": "How does the prototype keep the real LAMB repository untouched while still imitating its behavior?", + "expected_answer": "It runs in an isolated prototype folder and reimplements only the relevant collection, ingestion, retrieval, and chat behavior.", + "relevant_files": ["lamb_kb_server.md"], + "expected_concepts": ["prototype", "lamb", "kb-server"], + "kind": "single-hop", + }, + { + "id": "q8", + "question": "Which metric rewards placing the first relevant source as early as possible?", + "expected_answer": "Mean Reciprocal Rank rewards the first relevant result appearing early in the ranking.", + "relevant_files": ["evaluation_metrics.md"], + "expected_concepts": ["mean reciprocal rank", "mrr"], + "kind": "single-hop", + }, + { + "id": "q9", + "question": "What should the KG-RAG UI show to make retrieval less opaque?", + "expected_answer": "It should show entry concepts, traversed relationships, expanded chunks, and recent ChangeEvent metadata.", + "relevant_files": [ + "knowledge_graphs.md", + "traceability_and_curation.md", + ], + "expected_concepts": [ + "entry concepts", + "traversed relationships", + "changeevent", + ], + "kind": "multi-hop", + }, + { + "id": "q10", + "question": "Why is recall important for multi-hop educational questions?", + "expected_answer": "Recall measures whether all expected relevant sources were recovered, which matters when a correct answer needs evidence from multiple files.", + "relevant_files": ["evaluation_metrics.md", "rag_basics.md"], + "expected_concepts": ["recall", "multi-hop", "relevant sources"], + "kind": "multi-hop", + }, + ], + }, + "control": { + "name": "Control benchmark with no inter-document connections", + "description": "Single-document questions where every answer is contained in one standalone profile. KG-RAG should not materially outperform vector RAG.", + "expected_behavior": "parity: KG-RAG should be the same or nearly the same as vector baseline", + "recommended_top_k": 1, + "recommended_graph_depth": 1, + "questions": [ + { + "id": "c1", + "question": "What owner and inspection cadence are listed for Solar Kiln?", + "expected_answer": "Solar Kiln is owned by the Thermal Lab Desk and has an inspection cadence of every 14 days.", + "relevant_files": ["solar_kiln_profile.md"], + "expected_concepts": ["solar kiln", "thermal lab desk", "sk-14"], + "kind": "control-single-doc", + }, + { + "id": "c2", + "question": "What checksum code and owner are listed for Iris Archive?", + "expected_answer": "Iris Archive has checksum code IA-72 and is owned by the Records Integrity Desk.", + "relevant_files": ["iris_archive_profile.md"], + "expected_concepts": [ + "iris archive", + "ia-72", + "records integrity desk", + ], + "kind": "control-single-doc", + }, + { + "id": "c3", + "question": "What route code and inspection cadence are listed for Delta Ferry?", + "expected_answer": "Delta Ferry has route code DF-08 and an inspection cadence of every 9 days.", + "relevant_files": ["delta_ferry_profile.md"], + "expected_concepts": ["delta ferry", "df-08", "coastal transit desk"], + "kind": "control-single-doc", + }, + { + "id": "c4", + "question": "What protocol code and owner are listed for Orchid Lab?", + "expected_answer": "Orchid Lab has protocol code OL-31 and is owned by the Botanical Methods Desk.", + "relevant_files": ["orchid_lab_profile.md"], + "expected_concepts": ["orchid lab", "ol-31", "botanical methods desk"], + "kind": "control-single-doc", + }, + ], + }, + "paper": { + "name": "Paper-style MultiHop-RAG benchmark", + "description": "Small news-style benchmark inspired by MultiHop-RAG, HotpotQA, and MuSiQue. Questions require evidence across separated documents.", + "expected_behavior": "improvement: KG-RAG should improve recall and MRR on multi-hop chains", + "recommended_top_k": 5, + "recommended_graph_depth": 4, + "questions": [ + { + "id": "p1", + "question": "After the Aurora Port outage, which agency owns the project that mitigated the root failure?", + "expected_answer": "The Aurora Port outage was caused by TideNet router failure, which is mitigated by Harbor Battery Deployment; that project is owned by the Port Resilience Office.", + "relevant_files": [ + "city_incident_digest.md", + "failure_to_project_map.md", + "agency_directory.md", + ], + "expected_concepts": [ + "aurora port outage", + "tidenet router failure", + "harbor battery deployment", + "port resilience office", + ], + "kind": "paper-inference", + }, + { + "id": "p2", + "question": "Which mitigation had the larger budget: the project for the Aurora Port outage or the project for the Lantern Bridge closure?", + "expected_answer": "Aurora Port maps to Harbor Battery Deployment at 18.4 million euros, while Lantern Bridge maps to Eastbank Shuttle Loop at 12.1 million euros, so Harbor Battery Deployment is larger.", + "relevant_files": [ + "city_incident_digest.md", + "failure_to_project_map.md", + "budget_register.md", + ], + "expected_concepts": [ + "aurora port outage", + "lantern bridge closure", + "harbor battery deployment", + "eastbank shuttle loop", + ], + "kind": "paper-comparison", + }, + { + "id": "p3", + "question": "Which was completed earlier: the mitigation for the Lantern Bridge closure or the mitigation for the Cobalt Clinic evacuation?", + "expected_answer": "Lantern Bridge maps to Eastbank Shuttle Loop, completed on 2025-07-02. Cobalt Clinic maps to Clinic Microgrid Retrofit, completed on 2025-10-20. Eastbank Shuttle Loop was completed earlier.", + "relevant_files": [ + "city_incident_digest.md", + "failure_to_project_map.md", + "schedule_updates.md", + ], + "expected_concepts": [ + "lantern bridge closure", + "cobalt clinic evacuation", + "eastbank shuttle loop", + "clinic microgrid retrofit", + ], + "kind": "paper-temporal", + }, + { + "id": "p4", + "question": "What budget and owner correspond to the project that mitigated the MercyWing generator fault?", + "expected_answer": "MercyWing generator fault maps to Clinic Microgrid Retrofit, which has a budget of 9.6 million euros and is owned by the Health Facilities Bureau.", + "relevant_files": [ + "failure_to_project_map.md", + "budget_register.md", + "agency_directory.md", + ], + "expected_concepts": [ + "mercywing generator fault", + "clinic microgrid retrofit", + "health facilities bureau", + ], + "kind": "paper-inference", + }, + { + "id": "p5", + "question": "The Harbor Aquarium flood appears in a query. Which mitigation project, budget, and owner should be returned?", + "expected_answer": "No mitigation project, budget, or owner should be returned because the Harbor Aquarium flood is not connected to a registered mitigation project in this corpus.", + "relevant_files": ["null_registry.md"], + "expected_concepts": ["harbor aquarium flood"], + "kind": "paper-null", + "answerable": False, + }, + { + "id": "p6", + "question": "Which owner agency is responsible for the mitigation project that has the emergency mobility budget category?", + "expected_answer": "The emergency mobility budget category belongs to Eastbank Shuttle Loop, which is owned by the Transit Continuity Unit.", + "relevant_files": ["budget_register.md", "agency_directory.md"], + "expected_concepts": [ + "emergency mobility", + "eastbank shuttle loop", + "transit continuity unit", + ], + "kind": "paper-bridge", + }, + ], + }, + "extreme": { + "name": "Extreme KG-RAG multi-hop benchmark", + "description": "Adversarial questions where the surface clue lives in one file and the decisive answer lives several graph hops away.", + "expected_behavior": "improvement: adversarial multi-hop cases should favor KG-RAG over vector baseline", + "recommended_top_k": 6, + "recommended_graph_depth": 4, + "questions": [ + { + "id": "x1", + "question": "Case Orion-17 needs closure. What exact closure sequence should be applied?", + "expected_answer": "Follow protocol Glass Harbor: disable the quartz bypass, replace the cerulean latch, and run the cold-start assay.", + "relevant_files": [ + "anomaly_bridge_extreme.md", + "protocol_index_extreme.md", + "remediation_playbook_extreme.md", + ], + "expected_concepts": [ + "orion-17", + "aster valve drift", + "m-41", + "glass harbor", + ], + "kind": "extreme-multi-hop", + }, + { + "id": "x2", + "question": "Case Vega-03 needs closure. What exact closure sequence should be applied?", + "expected_answer": "Vega-03 has Lumen shard bloom, which maps to Q-Delta and then Night Orchard: isolate the amber bus, reseed the clock lattice, and run the midnight parity check.", + "relevant_files": [ + "anomaly_bridge_extreme.md", + "protocol_index_extreme.md", + "remediation_playbook_extreme.md", + ], + "expected_concepts": [ + "vega-03", + "lumen shard bloom", + "q-delta", + "night orchard", + ], + "kind": "extreme-multi-hop", + }, + { + "id": "x3", + "question": "Case Mira-22 needs closure. What exact closure sequence should be applied?", + "expected_answer": "Mira-22 has Sable checksum echo, which maps to R-9 and then Blue Thread: rotate the ivory token, rebuild the relay ledger, and run the archive handshake.", + "relevant_files": [ + "anomaly_bridge_extreme.md", + "protocol_index_extreme.md", + "remediation_playbook_extreme.md", + ], + "expected_concepts": [ + "mira-22", + "sable checksum echo", + "r-9", + "blue thread", + ], + "kind": "extreme-multi-hop", + }, + { + "id": "x4", + "question": "A HelioForge report only says Aster valve drift. What closure sequence follows from the chain?", + "expected_answer": "Aster valve drift indicates M-41; M-41 is governed by Glass Harbor; Glass Harbor requires disabling the quartz bypass, replacing the cerulean latch, and running the cold-start assay.", + "relevant_files": [ + "anomaly_bridge_extreme.md", + "protocol_index_extreme.md", + "remediation_playbook_extreme.md", + ], + "expected_concepts": ["aster valve drift", "m-41", "glass harbor"], + "kind": "extreme-multi-hop", + }, + { + "id": "x5", + "question": "A North Atrium report only says Lumen shard bloom. What closure sequence follows from the chain?", + "expected_answer": "Lumen shard bloom indicates Q-Delta; Q-Delta is governed by Night Orchard; Night Orchard requires isolating the amber bus, reseeding the clock lattice, and running the midnight parity check.", + "relevant_files": [ + "anomaly_bridge_extreme.md", + "protocol_index_extreme.md", + "remediation_playbook_extreme.md", + ], + "expected_concepts": ["lumen shard bloom", "q-delta", "night orchard"], + "kind": "extreme-multi-hop", + }, + { + "id": "x6", + "question": "An Archive Relay report only says Sable checksum echo. What closure sequence follows from the chain?", + "expected_answer": "Sable checksum echo indicates R-9; R-9 is governed by Blue Thread; Blue Thread requires rotating the ivory token, rebuilding the relay ledger, and running the archive handshake.", + "relevant_files": [ + "anomaly_bridge_extreme.md", + "protocol_index_extreme.md", + "remediation_playbook_extreme.md", + ], + "expected_concepts": ["sable checksum echo", "r-9", "blue thread"], + "kind": "extreme-multi-hop", + }, + ], + }, +} + +_ALIASES = { + "default": "educational", + "sample": "educational", + "education": "educational", + "multihop": "paper", + "multi-hop": "paper", + "paper-multihop": "paper", + "adversarial": "extreme", + "no-connections": "control", + "no_connections": "control", +} + + +class BenchmarkService: + """Run retrieval benchmarks against an existing LAMB collection.""" + + @classmethod + def list_datasets(cls) -> List[BenchmarkDatasetSummary]: + return [cls._dataset_summary(dataset_id) for dataset_id in _DATASETS] + + @classmethod + def get_dataset(cls, dataset_id: str) -> BenchmarkDataset: + canonical = cls._canonical_dataset_id(dataset_id) + config = _DATASETS[canonical] + return BenchmarkDataset( + id=canonical, + name=config["name"], + description=config["description"], + expected_behavior=config["expected_behavior"], + recommended_top_k=config["recommended_top_k"], + recommended_graph_depth=config["recommended_graph_depth"], + question_count=len(config["questions"]), + questions=[BenchmarkQuestion(**item) for item in config["questions"]], + ) + + @classmethod + def run( + cls, + db: Any, + collection_id: str, + request: BenchmarkRunRequest, + embedding_credentials: Optional[Dict[str, Any]] = None, + ) -> BenchmarkRunResponse: + from services.query_service import query_with_plugin + + creds = embedding_credentials or {} + dataset = cls.get_dataset(request.dataset_id or "educational") + questions = request.questions or dataset.questions + if not questions: + raise HTTPException( + status_code=400, detail="No benchmark questions supplied" + ) + + top_k = int(request.top_k or dataset.recommended_top_k or 5) + graph_depth = int(request.graph_depth or dataset.recommended_graph_depth or 2) + top_k = max(1, min(top_k, 50)) + graph_depth = max(1, min(graph_depth, 4)) + threshold = float(request.threshold or 0.0) + + baseline_scores: List[BenchmarkQueryScore] = [] + kg_scores: List[BenchmarkQueryScore] = [] + rows: List[BenchmarkQueryResult] = [] + + for question in questions: + baseline_response = query_with_plugin( + db=db, + collection_id=collection_id, + query_text=question.question, + plugin_name="simple_query", + plugin_params={"top_k": top_k, "threshold": threshold}, + embedding_credentials=creds, + ) + kg_params: Dict[str, Any] = { + "top_k": top_k, + "threshold": threshold, + "graph_depth": graph_depth, + "include_trace": True, + } + # Pass-through for optional tuning knobs (rrf_k, graph_weight, + # graph_limit_factor). These are set on the + # ``BenchmarkRunRequest`` via Pydantic ``model_extra`` (see + # ``BenchmarkRunRequest.model_config``); when absent the + # plugin falls back to its own defaults. + for extra in ("rrf_k", "graph_weight", "graph_limit_factor"): + if hasattr(request, extra): + val = getattr(request, extra, None) + if val is not None: + kg_params[extra] = val + # also pull from the model's extra fields if any + extras = getattr(request, "model_extra", None) or {} + if extra in extras and extras[extra] is not None: + kg_params[extra] = extras[extra] + kg_response = query_with_plugin( + db=db, + collection_id=collection_id, + query_text=question.question, + plugin_name="kg_rag_query", + plugin_params=kg_params, + embedding_credentials=creds, + ) + + baseline_score = cls._score_response( + question=question, + response=baseline_response, + top_k=top_k, + vector_ms=float( + baseline_response.get("timing", {}).get("total_ms", 0.0) + ), + graph_ms=0.0, + total_ms=float( + baseline_response.get("timing", {}).get("total_ms", 0.0) + ), + ) + trace = cls._trace_from_results(kg_response.get("results", [])) + kg_total_ms = float(kg_response.get("timing", {}).get("total_ms", 0.0)) + kg_graph_ms = float(trace.get("graph_latency_ms") or 0.0) + kg_vector_ms = float( + trace.get("vector_latency_ms") or max(kg_total_ms - kg_graph_ms, 0.0) + ) + kg_score = cls._score_response( + question=question, + response=kg_response, + top_k=top_k, + vector_ms=kg_vector_ms, + graph_ms=kg_graph_ms, + total_ms=kg_total_ms, + ) + + baseline_scores.append(baseline_score) + kg_scores.append(kg_score) + rows.append( + BenchmarkQueryResult( + question_id=question.id, + question=question.question, + kind=question.kind, + relevant_files=question.relevant_files, + expected_concepts=question.expected_concepts, + baseline=baseline_score, + kg_rag=kg_score, + ) + ) + + baseline_metrics = cls._aggregate(baseline_scores) + kg_metrics = cls._aggregate(kg_scores) + comparison = BenchmarkComparison( + delta_precision_at_k=kg_metrics.precision_at_k + - baseline_metrics.precision_at_k, + delta_recall_at_k=kg_metrics.recall_at_k - baseline_metrics.recall_at_k, + delta_mrr=kg_metrics.mrr - baseline_metrics.mrr, + graph_overhead_ms=kg_metrics.avg_total_ms - baseline_metrics.avg_total_ms, + expected_behavior=dataset.expected_behavior, + ) + + return BenchmarkRunResponse( + collection_id=collection_id, + dataset_id=dataset.id, + dataset_name=dataset.name, + top_k=top_k, + graph_depth=graph_depth, + baseline=baseline_metrics, + kg_rag=kg_metrics, + comparison=comparison, + results=rows, + ) + + @classmethod + def run_all( + cls, + db: Any, + collection_id: str, + dataset_ids: Sequence[str], + threshold: float = 0.0, + embedding_credentials: Optional[Dict[str, Any]] = None, + ) -> BenchmarkRunAllResponse: + seen = set() + runs: List[BenchmarkRunResponse] = [] + for dataset_id in dataset_ids or _DATASETS.keys(): + canonical = cls._canonical_dataset_id(dataset_id) + if canonical in seen: + continue + seen.add(canonical) + dataset = cls.get_dataset(canonical) + runs.append( + cls.run( + db=db, + collection_id=collection_id, + request=BenchmarkRunRequest( + dataset_id=canonical, + top_k=dataset.recommended_top_k, + graph_depth=dataset.recommended_graph_depth, + threshold=threshold, + ), + embedding_credentials=embedding_credentials, + ) + ) + + summary = { + "datasets": len(runs), + "avg_delta_recall_at_k": ( + mean(run.comparison.delta_recall_at_k for run in runs) if runs else 0.0 + ), + "avg_delta_mrr": ( + mean(run.comparison.delta_mrr for run in runs) if runs else 0.0 + ), + "avg_graph_overhead_ms": ( + mean(run.comparison.graph_overhead_ms for run in runs) if runs else 0.0 + ), + } + return BenchmarkRunAllResponse( + collection_id=collection_id, + runs=runs, + summary=summary, + ) + + @classmethod + def _dataset_summary(cls, dataset_id: str) -> BenchmarkDatasetSummary: + config = _DATASETS[dataset_id] + return BenchmarkDatasetSummary( + id=dataset_id, + name=config["name"], + description=config["description"], + expected_behavior=config["expected_behavior"], + recommended_top_k=config["recommended_top_k"], + recommended_graph_depth=config["recommended_graph_depth"], + question_count=len(config["questions"]), + ) + + @staticmethod + def _canonical_dataset_id(dataset_id: str) -> str: + normalized = (dataset_id or "educational").strip().lower().replace("_", "-") + canonical = _ALIASES.get(normalized, normalized) + if canonical not in _DATASETS: + valid = ", ".join(sorted(_DATASETS)) + raise HTTPException( + status_code=400, + detail=f"Unknown benchmark dataset '{dataset_id}'. Expected one of: {valid}", + ) + return canonical + + @classmethod + def _score_response( + cls, + *, + question: BenchmarkQuestion, + response: Dict[str, Any], + top_k: int, + vector_ms: float, + graph_ms: float, + total_ms: float, + ) -> BenchmarkQueryScore: + retrieved_files = cls._retrieved_files(response.get("results", []), top_k) + relevant = {cls._normalize_filename(name) for name in question.relevant_files} + relevant.discard("") + if not relevant: + return BenchmarkQueryScore( + vector_ms=vector_ms, + graph_ms=graph_ms, + total_ms=total_ms, + retrieved_files=retrieved_files, + ) + + hits = [filename in relevant for filename in retrieved_files] + precision = sum(1 for hit in hits if hit) / max(top_k, 1) + recall = len( + {filename for filename in retrieved_files if filename in relevant} + ) / len(relevant) + reciprocal = 0.0 + for index, hit in enumerate(hits, start=1): + if hit: + reciprocal = 1.0 / index + break + + return BenchmarkQueryScore( + precision_at_k=precision, + recall_at_k=recall, + mrr=reciprocal, + vector_ms=vector_ms, + graph_ms=graph_ms, + total_ms=total_ms, + retrieved_files=retrieved_files, + ) + + @staticmethod + def _aggregate(rows: Sequence[BenchmarkQueryScore]) -> BenchmarkMetrics: + if not rows: + return BenchmarkMetrics() + return BenchmarkMetrics( + precision_at_k=mean(row.precision_at_k for row in rows), + recall_at_k=mean(row.recall_at_k for row in rows), + mrr=mean(row.mrr for row in rows), + avg_vector_ms=mean(row.vector_ms for row in rows), + avg_graph_ms=mean(row.graph_ms for row in rows), + avg_total_ms=mean(row.total_ms for row in rows), + ) + + @classmethod + def _retrieved_files( + cls, results: Sequence[Dict[str, Any]], top_k: int + ) -> List[str]: + retrieved: List[str] = [] + seen = set() + for result in results: + metadata = result.get("metadata") or {} + filename = cls._result_filename(metadata) + if not filename or filename in seen: + continue + seen.add(filename) + retrieved.append(filename) + if len(retrieved) >= top_k: + break + return retrieved + + @classmethod + def _result_filename(cls, metadata: Dict[str, Any]) -> str: + for key in ("filename", "original_filename", "source", "file_path", "file_url"): + value = metadata.get(key) + filename = cls._normalize_filename(value) + if filename: + return filename + return "" + + @staticmethod + def _normalize_filename(value: Optional[Any]) -> str: + if not value: + return "" + text = str(value).strip().replace("\\", "/") + if not text: + return "" + return PurePosixPath(text).name.lower() + + @staticmethod + def _trace_from_results(results: Sequence[Dict[str, Any]]) -> Dict[str, Any]: + for result in results: + metadata = result.get("metadata") or {} + trace = metadata.get("kg_rag") + if isinstance(trace, dict): + return trace + return {} diff --git a/lamb-kb-server/backend/services/collection_service.py b/lamb-kb-server/backend/services/collection_service.py index f65b94231..22d06d39f 100644 --- a/lamb-kb-server/backend/services/collection_service.py +++ b/lamb-kb-server/backend/services/collection_service.py @@ -9,7 +9,12 @@ from config import STORAGE_DIR from database.models import Collection from fastapi import HTTPException, status -from plugins.base import ChunkingRegistry, EmbeddingRegistry, VectorDBRegistry +from plugins.base import ( + ChunkingRegistry, + EmbeddingRegistry, + LLMExtractionRegistry, + VectorDBRegistry, +) from plugins.chunking._common import validate_chunking_params from schemas.collection import CreateCollectionRequest, UpdateCollectionRequest from sqlalchemy.orm import Session @@ -36,6 +41,17 @@ def _validate_plugins(req: CreateCollectionRequest) -> None: f"Embedding vendor '{req.embedding.vendor}' is not registered. " f"Available: {[p['name'] for p in EmbeddingRegistry.list_plugins()]}" ) + # Validate extraction config when provided. Only enforced when + # graph_enabled=true; otherwise the field is irrelevant and ignored. + if getattr(req, "graph_enabled", False) and req.extraction is not None: + if req.extraction.vendor and not LLMExtractionRegistry.is_registered( + req.extraction.vendor + ): + errors.append( + f"LLM extraction vendor '{req.extraction.vendor}' is not " + f"registered. Available: " + f"{[p['name'] for p in LLMExtractionRegistry.list_plugins()]}" + ) if errors: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -162,6 +178,10 @@ def create_collection(db: Session, req: CreateCollectionRequest) -> Collection: backend_collection_id = backend_name # Persist metadata row. + # Extraction config is only meaningful when graph_enabled=true; we + # null it out otherwise so non-graph collections don't carry stale + # references to a vendor/model they never use. + extraction = req.extraction if bool(getattr(req, "graph_enabled", False)) else None collection = Collection( id=collection_id, organization_id=req.organization_id, @@ -175,6 +195,10 @@ def create_collection(db: Session, req: CreateCollectionRequest) -> Collection: vector_db_backend=req.vector_db_backend, backend_collection_id=backend_collection_id, storage_path=storage_path, + graph_enabled=bool(getattr(req, "graph_enabled", False)), + extraction_vendor=(extraction.vendor if extraction else None), + extraction_model=(extraction.model if extraction else None), + extraction_endpoint=(extraction.api_endpoint if extraction else None), status="ready", document_count=0, chunk_count=0, @@ -332,6 +356,7 @@ def delete_collection(db: Session, collection_id: str) -> None: """ collection = get_collection(db, collection_id) storage_path = collection.storage_path + graph_enabled = bool(getattr(collection, "graph_enabled", False)) # Step 2: drop vectors from the backend. Use the stored # backend_collection_id (with its "kb_" prefix) so the backend finds @@ -350,6 +375,26 @@ def delete_collection(db: Session, collection_id: str) -> None: collection_id, ) + # Step 2b: drop the matching subgraph in Neo4j if this collection + # had graph indexing enabled. Failure is non-fatal — the DB row is + # the source of truth and we'd rather orphan a few Neo4j nodes than + # block the user's delete. + if graph_enabled: + try: + import config as config_module # noqa: PLC0415 + + if config_module.get_kg_rag_config().get("enabled"): + from services.graph_store import get_graph_store # noqa: PLC0415 + + gs = get_graph_store() + if gs.is_configured() and gs.is_available(): + gs.delete_collection(collection_id) + except Exception: + logger.exception( + "Graph delete failed for collection %s — proceeding with DB delete", + collection_id, + ) + # Step 3: remove DB row first. db.delete(collection) db.commit() diff --git a/lamb-kb-server/backend/services/concept_extraction.py b/lamb-kb-server/backend/services/concept_extraction.py new file mode 100644 index 000000000..072630b05 --- /dev/null +++ b/lamb-kb-server/backend/services/concept_extraction.py @@ -0,0 +1,344 @@ +"""LLM-based concept and relationship extraction for KG-RAG indexing.""" + +from __future__ import annotations + +import json +import logging +import re +import unicodedata +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Set, Tuple + +import config as config_module +from plugins.base import LLMExtractionFunction, LLMExtractionRegistry + + +logger = logging.getLogger("lamb-kb") + + +@dataclass(frozen=True) +class TextChunk: + chunk_id: str + text: str + parent_text: str + metadata: Dict[str, Any] + + +@dataclass(frozen=True) +class ExtractedEntity: + name: str + display_name: str + entity_type: str = "concept" + description: str = "" + confidence: float = 1.0 + + +@dataclass(frozen=True) +class ExtractedRelationship: + source: str + target: str + relation: str + description: str = "" + evidence: str = "" + confidence: float = 1.0 + chunk_id: str = "" + + +@dataclass +class GraphExtraction: + concepts_by_chunk: Dict[str, List[str]] = field(default_factory=dict) + entities: Dict[str, ExtractedEntity] = field(default_factory=dict) + relationships: List[ExtractedRelationship] = field(default_factory=list) + + @property + def all_concepts(self) -> Set[str]: + return set(self.entities) + + +def normalize_concept(value: str) -> str: + value = unicodedata.normalize("NFKD", value) + value = "".join(char for char in value if not unicodedata.combining(char)) + value = re.sub(r"[^\w\- ]", " ", value.lower(), flags=re.UNICODE) + value = re.sub(r"\s+", " ", value).strip(" -_") + return value + + +def normalize_relation(value: str) -> str: + value = unicodedata.normalize("NFKD", value) + value = "".join(char for char in value if not unicodedata.combining(char)) + value = re.sub(r"[^\w]+", "_", value.lower(), flags=re.UNICODE) + value = re.sub(r"_+", "_", value).strip("_") + return value[:64] or "related_to" + + +def _clean_text(value: Any, limit: int = 500) -> str: + text = re.sub(r"\s+", " ", str(value or "")).strip() + return text[:limit] + + +def _confidence(value: Any) -> float: + try: + numeric = float(value) + except (TypeError, ValueError): + return 1.0 + return max(0.0, min(1.0, numeric)) + + +def _is_valid_entity_name(value: str) -> bool: + cleaned = _clean_text(value, limit=140) + normalized = normalize_concept(cleaned) + if len(normalized) < 2 or len(normalized) > 120: + return False + # Unreachable in practice: normalize_concept() strips to alphanumeric + # boundaries, so any normalized string of length >= 2 always contains an + # alphanumeric. Kept as a defensive guard. + if not any(char.isalpha() or char.isdigit() for char in normalized): # pragma: no cover + return False + if cleaned[0].isdigit(): + return False + if re.search(r"\.[a-z0-9]{1,8}$", cleaned.lower()): + return False + if normalized.endswith(" md"): + return False + if "\n" in cleaned or cleaned.count(".") > 1: + return False + if len(cleaned.split()) > 8: + return False + return True + + +class ConceptExtractor: + """Extract graph-ready concepts and typed relationships from text chunks.""" + + def __init__( + self, + kg_config: Optional[Dict[str, Any]] = None, + backend: Optional[LLMExtractionFunction] = None, + *, + vendor: Optional[str] = None, + model: Optional[str] = None, + api_endpoint: Optional[str] = None, + api_key: Optional[str] = None, + ): + """Build an extractor. + + Resolution order for the extraction backend: + 1. An explicit ``backend`` argument (used by tests). + 2. A per-collection ``vendor`` / ``model`` (passed from the + ingestion pipeline based on the collection's stored config). + 3. The server-level KG_RAG_* env defaults via ``kg_config``. + + ``api_key`` is request-scoped (ADR-4): when None, the chosen + backend falls back to its env-configured key. + """ + self.config = kg_config or config_module.get_kg_rag_config() + self.chat_model = self.config.get("chat_model") or "gpt-4o-mini" + self.model = model or self.config.get("extraction_model") or self.chat_model + configured_workers = int(self.config.get("extraction_max_workers") or 1) + self.max_workers = max(1, min(16, configured_workers)) + + timeout_seconds = float(self.config.get("openai_timeout_seconds") or 60.0) + + if backend is not None: + self.backend: Optional[LLMExtractionFunction] = backend + else: + # Default to OpenAI for back-compat with collections created + # before the vendor field existed. + resolved_vendor = vendor or "openai" + resolved_key = ( + api_key + if api_key is not None + else (self.config.get("openai_api_key") or "") + ) + resolved_endpoint = api_endpoint or "" + try: + self.backend = LLMExtractionRegistry.build( + resolved_vendor, + model=self.model, + api_key=resolved_key, + api_endpoint=resolved_endpoint, + timeout_seconds=timeout_seconds, + ) + except ValueError as exc: + logger.warning( + "KG-RAG: extraction vendor '%s' is not registered: %s", + resolved_vendor, + exc, + ) + self.backend = None + + def extract_for_chunks(self, chunks: List[TextChunk]) -> GraphExtraction: + if not chunks: + return GraphExtraction() + + extraction = GraphExtraction( + concepts_by_chunk={chunk.chunk_id: [] for chunk in chunks} + ) + if self.backend is None: + return extraction + + parent_groups: Dict[str, List[TextChunk]] = {} + for chunk in chunks: + parent_text = chunk.parent_text or chunk.text + parent_groups.setdefault(parent_text, []).append(chunk) + + groups = list(parent_groups.items()) + if self.max_workers > 1 and len(groups) > 1: + worker_count = min(self.max_workers, len(groups)) + with ThreadPoolExecutor(max_workers=worker_count) as executor: + futures = [ + executor.submit(self._extract_parent_group, parent_text, group) + for parent_text, group in groups + ] + group_results = [future.result() for future in futures] + else: + group_results = [ + self._extract_parent_group(parent_text, group) + for parent_text, group in groups + ] + + for group, parent_extraction in group_results: + + for entity_name, entity in parent_extraction.entities.items(): + extraction.entities.setdefault(entity_name, entity) + extraction.relationships.extend(parent_extraction.relationships) + + concept_names = sorted(parent_extraction.entities) + for chunk in group: + extraction.concepts_by_chunk[chunk.chunk_id] = concept_names + + return extraction + + def _extract_parent_group( + self, parent_text: str, group: List[TextChunk] + ) -> Tuple[List[TextChunk], GraphExtraction]: + source_labels = [ + str( + chunk.metadata.get("source_label") + or chunk.metadata.get("filename") + or chunk.chunk_id + ) + for chunk in group + ] + payload = self._extract_parent_text(parent_text, source_labels) + return group, self._parse_payload(payload, group[0].chunk_id) + + def _extract_parent_text( + self, text: str, source_labels: List[str] + ) -> Dict[str, Any]: + system = ( + "You extract knowledge-graph data for GraphRAG indexing. " + "Work for any domain and any document language. Do not use a fixed vocabulary. " + "Preserve entity names in the document language. Extract only entities or concepts that are explicit, specific, and useful for retrieval. " + "Do not output stopwords, generic verbs, generic adjectives, whole sentences, or vague phrases. " + "Extract persistent relationships only when the text states or strongly implies a typed connection. " + "Before returning, audit the graph and remove common nouns, example values, filenames, schema/property names, section labels, and entities that would not be meaningful outside this text. " + "Discard triples whose source or target is merely an example rank, a generic answer/evidence placeholder, a file/container word, or a schema field. " + "Prefer high-level named entities, technical terms, methods, systems, metrics, and domain concepts that participate in useful relationships. If unsure, omit the item. " + "Return only valid JSON matching the requested object shape." + ) + user = { + "task": "Extract entities and typed relationships from this text unit.", + "source_labels": source_labels, + "output_contract": { + "entities": [ + { + "name": "canonical entity or concept name from the text", + "type": "short type such as person, organization, concept, technology, metric, method, place, event", + "description": "one short description grounded in the text", + "confidence": "number from 0 to 1", + } + ], + "relationships": [ + { + "source": "entity name exactly as used in entities", + "target": "entity name exactly as used in entities", + "relation": "short verb phrase such as uses, stores, improves, depends_on, evaluates", + "description": "one short explanation grounded in the text", + "evidence": "short quote or paraphrase from the text", + "confidence": "number from 0 to 1", + } + ], + }, + # Doubled the per-chunk caps relative to the original + # legacy budget: small Wikipedia-style paragraphs routinely + # mention 10+ named entities, and the previous max=5 left + # bridging entities outside the graph. Cost stays bounded by + # the per-call ``max_tokens`` of the chat completion. + "limits": {"max_entities": 12, "max_relationships": 12}, + "text": text[:6000], + } + if self.backend is None: + return {"entities": [], "relationships": []} + fallback = self.chat_model if self.chat_model != self.model else None + parsed = self.backend.chat_json( + system=system, + user=json.dumps(user, ensure_ascii=False), + fallback_model=fallback, + ) + return parsed if isinstance(parsed, dict) else { + "entities": [], + "relationships": [], + } + + def _parse_payload(self, payload: Dict[str, Any], chunk_id: str) -> GraphExtraction: + entities: Dict[str, ExtractedEntity] = {} + raw_entities = payload.get("entities", []) + if not isinstance(raw_entities, list): + raw_entities = [] + + for item in raw_entities[:12]: + if not isinstance(item, dict): + continue + display_name = _clean_text(item.get("name"), limit=120) + if not _is_valid_entity_name(display_name): + continue + name = normalize_concept(display_name) + entities[name] = ExtractedEntity( + name=name, + display_name=display_name, + entity_type=normalize_relation(item.get("type") or "concept"), + description=_clean_text(item.get("description"), limit=600), + confidence=_confidence(item.get("confidence")), + ) + + relationships: List[ExtractedRelationship] = [] + raw_relationships = payload.get("relationships", []) + if not isinstance(raw_relationships, list): + raw_relationships = [] + + related_names: Set[str] = set() + for item in raw_relationships[:12]: + if not isinstance(item, dict): + continue + source = normalize_concept(_clean_text(item.get("source"), limit=120)) + target = normalize_concept(_clean_text(item.get("target"), limit=120)) + if source == target or source not in entities or target not in entities: + continue + relation = normalize_relation(item.get("relation") or "related_to") + relationships.append( + ExtractedRelationship( + source=source, + target=target, + relation=relation, + description=_clean_text(item.get("description"), limit=600), + evidence=_clean_text(item.get("evidence"), limit=500), + confidence=_confidence(item.get("confidence")), + chunk_id=chunk_id, + ) + ) + related_names.update({source, target}) + + if related_names: + entities = { + name: entity + for name, entity in entities.items() + if name in related_names + } + + return GraphExtraction( + concepts_by_chunk={chunk_id: sorted(entities)}, + entities=entities, + relationships=relationships, + ) diff --git a/lamb-kb-server/backend/services/graph_indexing.py b/lamb-kb-server/backend/services/graph_indexing.py new file mode 100644 index 000000000..21b58d184 --- /dev/null +++ b/lamb-kb-server/backend/services/graph_indexing.py @@ -0,0 +1,197 @@ +"""Shared helper that runs LLM concept extraction and writes graph data. + +Used by: + +* the ingestion path (``services/ingestion_service.py``) after Chroma + insertions complete, when the collection has ``graph_enabled=true``. +* the migration endpoint (``routers/graph.py`` ``/migrate``) when an + existing collection is opted into Graph RAG. + +The function fails *open* rather than rolling back vector ingestion: if +the LLM call or Neo4j write fails, vector retrieval still works — we +just don't get the graph augmentation. Returned dict contains ``error`` +when applicable so callers can surface it. +""" + +from __future__ import annotations + +import logging +import time +from typing import Any + +from database.models import Collection +from services.concept_extraction import ConceptExtractor, TextChunk + +logger = logging.getLogger(__name__) + + +def _build_chunks( + ids: list[str], + texts: list[str], + metadatas: list[dict[str, Any]], +) -> list[TextChunk]: + chunks: list[TextChunk] = [] + for i, chunk_id in enumerate(ids): + if not chunk_id or i >= len(texts): + continue + text = texts[i] or "" + if not text.strip(): + continue + metadata = metadatas[i] if i < len(metadatas) else {} + parent_text = "" + if isinstance(metadata, dict): + parent_text = str(metadata.get("parent_text") or "") + chunks.append( + TextChunk( + chunk_id=str(chunk_id), + text=text, + parent_text=parent_text or text, + metadata=dict(metadata or {}), + ) + ) + return chunks + + +def index_chunks_for_collection( + *, + collection: Collection, + ids: list[str], + texts: list[str], + metadatas: list[dict[str, Any]], + openai_api_key: str | None = None, + file_id: int | None = None, + filename: str | None = None, +) -> dict[str, Any]: + """Run extraction + Neo4j write for one batch of chunks. + + Args: + collection: ORM row, used for id + organization_id. + ids: Chunk IDs (must match what's in Chroma so KG-RAG expansion + can look them back up). + texts: Chunk text in the same order as ``ids``. + metadatas: Chunk metadata in the same order. The chunk's + ``permalink`` (when present) is preserved on the graph node. + openai_api_key: Per-request key. Falls back to + ``KG_RAG_OPENAI_API_KEY`` env when omitted. + file_id: Optional integer file registry ID (kept for parity with + the legacy graph schema; defaults to 0 in the new arch). + filename: Optional source filename. Falls back to the first chunk + metadata's ``filename`` / ``source_label`` / ``title``. + + Returns: + ``{"indexed": bool, "chunks": int, "extraction_ms": float, + "graph_ms": float, "error": str | None}``. + """ + import config as config_module + + kg_config = config_module.get_kg_rag_config() + if not kg_config.get("enabled"): + return {"indexed": False, "chunks": 0, "reason": "kg_rag_disabled"} + + chunks = _build_chunks(ids, texts, metadatas) + if not chunks: + return {"indexed": False, "chunks": 0, "reason": "no_chunks"} + + # Resolve the extraction vendor/model/endpoint from the collection + # (locked at creation), falling back to server-level env defaults for + # collections created before this surface existed. + coll_vendor = getattr(collection, "extraction_vendor", None) + coll_model = getattr(collection, "extraction_model", None) + coll_endpoint = getattr(collection, "extraction_endpoint", None) + resolved_vendor = coll_vendor or "openai" + + # Per-request key handling depends on vendor: OpenAI requires a key + # (returned-key fallback handled by the plugin); Ollama doesn't unless + # the operator put it behind an auth proxy. + api_key = (openai_api_key or kg_config.get("openai_api_key") or "").strip() + if resolved_vendor == "openai" and not api_key: + return { + "indexed": False, + "chunks": len(chunks), + "reason": "no_openai_api_key", + "error": ( + "OpenAI API key not supplied. Pass via X-OpenAI-Api-Key header " + "or set KG_RAG_OPENAI_API_KEY in the KB server env." + ), + } + + extractor_config = dict(kg_config) + extractor_config["openai_api_key"] = api_key + extractor = ConceptExtractor( + kg_config=extractor_config, + vendor=resolved_vendor, + model=coll_model, + api_endpoint=coll_endpoint, + api_key=api_key if resolved_vendor == "openai" else "", + ) + + extraction_start = time.perf_counter() + try: + extraction = extractor.extract_for_chunks(chunks) + except Exception as exc: # noqa: BLE001 + logger.warning("Concept extraction failed: %s", exc) + return { + "indexed": False, + "chunks": len(chunks), + "error": f"concept_extraction_failed: {exc}", + } + extraction_ms = (time.perf_counter() - extraction_start) * 1000 + + # Pick a filename if we weren't given one. + if not filename: + for ck in chunks: + for key in ("filename", "source_label", "title"): + value = ck.metadata.get(key) + if value: + filename = str(value) + break + if filename: + break + filename = filename or "collection_chunks" + + collection_payload = { + "id": collection.id, + "name": collection.name, + "organization_id": collection.organization_id, + "owner": collection.organization_id, + } + + from services.graph_store import get_graph_store + + graph_store = get_graph_store() + if not graph_store.is_configured() or not graph_store.is_available(): + return { + "indexed": False, + "chunks": len(chunks), + "error": "neo4j_unavailable", + } + + graph_start = time.perf_counter() + try: + graph_store.ingest_chunks( + collection=collection_payload, + file_id=file_id, + filename=filename, + chunks=chunks, + concepts_by_chunk=extraction.concepts_by_chunk, + entities=extraction.entities, + relationships=extraction.relationships, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("Graph write failed: %s", exc) + return { + "indexed": False, + "chunks": len(chunks), + "extraction_ms": extraction_ms, + "error": f"graph_write_failed: {exc}", + } + graph_ms = (time.perf_counter() - graph_start) * 1000 + + return { + "indexed": True, + "chunks": len(chunks), + "entities": len(extraction.entities), + "relationships": len(extraction.relationships), + "extraction_ms": extraction_ms, + "graph_ms": graph_ms, + } diff --git a/lamb-kb-server/backend/services/graph_store.py b/lamb-kb-server/backend/services/graph_store.py new file mode 100644 index 000000000..291a9bca6 --- /dev/null +++ b/lamb-kb-server/backend/services/graph_store.py @@ -0,0 +1,2584 @@ +"""Neo4j storage and traversal service for optional KG-RAG.""" + +from __future__ import annotations + +import itertools +import json +import logging +import time +from datetime import datetime, timezone +from typing import Any, Dict, Iterable, List, Optional, Set + +import config as config_module +from services.concept_extraction import ( + ExtractedEntity, + ExtractedRelationship, + TextChunk, + normalize_concept, + normalize_relation, +) + +try: + from neo4j import GraphDatabase +except Exception: # pragma: no cover - handled when dependency is not installed yet + GraphDatabase = None + + +logger = logging.getLogger("lamb-kb") +_GRAPH_STORE: Optional["GraphStore"] = None + + +def utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def get_graph_store() -> "GraphStore": + global _GRAPH_STORE + if _GRAPH_STORE is None: + _GRAPH_STORE = GraphStore() + return _GRAPH_STORE + + +class GraphStore: + """Small Neo4j wrapper used by ingestion and KG-RAG query plugins.""" + + def __init__(self, kg_config: Optional[Dict[str, Any]] = None): + self.config = kg_config or config_module.get_kg_rag_config() + self.enabled = bool(self.config.get("enabled")) + self.uri = self.config.get("neo4j_uri") or "" + self.user = self.config.get("neo4j_user") or "neo4j" + self.password = self.config.get("neo4j_password") or "" + self.driver = None + self._schema_ready = False + + if self.is_configured(): + try: + self.driver = GraphDatabase.driver( + self.uri, + auth=(self.user, self.password), + ) + except Exception as exc: + logger.warning("KG-RAG Neo4j driver could not be created: %s", exc) + + def is_configured(self) -> bool: + return bool( + self.enabled and GraphDatabase and self.uri and self.user and self.password + ) + + def close(self) -> None: + if self.driver is not None: + self.driver.close() + + def is_available(self) -> bool: + if self.driver is None: + return False + try: + self.driver.verify_connectivity() + return True + except Exception as exc: + logger.warning("KG-RAG Neo4j is unavailable: %s", exc) + return False + + def ensure_schema(self) -> bool: + if self._schema_ready: + return True + if not self.is_available(): + return False + + statements = [ + "CREATE CONSTRAINT org_id IF NOT EXISTS FOR (o:Organization) REQUIRE o.org_id IS UNIQUE", + "CREATE CONSTRAINT collection_id IF NOT EXISTS FOR (c:Collection) REQUIRE c.collection_id IS UNIQUE", + "CREATE CONSTRAINT document_id IF NOT EXISTS FOR (d:Document) REQUIRE d.document_id IS UNIQUE", + "CREATE CONSTRAINT chunk_id IF NOT EXISTS FOR (c:Chunk) REQUIRE c.chunk_id IS UNIQUE", + "CREATE CONSTRAINT concept_key IF NOT EXISTS FOR (c:Concept) REQUIRE (c.org_id, c.name) IS UNIQUE", + "CREATE INDEX concept_collection IF NOT EXISTS FOR (c:Concept) ON (c.org_id)", + "CREATE INDEX chunk_collection IF NOT EXISTS FOR (c:Chunk) ON (c.collection_id)", + "CREATE INDEX change_collection IF NOT EXISTS FOR (e:ChangeEvent) ON (e.collection_id)", + ] + try: + with self.driver.session() as session: + for statement in statements: + session.run(statement) + self._schema_ready = True + return True + except Exception as exc: + logger.warning("KG-RAG Neo4j schema setup failed: %s", exc) + return False + + def delete_collection(self, collection_id: str) -> None: + if not self.ensure_schema(): + return + with self.driver.session() as session: + session.run( + """ + MATCH ()-[rel]-() + WHERE rel.collection_id = $collection_id + DELETE rel + """, + collection_id=collection_id, + ) + session.run( + """ + MATCH (node) + WHERE node.collection_id = $collection_id + DETACH DELETE node + """, + collection_id=collection_id, + ) + session.run(""" + MATCH (concept:Concept) + WHERE NOT EXISTS { MATCH (:Chunk)-[:MENTIONS]->(concept) } + DETACH DELETE concept + """) + + def delete_document( + self, collection_id: str, org_id: str, source_item_id: str + ) -> None: + """Remove one document's graph data when its vectors are deleted. + + Deletes the Document node, all its Chunk nodes (and their MENTIONS / + CONTAINS edges via DETACH DELETE), and any Concept nodes that become + fully orphaned afterwards (no remaining MENTIONS edges in any + collection for this org). + """ + if not self.is_configured(): + return + if not self.ensure_schema(): + return + with self.driver.session() as session: + session.run( + """ + MATCH (doc:Document {collection_id: $collection_id, filename: $filename}) + -[:CONTAINS]->(chunk:Chunk) + DETACH DELETE chunk + """, + collection_id=collection_id, + filename=source_item_id, + ) + session.run( + """ + MATCH (doc:Document {collection_id: $collection_id, filename: $filename}) + DETACH DELETE doc + """, + collection_id=collection_id, + filename=source_item_id, + ) + session.run( + """ + MATCH (concept:Concept {org_id: $org_id}) + WHERE NOT EXISTS { MATCH (:Chunk)-[:MENTIONS]->(concept) } + DETACH DELETE concept + """, + org_id=org_id, + ) + + def list_changes( + self, + collection_id: str, + org_id: str, + *, + concept: Optional[str] = None, + relationship_source: Optional[str] = None, + relationship_target: Optional[str] = None, + relationship_relation: Optional[str] = None, + document_id: Optional[str] = None, + filename: Optional[str] = None, + operation: Optional[str] = None, + limit: int = 25, + ) -> List[Dict[str, Any]]: + if not self.ensure_schema(): + return [] + limit = max(1, min(int(limit or 25), 200)) + concept_filter = normalize_concept(concept or "") or None + source_filter = normalize_concept(relationship_source or "") or None + target_filter = normalize_concept(relationship_target or "") or None + relation_filter = ( + normalize_relation(relationship_relation) if relationship_relation else None + ) + fetch_limit = min(max(limit * 10, 100), 1000) + with self.driver.session() as session: + rows = session.run( + """ + MATCH (event:ChangeEvent {collection_id: $collection_id, org_id: $org_id}) + OPTIONAL MATCH (event)-[:RECORDED_CHANGE]->(doc:Document) + WHERE ($concept IS NULL OR $concept IN coalesce(event.concepts, [])) + AND ($relationship_source IS NULL OR $relationship_source IN coalesce(event.concepts, [])) + AND ($relationship_target IS NULL OR $relationship_target IN coalesce(event.concepts, [])) + AND ($document_id IS NULL OR doc.document_id = $document_id) + AND ($filename IS NULL OR event.filename = $filename OR doc.filename = $filename) + AND ($operation IS NULL OR event.operation = $operation) + RETURN event.event_id AS event_id, + event.collection_id AS collection_id, + event.org_id AS org_id, + event.operation AS operation, + event.actor AS actor, + event.timestamp AS timestamp, + event.filename AS filename, + coalesce(event.concepts, []) AS concepts, + event.payload_json AS payload_json, + doc.document_id AS document_id, + doc.file_id AS file_id + ORDER BY event.timestamp DESC + LIMIT $fetch_limit + """, + collection_id=collection_id, + org_id=org_id, + concept=concept_filter, + relationship_source=source_filter, + relationship_target=target_filter, + document_id=document_id, + filename=filename, + operation=operation, + fetch_limit=fetch_limit, + ).data() + if source_filter or target_filter or relation_filter: + rows = [ + row + for row in rows + if GraphStore._event_matches_relationship( + row, + source_filter, + target_filter, + relation_filter, + ) + ] + elif concept_filter: + rows = [ + row + for row in rows + if row.get("operation") + not in { + "manual_edit_relationship", + "manual_curate_relationship", + "manual_expunge_relationship", + } + ] + return rows[:limit] + + @staticmethod + def _event_payload(row: Dict[str, Any]) -> Dict[str, Any]: + try: + payload = json.loads(row.get("payload_json") or "{}") + except (TypeError, json.JSONDecodeError): + payload = {} + return payload if isinstance(payload, dict) else {} + + @staticmethod + def _event_matches_relationship( + row: Dict[str, Any], + source: Optional[str], + target: Optional[str], + relation: Optional[str], + ) -> bool: + payload = GraphStore._event_payload(row) + + def relationship_matches(candidate: Dict[str, Any]) -> bool: + candidate_source = normalize_concept(str(candidate.get("source") or "")) + candidate_target = normalize_concept(str(candidate.get("target") or "")) + candidate_relations = { + normalize_relation(str(candidate.get("relation") or "related_to")) + } + for key in ("new_relation", "old_relation", "target_relation"): + value = candidate.get(key) + if value: + candidate_relations.add(normalize_relation(str(value))) + if source and candidate_source != source: + return False + if target and candidate_target != target: + return False + return not relation or relation in candidate_relations + + details = payload.get("relationship_details") + if isinstance(details, list) and any( + isinstance(item, dict) and relationship_matches(item) for item in details + ): + return True + + removed = payload.get("removed_relationships") + if isinstance(removed, list) and any( + isinstance(item, dict) and relationship_matches(item) for item in removed + ): + return True + + if payload.get("source") or payload.get("target"): + return relationship_matches(payload) + return False + + def get_change( + self, collection_id: str, org_id: str, event_id: str + ) -> Optional[Dict[str, Any]]: + if not self.ensure_schema(): + return None + with self.driver.session() as session: + row = session.run( + """ + MATCH (event:ChangeEvent { + event_id: $event_id, + collection_id: $collection_id, + org_id: $org_id + }) + OPTIONAL MATCH (event)-[:RECORDED_CHANGE]->(doc:Document) + OPTIONAL MATCH (doc)-[:CONTAINS]->(chunk:Chunk) + RETURN event.event_id AS event_id, + event.collection_id AS collection_id, + event.org_id AS org_id, + event.operation AS operation, + event.actor AS actor, + event.timestamp AS timestamp, + event.filename AS filename, + coalesce(event.concepts, []) AS concepts, + event.payload_json AS payload_json, + doc.document_id AS document_id, + doc.file_id AS file_id, + collect(DISTINCT chunk.chunk_id) AS chunk_ids + """, + event_id=event_id, + collection_id=collection_id, + org_id=org_id, + ).single() + return dict(row) if row else None + + def get_collection_graph( + self, + collection_id: str, + org_id: str, + *, + concept: Optional[str] = None, + document_id: Optional[str] = None, + chunk_id: Optional[str] = None, + filename: Optional[str] = None, + include_chunks: bool = True, + limit: int = 60, + ) -> Dict[str, Any]: + if not self.ensure_schema(): + return { + "collection_id": collection_id, + "nodes": [], + "edges": [], + "filters": { + "concept": concept or "", + "document_id": document_id or "", + "chunk_id": chunk_id or "", + "filename": filename or "", + "include_chunks": include_chunks, + "limit": limit, + }, + "counts": {"concepts": 0, "documents": 0, "chunks": 0, "edges": 0}, + } + + limit = max(1, min(int(limit or 60), 200)) + concept_filter = normalize_concept(concept or "") or None + filename_filter = (filename or "").strip().lower() or None + + with self.driver.session() as session: + concept_rows = session.run( + """ + MATCH (concept:Concept {org_id: $org_id}) + WHERE ( + EXISTS { MATCH (:Chunk {collection_id: $collection_id})-[:MENTIONS]->(concept) } + OR EXISTS { MATCH (concept)-[rel:RELATES_TO]-(:Concept) WHERE rel.collection_id = $collection_id } + ) + AND ( + $concept_filter IS NULL + OR concept.name CONTAINS $concept_filter + OR toLower(coalesce(concept.display_name, '')) CONTAINS $concept_filter + ) + AND ( + $document_id IS NULL + OR EXISTS { + MATCH (:Document {document_id: $document_id})-[:CONTAINS]->(:Chunk {collection_id: $collection_id})-[:MENTIONS]->(concept) + } + ) + AND ( + $chunk_id IS NULL + OR EXISTS { + MATCH (:Chunk {collection_id: $collection_id, chunk_id: $chunk_id})-[:MENTIONS]->(concept) + } + ) + AND ( + $filename_filter IS NULL + OR EXISTS { + MATCH (doc:Document {collection_id: $collection_id})-[:CONTAINS]->(:Chunk {collection_id: $collection_id})-[:MENTIONS]->(concept) + WHERE toLower(coalesce(doc.filename, '')) CONTAINS $filename_filter + OR toLower(coalesce(doc.document_id, '')) CONTAINS $filename_filter + } + ) + OPTIONAL MATCH (concept)<-[:MENTIONS]-(chunk:Chunk {collection_id: $collection_id}) + RETURN concept.name AS name, + coalesce(concept.display_name, concept.name) AS display_name, + coalesce(concept.entity_type, 'concept') AS entity_type, + concept.description AS description, + concept.confidence AS confidence, + concept.notes AS notes, + coalesce(concept.tags, []) AS tags, + count(DISTINCT chunk) AS chunk_count + ORDER BY chunk_count DESC, display_name ASC + LIMIT $limit + """, + collection_id=collection_id, + org_id=org_id, + concept_filter=concept_filter, + document_id=document_id, + chunk_id=chunk_id, + filename_filter=filename_filter, + limit=limit, + ).data() + + concept_names = [row["name"] for row in concept_rows] + if not concept_names: + return { + "collection_id": collection_id, + "nodes": [], + "edges": [], + "filters": { + "concept": concept or "", + "document_id": document_id or "", + "chunk_id": chunk_id or "", + "filename": filename or "", + "include_chunks": include_chunks, + "limit": limit, + }, + "counts": {"concepts": 0, "documents": 0, "chunks": 0, "edges": 0}, + } + + # Per-collection verification: read from MENTIONS relationships + # (scoped to this collection) rather than the org-level Concept node. + # This ensures a fresh KS starts with all concepts unverified even if + # the same concepts were approved in a different KS. + coll_vs_rows = session.run( + """ + MATCH (:Chunk {collection_id: $collection_id})-[m:MENTIONS]->(concept:Concept {org_id: $org_id}) + WHERE concept.name IN $concept_names + WITH concept.name AS name, collect(m.verification_state)[0] AS vs + RETURN name, vs + """, + collection_id=collection_id, + org_id=org_id, + concept_names=concept_names, + ).data() + coll_vs: Dict[str, Optional[str]] = {row["name"]: row["vs"] for row in coll_vs_rows} + # Filter out concepts rejected in this collection + concept_rows = [r for r in concept_rows if coll_vs.get(r["name"]) != "rejected"] + concept_names = [r["name"] for r in concept_rows] + + document_rows = session.run( + """ + MATCH (doc:Document {collection_id: $collection_id})-[:CONTAINS]->(chunk:Chunk {collection_id: $collection_id})-[:MENTIONS]->(concept:Concept {org_id: $org_id}) + WHERE concept.name IN $concept_names + AND ($document_id IS NULL OR doc.document_id = $document_id) + AND ($chunk_id IS NULL OR chunk.chunk_id = $chunk_id) + AND ( + $filename_filter IS NULL + OR toLower(coalesce(doc.filename, '')) CONTAINS $filename_filter + OR toLower(coalesce(doc.document_id, '')) CONTAINS $filename_filter + ) + WITH doc, count(DISTINCT chunk) AS chunk_count, collect(DISTINCT concept.name) AS concepts + RETURN doc.document_id AS document_id, + doc.filename AS filename, + doc.file_id AS file_id, + chunk_count AS chunk_count, + concepts[0..10] AS concepts + ORDER BY filename ASC, document_id ASC + LIMIT $document_limit + """, + collection_id=collection_id, + org_id=org_id, + concept_names=concept_names, + document_id=document_id, + chunk_id=chunk_id, + filename_filter=filename_filter, + document_limit=limit, + ).data() + + relationship_rows = session.run( + """ + MATCH (source:Concept {org_id: $org_id})-[rel:RELATES_TO]->(target:Concept {org_id: $org_id}) + WHERE rel.collection_id = $collection_id + AND source.name IN $concept_names + AND target.name IN $concept_names + AND coalesce(rel.verification_state, '') <> 'rejected' + RETURN source.name AS source, + target.name AS target, + type(rel) AS type, + coalesce(rel.relation, type(rel)) AS relation, + coalesce(rel.weight, 1.0) AS weight, + rel.description AS description, + rel.evidence AS evidence, + rel.chunk_id AS chunk_id, + rel.notes AS notes, + coalesce(rel.tags, []) AS tags, + rel.verification_state AS verification_state + ORDER BY weight DESC, source ASC, target ASC + LIMIT $edge_limit + """, + collection_id=collection_id, + org_id=org_id, + concept_names=concept_names, + edge_limit=limit * 2, + ).data() + + chunk_rows: List[Dict[str, Any]] = [] + if include_chunks: + chunk_rows = session.run( + """ + MATCH (doc:Document {collection_id: $collection_id})-[:CONTAINS]->(chunk:Chunk {collection_id: $collection_id})-[:MENTIONS]->(concept:Concept {org_id: $org_id}) + WHERE concept.name IN $concept_names + AND ($document_id IS NULL OR doc.document_id = $document_id) + AND ($chunk_id IS NULL OR chunk.chunk_id = $chunk_id) + AND ( + $filename_filter IS NULL + OR toLower(coalesce(doc.filename, '')) CONTAINS $filename_filter + OR toLower(coalesce(doc.document_id, '')) CONTAINS $filename_filter + ) + WITH chunk, doc, collect(DISTINCT concept.name) AS concepts + RETURN chunk.chunk_id AS chunk_id, + coalesce(chunk.source_label, chunk.chunk_id) AS source_label, + chunk.filename AS filename, + doc.document_id AS document_id, + left(coalesce(chunk.text, ''), 240) AS text_preview, + coalesce(chunk.permalink_original, '') AS permalink_original, + coalesce(chunk.permalink_full_markdown, '') AS permalink_full_markdown, + coalesce(chunk.permalink_page, '') AS permalink_page, + concepts AS concepts + ORDER BY filename ASC, source_label ASC + LIMIT $chunk_limit + """, + collection_id=collection_id, + org_id=org_id, + concept_names=concept_names, + document_id=document_id, + chunk_id=chunk_id, + filename_filter=filename_filter, + chunk_limit=limit * 3, + ).data() + + nodes: List[Dict[str, Any]] = [] + edges: List[Dict[str, Any]] = [] + + for row in concept_rows: + nodes.append( + { + "id": f"concept:{row['name']}", + "type": "concept", + "label": row.get("display_name") or row["name"], + "data": { + "name": row["name"], + "entity_type": row.get("entity_type") or "concept", + "description": row.get("description") or "", + "confidence": row.get("confidence"), + "notes": row.get("notes") or "", + "tags": row.get("tags") or [], + "verification_state": coll_vs.get(row["name"]) or "unverified", + "chunk_count": int(row.get("chunk_count") or 0), + }, + } + ) + + for row in document_rows: + row_document_id = row.get("document_id") + if not row_document_id: + continue + nodes.append( + { + "id": f"document:{row_document_id}", + "type": "document", + "label": row.get("filename") or row_document_id, + "data": { + "document_id": row_document_id, + "filename": row.get("filename") or "", + "file_id": row.get("file_id"), + "chunk_count": int(row.get("chunk_count") or 0), + "concepts": row.get("concepts") or [], + }, + } + ) + + for row in relationship_rows: + edge_type = row.get("type") or "RELATES_TO" + relation = row.get("relation") or edge_type + edges.append( + { + "id": f"relationship:{row['source']}:{relation}:{row['target']}", + "type": edge_type, + "source": f"concept:{row['source']}", + "target": f"concept:{row['target']}", + "label": relation, + "weight": float(row.get("weight") or 1.0), + "data": { + "source": row["source"], + "target": row["target"], + "relation": relation, + "description": row.get("description") or "", + "evidence": row.get("evidence") or "", + "chunk_id": row.get("chunk_id") or "", + "notes": row.get("notes") or "", + "tags": row.get("tags") or [], + "verification_state": row.get("verification_state") + or "unverified", + }, + } + ) + + for row in chunk_rows: + chunk_id = row.get("chunk_id") + if not chunk_id: + continue + nodes.append( + { + "id": f"chunk:{chunk_id}", + "type": "chunk", + "label": row.get("source_label") or chunk_id, + "data": { + "chunk_id": chunk_id, + "filename": row.get("filename") or "", + "document_id": row.get("document_id") or "", + "text_preview": row.get("text_preview") or "", + "permalink_original": row.get("permalink_original") or "", + "permalink_full_markdown": row.get( + "permalink_full_markdown" + ) + or "", + "permalink_page": row.get("permalink_page") or "", + "concepts": row.get("concepts") or [], + }, + } + ) + for mentioned_concept in row.get("concepts") or []: + edges.append( + { + "id": f"mention:{chunk_id}:{mentioned_concept}", + "type": "MENTIONS", + "source": f"chunk:{chunk_id}", + "target": f"concept:{mentioned_concept}", + "label": "mentions", + "weight": 1.0, + "data": {"chunk_id": chunk_id, "concept": mentioned_concept}, + } + ) + document_id_value = row.get("document_id") or "" + if document_id_value: + edges.append( + { + "id": f"contains:{document_id_value}:{chunk_id}", + "type": "CONTAINS", + "source": f"document:{document_id_value}", + "target": f"chunk:{chunk_id}", + "label": "contains", + "weight": 1.0, + "data": { + "document_id": document_id_value, + "chunk_id": chunk_id, + }, + } + ) + + if not include_chunks: + for row in document_rows: + document_id_value = row.get("document_id") or "" + if not document_id_value: + continue + for mentioned_concept in row.get("concepts") or []: + edges.append( + { + "id": f"document-mention:{document_id_value}:{mentioned_concept}", + "type": "DOCUMENT_MENTIONS", + "source": f"document:{document_id_value}", + "target": f"concept:{mentioned_concept}", + "label": "mentions", + "weight": 1.0, + "data": { + "document_id": document_id_value, + "concept": mentioned_concept, + }, + } + ) + + return { + "collection_id": collection_id, + "nodes": nodes, + "edges": edges, + "filters": { + "concept": concept or "", + "document_id": document_id or "", + "chunk_id": chunk_id or "", + "filename": filename or "", + "include_chunks": include_chunks, + "limit": limit, + }, + "counts": { + "concepts": len(concept_rows), + "documents": len(document_rows), + "chunks": len(chunk_rows), + "edges": len(edges), + }, + } + + def revert_change( + self, + collection_id: str, + org_id: str, + event_id: str, + *, + actor: str = "graph-traceability-api", + reason: str = "", + ) -> Dict[str, Any]: + if not self.ensure_schema(): + return {"reverted": False, "reason": "neo4j_not_available"} + timestamp = utc_now() + with self.driver.session() as session: + return session.execute_write( + self._revert_change_tx, + collection_id, + org_id, + event_id, + actor, + reason, + timestamp, + ) + + @staticmethod + def _revert_change_tx( + tx, + collection_id: str, + org_id: str, + event_id: str, + actor: str, + reason: str, + timestamp: str, + ) -> Dict[str, Any]: + event_row = tx.run( + """ + MATCH (event:ChangeEvent { + event_id: $event_id, + collection_id: $collection_id, + org_id: $org_id + }) + OPTIONAL MATCH (event)-[:RECORDED_CHANGE]->(doc:Document) + RETURN event.operation AS operation, + event.filename AS filename, + coalesce(event.concepts, []) AS concepts, + event.payload_json AS payload_json, + doc.document_id AS document_id + """, + event_id=event_id, + collection_id=collection_id, + org_id=org_id, + ).single() + if not event_row: + return { + "reverted": False, + "reason": "change_not_found", + "event_id": event_id, + } + + operation = event_row.get("operation") + try: + payload = json.loads(event_row.get("payload_json") or "{}") + except (TypeError, json.JSONDecodeError): + payload = {} + + if operation == "manual_expunge_relationship": + return GraphStore._restore_expunged_relationship_tx( + tx, + collection_id, + org_id, + event_id, + event_row, + payload, + actor, + reason, + timestamp, + ) + if operation == "manual_expunge_concept": + return GraphStore._restore_expunged_concept_tx( + tx, + collection_id, + org_id, + event_id, + event_row, + payload, + actor, + reason, + timestamp, + ) + + if operation != "automatic_ingestion": + return { + "reverted": False, + "reason": "unsupported_operation", + "event_id": event_id, + "operation": operation, + } + + document_id = event_row.get("document_id") + if not document_id: + return { + "reverted": False, + "reason": "change_has_no_document", + "event_id": event_id, + } + + chunk_row = tx.run( + """ + MATCH (doc:Document {document_id: $document_id, collection_id: $collection_id})-[:CONTAINS]->(chunk:Chunk) + RETURN collect(chunk.chunk_id) AS chunk_ids + """, + document_id=document_id, + collection_id=collection_id, + ).single() + chunk_ids = list(chunk_row.get("chunk_ids") or []) if chunk_row else [] + + relationship_details = payload.get("relationship_details") + if not isinstance(relationship_details, list): + relationship_details = payload.get("relationships") + if not isinstance(relationship_details, list): + relationship_details = [] + + for relationship in relationship_details: + if not isinstance(relationship, dict): + continue + tx.run( + """ + MATCH (source:Concept {org_id: $org_id, name: $source}) + MATCH (target:Concept {org_id: $org_id, name: $target}) + MATCH (source)-[rel:RELATES_TO {collection_id: $collection_id, relation: $relation}]->(target) + SET rel.weight = coalesce(rel.weight, 0) - $confidence + WITH rel + WHERE coalesce(rel.weight, 0) <= 0 + DELETE rel + """, + org_id=org_id, + collection_id=collection_id, + source=relationship.get("source") or "", + target=relationship.get("target") or "", + relation=relationship.get("relation") or "related_to", + confidence=float(relationship.get("confidence") or 1.0), + ) + + tx.run( + """ + MATCH (doc:Document {document_id: $document_id, collection_id: $collection_id}) + OPTIONAL MATCH (doc)-[:CONTAINS]->(chunk:Chunk) + DETACH DELETE chunk + WITH doc + DETACH DELETE doc + """, + document_id=document_id, + collection_id=collection_id, + ) + tx.run( + """ + MATCH (concept:Concept {org_id: $org_id}) + WHERE NOT EXISTS { MATCH (:Chunk)-[:MENTIONS]->(concept) } + DETACH DELETE concept + """, + org_id=org_id, + ) + revert_row = tx.run( + """ + CREATE (event:ChangeEvent { + event_id: randomUUID(), + collection_id: $collection_id, + org_id: $org_id, + operation: 'revert_change', + actor: $actor, + timestamp: $timestamp, + filename: $filename, + concepts: $concepts, + payload_json: $payload_json + }) + WITH event + MATCH (original:ChangeEvent {event_id: $event_id, collection_id: $collection_id, org_id: $org_id}) + MERGE (event)-[:REVERTS]->(original) + RETURN event.event_id AS revert_event_id + """, + collection_id=collection_id, + org_id=org_id, + event_id=event_id, + actor=actor, + timestamp=timestamp, + filename=event_row.get("filename") or "", + concepts=event_row.get("concepts") or [], + payload_json=json.dumps( + { + "reverted_event_id": event_id, + "reverted_operation": operation, + "document_id": document_id, + "chunk_ids": chunk_ids, + "reason": reason, + }, + ensure_ascii=False, + ), + ).single() + return { + "reverted": True, + "event_id": event_id, + "revert_event_id": ( + revert_row.get("revert_event_id") if revert_row else None + ), + "operation": operation, + "document_id": document_id, + "chunk_ids": chunk_ids, + } + + @staticmethod + def _create_revert_event_tx( + tx, + collection_id: str, + org_id: str, + event_id: str, + actor: str, + timestamp: str, + filename: str, + concepts: List[str], + payload: Dict[str, Any], + ) -> Optional[str]: + revert_row = tx.run( + """ + CREATE (event:ChangeEvent { + event_id: randomUUID(), + collection_id: $collection_id, + org_id: $org_id, + operation: 'revert_change', + actor: $actor, + timestamp: $timestamp, + filename: $filename, + concepts: $concepts, + payload_json: $payload_json + }) + WITH event + MATCH (original:ChangeEvent {event_id: $event_id, collection_id: $collection_id, org_id: $org_id}) + MERGE (event)-[:REVERTS]->(original) + RETURN event.event_id AS revert_event_id + """, + collection_id=collection_id, + org_id=org_id, + event_id=event_id, + actor=actor, + timestamp=timestamp, + filename=filename, + concepts=sorted({concept for concept in concepts if concept}), + payload_json=json.dumps(payload, ensure_ascii=False), + ).single() + return revert_row.get("revert_event_id") if revert_row else None + + @staticmethod + def _restore_concept_node_tx( + tx, + collection_id: str, + org_id: str, + concept: str, + timestamp: str, + *, + notes: Optional[str] = None, + tags: Optional[List[str]] = None, + ) -> str: + normalized = normalize_concept(concept or "") + if not normalized: + return "" + tx.run( + """ + MERGE (concept:Concept {org_id: $org_id, name: $concept}) + ON CREATE SET concept.created_at = $timestamp, + concept.sources = [] + SET concept.updated_at = $timestamp, + concept.collection_hint = $collection_id, + concept.display_name = coalesce(concept.display_name, $concept), + concept.entity_type = coalesce(concept.entity_type, 'concept'), + concept.notes = CASE WHEN $notes IS NULL THEN concept.notes ELSE $notes END, + concept.tags = CASE WHEN $tags IS NULL THEN concept.tags ELSE $tags END, + concept.verification_state = 'verified' + """, + collection_id=collection_id, + org_id=org_id, + concept=normalized, + notes=notes, + tags=tags if isinstance(tags, list) else None, + timestamp=timestamp, + ) + return normalized + + @staticmethod + def _relationship_restore_weight(value: Any) -> float: + try: + return float(value if value is not None else 1.0) + except (TypeError, ValueError): + return 1.0 + + @staticmethod + def _restore_relationship_payload_tx( + tx, + collection_id: str, + org_id: str, + relationship: Dict[str, Any], + timestamp: str, + ) -> Optional[Dict[str, str]]: + source = GraphStore._restore_concept_node_tx( + tx, + collection_id, + org_id, + str(relationship.get("source") or ""), + timestamp, + ) + target = GraphStore._restore_concept_node_tx( + tx, + collection_id, + org_id, + str(relationship.get("target") or ""), + timestamp, + ) + relation = normalize_relation(str(relationship.get("relation") or "related_to")) + if not source or not target or not relation: + return None + tags = relationship.get("tags") + tx.run( + """ + MATCH (source:Concept {org_id: $org_id, name: $source}) + MATCH (target:Concept {org_id: $org_id, name: $target}) + MERGE (source)-[rel:RELATES_TO {collection_id: $collection_id, relation: $relation}]->(target) + ON CREATE SET rel.created_at = $timestamp + SET rel.updated_at = $timestamp, + rel.weight = $weight, + rel.description = $description, + rel.evidence = $evidence, + rel.chunk_id = $chunk_id, + rel.notes = $notes, + rel.tags = $tags, + rel.verification_state = 'verified' + """, + collection_id=collection_id, + org_id=org_id, + source=source, + target=target, + relation=relation, + weight=GraphStore._relationship_restore_weight(relationship.get("weight")), + description=relationship.get("description") or "", + evidence=relationship.get("evidence") or "", + chunk_id=relationship.get("chunk_id") or "", + notes=relationship.get("notes"), + tags=tags if isinstance(tags, list) else [], + timestamp=timestamp, + ) + return {"source": source, "target": target, "relation": relation} + + @staticmethod + def _restore_expunged_relationship_tx( + tx, + collection_id: str, + org_id: str, + event_id: str, + event_row: Dict[str, Any], + payload: Dict[str, Any], + actor: str, + reason: str, + timestamp: str, + ) -> Dict[str, Any]: + restored = GraphStore._restore_relationship_payload_tx( + tx, + collection_id, + org_id, + { + "source": payload.get("source"), + "target": payload.get("target"), + "relation": payload.get("relation") or payload.get("new_relation"), + "weight": payload.get("old_weight"), + "description": payload.get("old_description"), + "evidence": payload.get("old_evidence"), + "chunk_id": payload.get("old_chunk_id"), + "notes": payload.get("old_notes"), + "tags": payload.get("old_tags"), + }, + timestamp, + ) + if not restored: + return { + "reverted": False, + "reason": "invalid_expunge_payload", + "event_id": event_id, + "operation": event_row.get("operation"), + } + revert_event_id = GraphStore._create_revert_event_tx( + tx, + collection_id, + org_id, + event_id, + actor, + timestamp, + event_row.get("filename") or "", + [restored["source"], restored["target"]], + { + "reverted_event_id": event_id, + "reverted_operation": event_row.get("operation"), + "source": restored["source"], + "target": restored["target"], + "relation": restored["relation"], + "verification_state": "verified", + "reason": reason, + }, + ) + return { + "reverted": True, + "event_id": event_id, + "revert_event_id": revert_event_id, + "operation": event_row.get("operation"), + "chunk_ids": [], + } + + @staticmethod + def _restore_expunged_concept_tx( + tx, + collection_id: str, + org_id: str, + event_id: str, + event_row: Dict[str, Any], + payload: Dict[str, Any], + actor: str, + reason: str, + timestamp: str, + ) -> Dict[str, Any]: + concept = GraphStore._restore_concept_node_tx( + tx, + collection_id, + org_id, + str(payload.get("concept") or (event_row.get("concepts") or [""])[0]), + timestamp, + notes=payload.get("old_notes"), + tags=payload.get("old_tags"), + ) + if not concept: + return { + "reverted": False, + "reason": "invalid_expunge_payload", + "event_id": event_id, + "operation": event_row.get("operation"), + } + + chunk_ids = [ + str(chunk_id) + for chunk_id in payload.get("removed_chunk_mentions") or [] + if chunk_id + ] + for chunk_id in chunk_ids: + tx.run( + """ + MATCH (chunk:Chunk {collection_id: $collection_id, chunk_id: $chunk_id}) + MATCH (concept:Concept {org_id: $org_id, name: $concept}) + MERGE (chunk)-[mention:MENTIONS]->(concept) + ON CREATE SET mention.created_at = $timestamp + SET mention.collection_id = $collection_id + """, + collection_id=collection_id, + org_id=org_id, + chunk_id=chunk_id, + concept=concept, + timestamp=timestamp, + ) + + restored_relationships: List[Dict[str, str]] = [] + for relationship in payload.get("removed_relationships") or []: + if not isinstance(relationship, dict): + continue + restored = GraphStore._restore_relationship_payload_tx( + tx, + collection_id, + org_id, + relationship, + timestamp, + ) + if restored: + restored_relationships.append(restored) + + touched_concepts = {concept} + for relationship in restored_relationships: + touched_concepts.add(relationship["source"]) + touched_concepts.add(relationship["target"]) + revert_event_id = GraphStore._create_revert_event_tx( + tx, + collection_id, + org_id, + event_id, + actor, + timestamp, + event_row.get("filename") or "", + list(touched_concepts), + { + "reverted_event_id": event_id, + "reverted_operation": event_row.get("operation"), + "concept": concept, + "verification_state": "verified", + "restored_chunk_mentions": chunk_ids, + "restored_relationships": restored_relationships, + "reason": reason, + }, + ) + return { + "reverted": True, + "event_id": event_id, + "revert_event_id": revert_event_id, + "operation": event_row.get("operation"), + "chunk_ids": chunk_ids, + } + + def rename_concept( + self, + collection_id: str, + org_id: str, + old_name: str, + new_name: str, + *, + actor: str = "graph-curation-api", + reason: str = "", + ) -> Dict[str, Any]: + if not self.ensure_schema(): + return {"ok": False, "reason": "neo4j_not_available"} + timestamp = utc_now() + with self.driver.session() as session: + return session.execute_write( + self._rename_concept_tx, + collection_id, + org_id, + old_name, + new_name, + actor, + reason, + timestamp, + ) + + def merge_concepts( + self, + collection_id: str, + org_id: str, + source_names: List[str], + target_name: str, + *, + actor: str = "graph-curation-api", + reason: str = "", + ) -> Dict[str, Any]: + if not self.ensure_schema(): + return {"ok": False, "reason": "neo4j_not_available"} + timestamp = utc_now() + with self.driver.session() as session: + return session.execute_write( + self._merge_concepts_tx, + collection_id, + org_id, + source_names, + target_name, + actor, + reason, + timestamp, + ) + + def edit_relationship( + self, + collection_id: str, + org_id: str, + *, + source_name: str, + target_name: str, + relation: str, + new_relation: Optional[str] = None, + weight: Optional[float] = None, + description: Optional[str] = None, + evidence: Optional[str] = None, + notes: Optional[str] = None, + tags: Optional[List[str]] = None, + verification_state: Optional[str] = None, + actor: str = "graph-curation-api", + reason: str = "", + operation: str = "manual_edit_relationship", + ) -> Dict[str, Any]: + if not self.ensure_schema(): + return {"ok": False, "reason": "neo4j_not_available"} + timestamp = utc_now() + with self.driver.session() as session: + return session.execute_write( + self._edit_relationship_tx, + collection_id, + org_id, + source_name, + target_name, + relation, + new_relation, + weight, + description, + evidence, + notes, + tags, + verification_state, + actor, + reason, + operation, + timestamp, + ) + + def update_concept_curation( + self, + collection_id: str, + org_id: str, + concept_name: str, + *, + notes: Optional[str] = None, + tags: Optional[List[str]] = None, + verification_state: Optional[str] = None, + actor: str = "graph-curation-api", + reason: str = "", + ) -> Dict[str, Any]: + if not self.ensure_schema(): + return {"ok": False, "reason": "neo4j_not_available"} + timestamp = utc_now() + with self.driver.session() as session: + return session.execute_write( + self._update_concept_curation_tx, + collection_id, + org_id, + concept_name, + notes, + tags, + verification_state, + actor, + reason, + timestamp, + ) + + @staticmethod + def _manual_change_event_tx( + tx, + collection_id: str, + org_id: str, + operation: str, + actor: str, + timestamp: str, + concepts: List[str], + payload: Dict[str, Any], + ) -> Optional[str]: + row = tx.run( + """ + CREATE (event:ChangeEvent { + event_id: randomUUID(), + collection_id: $collection_id, + org_id: $org_id, + operation: $operation, + actor: $actor, + timestamp: $timestamp, + filename: '', + concepts: $concepts, + payload_json: $payload_json + }) + RETURN event.event_id AS event_id + """, + collection_id=collection_id, + org_id=org_id, + operation=operation, + actor=actor, + timestamp=timestamp, + concepts=sorted({concept for concept in concepts if concept}), + payload_json=json.dumps(payload, ensure_ascii=False), + ).single() + return row.get("event_id") if row else None + + @staticmethod + def _concept_is_used_query() -> str: + return """ + MATCH (concept:Concept {org_id: $org_id, name: $concept}) + WHERE EXISTS { MATCH (:Chunk {collection_id: $collection_id})-[:MENTIONS]->(concept) } + OR EXISTS { MATCH (concept)-[rel:RELATES_TO]-(:Concept) WHERE rel.collection_id = $collection_id } + RETURN concept.name AS name, + concept.notes AS old_notes, + concept.tags AS old_tags, + concept.verification_state AS old_verification_state + """ + + @staticmethod + def _move_concept_in_collection_tx( + tx, + collection_id: str, + org_id: str, + source_name: str, + target_name: str, + target_display_name: str, + timestamp: str, + ) -> Dict[str, Any]: + source_row = tx.run( + GraphStore._concept_is_used_query(), + collection_id=collection_id, + org_id=org_id, + concept=source_name, + ).single() + if not source_row: + return {"moved": False, "reason": "source_concept_not_found"} + + tx.run( + """ + MERGE (target:Concept {org_id: $org_id, name: $target}) + ON CREATE SET target.created_at = $timestamp, + target.entity_type = 'concept', + target.sources = [] + SET target.updated_at = $timestamp, + target.display_name = $target_display_name, + target.collection_hint = $collection_id + """, + org_id=org_id, + target=target_name, + target_display_name=target_display_name, + collection_id=collection_id, + timestamp=timestamp, + ) + + removed_between = tx.run( + """ + MATCH (source:Concept {org_id: $org_id, name: $source})-[rel]-(target:Concept {org_id: $org_id, name: $target}) + WHERE rel.collection_id = $collection_id + DELETE rel + RETURN count(rel) AS count + """, + org_id=org_id, + source=source_name, + target=target_name, + collection_id=collection_id, + ).single() + + mentions = tx.run( + """ + MATCH (target:Concept {org_id: $org_id, name: $target}) + MATCH (chunk:Chunk {collection_id: $collection_id})-[mention:MENTIONS]->(source:Concept {org_id: $org_id, name: $source}) + MERGE (chunk)-[newMention:MENTIONS]->(target) + ON CREATE SET newMention.created_at = $timestamp + SET newMention.collection_id = $collection_id + DELETE mention + RETURN count(mention) AS count + """, + org_id=org_id, + source=source_name, + target=target_name, + collection_id=collection_id, + timestamp=timestamp, + ).single() + + outgoing = tx.run( + """ + MATCH (target:Concept {org_id: $org_id, name: $target}) + MATCH (source:Concept {org_id: $org_id, name: $source})-[rel:RELATES_TO {collection_id: $collection_id}]->(other:Concept {org_id: $org_id}) + WHERE other.name <> $target + WITH target, other, rel, coalesce(rel.relation, 'related_to') AS relation + MERGE (target)-[newRel:RELATES_TO {collection_id: $collection_id, relation: relation}]->(other) + ON CREATE SET newRel.created_at = $timestamp, + newRel.weight = 0 + SET newRel.weight = coalesce(newRel.weight, 0) + coalesce(rel.weight, 1), + newRel.updated_at = $timestamp, + newRel.description = coalesce(rel.description, newRel.description, ''), + newRel.evidence = coalesce(rel.evidence, newRel.evidence, ''), + newRel.chunk_id = coalesce(rel.chunk_id, newRel.chunk_id, '') + DELETE rel + RETURN count(rel) AS count + """, + org_id=org_id, + source=source_name, + target=target_name, + collection_id=collection_id, + timestamp=timestamp, + ).single() + + incoming = tx.run( + """ + MATCH (target:Concept {org_id: $org_id, name: $target}) + MATCH (other:Concept {org_id: $org_id})-[rel:RELATES_TO {collection_id: $collection_id}]->(source:Concept {org_id: $org_id, name: $source}) + WHERE other.name <> $target + WITH target, other, rel, coalesce(rel.relation, 'related_to') AS relation + MERGE (other)-[newRel:RELATES_TO {collection_id: $collection_id, relation: relation}]->(target) + ON CREATE SET newRel.created_at = $timestamp, + newRel.weight = 0 + SET newRel.weight = coalesce(newRel.weight, 0) + coalesce(rel.weight, 1), + newRel.updated_at = $timestamp, + newRel.description = coalesce(rel.description, newRel.description, ''), + newRel.evidence = coalesce(rel.evidence, newRel.evidence, ''), + newRel.chunk_id = coalesce(rel.chunk_id, newRel.chunk_id, '') + DELETE rel + RETURN count(rel) AS count + """, + org_id=org_id, + source=source_name, + target=target_name, + collection_id=collection_id, + timestamp=timestamp, + ).single() + + deleted_source = tx.run( + """ + MATCH (source:Concept {org_id: $org_id, name: $source}) + WHERE NOT EXISTS { MATCH (:Chunk)-[:MENTIONS]->(source) } + AND NOT EXISTS { MATCH (source)-[:RELATES_TO]-(:Concept) } + DETACH DELETE source + RETURN count(source) AS count + """, + org_id=org_id, + source=source_name, + ).single() + + return { + "moved": True, + "source": source_name, + "target": target_name, + "mentions": mentions.get("count", 0) if mentions else 0, + "removed_between": ( + removed_between.get("count", 0) if removed_between else 0 + ), + "outgoing_relationships": outgoing.get("count", 0) if outgoing else 0, + "incoming_relationships": incoming.get("count", 0) if incoming else 0, + "deleted_source": deleted_source.get("count", 0) if deleted_source else 0, + } + + @staticmethod + def _rename_concept_tx( + tx, + collection_id: str, + org_id: str, + old_name: str, + new_name: str, + actor: str, + reason: str, + timestamp: str, + ) -> Dict[str, Any]: + source_name = normalize_concept(old_name) + target_name = normalize_concept(new_name) + if not source_name or not target_name: + return {"ok": False, "reason": "invalid_concept_name"} + if source_name == target_name: + return {"ok": False, "reason": "concept_names_are_equal"} + + move = GraphStore._move_concept_in_collection_tx( + tx, + collection_id, + org_id, + source_name, + target_name, + new_name.strip() or target_name, + timestamp, + ) + if not move.get("moved"): + return {"ok": False, **move} + + event_id = GraphStore._manual_change_event_tx( + tx, + collection_id, + org_id, + "manual_rename_concept", + actor, + timestamp, + [source_name, target_name], + { + "old_name": old_name, + "new_name": new_name, + "normalized_old_name": source_name, + "normalized_new_name": target_name, + "reason": reason, + "move": move, + }, + ) + return { + "ok": True, + "operation": "manual_rename_concept", + "event_id": event_id, + "details": move, + } + + @staticmethod + def _merge_concepts_tx( + tx, + collection_id: str, + org_id: str, + source_names: List[str], + target_name: str, + actor: str, + reason: str, + timestamp: str, + ) -> Dict[str, Any]: + normalized_target = normalize_concept(target_name) + normalized_sources = [] + for source_name in source_names: + normalized_source = normalize_concept(source_name) + if normalized_source and normalized_source != normalized_target: + normalized_sources.append(normalized_source) + normalized_sources = sorted(set(normalized_sources)) + + if not normalized_target or not normalized_sources: + return {"ok": False, "reason": "invalid_merge_request"} + + moved = [] + missing = [] + for normalized_source in normalized_sources: + move = GraphStore._move_concept_in_collection_tx( + tx, + collection_id, + org_id, + normalized_source, + normalized_target, + target_name.strip() or normalized_target, + timestamp, + ) + if move.get("moved"): + moved.append(move) + else: + missing.append(normalized_source) + + if not moved: + return { + "ok": False, + "reason": "source_concepts_not_found", + "missing": missing, + } + + event_id = GraphStore._manual_change_event_tx( + tx, + collection_id, + org_id, + "manual_merge_concepts", + actor, + timestamp, + [normalized_target, *normalized_sources], + { + "target_name": target_name, + "normalized_target_name": normalized_target, + "source_names": source_names, + "normalized_source_names": normalized_sources, + "missing_source_names": missing, + "reason": reason, + "moves": moved, + }, + ) + return { + "ok": True, + "operation": "manual_merge_concepts", + "event_id": event_id, + "details": { + "target": normalized_target, + "moved": moved, + "missing": missing, + }, + } + + @staticmethod + def _optional_text_changed( + current: Optional[str], requested: Optional[str] + ) -> bool: + if requested is None: + return False + return (current or "") != (requested or "") + + @staticmethod + def _optional_tags_changed( + current: Optional[List[str]], requested: Optional[List[str]] + ) -> bool: + if requested is None: + return False + return list(current or []) != list(requested or []) + + @staticmethod + def _optional_weight_changed( + current: Optional[float], requested: Optional[float] + ) -> bool: + if requested is None: + return False + try: + return ( + abs(float(current if current is not None else 1.0) - float(requested)) + > 1e-9 + ) + except (TypeError, ValueError): + return current != requested + + @staticmethod + def _optional_state_changed( + current: Optional[str], requested: Optional[str] + ) -> bool: + if requested is None: + return False + return str(current or "unverified") != requested + + @staticmethod + def _relationship_update_has_changes( + rel_row: Dict[str, Any], + *, + current_relation: str, + target_relation: str, + weight: Optional[float], + description: Optional[str], + evidence: Optional[str], + notes: Optional[str], + tags: Optional[List[str]], + verification_state: Optional[str], + ) -> bool: + if target_relation != current_relation: + return True + return any( + ( + GraphStore._optional_weight_changed(rel_row.get("old_weight"), weight), + GraphStore._optional_text_changed( + rel_row.get("old_description"), description + ), + GraphStore._optional_text_changed( + rel_row.get("old_evidence"), evidence + ), + GraphStore._optional_text_changed(rel_row.get("old_notes"), notes), + GraphStore._optional_tags_changed(rel_row.get("old_tags"), tags), + GraphStore._optional_state_changed( + rel_row.get("old_verification_state"), verification_state + ), + ) + ) + + @staticmethod + def _concept_update_has_changes( + concept_row: Dict[str, Any], + *, + notes: Optional[str], + tags: Optional[List[str]], + verification_state: Optional[str], + ) -> bool: + return any( + ( + GraphStore._optional_text_changed(concept_row.get("old_notes"), notes), + GraphStore._optional_tags_changed(concept_row.get("old_tags"), tags), + GraphStore._optional_state_changed( + concept_row.get("old_verification_state"), verification_state + ), + ) + ) + + @staticmethod + def _edit_relationship_tx( + tx, + collection_id: str, + org_id: str, + source_name: str, + target_name: str, + relation: str, + new_relation: Optional[str], + weight: Optional[float], + description: Optional[str], + evidence: Optional[str], + notes: Optional[str], + tags: Optional[List[str]], + verification_state: Optional[str], + actor: str, + reason: str, + operation: str, + timestamp: str, + ) -> Dict[str, Any]: + source = normalize_concept(source_name) + target = normalize_concept(target_name) + current_relation = normalize_relation(relation) + target_relation = normalize_relation(new_relation or relation) + if not source or not target or not current_relation: + return {"ok": False, "reason": "invalid_relationship_identity"} + + rel_row = tx.run( + """ + MATCH (source:Concept {org_id: $org_id, name: $source})-[rel:RELATES_TO {collection_id: $collection_id, relation: $relation}]->(target:Concept {org_id: $org_id, name: $target}) + RETURN rel.weight AS old_weight, + rel.description AS old_description, + rel.evidence AS old_evidence, + rel.chunk_id AS old_chunk_id, + rel.notes AS old_notes, + rel.tags AS old_tags, + rel.verification_state AS old_verification_state + """, + org_id=org_id, + collection_id=collection_id, + source=source, + target=target, + relation=current_relation, + ).single() + if not rel_row: + return {"ok": False, "reason": "relationship_not_found"} + + if verification_state == "rejected": + event_id = GraphStore._manual_change_event_tx( + tx, + collection_id, + org_id, + "manual_expunge_relationship", + actor, + timestamp, + [source, target], + { + "source": source, + "target": target, + "relation": current_relation, + "old_weight": rel_row.get("old_weight"), + "old_description": rel_row.get("old_description"), + "old_evidence": rel_row.get("old_evidence"), + "old_chunk_id": rel_row.get("old_chunk_id"), + "old_notes": rel_row.get("old_notes"), + "old_tags": rel_row.get("old_tags"), + "old_verification_state": rel_row.get("old_verification_state"), + "verification_state": "rejected", + "reason": reason, + }, + ) + tx.run( + """ + MATCH (source:Concept {org_id: $org_id, name: $source})-[rel:RELATES_TO {collection_id: $collection_id, relation: $relation}]->(target:Concept {org_id: $org_id, name: $target}) + DELETE rel + """, + org_id=org_id, + collection_id=collection_id, + source=source, + target=target, + relation=current_relation, + ) + return { + "ok": True, + "operation": "manual_expunge_relationship", + "event_id": event_id, + "details": { + "source": source, + "target": target, + "relation": current_relation, + "expunged": True, + }, + } + + if not GraphStore._relationship_update_has_changes( + rel_row, + current_relation=current_relation, + target_relation=target_relation, + weight=weight, + description=description, + evidence=evidence, + notes=notes, + tags=tags, + verification_state=verification_state, + ): + return { + "ok": True, + "operation": None, + "event_id": None, + "reason": "no_change", + "details": { + "source": source, + "target": target, + "old_relation": current_relation, + "new_relation": target_relation, + "changed": False, + }, + } + + if target_relation == current_relation: + tx.run( + """ + MATCH (source:Concept {org_id: $org_id, name: $source})-[rel:RELATES_TO {collection_id: $collection_id, relation: $relation}]->(target:Concept {org_id: $org_id, name: $target}) + SET rel.updated_at = $timestamp, + rel.weight = CASE WHEN $weight IS NULL THEN rel.weight ELSE $weight END, + rel.description = CASE WHEN $description IS NULL THEN rel.description ELSE $description END, + rel.evidence = CASE WHEN $evidence IS NULL THEN rel.evidence ELSE $evidence END, + rel.notes = CASE WHEN $notes IS NULL THEN rel.notes ELSE $notes END, + rel.tags = CASE WHEN $tags IS NULL THEN rel.tags ELSE $tags END, + rel.verification_state = CASE WHEN $verification_state IS NULL THEN rel.verification_state ELSE $verification_state END + """, + org_id=org_id, + collection_id=collection_id, + source=source, + target=target, + relation=current_relation, + weight=weight, + description=description, + evidence=evidence, + notes=notes, + tags=tags, + verification_state=verification_state, + timestamp=timestamp, + ) + else: + tx.run( + """ + MATCH (source:Concept {org_id: $org_id, name: $source})-[oldRel:RELATES_TO {collection_id: $collection_id, relation: $current_relation}]->(target:Concept {org_id: $org_id, name: $target}) + WITH source, target, oldRel, + coalesce(oldRel.weight, 1.0) AS old_weight, + oldRel.description AS old_description, + oldRel.evidence AS old_evidence, + oldRel.notes AS old_notes, + oldRel.tags AS old_tags, + oldRel.verification_state AS old_verification_state + MERGE (source)-[rel:RELATES_TO {collection_id: $collection_id, relation: $target_relation}]->(target) + ON CREATE SET rel.created_at = $timestamp + SET rel.updated_at = $timestamp, + rel.weight = CASE WHEN $weight IS NULL THEN old_weight ELSE $weight END, + rel.description = CASE WHEN $description IS NULL THEN old_description ELSE $description END, + rel.evidence = CASE WHEN $evidence IS NULL THEN old_evidence ELSE $evidence END, + rel.notes = CASE WHEN $notes IS NULL THEN old_notes ELSE $notes END, + rel.tags = CASE WHEN $tags IS NULL THEN old_tags ELSE $tags END, + rel.verification_state = CASE WHEN $verification_state IS NULL THEN old_verification_state ELSE $verification_state END + DELETE oldRel + """, + org_id=org_id, + collection_id=collection_id, + source=source, + target=target, + current_relation=current_relation, + target_relation=target_relation, + weight=weight, + description=description, + evidence=evidence, + notes=notes, + tags=tags, + verification_state=verification_state, + timestamp=timestamp, + ) + + event_id = GraphStore._manual_change_event_tx( + tx, + collection_id, + org_id, + operation, + actor, + timestamp, + [source, target], + { + "source": source, + "target": target, + "relation": current_relation, + "new_relation": target_relation, + "old_weight": rel_row.get("old_weight"), + "new_weight": weight, + "old_notes": rel_row.get("old_notes"), + "notes": notes, + "old_tags": rel_row.get("old_tags"), + "tags": tags, + "old_verification_state": rel_row.get("old_verification_state"), + "verification_state": verification_state, + "reason": reason, + }, + ) + return { + "ok": True, + "operation": operation, + "event_id": event_id, + "details": { + "source": source, + "target": target, + "old_relation": current_relation, + "new_relation": target_relation, + }, + } + + @staticmethod + def _update_concept_curation_tx( + tx, + collection_id: str, + org_id: str, + concept_name: str, + notes: Optional[str], + tags: Optional[List[str]], + verification_state: Optional[str], + actor: str, + reason: str, + timestamp: str, + ) -> Dict[str, Any]: + concept = normalize_concept(concept_name) + if not concept: + return {"ok": False, "reason": "invalid_concept_name"} + + concept_row = tx.run( + GraphStore._concept_is_used_query(), + collection_id=collection_id, + org_id=org_id, + concept=concept, + ).single() + if not concept_row: + return {"ok": False, "reason": "concept_not_found"} + + # Change-detection for verification_state must use the per-collection + # MENTIONS.verification_state, not the org-level concept.verification_state. + # The org-level value may be 'verified' from another KS; for this KS it may + # still be null (unverified). Fetch the collection-scoped value and override + # the comparison baseline so the write is not incorrectly skipped. + if verification_state is not None: + mentions_row = tx.run( + """ + MATCH (:Chunk {collection_id: $collection_id})-[m:MENTIONS]-> + (:Concept {org_id: $org_id, name: $concept}) + RETURN m.verification_state AS vs + LIMIT 1 + """, + collection_id=collection_id, + org_id=org_id, + concept=concept, + ).single() + concept_row = dict(concept_row) + concept_row["old_verification_state"] = ( + mentions_row.get("vs") if mentions_row else None + ) + + if verification_state == "rejected": + chunk_rows = tx.run( + """ + MATCH (chunk:Chunk {collection_id: $collection_id})-[mention:MENTIONS]->(concept:Concept {org_id: $org_id, name: $concept}) + RETURN collect(DISTINCT chunk.chunk_id) AS chunk_ids + """, + collection_id=collection_id, + org_id=org_id, + concept=concept, + ).single() + relationship_rows = tx.run( + """ + MATCH (concept:Concept {org_id: $org_id, name: $concept})-[rel:RELATES_TO {collection_id: $collection_id}]-(other:Concept {org_id: $org_id}) + RETURN startNode(rel).name AS source, + endNode(rel).name AS target, + coalesce(rel.relation, type(rel)) AS relation, + type(rel) AS type, + rel.weight AS weight, + rel.description AS description, + rel.evidence AS evidence, + rel.chunk_id AS chunk_id, + rel.notes AS notes, + rel.tags AS tags, + rel.verification_state AS verification_state + """, + collection_id=collection_id, + org_id=org_id, + concept=concept, + ).data() + event_id = GraphStore._manual_change_event_tx( + tx, + collection_id, + org_id, + "manual_expunge_concept", + actor, + timestamp, + [concept], + { + "concept": concept, + "old_notes": concept_row.get("old_notes"), + "old_tags": concept_row.get("old_tags"), + "old_verification_state": concept_row.get("old_verification_state"), + "verification_state": "rejected", + "removed_chunk_mentions": (chunk_rows or {}).get("chunk_ids", []), + "removed_relationships": relationship_rows, + "reason": reason, + }, + ) + tx.run( + """ + MATCH (:Chunk {collection_id: $collection_id})-[mention:MENTIONS]->(:Concept {org_id: $org_id, name: $concept}) + DELETE mention + """, + collection_id=collection_id, + org_id=org_id, + concept=concept, + ) + tx.run( + """ + MATCH (:Concept {org_id: $org_id, name: $concept})-[rel:RELATES_TO {collection_id: $collection_id}]-(:Concept {org_id: $org_id}) + DELETE rel + """, + collection_id=collection_id, + org_id=org_id, + concept=concept, + ) + tx.run( + """ + MATCH (concept:Concept {org_id: $org_id, name: $concept}) + WHERE NOT EXISTS { MATCH (:Chunk)-[:MENTIONS]->(concept) } + AND NOT EXISTS { MATCH (concept)-[:RELATES_TO]-(:Concept) } + DETACH DELETE concept + """, + org_id=org_id, + concept=concept, + ) + return { + "ok": True, + "operation": "manual_expunge_concept", + "event_id": event_id, + "details": { + "concept": concept, + "expunged": True, + "removed_chunk_mentions": len( + (chunk_rows or {}).get("chunk_ids", []) + ), + "removed_relationships": len(relationship_rows), + }, + } + + if not GraphStore._concept_update_has_changes( + concept_row, + notes=notes, + tags=tags, + verification_state=verification_state, + ): + return { + "ok": True, + "operation": None, + "event_id": None, + "reason": "no_change", + "details": {"concept": concept, "changed": False}, + } + + # Write per-collection verification state onto the MENTIONS relationships + # (scoped to this collection) so each KS tracks its own approval independently + # of the org-level Concept node. + if verification_state is not None: + mentions_vs = None if verification_state == "unverified" else verification_state + tx.run( + """ + MATCH (:Chunk {collection_id: $collection_id})-[m:MENTIONS]->(:Concept {org_id: $org_id, name: $concept}) + SET m.verification_state = $vs + """, + collection_id=collection_id, + org_id=org_id, + concept=concept, + vs=mentions_vs, + ) + # Also promote to the org-level Concept node when verifying so that + # RELATES_TO path traversal (which uses concept.verification_state) + # can cross this concept. Deliberately NOT cleared on unverify — + # another KS in the same org might still have it approved. + if verification_state == "verified": + tx.run( + """ + MATCH (concept:Concept {org_id: $org_id, name: $concept}) + SET concept.verification_state = 'verified' + """, + org_id=org_id, + concept=concept, + ) + + tx.run( + """ + MATCH (concept:Concept {org_id: $org_id, name: $concept}) + SET concept.updated_at = $timestamp, + concept.notes = CASE WHEN $notes IS NULL THEN concept.notes ELSE $notes END, + concept.tags = CASE WHEN $tags IS NULL THEN concept.tags ELSE $tags END + """, + org_id=org_id, + concept=concept, + notes=notes, + tags=tags, + timestamp=timestamp, + ) + + event_id = GraphStore._manual_change_event_tx( + tx, + collection_id, + org_id, + "manual_curate_concept", + actor, + timestamp, + [concept], + { + "concept": concept, + "old_notes": concept_row.get("old_notes"), + "notes": notes, + "old_tags": concept_row.get("old_tags"), + "tags": tags, + "old_verification_state": concept_row.get("old_verification_state"), + "verification_state": verification_state, + "reason": reason, + }, + ) + return { + "ok": True, + "operation": "manual_curate_concept", + "event_id": event_id, + "details": {"concept": concept}, + } + + def ingest_chunks( + self, + *, + collection: Dict[str, Any], + file_id: Optional[int], + filename: str, + chunks: Iterable[TextChunk], + concepts_by_chunk: Dict[str, List[str]], + entities: Dict[str, ExtractedEntity], + relationships: List[ExtractedRelationship], + actor: str = "lamb-ingestion-pipeline", + ) -> int: + chunk_list = list(chunks) + if not chunk_list: + return 0 + if not self.ensure_schema(): + return 0 + + entity_map = dict(entities) + for concept in itertools.chain.from_iterable(concepts_by_chunk.values()): + entity_map.setdefault( + concept, + ExtractedEntity( + name=concept, + display_name=concept, + entity_type="concept", + ), + ) + for relationship in relationships: + entity_map.setdefault( + relationship.source, + ExtractedEntity( + name=relationship.source, + display_name=relationship.source, + entity_type="concept", + ), + ) + entity_map.setdefault( + relationship.target, + ExtractedEntity( + name=relationship.target, + display_name=relationship.target, + entity_type="concept", + ), + ) + + relationship_payloads = [ + relationship.__dict__ + for relationship in relationships + if relationship.source in entity_map and relationship.target in entity_map + ] + entity_payloads = [entity.__dict__ for entity in entity_map.values()] + + with self.driver.session() as session: + session.execute_write( + self._ingest_tx, + collection, + int(file_id or 0), + filename, + chunk_list, + concepts_by_chunk, + sorted(entity_payloads, key=lambda item: item["name"]), + relationship_payloads, + actor, + ) + + return len(chunk_list) + len(entity_map) + len(relationship_payloads) + 1 + + @staticmethod + def _ingest_tx( + tx, + collection: Dict[str, Any], + file_id: int, + filename: str, + chunks: List[TextChunk], + concepts_by_chunk: Dict[str, List[str]], + entities: List[Dict[str, Any]], + relationships: List[Dict[str, Any]], + actor: str, + ) -> None: + collection_id = str(collection["id"]) + org_id = str( + collection.get("organization_id") or collection.get("owner") or "default" + ) + document_id = f"{collection_id}:{file_id}:{filename}" + timestamp = utc_now() + + tx.run( + """ + MERGE (org:Organization {org_id: $org_id}) + ON CREATE SET org.created_at = $timestamp + MERGE (collection:Collection {collection_id: $collection_id}) + ON CREATE SET collection.created_at = $timestamp + SET collection.name = $name, + collection.description = $description, + collection.owner = $org_id, + collection.collection_id = $collection_id + MERGE (org)-[:OWNS]->(collection) + MERGE (doc:Document {document_id: $document_id}) + ON CREATE SET doc.created_at = $timestamp + SET doc.collection_id = $collection_id, + doc.file_id = $file_id, + doc.filename = $filename, + doc.org_id = $org_id + MERGE (collection)-[:CONTAINS]->(doc) + """, + org_id=org_id, + collection_id=collection_id, + name=collection.get("name", ""), + description=collection.get("description") or "", + document_id=document_id, + file_id=file_id, + filename=filename, + timestamp=timestamp, + ) + + for entity in entities: + tx.run( + """ + MERGE (concept:Concept {org_id: $org_id, name: $name}) + ON CREATE SET concept.created_at = $timestamp, + concept.sources = [] + SET concept.updated_at = $timestamp, + concept.collection_hint = $collection_id, + concept.display_name = $display_name, + concept.entity_type = $entity_type, + concept.description = $description, + concept.confidence = $confidence + """, + org_id=org_id, + name=entity["name"], + display_name=entity.get("display_name") or entity["name"], + entity_type=entity.get("entity_type") or "concept", + description=entity.get("description") or "", + confidence=float(entity.get("confidence") or 1.0), + collection_id=collection_id, + timestamp=timestamp, + ) + + for chunk in chunks: + metadata = chunk.metadata or {} + # Carry permalinks from the chunk metadata onto the Chunk node so + # graph-driven citations can link back to source content. The + # ingestion path puts these under top-level metadata keys + # (``permalink_original``, ``permalink_full_markdown``, + # ``permalink_page``) via the chunking strategies. Anything + # missing becomes empty string so Neo4j stays typed. + permalink_original = str( + metadata.get("permalink_original") + or metadata.get("permalink") + or "" + ) + permalink_full_markdown = str( + metadata.get("permalink_full_markdown") or "" + ) + permalink_page = str(metadata.get("permalink_page") or "") + tx.run( + """ + MATCH (doc:Document {document_id: $document_id}) + MERGE (chunk:Chunk {chunk_id: $chunk_id}) + ON CREATE SET chunk.created_at = $timestamp + SET chunk.collection_id = $collection_id, + chunk.org_id = $org_id, + chunk.file_id = $file_id, + chunk.filename = $filename, + chunk.text = $text, + chunk.parent_text = $parent_text, + chunk.section_title = $section_title, + chunk.source_label = $source_label, + chunk.permalink_original = $permalink_original, + chunk.permalink_full_markdown = $permalink_full_markdown, + chunk.permalink_page = $permalink_page + MERGE (doc)-[:CONTAINS]->(chunk) + """, + document_id=document_id, + chunk_id=chunk.chunk_id, + collection_id=collection_id, + org_id=org_id, + file_id=file_id, + # Prefer the per-chunk filename from metadata so each + # chunk keeps the source document it came from. Falls + # back to the batch-level filename only if the chunk + # didn't carry its own (e.g. legacy ingestion paths). + filename=str(metadata.get("filename") or filename), + text=chunk.text, + parent_text=chunk.parent_text, + section_title=str(metadata.get("section_title") or "Document"), + source_label=str(metadata.get("source_label") or chunk.chunk_id), + permalink_original=permalink_original, + permalink_full_markdown=permalink_full_markdown, + permalink_page=permalink_page, + timestamp=timestamp, + ) + for concept in concepts_by_chunk.get(chunk.chunk_id, []): + tx.run( + """ + MATCH (chunk:Chunk {chunk_id: $chunk_id}) + MATCH (concept:Concept {org_id: $org_id, name: $concept}) + MERGE (chunk)-[mention:MENTIONS]->(concept) + ON CREATE SET mention.created_at = $timestamp + SET mention.collection_id = $collection_id + """, + chunk_id=chunk.chunk_id, + org_id=org_id, + concept=concept, + collection_id=collection_id, + timestamp=timestamp, + ) + + for relationship in relationships: + tx.run( + """ + MATCH (source:Concept {org_id: $org_id, name: $source}) + MATCH (target:Concept {org_id: $org_id, name: $target}) + MERGE (source)-[rel:RELATES_TO {collection_id: $collection_id, relation: $relation}]->(target) + ON CREATE SET rel.created_at = $timestamp, + rel.weight = 0 + SET rel.weight = coalesce(rel.weight, 0) + $confidence, + rel.updated_at = $timestamp, + rel.description = $description, + rel.evidence = $evidence, + rel.chunk_id = $chunk_id + """, + org_id=org_id, + collection_id=collection_id, + source=relationship["source"], + target=relationship["target"], + relation=relationship.get("relation") or "related_to", + description=relationship.get("description") or "", + evidence=relationship.get("evidence") or "", + chunk_id=relationship.get("chunk_id") or "", + confidence=float(relationship.get("confidence") or 1.0), + timestamp=timestamp, + ) + + tx.run( + """ + CREATE (event:ChangeEvent { + event_id: randomUUID(), + collection_id: $collection_id, + org_id: $org_id, + operation: 'automatic_ingestion', + actor: $actor, + timestamp: $timestamp, + filename: $filename, + concepts: $concepts, + payload_json: $payload_json + }) + WITH event + MATCH (doc:Document {document_id: $document_id}) + MERGE (event)-[:RECORDED_CHANGE]->(doc) + """, + collection_id=collection_id, + org_id=org_id, + actor=actor, + timestamp=timestamp, + filename=filename, + concepts=[entity["name"] for entity in entities], + payload_json=json.dumps( + { + "chunks": len(chunks), + "chunk_ids": [chunk.chunk_id for chunk in chunks], + "concepts": len(entities), + "relationships": len(relationships), + "relationship_details": relationships, + }, + ensure_ascii=False, + ), + document_id=document_id, + ) + + def expand_from_concept_names( + self, + collection_id: str, + org_id: str, + concept_names: List[str], + depth: int, + limit: int, + ) -> Dict[str, Any]: + """Local-search-style expansion seeded directly by concept names. + + This method takes concept names directly (after + :func:`normalize_concept`) and returns the chunks that mention + any of them or any concept reachable via + ``RELATES_TO*1..depth``. + + This is the entry point for question-entity seeding: the caller + extracts named-entity-like tokens from the question text and + passes them here so the graph can contribute results even when + the vector baseline missed every gold chunk. + """ + depth = max(1, min(int(depth or 2), 4)) + limit = max(1, int(limit or 10)) + start = time.perf_counter() + normalized = sorted({normalize_concept(n) for n in concept_names if n}) + normalized = [n for n in normalized if n] + if not normalized: + return self._empty_expansion(start) + if not self.ensure_schema(): + return self._empty_expansion( + start, + warning="Neo4j is not configured or available; KG expansion skipped", + ) + + with self.driver.session() as session: + # Filter out highly-mentioned concepts: a question entity that + # is mentioned by N+ chunks in the collection is too generic + # to be a useful seed (e.g. "radar station" in a HotPotQA + # corpus of Wikipedia paragraphs). Specific named entities + # typically appear in ≤10 chunks. + MAX_MENTIONS_PER_CONCEPT = 25 + + matched_rows = session.run( + """ + MATCH (concept:Concept {org_id: $org_id})<-[m:MENTIONS]-(chunk:Chunk {collection_id: $collection_id}) + WHERE concept.name IN $names + AND m.verification_state = 'verified' + WITH concept, count(chunk) AS mentions + WHERE mentions <= $max_mentions + RETURN concept.name AS name, mentions + ORDER BY mentions ASC + """, + org_id=org_id, + names=normalized, + collection_id=collection_id, + max_mentions=MAX_MENTIONS_PER_CONCEPT, + ).data() + entry_concepts = [r["name"] for r in matched_rows] + if not entry_concepts: + return self._empty_expansion(start) + + chunk_ids: List[str] = [] + seen: Set[str] = set() + + def _add(cid: str) -> None: + if cid and cid not in seen: + seen.add(cid) + chunk_ids.append(cid) + + direct = session.run( + """ + MATCH (concept:Concept {org_id: $org_id})<-[:MENTIONS]-(chunk:Chunk {collection_id: $collection_id}) + WHERE concept.name IN $entry + RETURN chunk.chunk_id AS chunk_id, count(*) AS mentions + ORDER BY mentions DESC + LIMIT $limit + """, + org_id=org_id, + collection_id=collection_id, + entry=entry_concepts, + limit=limit, + ).data() + for row in direct: + _add(row["chunk_id"]) + + related_query = f""" + MATCH (entry:Concept {{org_id: $org_id}}) + WHERE entry.name IN $entry + MATCH path = (entry)-[:RELATES_TO*1..{depth}]-(related:Concept {{org_id: $org_id}}) + WHERE related.name <> entry.name + AND all(rel IN relationships(path) WHERE rel.collection_id = $collection_id) + AND all(node IN nodes(path) WHERE node.verification_state = 'verified') + AND all(rel IN relationships(path) WHERE rel.verification_state = 'verified') + WITH related, length(path) AS hops + ORDER BY hops ASC + LIMIT $limit + OPTIONAL MATCH (related)<-[:MENTIONS]-(chunk:Chunk {{collection_id: $collection_id}}) + RETURN DISTINCT chunk.chunk_id AS chunk_id + """ + related = session.run( + related_query, + org_id=org_id, + collection_id=collection_id, + entry=entry_concepts, + limit=limit, + ).data() + for row in related: + _add(row.get("chunk_id")) + + return { + "entry_concepts": entry_concepts, + "expanded_chunk_ids": chunk_ids[:limit], + "traversed_edges": [], + "latest_changes": [], + "graph_latency_ms": (time.perf_counter() - start) * 1000, + } + + @staticmethod + def _empty_expansion(start: float, warning: Optional[str] = None) -> Dict[str, Any]: + changes: List[Dict[str, Any]] = [] + if warning: + changes.append({"warning": warning}) + return { + "entry_concepts": [], + "expanded_chunk_ids": [], + "traversed_edges": [], + "latest_changes": changes, + "graph_latency_ms": (time.perf_counter() - start) * 1000, + } diff --git a/lamb-kb-server/backend/services/ingestion_service.py b/lamb-kb-server/backend/services/ingestion_service.py index 4d0671b48..922cfb503 100644 --- a/lamb-kb-server/backend/services/ingestion_service.py +++ b/lamb-kb-server/backend/services/ingestion_service.py @@ -285,6 +285,97 @@ def execute_ingestion_job( total_chunks_added, ) + _maybe_index_graph( + collection=collection, + docs_list=docs_list, + backend=backend, + embedding_function=embedding_function, + openai_api_key=credentials.get("kg_rag_openai_api_key", "") + or credentials.get("api_key", ""), + ) + + +def _maybe_index_graph( + *, + collection: Collection, + docs_list: list[dict], + backend, + embedding_function, + openai_api_key: str = "", +) -> None: + """Run graph indexing for this batch if the collection opted in. + + Failures here are logged and swallowed: vector ingestion already + committed, and we don't want a Neo4j hiccup to roll back successful + chunk insertions. + + Pulls chunks back from the vector backend (via the public + ``get_chunks_by_source`` surface) so the graph indexer has stable + chunk IDs that match what's searchable. Backends that don't + implement that method silently return empty lists, which we treat as + "no chunks to index" and log a one-line warning. + """ + if not getattr(collection, "graph_enabled", False): + return + + import config as config_module # noqa: PLC0415 + + kg_config = config_module.get_kg_rag_config() + if not kg_config.get("enabled") or not kg_config.get("index_on_ingest", True): + return + + from services.graph_indexing import ( # noqa: PLC0415 + index_chunks_for_collection, + ) + + source_ids = [doc["source_item_id"] for doc in docs_list] + backend_collection_id = collection.backend_collection_id or collection.id + for source_id in source_ids: + try: + fetched = backend.get_chunks_by_source( + collection_id=backend_collection_id, + storage_path=collection.storage_path, + source_item_id=source_id, + embedding_function=embedding_function, + ) + except Exception as exc: # noqa: BLE001 + logger.warning( + "Graph indexing: backend %s failed to fetch source %s: %s", + collection.vector_db_backend, + source_id, + exc, + ) + continue + + if not fetched: + logger.debug( + "Graph indexing: no chunks returned for source %s on backend %s; " + "skipping. (Most likely the backend does not support " + "get_chunks_by_source.)", + source_id, + collection.vector_db_backend, + ) + continue + + result = index_chunks_for_collection( + collection=collection, + ids=[ + str(item.metadata.get("chunk_id") or item.metadata.get("document_id") or "") + for item in fetched + ], + texts=[item.text for item in fetched], + metadatas=[dict(item.metadata or {}) for item in fetched], + openai_api_key=openai_api_key, + filename=source_id, + ) + if result.get("error"): + logger.warning( + "Graph indexing for source %s in collection %s reported: %s", + source_id, + collection.id, + result["error"], + ) + def delete_vectors( db: Session, collection_id: str, source_item_id: str @@ -346,6 +437,8 @@ def delete_vectors( ) db.commit() + _maybe_delete_graph_document(collection, source_item_id) + logger.info( "Deleted %d vectors for source_item_id '%s' from collection %s", deleted_count, @@ -355,6 +448,35 @@ def delete_vectors( return deleted_count +def _maybe_delete_graph_document(collection: Collection, source_item_id: str) -> None: + """Remove graph data for a deleted source item if the collection has graph_enabled.""" + if not getattr(collection, "graph_enabled", False): + return + + import config as config_module # noqa: PLC0415 + + if not config_module.get_kg_rag_config().get("enabled"): + return + + from services.graph_store import get_graph_store # noqa: PLC0415 + + org_id = str(collection.organization_id or "") + try: + get_graph_store().delete_document(collection.id, org_id, source_item_id) + logger.info( + "Graph data removed for source_item_id '%s' in collection %s", + source_item_id, + collection.id, + ) + except Exception as exc: # noqa: BLE001 + logger.warning( + "Graph cleanup for source '%s' in collection %s failed: %s", + source_item_id, + collection.id, + exc, + ) + + def cancel_job(db: Session, job_id: str) -> IngestionJob: """Cancel a pending or in-flight ingestion job. diff --git a/lamb-kb-server/backend/services/query_service.py b/lamb-kb-server/backend/services/query_service.py index ade1dbbf7..e12d4f424 100644 --- a/lamb-kb-server/backend/services/query_service.py +++ b/lamb-kb-server/backend/services/query_service.py @@ -1,6 +1,7 @@ """Business logic for vector similarity queries.""" import logging +from typing import Any from database.models import Collection from fastapi import HTTPException, status @@ -62,6 +63,54 @@ def query_collection( embedding_function=embedding_function, ) + # Auto-route through the KG-RAG plugin when the collection was built + # with a graph. This makes the regular /query endpoint (used by the + # "Test Query" affordance and by ``knowledge_store_rag.py`` at chat + # time) actually exercise question-entity extraction + graph + # expansion — without any caller-side opt-in. The plugin gracefully + # degrades to the vector baseline when Neo4j isn't reachable or the + # graph returns nothing, so this is safe to always-on for + # graph-enabled collections. + if getattr(collection, "graph_enabled", False): + import config as config_module # noqa: PLC0415 + + if config_module.KG_RAG_ENABLED: + from plugins.kg_rag_query import KGRAGQueryPlugin # noqa: PLC0415 + + baseline_dicts = [ + { + "similarity": r.score, + "data": r.text, + "metadata": dict(r.metadata or {}), + } + for r in results + ] + try: + augmented = KGRAGQueryPlugin().augment( + db=db, + collection=collection, + backend=backend, + embedding_function=embedding_function, + query_text=req.query_text, + baseline_results=baseline_dicts, + params={"top_k": req.top_k, "include_trace": True}, + ) + results = [ + QueryResult( + text=item.get("data", "") or "", + score=float(item.get("similarity") or 0.0), + metadata=dict(item.get("metadata") or {}), + ) + for item in augmented + ] + except Exception as exc: # noqa: BLE001 — degrade to baseline on any failure + logger.warning( + "KG-RAG augmentation failed for collection %s, " + "returning vector baseline: %s", + collection_id, + exc, + ) + logger.debug( "Query on collection %s returned %d results for '%s'", collection_id, @@ -69,3 +118,117 @@ def query_collection( req.query_text[:80], ) return results + + +def query_with_plugin( + *, + db: Session, + collection_id: str, + query_text: str, + plugin_name: str = "simple_query", + plugin_params: dict[str, Any] | None = None, + embedding_credentials: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Run a query through a named plugin (``simple_query`` or ``kg_rag_query``). + + This is the path used by KG-RAG / benchmark callers that need both + results and side-channel metadata (graph trace, latency timings). The + return shape is a dict with ``{results, query, top_k, timing}`` so + benchmark code can read ``timing.total_ms`` and walk the result-attached + ``metadata.kg_rag`` trace. + + Args: + db: Database session. + collection_id: Target collection ID. + query_text: Free-text query. + plugin_name: ``simple_query`` (baseline vector retrieval) or + ``kg_rag_query`` (vector seed + Neo4j graph expansion). + plugin_params: Plugin-specific tuning (``top_k``, ``threshold``, + ``graph_depth``, ``include_trace``, ...). + embedding_credentials: Per-request embedding credentials + (``{"api_key": ..., "api_endpoint": ...}``). Falls back to the + collection-level endpoint when no key is supplied. + + Returns: + ``{"results": [...], "query": str, "top_k": int, "timing": {...}}``. + """ + import time + + params = dict(plugin_params or {}) + top_k = int(params.get("top_k", 5) or 5) + threshold = float(params.get("threshold", 0.0) or 0.0) + creds = embedding_credentials or {} + + collection = ( + db.query(Collection).filter(Collection.id == collection_id).first() + ) + if collection is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Collection '{collection_id}' not found.", + ) + + embedding_function = EmbeddingRegistry.build( + collection.embedding_vendor, + model=collection.embedding_model, + api_key=creds.get("api_key", ""), + api_endpoint=creds.get("api_endpoint") or collection.embedding_endpoint or "", + ) + + backend = VectorDBRegistry.get(collection.vector_db_backend) + if backend is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=( + f"Vector DB backend '{collection.vector_db_backend}' is not available." + ), + ) + + start = time.perf_counter() + raw_results = backend.query( + collection_id=collection.backend_collection_id or collection_id, + storage_path=collection.storage_path, + query_text=query_text, + top_k=top_k, + embedding_function=embedding_function, + ) + + formatted = [ + { + "similarity": r.score, + "data": r.text, + "metadata": dict(r.metadata or {}), + } + for r in raw_results + if r.score >= threshold + ] + + # Optional KG-RAG augmentation via the registered query plugin. + if plugin_name == "kg_rag_query": + try: + from plugins.kg_rag_query import KGRAGQueryPlugin + + plugin = KGRAGQueryPlugin() + formatted = plugin.augment( + db=db, + collection=collection, + backend=backend, + embedding_function=embedding_function, + query_text=query_text, + baseline_results=formatted, + params=params, + ) + except Exception as exc: # noqa: BLE001 — degrade gracefully + logger.warning( + "KG-RAG augmentation failed for collection %s: %s", + collection_id, + exc, + ) + elapsed_ms = (time.perf_counter() - start) * 1000 + + return { + "results": formatted, + "query": query_text, + "top_k": top_k, + "timing": {"total_ms": elapsed_ms}, + } diff --git a/lamb-kb-server/backend/tasks/worker.py b/lamb-kb-server/backend/tasks/worker.py index 01063ada0..87a4c54ce 100644 --- a/lamb-kb-server/backend/tasks/worker.py +++ b/lamb-kb-server/backend/tasks/worker.py @@ -16,9 +16,14 @@ import asyncio import logging from concurrent.futures import ThreadPoolExecutor -from datetime import UTC, datetime - -from config import INGESTION_TASK_TIMEOUT_SECONDS, MAX_CONCURRENT_INGESTIONS, MAX_JOB_ATTEMPTS +from datetime import UTC, datetime, timedelta + +from config import ( + INGESTION_TASK_TIMEOUT_SECONDS, + MAX_CONCURRENT_INGESTIONS, + MAX_JOB_ATTEMPTS, + RETRY_CACHE_TTL_MINUTES, +) from database.connection import get_session_direct from database.models import Collection, IngestionJob from sqlalchemy.orm import Session @@ -32,22 +37,31 @@ _running = False # In-memory store for embedding credentials — never written to disk (ADR-4). -# Maps job_id → credentials dict. Entries are removed once the worker picks -# them up. If the service restarts before the worker picks up a job, the -# credentials are lost and the job fails cleanly — exactly the behavior the -# Library Manager uses for its own API keys. +# Maps job_id → credentials dict. The document payload itself is persisted in +# the ``ingestion_jobs`` row, so only the credentials need to survive in +# memory. Unlike the original design (which popped credentials on first +# pickup), entries are RETAINED across attempts so a failed job can be retried +# without the client re-sending credentials. They are discarded when the job +# succeeds, exhausts its attempts, or ages past ``RETRY_CACHE_TTL_MINUTES`` +# (see ``_purge_expired_credentials``). A restart still loses them, in which +# case a retry requires a fresh add-content request. _job_credentials: dict[str, dict[str, str]] = {} +# Parallel map of job_id → when the credentials were stored, for TTL expiry. +_job_cred_stored_at: dict[str, datetime] = {} # How often (seconds) the worker checks for new pending jobs. _POLL_INTERVAL = 2.0 +# How often (seconds) expired credentials are purged from the retry cache. +_CLEANUP_INTERVAL = 600.0 def store_credentials(job_id: str, credentials: dict[str, str] | None) -> None: - """Hold embedding credentials in memory for a job until the worker runs it. + """Hold embedding credentials in memory for a job and its retries. Called by ``ingestion_service`` immediately after committing the job row - to SQLite. Credentials live only in the module-level dict and are popped - by the worker when processing starts. + to SQLite. Credentials live only in the module-level dict and are retained + until the job reaches a terminal state, exhausts its attempts, or the + retention window elapses. Args: job_id: The ingestion job ID. @@ -55,6 +69,40 @@ def store_credentials(job_id: str, credentials: dict[str, str] | None) -> None: """ if credentials: _job_credentials[job_id] = credentials + _job_cred_stored_at[job_id] = datetime.now(UTC) + + +def _discard_credentials(job_id: str) -> None: + """Forget a job's cached credentials (terminal state or exhausted retries).""" + _job_credentials.pop(job_id, None) + _job_cred_stored_at.pop(job_id, None) + + +def retry_available(job_id: str) -> bool: + """Whether a failed job can still be retried with its cached credentials. + + True only while the credentials remain in memory (i.e. within the + retention window and before a restart). The HTTP layer combines this with + the job's status and attempt count to decide whether to offer a retry. + """ + return job_id in _job_credentials + + +def _purge_expired_credentials() -> None: + """Drop cached credentials older than ``RETRY_CACHE_TTL_MINUTES``. + + Bounds the in-memory footprint and enforces the retention window: once a + failed job's credentials age out, a retry must re-send them. Called + periodically by the cleanup loop and once during stale-job recovery. + """ + cutoff = datetime.now(UTC) - timedelta(minutes=RETRY_CACHE_TTL_MINUTES) + expired = [jid for jid, ts in _job_cred_stored_at.items() if ts < cutoff] + for jid in expired: + _discard_credentials(jid) + if expired: + logger.info( + "Purged credentials for %d job(s) past the retry window", len(expired) + ) def is_worker_running() -> bool: @@ -102,10 +150,11 @@ def _process_job_sync(job_id: str) -> None: # respect that and return without flipping back to processing. if job.status == "cancelled": logger.info("Job %s was cancelled before pickup — skipping", job_id) - _job_credentials.pop(job_id, None) + _discard_credentials(job_id) return - credentials = _job_credentials.pop(job_id, {}) + # Read (do not pop) so the credentials survive for a possible retry. + credentials = _job_credentials.get(job_id, {}) collection = ( db.query(Collection).filter(Collection.id == job.collection_id).first() @@ -119,6 +168,8 @@ def _process_job_sync(job_id: str) -> None: job.error_message = error_msg job.completed_at = datetime.now(UTC) db.commit() + # The collection is gone — retrying cannot help, so free creds. + _discard_credentials(job_id) logger.error("Job %s aborted — collection missing", job_id) return @@ -141,11 +192,13 @@ def _process_job_sync(job_id: str) -> None: # noticed and bailed out. Leave the row alone so the cancellation # timestamp / status survive. logger.info("Job %s cancelled cooperatively: %s", job_id, exc) + _discard_credentials(job_id) return job.status = "completed" job.completed_at = datetime.now(UTC) db.commit() + _discard_credentials(job_id) logger.info( "Job %s completed — %d documents, %d chunks", @@ -167,6 +220,10 @@ def _process_job_sync(job_id: str) -> None: job.error_message = error_msg job.completed_at = datetime.now(UTC) db.commit() + # Keep the credentials so the user can retry, unless the job + # has exhausted its attempts — then a retry is not allowed. + if job.attempts >= _MAX_ATTEMPTS: + _discard_credentials(job_id) except Exception: logger.exception("Failed to record error for job %s", job_id) finally: @@ -196,6 +253,8 @@ async def _process_job_async(job_id: str) -> None: job.error_message = timeout_msg job.completed_at = datetime.now(UTC) db.commit() + if job.attempts >= _MAX_ATTEMPTS: + _discard_credentials(job_id) finally: db.close() @@ -267,6 +326,17 @@ async def start_worker() -> None: ) asyncio.create_task(_poll_loop()) + asyncio.create_task(_cleanup_loop()) + + +async def _cleanup_loop() -> None: + """Periodically purge credentials of jobs past the retry window.""" + while _running: + await asyncio.sleep(_CLEANUP_INTERVAL) + try: + _purge_expired_credentials() + except Exception: # noqa: BLE001 — never let cleanup kill the loop + logger.exception("Credential purge failed") async def stop_worker() -> None: @@ -290,6 +360,7 @@ def recover_stale_jobs() -> None: Jobs exceeding ``_MAX_ATTEMPTS`` are marked failed instead of being retried. Called once at startup, before the worker begins polling. """ + _purge_expired_credentials() db = _get_db() try: stale = ( @@ -305,6 +376,7 @@ def recover_stale_jobs() -> None: ) job.status = "failed" job.error_message = error_msg + _discard_credentials(job.id) logger.warning( "Job %s exceeded max attempts, marked failed", job.id ) diff --git a/lamb-kb-server/docker-compose.test.yml b/lamb-kb-server/docker-compose.test.yml index 24d328b94..2ab8a0576 100644 --- a/lamb-kb-server/docker-compose.test.yml +++ b/lamb-kb-server/docker-compose.test.yml @@ -11,26 +11,3 @@ services: retries: 30 tmpfs: - /qdrant/storage - - ollama-test: - image: ollama/ollama:latest - container_name: kbs-test-ollama - ports: - - "${OLLAMA_TEST_PORT:-11434}:11434" - volumes: - - ollama-models:/root/.ollama - healthcheck: - test: ["CMD-SHELL", "ollama list >/dev/null 2>&1 || exit 1"] - interval: 2s - timeout: 3s - retries: 60 - entrypoint: ["/bin/sh", "-c"] - command: - - | - ollama serve & - sleep 3 - ollama pull nomic-embed-text || true - wait - -volumes: - ollama-models: diff --git a/lamb-kb-server/pyproject.toml b/lamb-kb-server/pyproject.toml index c3b4e9e11..f60c3c111 100644 --- a/lamb-kb-server/pyproject.toml +++ b/lamb-kb-server/pyproject.toml @@ -11,6 +11,8 @@ dependencies = [ # Database "sqlalchemy>=2.0.30", + # Schema migrations (run at startup via init_db). + "alembic>=1.13", # File upload handling (multipart bodies) "python-multipart>=0.0.9", @@ -47,8 +49,14 @@ openai = [ "openai>=1.0.0", ] +# Optional KG-RAG: Neo4j graph store + OpenAI extractor. +kg-rag = [ + "neo4j>=5.0.0", + "openai>=1.0.0", +] + all = [ - "lamb-kb-server[qdrant,local,openai]", + "lamb-kb-server[qdrant,local,openai,kg-rag]", ] dev = [ @@ -79,6 +87,11 @@ select = ["E", "F", "W", "I", "UP", "B", "SIM"] # B905: zip strict — prefer concise zip() over strict=False noise everywhere. ignore = ["B008", "B904", "B905"] +[tool.ruff.lint.per-file-ignores] +# Autogenerated Alembic migrations: long DDL lines are unavoidable and the +# Sequence import is only used in some revisions. +"backend/migrations/versions/*" = ["E501", "F401", "I001"] + [tool.pytest.ini_options] asyncio_mode = "auto" testpaths = ["tests"] diff --git a/lamb-kb-server/tests/README.md b/lamb-kb-server/tests/README.md index 27ea40c5c..d3d6364c1 100644 --- a/lamb-kb-server/tests/README.md +++ b/lamb-kb-server/tests/README.md @@ -1,6 +1,6 @@ # Test suite -Three tiers, each with its own scope, fixtures, and runtime profile. The combined run hits **99% line + branch coverage** on `backend/`. +Three tiers, each with its own scope, fixtures, and runtime profile. The combined run reaches **~94% line + branch coverage** on `backend/` (the e2e tier's server runs in a subprocess the parent coverage process cannot instrument, so its server-side paths are not counted). | Tier | Where | Tests | Runtime | What it proves | |---|---|---|---|---| diff --git a/lamb-kb-server/tests/conftest.py b/lamb-kb-server/tests/conftest.py index 96e3cae7a..f413cb701 100644 --- a/lamb-kb-server/tests/conftest.py +++ b/lamb-kb-server/tests/conftest.py @@ -30,6 +30,14 @@ os.environ.setdefault("EMBEDDING_LOCAL", "DISABLE") os.environ.setdefault("VECTOR_DB_QDRANT", "DISABLE") +# Canonical Ollama embedding endpoint the suite asserts against. Read at +# import time by ``plugins.embedding.ollama``; pin it here (before any plugin +# import) so the default-endpoint tests are deterministic in every +# environment. A real CI/Docker export still wins via ``setdefault``. +os.environ.setdefault( + "OLLAMA_DEFAULT_ENDPOINT", "http://host.docker.internal:11435/api/embeddings" +) + # Make backend & tests packages importable without editable install. _ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(_ROOT / "backend")) diff --git a/lamb-kb-server/tests/e2e/conftest.py b/lamb-kb-server/tests/e2e/conftest.py index d1bf306f6..3a5c910ca 100644 --- a/lamb-kb-server/tests/e2e/conftest.py +++ b/lamb-kb-server/tests/e2e/conftest.py @@ -1,8 +1,11 @@ """E2E-tier fixtures: real uvicorn subprocess + docker stack + VCR. -The docker-compose stack (Qdrant + Ollama) is brought up at session start +The docker-compose stack (Qdrant only) is brought up at session start and torn down at session end via ``tests/e2e/_compose.py``. If Docker isn't available the entire e2e tier is skipped with a clear message. + +Embeddings are produced by an OpenAI-compatible endpoint (LM Studio running +on the host). No Ollama is used or required anywhere in the e2e tier. """ from __future__ import annotations @@ -25,6 +28,21 @@ _E2E_ROOT = Path(__file__).resolve().parent _KB_ROOT = _E2E_ROOT.parent.parent +# --- LM Studio (OpenAI-compatible) embedding configuration ----------------- +# LM Studio runs on the host. From inside a --network host container, +# ``localhost:1234`` reaches it directly; ``host.docker.internal`` works as a +# fallback for bridge-networked runs. Both are tried by ``_resolve_embedding``. +_LMSTUDIO_HOSTS = ( + os.environ.get("LMSTUDIO_BASE_URL"), + "http://localhost:1234/v1", + "http://host.docker.internal:1234/v1", +) +EMBEDDING_VENDOR = "openai" +EMBEDDING_MODEL = os.environ.get( + "LMSTUDIO_EMBEDDING_MODEL", "text-embedding-nomic-embed-text-v1.5" +) +EMBEDDING_API_KEY = "lm-studio" + def _docker_available() -> bool: if shutil.which("docker") is None: @@ -89,81 +107,91 @@ def _wait_for_http(url: str, timeout: float = 30.0) -> bool: return False -def _wait_for_ollama_model(ollama_url: str, model: str, timeout: float = 300.0) -> bool: - """Poll ``/api/tags`` until *model* is listed (pull complete).""" - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: +def _resolve_embedding_base_url() -> str | None: + """Return the first reachable LM Studio base URL (``.../v1``), or None. + + Probes ``GET {base}/models`` on each candidate host. Returns the base URL + (suitable for appending ``/embeddings``) of the first that responds. + """ + for base in _LMSTUDIO_HOSTS: + if not base: + continue + base = base.rstrip("/") try: - r = httpx.get(f"{ollama_url}/api/tags", timeout=5.0) + r = httpx.get(f"{base}/models", timeout=3.0) if r.status_code == 200: - models = [m.get("name", "") for m in r.json().get("models", [])] - if any(name.startswith(model) for name in models): - return True + return base except Exception: - pass - time.sleep(2.0) - return False + continue + return None @pytest.fixture(scope="session") def docker_stack() -> Iterator[dict]: - """Bring up Qdrant + Ollama containers for the e2e tier. + """Bring up the Qdrant container for the e2e tier and resolve LM Studio. - If ``QDRANT_TEST_PORT`` and ``OLLAMA_TEST_PORT`` environment variables are - already set (i.e. the stack was started externally before the test session), - the fixture skips compose_up/down and simply verifies the containers are - reachable at those ports. This allows running the e2e tier against a - pre-started stack without port conflicts or container name collisions. - """ - if not _docker_available(): - pytest.skip("Docker not available; skipping e2e tier") + If ``QDRANT_TEST_PORT`` is already set (i.e. the stack was started + externally before the test session), the fixture skips compose_up/down and + simply verifies Qdrant is reachable at that port. This allows running the + e2e tier against a pre-started stack without port conflicts or container + name collisions. - from tests.e2e._compose import compose_down, compose_up + Embeddings are served by LM Studio (OpenAI-compatible) on the host. The + fixture probes for a reachable endpoint and skips the tier if none is found. + """ + embedding_base = _resolve_embedding_base_url() + if embedding_base is None: + pytest.skip( + "LM Studio (OpenAI-compatible) endpoint not reachable at " + f"{[h for h in _LMSTUDIO_HOSTS if h]}; start LM Studio with an " + "embedding model loaded" + ) + embedding = { + "vendor": EMBEDDING_VENDOR, + "model": EMBEDDING_MODEL, + # Plugins/tests append nothing — they receive a full /embeddings URL. + "api_endpoint": f"{embedding_base}/embeddings", + "api_key": EMBEDDING_API_KEY, + } # --- Pre-started stack mode ------------------------------------------- - # When the orchestrator (or CI) has already brought up the stack, detect - # it via env vars OR by inspecting the well-known container names. This - # avoids port conflicts and container name collisions when re-running tests. + # When the orchestrator (or CI) has already brought up Qdrant, detect it + # via env var OR by inspecting the well-known container name. This avoids + # port conflicts and container name collisions when re-running tests. + # + # This mode requires NO local docker CLI (the test runner may itself be a + # container with Qdrant reachable over the network), so it is checked before + # any docker-availability gating. pre_qdrant_port = os.environ.get("QDRANT_TEST_PORT") - pre_ollama_port = os.environ.get("OLLAMA_TEST_PORT") - - # Fall back to docker inspect if env vars not set but containers exist. - if not pre_qdrant_port: + if not pre_qdrant_port and _docker_available(): discovered = _container_host_port("kbs-test-qdrant", 6333) if discovered: pre_qdrant_port = str(discovered) - if not pre_ollama_port: - discovered = _container_host_port("kbs-test-ollama", 11434) - if discovered: - pre_ollama_port = str(discovered) - if pre_qdrant_port and pre_ollama_port: + if pre_qdrant_port: qdrant_url = f"http://127.0.0.1:{pre_qdrant_port}" - ollama_url = f"http://127.0.0.1:{pre_ollama_port}" if not _wait_for_http(f"{qdrant_url}/", timeout=10): pytest.skip(f"Pre-started Qdrant not reachable at {qdrant_url}") - if not _wait_for_http(f"{ollama_url}/api/tags", timeout=10): - pytest.skip(f"Pre-started Ollama not reachable at {ollama_url}") - if not _wait_for_ollama_model(ollama_url, "nomic-embed-text", timeout=300): - pytest.skip( - f"Pre-started Ollama at {ollama_url} does not have " - f"nomic-embed-text pulled (run: ollama pull nomic-embed-text)" - ) yield { "qdrant_url": qdrant_url, - "ollama_url": ollama_url, "qdrant_port": int(pre_qdrant_port), - "ollama_port": int(pre_ollama_port), + "embedding": embedding, } return # do NOT tear down a pre-started stack # --- Self-managed stack mode ------------------------------------------ + # No pre-started Qdrant; bring one up via docker compose. This path needs a + # local docker CLI. + if not _docker_available(): + pytest.skip( + "Docker not available and no pre-started Qdrant (set " + "QDRANT_TEST_PORT); skipping e2e tier" + ) + + from tests.e2e._compose import compose_down, compose_up + qdrant_port = _free_port() - ollama_port = _free_port() - env = { - "QDRANT_TEST_PORT": str(qdrant_port), - "OLLAMA_TEST_PORT": str(ollama_port), - } + env = {"QDRANT_TEST_PORT": str(qdrant_port)} try: compose_up(env) @@ -171,25 +199,15 @@ def docker_stack() -> Iterator[dict]: pytest.skip(f"docker compose up failed: {exc}") qdrant_url = f"http://127.0.0.1:{qdrant_port}" - ollama_url = f"http://127.0.0.1:{ollama_port}" if not _wait_for_http(f"{qdrant_url}/", timeout=30): compose_down(env) pytest.skip("Qdrant container failed to become ready") - if not _wait_for_http(f"{ollama_url}/api/tags", timeout=60): - compose_down(env) - pytest.skip("Ollama container failed to become ready") - # Wait for the embedding model to finish pulling — /api/tags returns 200 - # before the model is ready, so tests racing the pull see 404s. - if not _wait_for_ollama_model(ollama_url, "nomic-embed-text", timeout=300): - compose_down(env) - pytest.skip("Ollama failed to pull nomic-embed-text within 300s") info = { "qdrant_url": qdrant_url, - "ollama_url": ollama_url, "qdrant_port": qdrant_port, - "ollama_port": ollama_port, + "embedding": embedding, } try: yield info @@ -214,6 +232,7 @@ def _spawn_kb_server(data_dir: str, env_overrides: dict[str, str] | None = None) "MAX_CONCURRENT_INGESTIONS": "2", "INGESTION_TASK_TIMEOUT_SECONDS": "30", "VECTOR_DB_QDRANT": "ENABLE", + "EMBEDDING_OPENAI": "ENABLE", "EMBEDDING_LOCAL": "DISABLE", "QDRANT_URL": "", # local on-disk mode by default } @@ -270,7 +289,7 @@ def kb_server_process_standalone() -> Iterator[dict]: """Launch the KB server in a subprocess on a free port — no Docker required. Use this fixture for tests that only need a live KB server (auth, routing, - error paths, capability listings) and do NOT need real Ollama or Qdrant + error paths, capability listings) and do NOT need real embeddings or Qdrant services. Tests that perform actual ingestion with real embeddings must use ``kb_server_process`` (which depends on ``docker_stack``). """ @@ -295,14 +314,14 @@ def kb_server_process(docker_stack: dict) -> Iterator[dict]: shutil.rmtree(data_dir, ignore_errors=True) -def _make_no_chromadb_fixture(ollama_url: str) -> Iterator[dict]: +def _make_no_chromadb_fixture(embedding_endpoint: str) -> Iterator[dict]: """Shared implementation for the two-phase 503-backend-unavailable fixture. Phase 1: spawn a default-config server, create a chromadb-backed collection - (using *ollama_url* for the embedding endpoint — Ollama is never contacted), - stop the server. Phase 2: spawn a second server against the same DATA_DIR - with ``VECTOR_DB_CHROMADB=DISABLE`` so the persisted collection's backend is - no longer registered. Yields ``{"info": phase2_info, "collection_id": str}``. + (using *embedding_endpoint* — the endpoint is never contacted), stop the + server. Phase 2: spawn a second server against the same DATA_DIR with + ``VECTOR_DB_CHROMADB=DISABLE`` so the persisted collection's backend is no + longer registered. Yields ``{"info": phase2_info, "collection_id": str}``. """ data_dir = tempfile.mkdtemp(prefix="kbs-e2e-503-") @@ -316,9 +335,9 @@ def _make_no_chromadb_fixture(ollama_url: str) -> Iterator[dict]: "chunking_strategy": "simple", "chunking_params": {"chunk_size": 400, "chunk_overlap": 0}, "embedding": { - "vendor": "ollama", - "model": "nomic-embed-text", - "api_endpoint": f"{ollama_url}/api/embeddings", + "vendor": EMBEDDING_VENDOR, + "model": EMBEDDING_MODEL, + "api_endpoint": embedding_endpoint, }, "vector_db_backend": "chromadb", } @@ -343,11 +362,11 @@ def _make_no_chromadb_fixture(ollama_url: str) -> Iterator[dict]: def kb_server_no_chromadb_standalone() -> Iterator[dict]: """Two-phase 503-backend-unavailable fixture — no Docker required. - Ollama is listed as the embedding vendor in the collection payload but is - never contacted: 503 fires before the embedding callable is invoked, so - a dummy (unreachable) endpoint suffices. + The embedding endpoint is listed in the collection payload but is never + contacted: 503 fires before the embedding callable is invoked, so a dummy + (unreachable) endpoint suffices. """ - yield from _make_no_chromadb_fixture("http://127.0.0.1:19999") + yield from _make_no_chromadb_fixture("http://127.0.0.1:19999/v1/embeddings") @pytest.fixture @@ -359,12 +378,10 @@ def kb_server_no_chromadb(docker_stack: dict) -> Iterator[dict]: with ``VECTOR_DB_CHROMADB=DISABLE`` so the persisted collection's backend is no longer registered. Yields ``{"info": phase2_info, "collection_id": str}``. - Depends on docker_stack only because Ollama is the simplest registered - embedding vendor available to a fresh subprocess (the test fake plugin - only registers in the test process). Ollama is not actually contacted — - 503 fires before the embedding callable is invoked. + The embedding endpoint is the LM Studio endpoint but is not actually + contacted — 503 fires before the embedding callable is invoked. """ - yield from _make_no_chromadb_fixture(docker_stack["ollama_url"]) + yield from _make_no_chromadb_fixture(docker_stack["embedding"]["api_endpoint"]) @pytest.fixture diff --git a/lamb-kb-server/tests/e2e/test_concurrency.py b/lamb-kb-server/tests/e2e/test_concurrency.py index 3b8baf958..05aba3233 100644 --- a/lamb-kb-server/tests/e2e/test_concurrency.py +++ b/lamb-kb-server/tests/e2e/test_concurrency.py @@ -6,7 +6,7 @@ - MAX_CONCURRENT_INGESTIONS=2 is respected (≤2 jobs in "processing" at once). - SQLite WAL mode handles concurrent readers/writers without lock errors. -Tests are SLOW (~30–90 s each due to Ollama latency × multiple jobs). +Tests are SLOW (~30–90 s each due to embedding latency × multiple jobs). Each test creates its own ``httpx.Client`` with a generous timeout rather than using the session ``http`` fixture (which has ``timeout=30s``). Concurrent @@ -129,9 +129,9 @@ def _poll_until_complete( def _make_collection( base_url: str, org_id: str, - vendor: str = "ollama", + vendor: str = "openai", api_endpoint: str = "", - model: str = "nomic-embed-text", + model: str = "text-embedding-nomic-embed-text-v1.5", name_suffix: str = "", ) -> dict: """Create a collection using the given embedding vendor. @@ -208,13 +208,12 @@ def test_20_concurrent_add_content_requests_succeed( # Drain stdout to prevent pipe buffer deadlock with LOG_LEVEL=DEBUG. _drain_server_stdout(kb_server_process["process"]) - ollama_url = docker_stack["ollama_url"] - api_endpoint = f"{ollama_url}/api/embeddings" + api_endpoint = docker_stack["embedding"]["api_endpoint"] base_url = kb_server_process["base_url"] org_id = f"org-conc-{uuid4().hex[:8]}" collection = _make_collection( - base_url, org_id, vendor="ollama", api_endpoint=api_endpoint + base_url, org_id, vendor="openai", api_endpoint=api_endpoint ) collection_id = collection["id"] @@ -299,13 +298,12 @@ def test_concurrent_ingest_respects_semaphore( # Drain stdout to prevent pipe buffer deadlock with LOG_LEVEL=DEBUG. _drain_server_stdout(kb_server_process["process"]) - ollama_url = docker_stack["ollama_url"] - api_endpoint = f"{ollama_url}/api/embeddings" + api_endpoint = docker_stack["embedding"]["api_endpoint"] base_url = kb_server_process["base_url"] org_id = f"org-sem-{uuid4().hex[:8]}" collection = _make_collection( - base_url, org_id, vendor="ollama", api_endpoint=api_endpoint + base_url, org_id, vendor="openai", api_endpoint=api_endpoint ) collection_id = collection["id"] @@ -372,8 +370,7 @@ def test_concurrent_collection_creation_no_conflicts( Verifies that concurrent SQLite writes for collection creation do not cause lock errors or duplicate-ID collisions. """ - ollama_url = docker_stack["ollama_url"] - api_endpoint = f"{ollama_url}/api/embeddings" + api_endpoint = docker_stack["embedding"]["api_endpoint"] base_url = kb_server_process["base_url"] n = 10 @@ -394,8 +391,8 @@ def _create_one(i: int) -> dict: "chunking_strategy": "simple", "chunking_params": {"chunk_size": 300, "chunk_overlap": 30}, "embedding": { - "vendor": "ollama", - "model": "nomic-embed-text", + "vendor": "openai", + "model": "text-embedding-nomic-embed-text-v1.5", "api_endpoint": api_endpoint, }, "vector_db_backend": "chromadb", @@ -438,13 +435,12 @@ def test_concurrent_queries_after_ingest( # Drain stdout to prevent pipe buffer deadlock with LOG_LEVEL=DEBUG. _drain_server_stdout(kb_server_process["process"]) - ollama_url = docker_stack["ollama_url"] - api_endpoint = f"{ollama_url}/api/embeddings" + api_endpoint = docker_stack["embedding"]["api_endpoint"] base_url = kb_server_process["base_url"] org_id = f"org-qc-{uuid4().hex[:8]}" collection = _make_collection( - base_url, org_id, vendor="ollama", api_endpoint=api_endpoint + base_url, org_id, vendor="openai", api_endpoint=api_endpoint ) collection_id = collection["id"] @@ -525,8 +521,7 @@ def test_high_concurrency_no_db_lock_errors( # Drain stdout to prevent pipe buffer deadlock with LOG_LEVEL=DEBUG. _drain_server_stdout(kb_server_process["process"]) - ollama_url = docker_stack["ollama_url"] - api_endpoint = f"{ollama_url}/api/embeddings" + api_endpoint = docker_stack["embedding"]["api_endpoint"] base_url = kb_server_process["base_url"] n_collections = 5 @@ -539,7 +534,7 @@ def test_high_concurrency_no_db_lock_errors( coll = _make_collection( base_url, org_id, - vendor="ollama", + vendor="openai", api_endpoint=api_endpoint, name_suffix=f"wal{ci}", ) diff --git a/lamb-kb-server/tests/e2e/test_crash_recovery.py b/lamb-kb-server/tests/e2e/test_crash_recovery.py index 2edd6a53a..aacef0b59 100644 --- a/lamb-kb-server/tests/e2e/test_crash_recovery.py +++ b/lamb-kb-server/tests/e2e/test_crash_recovery.py @@ -5,7 +5,7 @@ The session-scoped ``kb_server_process`` fixture is intentionally NOT used here because these tests require independent start/kill/restart cycles. -``docker_stack`` is the only session fixture used (Ollama container for +``docker_stack`` is the only session fixture used (embeddings/Qdrant for real embeddings). Vector storage uses Qdrant **local on-disk mode** so each test's ``data_dir`` is fully self-contained and isolated — no shared remote Qdrant instance is involved. @@ -153,9 +153,8 @@ def _poll_job_status( # --------------------------------------------------------------------------- -def _make_collection_payload(ollama_url: str, name: str = "crash-test") -> dict: - """Return a ``POST /collections`` body using Ollama embeddings + Qdrant local.""" - ollama_endpoint = f"{ollama_url}/api/embeddings" +def _make_collection_payload(emb_endpoint: str, name: str = "crash-test") -> dict: + """Return a ``POST /collections`` body using LM Studio embeddings + Qdrant local.""" return { "organization_id": "org-crash-test", "name": name, @@ -163,17 +162,16 @@ def _make_collection_payload(ollama_url: str, name: str = "crash-test") -> dict: # chunk_overlap must be < chunk_size — use chunk_overlap (not 'overlap'). "chunking_params": {"chunk_size": 500, "chunk_overlap": 50}, "embedding": { - "vendor": "ollama", - "model": "nomic-embed-text", - "api_endpoint": ollama_endpoint, + "vendor": "openai", + "model": "text-embedding-nomic-embed-text-v1.5", + "api_endpoint": emb_endpoint, }, "vector_db_backend": "qdrant", } -def _make_add_content_payload(ollama_url: str, text: str, source_id: str = "item-1") -> dict: +def _make_add_content_payload(emb_endpoint: str, text: str, source_id: str = "item-1") -> dict: """Return a ``POST /collections/{id}/add-content`` body.""" - ollama_endpoint = f"{ollama_url}/api/embeddings" return { "documents": [ { @@ -184,7 +182,7 @@ def _make_add_content_payload(ollama_url: str, text: str, source_id: str = "item ], "embedding_credentials": { "api_key": "", - "api_endpoint": ollama_endpoint, + "api_endpoint": emb_endpoint, }, } @@ -198,7 +196,7 @@ class TestSigkillThenRestartResumesPendingJobs: """Queue several jobs, SIGKILL mid-flight, restart, confirm all terminal.""" def test_sigkill_then_restart_resumes_pending_jobs(self, docker_stack: dict) -> None: - ollama_url = docker_stack["ollama_url"] + emb_endpoint = docker_stack["embedding"]["api_endpoint"] data_dir = tempfile.mkdtemp(prefix="kbs-crash-") port1 = _free_port() @@ -210,13 +208,12 @@ def test_sigkill_then_restart_resumes_pending_jobs(self, docker_stack: dict) -> base1 = f"http://127.0.0.1:{port1}" assert _wait_health(base1, timeout=30), "Server #1 failed to start" - ollama_endpoint = f"{ollama_url}/api/embeddings" with httpx.Client(base_url=base1, timeout=30.0) as client: - # Create a collection that uses Ollama (real embedding, slower than fake). + # Create a collection that uses LM Studio (real embedding, slower than fake). col_r = client.post( "/collections", - json=_make_collection_payload(ollama_url, name="resume-test"), + json=_make_collection_payload(emb_endpoint, name="resume-test"), headers=_AUTH_HEADERS, ) assert col_r.status_code == 201, col_r.text @@ -243,7 +240,7 @@ def test_sigkill_then_restart_resumes_pending_jobs(self, docker_stack: dict) -> ], "embedding_credentials": { "api_key": "", - "api_endpoint": ollama_endpoint, + "api_endpoint": emb_endpoint, }, }, headers=_AUTH_HEADERS, @@ -263,7 +260,7 @@ def test_sigkill_then_restart_resumes_pending_jobs(self, docker_stack: dict) -> break if not found_processing: time.sleep(0.3) - # Even if we didn't observe 'processing' (fast Ollama on GPU), + # Even if we didn't observe 'processing' (fast embeddings), # the SIGKILL still exercises the recovery path. # --- SIGKILL server #1 --- @@ -303,7 +300,7 @@ def test_sigkill_then_restart_resumes_pending_jobs(self, docker_stack: dict) -> def test_sigkill_preserves_completed_jobs(docker_stack: dict) -> None: """A completed job is still retrievable after SIGKILL + restart.""" - ollama_url = docker_stack["ollama_url"] + emb_endpoint = docker_stack["embedding"]["api_endpoint"] data_dir = tempfile.mkdtemp(prefix="kbs-crash-") port = _free_port() @@ -318,7 +315,7 @@ def test_sigkill_preserves_completed_jobs(docker_stack: dict) -> None: # Create collection. col_r = client.post( "/collections", - json=_make_collection_payload(ollama_url, name="preserve-jobs-test"), + json=_make_collection_payload(emb_endpoint, name="preserve-jobs-test"), headers=_AUTH_HEADERS, ) assert col_r.status_code == 201, col_r.text @@ -328,7 +325,7 @@ def test_sigkill_preserves_completed_jobs(docker_stack: dict) -> None: add_r = client.post( f"/collections/{col_id}/add-content", json=_make_add_content_payload( - ollama_url, + emb_endpoint, "LAMB helps teachers build AI learning assistants easily.", source_id="item-preserve", ), @@ -379,7 +376,7 @@ def test_sigkill_preserves_completed_jobs(docker_stack: dict) -> None: def test_sigkill_preserves_collection_metadata(docker_stack: dict) -> None: """Collections created before SIGKILL are still listed and accessible after restart.""" - ollama_url = docker_stack["ollama_url"] + emb_endpoint = docker_stack["embedding"]["api_endpoint"] data_dir = tempfile.mkdtemp(prefix="kbs-crash-") port = _free_port() @@ -395,7 +392,7 @@ def test_sigkill_preserves_collection_metadata(docker_stack: dict) -> None: for i in range(3): r = client.post( "/collections", - json=_make_collection_payload(ollama_url, name=f"meta-test-{i}"), + json=_make_collection_payload(emb_endpoint, name=f"meta-test-{i}"), headers=_AUTH_HEADERS, ) assert r.status_code == 201, r.text @@ -444,7 +441,7 @@ def test_sigkill_preserves_chromadb_storage(docker_stack: dict) -> None: alongside the SQLite database. Both survive the crash and the restarted server should serve identical query results. """ - ollama_url = docker_stack["ollama_url"] + emb_endpoint = docker_stack["embedding"]["api_endpoint"] data_dir = tempfile.mkdtemp(prefix="kbs-crash-") port = _free_port() @@ -455,7 +452,6 @@ def test_sigkill_preserves_chromadb_storage(docker_stack: dict) -> None: base = f"http://127.0.0.1:{port}" assert _wait_health(base, timeout=30), "Server failed to start" - ollama_endpoint = f"{ollama_url}/api/embeddings" query_text = "LAMB open-source educators AI learning" # 5 short documents → 5 vectors after simple chunking (each fits in 500 chars). @@ -475,7 +471,7 @@ def test_sigkill_preserves_chromadb_storage(docker_stack: dict) -> None: # Create collection. col_r = client.post( "/collections", - json=_make_collection_payload(ollama_url, name="vector-persist-test"), + json=_make_collection_payload(emb_endpoint, name="vector-persist-test"), headers=_AUTH_HEADERS, ) assert col_r.status_code == 201, col_r.text @@ -488,7 +484,7 @@ def test_sigkill_preserves_chromadb_storage(docker_stack: dict) -> None: "documents": docs, "embedding_credentials": { "api_key": "", - "api_endpoint": ollama_endpoint, + "api_endpoint": emb_endpoint, }, }, headers=_AUTH_HEADERS, @@ -509,7 +505,7 @@ def test_sigkill_preserves_chromadb_storage(docker_stack: dict) -> None: "top_k": 5, "embedding_credentials": { "api_key": "", - "api_endpoint": ollama_endpoint, + "api_endpoint": emb_endpoint, }, }, headers=_AUTH_HEADERS, @@ -538,7 +534,7 @@ def test_sigkill_preserves_chromadb_storage(docker_stack: dict) -> None: "top_k": 5, "embedding_credentials": { "api_key": "", - "api_endpoint": ollama_endpoint, + "api_endpoint": emb_endpoint, }, }, headers=_AUTH_HEADERS, diff --git a/lamb-kb-server/tests/e2e/test_error_paths.py b/lamb-kb-server/tests/e2e/test_error_paths.py index 8c00dd014..7a9de9394 100644 --- a/lamb-kb-server/tests/e2e/test_error_paths.py +++ b/lamb-kb-server/tests/e2e/test_error_paths.py @@ -9,7 +9,7 @@ **Embedding vendor in helpers:** The e2e subprocess does NOT register the ``FakeEmbedding`` test double — that is only available in unit/integration tiers. -We use ``ollama`` with a dummy endpoint for error-path tests that only need a +We use ``openai`` with a dummy endpoint for error-path tests that only need a collection to *exist* (they fail before any embedding is invoked). """ @@ -46,14 +46,14 @@ def _minimal_collection_payload( name: str | None = None, *, chunking_strategy: str = "simple", - embedding_vendor: str = "ollama", - embedding_model: str = "nomic-embed-text", + embedding_vendor: str = "openai", + embedding_model: str = "text-embedding-nomic-embed-text-v1.5", embedding_endpoint: str = "http://127.0.0.1:19999", # unreachable — not called vector_db_backend: str = "chromadb", ) -> dict: """Build a minimal CreateCollectionRequest payload. - Defaults use ``ollama`` with an intentionally unreachable endpoint because + Defaults use ``openai`` with an intentionally unreachable endpoint because error-path tests that only create/list/delete collections never actually invoke the embedding backend. The endpoint will not be contacted. """ @@ -493,8 +493,8 @@ def test_422_missing_required_field_name(http_standalone: httpx.Client) -> None: "organization_id": org_id, "chunking_strategy": "simple", "embedding": { - "vendor": "ollama", - "model": "nomic-embed-text", + "vendor": "openai", + "model": "text-embedding-nomic-embed-text-v1.5", }, "vector_db_backend": "chromadb", } diff --git a/lamb-kb-server/tests/e2e/test_multitenancy.py b/lamb-kb-server/tests/e2e/test_multitenancy.py index 1049c3fe6..a21aaa3b0 100644 --- a/lamb-kb-server/tests/e2e/test_multitenancy.py +++ b/lamb-kb-server/tests/e2e/test_multitenancy.py @@ -1,7 +1,7 @@ """E2E multi-tenancy isolation tests. Verifies per-org isolation over real HTTP against a real uvicorn subprocess -with a real ChromaDB / Qdrant backend and real Ollama embeddings. +with a real ChromaDB / Qdrant backend and real LM Studio (OpenAI-compatible) embeddings. ADR-6: LAMB owns ACL; KB Server does not enforce per-org access on GET by id. ADR-9: Per-org filesystem isolation at data/storage/{org_id}/{collection_id}/. @@ -31,12 +31,12 @@ def _create_collection( http: httpx.Client, org_id: str, name: str, - vendor: str = "ollama", + vendor: str = "openai", api_endpoint: str = "", ) -> dict: """POST /collections and return the parsed JSON body. - Uses Ollama as the embedding vendor so real vectors are produced, which + Uses LM Studio (OpenAI-compatible) as the embedding vendor so real vectors are produced, which makes the cross-org leak test (#5) meaningful. """ payload = { @@ -47,7 +47,7 @@ def _create_collection( "chunking_params": {"chunk_size": 500, "chunk_overlap": 50}, "embedding": { "vendor": vendor, - "model": "nomic-embed-text", + "model": "text-embedding-nomic-embed-text-v1.5", "api_endpoint": api_endpoint, }, "vector_db_backend": "chromadb", @@ -118,12 +118,12 @@ def test_two_orgs_same_collection_name( Name uniqueness is scoped per (organization_id, name) — the same name in different orgs must succeed. """ - ollama_url = docker_stack["ollama_url"] + emb_endpoint = docker_stack["embedding"]["api_endpoint"] org_a = _org_id() org_b = _org_id() - col_a = _create_collection(http, org_a, "kb1", api_endpoint=ollama_url) - col_b = _create_collection(http, org_b, "kb1", api_endpoint=ollama_url) + col_a = _create_collection(http, org_a, "kb1", api_endpoint=emb_endpoint) + col_b = _create_collection(http, org_b, "kb1", api_endpoint=emb_endpoint) assert col_a["organization_id"] == org_a assert col_b["organization_id"] == org_b @@ -142,13 +142,13 @@ def test_filesystem_isolation_per_org( After creating and ingesting into collections for org-A and org-B, each org's storage directory must exist and be completely separate. """ - ollama_url = docker_stack["ollama_url"] + emb_endpoint = docker_stack["embedding"]["api_endpoint"] data_dir = Path(kb_server_process["data_dir"]) org_a = _org_id() org_b = _org_id() - col_a = _create_collection(http, org_a, "fs-test", api_endpoint=ollama_url) - col_b = _create_collection(http, org_b, "fs-test", api_endpoint=ollama_url) + col_a = _create_collection(http, org_a, "fs-test", api_endpoint=emb_endpoint) + col_b = _create_collection(http, org_b, "fs-test", api_endpoint=emb_endpoint) # Ingest something to ensure the storage directories are populated. job_a = _ingest( @@ -156,14 +156,14 @@ def test_filesystem_isolation_per_org( col_a["id"], "doc-a", "Filesystem isolation test content for org A.", - api_endpoint=ollama_url, + api_endpoint=emb_endpoint, ) job_b = _ingest( http, col_b["id"], "doc-b", "Filesystem isolation test content for org B.", - api_endpoint=ollama_url, + api_endpoint=emb_endpoint, ) result_a = _poll_job(http, job_a) result_b = _poll_job(http, job_b) @@ -205,10 +205,10 @@ def test_cross_org_collection_access_via_id( This documents the actual behavior: the collection is returned with a 200. """ - ollama_url = docker_stack["ollama_url"] + emb_endpoint = docker_stack["embedding"]["api_endpoint"] org_a = _org_id() - col_a = _create_collection(http, org_a, "acl-test", api_endpoint=ollama_url) + col_a = _create_collection(http, org_a, "acl-test", api_endpoint=emb_endpoint) collection_id = col_a["id"] # Retrieve the collection by ID — no org filter is required or checked. @@ -231,12 +231,12 @@ def test_list_with_org_filter_excludes_other_orgs( Create kb-1 in org-A and kb-2 in org-B. Listing with org-A's filter must contain kb-1 but not kb-2, and vice versa. """ - ollama_url = docker_stack["ollama_url"] + emb_endpoint = docker_stack["embedding"]["api_endpoint"] org_a = _org_id() org_b = _org_id() - col_a = _create_collection(http, org_a, "kb-1", api_endpoint=ollama_url) - col_b = _create_collection(http, org_b, "kb-2", api_endpoint=ollama_url) + col_a = _create_collection(http, org_a, "kb-1", api_endpoint=emb_endpoint) + col_b = _create_collection(http, org_b, "kb-2", api_endpoint=emb_endpoint) # List org-A. r_a = http.get(f"/collections?organization_id={org_a}") @@ -270,17 +270,17 @@ def test_query_isolation_between_orgs( ) -> None: """Querying org-A's collection must not surface org-B's chunks. - Uses real Ollama embeddings so each phrase produces a real semantic vector. + Uses real LM Studio embeddings so each phrase produces a real semantic vector. Ingest "secret-org-A-data" into org-A's collection and "secret-org-B-data" into org-B's collection. Query org-A's collection for "secret-org-B-data" — no result with source_item_id from org-B must appear. """ - ollama_url = docker_stack["ollama_url"] + emb_endpoint = docker_stack["embedding"]["api_endpoint"] org_a = _org_id() org_b = _org_id() - col_a = _create_collection(http, org_a, "query-iso-a", api_endpoint=ollama_url) - col_b = _create_collection(http, org_b, "query-iso-b", api_endpoint=ollama_url) + col_a = _create_collection(http, org_a, "query-iso-a", api_endpoint=emb_endpoint) + col_b = _create_collection(http, org_b, "query-iso-b", api_endpoint=emb_endpoint) # Ingest into org-A. job_a = _ingest( @@ -288,7 +288,7 @@ def test_query_isolation_between_orgs( col_a["id"], "src-a", "secret-org-A-data: confidential information belonging to organization A.", - api_endpoint=ollama_url, + api_endpoint=emb_endpoint, ) # Ingest into org-B. job_b = _ingest( @@ -296,7 +296,7 @@ def test_query_isolation_between_orgs( col_b["id"], "src-b", "secret-org-B-data: confidential information belonging to organization B.", - api_endpoint=ollama_url, + api_endpoint=emb_endpoint, ) result_a = _poll_job(http, job_a) @@ -312,7 +312,7 @@ def test_query_isolation_between_orgs( "top_k": 10, "embedding_credentials": { "api_key": "", - "api_endpoint": ollama_url, + "api_endpoint": emb_endpoint, }, }, ) @@ -338,13 +338,13 @@ def test_delete_one_orgs_collection_doesnt_affect_other( - org-A's storage dir is gone - org-B's storage dir still exists """ - ollama_url = docker_stack["ollama_url"] + emb_endpoint = docker_stack["embedding"]["api_endpoint"] data_dir = Path(kb_server_process["data_dir"]) org_a = _org_id() org_b = _org_id() - col_a = _create_collection(http, org_a, "del-test-a", api_endpoint=ollama_url) - col_b = _create_collection(http, org_b, "del-test-b", api_endpoint=ollama_url) + col_a = _create_collection(http, org_a, "del-test-a", api_endpoint=emb_endpoint) + col_b = _create_collection(http, org_b, "del-test-b", api_endpoint=emb_endpoint) # Ingest into both so storage directories are created. job_a = _ingest( @@ -352,14 +352,14 @@ def test_delete_one_orgs_collection_doesnt_affect_other( col_a["id"], "doc-del-a", "Org A content for deletion test.", - api_endpoint=ollama_url, + api_endpoint=emb_endpoint, ) job_b = _ingest( http, col_b["id"], "doc-del-b", "Org B content for deletion test.", - api_endpoint=ollama_url, + api_endpoint=emb_endpoint, ) assert _poll_job(http, job_a)["status"] == "completed" assert _poll_job(http, job_b)["status"] == "completed" @@ -408,12 +408,12 @@ def test_concurrent_orgs_no_id_collision( ThreadPoolExecutor to issue creates concurrently, verifying that the server's ID-generation has no collision under concurrent load. """ - ollama_url = docker_stack["ollama_url"] + emb_endpoint = docker_stack["embedding"]["api_endpoint"] orgs = [_org_id() for _ in range(5)] def _create(org_id: str) -> dict: return _create_collection( - http, org_id, "shared-name", api_endpoint=ollama_url + http, org_id, "shared-name", api_endpoint=emb_endpoint ) # ThreadPoolExecutor for concurrent HTTP requests. diff --git a/lamb-kb-server/tests/e2e/test_pipeline_matrix.py b/lamb-kb-server/tests/e2e/test_pipeline_matrix.py index 63f9e973a..5fba2813d 100644 --- a/lamb-kb-server/tests/e2e/test_pipeline_matrix.py +++ b/lamb-kb-server/tests/e2e/test_pipeline_matrix.py @@ -1,7 +1,7 @@ """E2E pipeline matrix: create → ingest → query → delete over real HTTP. Parametrized over (vector_db, embedding_vendor) combinations using real -backend containers (Qdrant on :16333, Ollama on :11435 with nomic-embed-text). +backend containers (Qdrant on :16333, LM Studio on host (OpenAI-compatible)). OpenAI is excluded from this matrix because no real API key is available for recording VCR cassettes, and hand-crafted cassettes are fragile under @@ -11,7 +11,7 @@ Stack requirements (provided by the session-scope docker_stack fixture): - Qdrant : http://127.0.0.1:{qdrant_port} - - Ollama : http://127.0.0.1:{ollama_port} (nomic-embed-text must be pulled) + - Embeddings : LM Studio on the host (OpenAI-compatible, /v1) """ from __future__ import annotations @@ -30,7 +30,7 @@ # --------------------------------------------------------------------------- _POLL_INTERVAL = 1.0 # seconds between polls -_POLL_TIMEOUT = 120.0 # Ollama embeddings can be slow on first call +_POLL_TIMEOUT = 120.0 # Embeddings can be slow on first call def _poll_job_sync( @@ -59,9 +59,9 @@ def _unique_id(prefix: str = "") -> str: return f"{prefix}{uuid.uuid4().hex[:8]}" -def _ollama_endpoint(docker_stack: dict) -> str: - """Return the Ollama /api/embeddings URL for the running container.""" - return f"{docker_stack['ollama_url']}/api/embeddings" +def _embedding_endpoint(docker_stack: dict) -> str: + """Return the OpenAI-compatible (LM Studio) /embeddings URL.""" + return docker_stack["embedding"]["api_endpoint"] # --------------------------------------------------------------------------- @@ -69,8 +69,8 @@ def _ollama_endpoint(docker_stack: dict) -> str: # --------------------------------------------------------------------------- _MATRIX = [ - pytest.param("chromadb", "ollama", id="chromadb-ollama"), - pytest.param("qdrant", "ollama", id="qdrant-ollama"), + pytest.param("chromadb", "openai", id="chromadb-openai"), + pytest.param("qdrant", "openai", id="qdrant-openai"), ] @@ -90,7 +90,7 @@ def test_full_pipeline( 1. POST /collections with the given (vector_db, embedding_vendor) combo. 2. POST /collections/{id}/add-content with 3 short documents using the "simple" chunking strategy. - 3. Poll /jobs/{id} until completed (up to 120 s for Ollama). + 3. Poll /jobs/{id} until completed (up to 120 s for embeddings). 4. POST /collections/{id}/query with text known to be in doc-2. 5. Assert top result has the expected source_item_id and score > 0.3. 6. DELETE /collections/{id}/content/{source_item_id} for doc-2. @@ -100,7 +100,7 @@ def test_full_pipeline( """ org_id = _unique_id("org-") collection_name = _unique_id(f"mat-{vector_db[:4]}-") - ollama_ep = _ollama_endpoint(docker_stack) + embedding_ep = _embedding_endpoint(docker_stack) # ------------------------------------------------------------------ # Step 1: Create collection @@ -113,8 +113,8 @@ def test_full_pipeline( "chunking_params": {"chunk_size": 512, "chunk_overlap": 50}, "embedding": { "vendor": embedding_vendor, - "model": "nomic-embed-text", - "api_endpoint": ollama_ep, + "model": "text-embedding-nomic-embed-text-v1.5", + "api_endpoint": embedding_ep, }, "vector_db_backend": vector_db, } @@ -172,7 +172,7 @@ def test_full_pipeline( ], "embedding_credentials": { "api_key": "", - "api_endpoint": ollama_ep, + "api_endpoint": embedding_ep, }, } r_ingest = http.post( @@ -203,7 +203,7 @@ def test_full_pipeline( "top_k": 3, "embedding_credentials": { "api_key": "", - "api_endpoint": ollama_ep, + "api_endpoint": embedding_ep, }, } r_query = http.post( @@ -282,7 +282,7 @@ def test_full_pipeline( # --------------------------------------------------------------------------- -# Chunking variety tests (chromadb + ollama, non-parametrized) +# Chunking variety tests (chromadb + openai, non-parametrized) # --------------------------------------------------------------------------- @@ -293,13 +293,13 @@ def test_pipeline_with_hierarchical_chunking( kb_server_process: dict, docker_stack: dict, ) -> None: - """Full pipeline using hierarchical chunking strategy with real Ollama embeddings. + """Full pipeline using hierarchical chunking strategy with real LM Studio embeddings. Verifies that hierarchical parent/child chunks are produced and that query results carry section_title metadata from the markdown headers. """ org_id = _unique_id("org-") - ollama_ep = _ollama_endpoint(docker_stack) + embedding_ep = _embedding_endpoint(docker_stack) # Create collection r_create = http.post( @@ -315,9 +315,9 @@ def test_pipeline_with_hierarchical_chunking( "child_chunk_overlap": 30, }, "embedding": { - "vendor": "ollama", - "model": "nomic-embed-text", - "api_endpoint": ollama_ep, + "vendor": "openai", + "model": "text-embedding-nomic-embed-text-v1.5", + "api_endpoint": embedding_ep, }, "vector_db_backend": "chromadb", }, @@ -360,7 +360,7 @@ def test_pipeline_with_hierarchical_chunking( "text": markdown_doc, } ], - "embedding_credentials": {"api_key": "", "api_endpoint": ollama_ep}, + "embedding_credentials": {"api_key": "", "api_endpoint": embedding_ep}, }, ) assert r_ingest.status_code == 202, r_ingest.text @@ -374,7 +374,7 @@ def test_pipeline_with_hierarchical_chunking( json={ "query_text": "retrieval augmented generation RAG language model", "top_k": 5, - "embedding_credentials": {"api_key": "", "api_endpoint": ollama_ep}, + "embedding_credentials": {"api_key": "", "api_endpoint": embedding_ep}, }, ) assert r_query.status_code == 200, r_query.text @@ -404,13 +404,13 @@ def test_pipeline_with_by_page_chunking( kb_server_process: dict, docker_stack: dict, ) -> None: - """Full pipeline using by_page chunking with pre-split pages and real Ollama. + """Full pipeline using by_page chunking with pre-split pages and real LM Studio. Verifies page_range metadata is present in query results and that each page produces a distinct chunk with the correct page number. """ org_id = _unique_id("org-") - ollama_ep = _ollama_endpoint(docker_stack) + embedding_ep = _embedding_endpoint(docker_stack) # Create collection r_create = http.post( @@ -422,9 +422,9 @@ def test_pipeline_with_by_page_chunking( "chunking_strategy": "by_page", "chunking_params": {"pages_per_chunk": 1}, "embedding": { - "vendor": "ollama", - "model": "nomic-embed-text", - "api_endpoint": ollama_ep, + "vendor": "openai", + "model": "text-embedding-nomic-embed-text-v1.5", + "api_endpoint": embedding_ep, }, "vector_db_backend": "chromadb", }, @@ -479,7 +479,7 @@ def test_pipeline_with_by_page_chunking( }, } ], - "embedding_credentials": {"api_key": "", "api_endpoint": ollama_ep}, + "embedding_credentials": {"api_key": "", "api_endpoint": embedding_ep}, }, ) assert r_ingest.status_code == 202, r_ingest.text @@ -496,7 +496,7 @@ def test_pipeline_with_by_page_chunking( json={ "query_text": "convolutional neural networks image recognition filters", "top_k": 3, - "embedding_credentials": {"api_key": "", "api_endpoint": ollama_ep}, + "embedding_credentials": {"api_key": "", "api_endpoint": embedding_ep}, }, ) assert r_query.status_code == 200, r_query.text @@ -527,13 +527,13 @@ def test_pipeline_with_by_section_chunking( kb_server_process: dict, docker_stack: dict, ) -> None: - """Full pipeline using by_section chunking with real Ollama embeddings. + """Full pipeline using by_section chunking with real LM Studio embeddings. Verifies that sections are split on H2 headings and that each chunk carries a section_title in its metadata. """ org_id = _unique_id("org-") - ollama_ep = _ollama_endpoint(docker_stack) + embedding_ep = _embedding_endpoint(docker_stack) # Create collection r_create = http.post( @@ -548,9 +548,9 @@ def test_pipeline_with_by_section_chunking( "headings_per_chunk": 1, }, "embedding": { - "vendor": "ollama", - "model": "nomic-embed-text", - "api_endpoint": ollama_ep, + "vendor": "openai", + "model": "text-embedding-nomic-embed-text-v1.5", + "api_endpoint": embedding_ep, }, "vector_db_backend": "chromadb", }, @@ -589,7 +589,7 @@ def test_pipeline_with_by_section_chunking( "text": section_doc, } ], - "embedding_credentials": {"api_key": "", "api_endpoint": ollama_ep}, + "embedding_credentials": {"api_key": "", "api_endpoint": embedding_ep}, }, ) assert r_ingest.status_code == 202, r_ingest.text @@ -606,7 +606,7 @@ def test_pipeline_with_by_section_chunking( json={ "query_text": "ocean acidification carbon dioxide seawater carbonic acid", "top_k": 3, - "embedding_credentials": {"api_key": "", "api_endpoint": ollama_ep}, + "embedding_credentials": {"api_key": "", "api_endpoint": embedding_ep}, }, ) assert r_query.status_code == 200, r_query.text diff --git a/lamb-kb-server/tests/integration/test_collections_branches.py b/lamb-kb-server/tests/integration/test_collections_branches.py new file mode 100644 index 000000000..5b9e9f597 --- /dev/null +++ b/lamb-kb-server/tests/integration/test_collections_branches.py @@ -0,0 +1,188 @@ +"""Branch tests for ``services.collection_service`` via the collections router. + +Covers the graph/extraction-validation branch, the plugin-params passthrough +logging, the update-time chunking_params validation, and deletion of a +graph-enabled collection. +""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest +from httpx import AsyncClient + +from tests._helpers import AUTH_HEADERS + + +def _payload(org_id: str, name: str | None = None, **over) -> dict: + base = { + "organization_id": org_id, + "name": name or f"kb-{uuid4().hex[:6]}", + "description": "pytest", + "chunking_strategy": "simple", + "chunking_params": {"chunk_size": 400, "chunk_overlap": 50}, + "embedding": {"vendor": "fake", "model": "fake-model"}, + "vector_db_backend": "chromadb", + } + base.update(over) + return base + + +@pytest.mark.asyncio +async def test_create_graph_enabled_invalid_extraction_vendor( + client: AsyncClient, org_id: str +) -> None: + payload = _payload(org_id, graph_enabled=True, extraction={ + "vendor": "no-such-vendor", "model": "x", + }) + resp = await client.post("/collections", json=payload, headers=AUTH_HEADERS) + assert resp.status_code == 400 + assert "extraction vendor" in resp.text.lower() + + +@pytest.mark.asyncio +async def test_create_with_plugin_params_passthrough( + client: AsyncClient, org_id: str +) -> None: + # Non-empty embedding_params / vector_db_params exercise the + # "received but not yet wired" logging branches. + payload = _payload( + org_id, + embedding_params={"extra_knob": "v"}, + vector_db_params={"hnsw_ef": 64}, + ) + resp = await client.post("/collections", json=payload, headers=AUTH_HEADERS) + assert resp.status_code == 201 + + +@pytest.mark.asyncio +async def test_update_chunking_params_valid_and_invalid( + client: AsyncClient, org_id: str +) -> None: + created = await client.post( + "/collections", json=_payload(org_id), headers=AUTH_HEADERS + ) + cid = created.json()["id"] + + # Valid chunking_params update applies. + ok = await client.put( + f"/collections/{cid}", + json={"chunking_params": {"chunk_size": 800, "chunk_overlap": 80}}, + headers=AUTH_HEADERS, + ) + assert ok.status_code == 200 + assert ok.json()["chunking_params"]["chunk_size"] == 800 + + # Invalid chunking_params (typo'd key) is rejected at update time. + bad = await client.put( + f"/collections/{cid}", + json={"chunking_params": {"chunk_size": 800, "bogus_key": 1}}, + headers=AUTH_HEADERS, + ) + assert bad.status_code == 422 + + +@pytest.mark.asyncio +async def test_delete_graph_enabled_collection( + client: AsyncClient, org_id: str +) -> None: + # Creating + deleting a graph_enabled collection exercises the graph + # teardown branch in delete (KG-RAG disabled in tests -> Neo4j call is + # skipped, but the branch + config guard run). + created = await client.post( + "/collections", json=_payload(org_id, graph_enabled=True), headers=AUTH_HEADERS + ) + cid = created.json()["id"] + resp = await client.delete(f"/collections/{cid}", headers=AUTH_HEADERS) + assert resp.status_code in (200, 204) + + +def test_delete_graph_enabled_calls_neo4j_teardown(monkeypatch): + """Direct service call with KG-RAG enabled + an available fake graph store + exercises the Neo4j subgraph-delete branch on collection delete.""" + import config as config_module + import services.collection_service as cs + import services.graph_store as gs_module + from database.connection import get_session_direct + from database.models import Collection + + cid = f"col-graphdel-{uuid4().hex[:8]}" + session = get_session_direct() + try: + session.add(Collection( + id=cid, organization_id="org-z", name=f"graphdel-{uuid4().hex[:6]}", + chunking_strategy="simple", embedding_vendor="fake", + embedding_model="fake-model", vector_db_backend="chromadb", + storage_path="/tmp/does-not-matter", graph_enabled=True, + )) + session.commit() + finally: + session.close() + + monkeypatch.setattr(config_module, "get_kg_rag_config", lambda: {"enabled": True}) + + deleted = {"called": False} + + class _FakeStore: + def is_configured(self): + return True + + def is_available(self): + return True + + def delete_collection(self, collection_id): + deleted["called"] = collection_id == cid + + monkeypatch.setattr(gs_module, "get_graph_store", lambda: _FakeStore()) + + session2 = get_session_direct() + try: + cs.delete_collection(session2, cid) + finally: + session2.close() + assert deleted["called"] is True + + +def test_delete_graph_teardown_error_is_swallowed(monkeypatch): + """If the Neo4j teardown raises, delete_collection swallows it and still + deletes the DB row (covers the graph except-handler branch).""" + import config as config_module + import services.collection_service as cs + import services.graph_store as gs_module + from database.connection import get_session_direct + from database.models import Collection + + cid = f"col-graphdelerr-{uuid4().hex[:8]}" + session = get_session_direct() + try: + session.add(Collection( + id=cid, organization_id="org-z", name=f"gderr-{uuid4().hex[:6]}", + chunking_strategy="simple", embedding_vendor="fake", + embedding_model="fake-model", vector_db_backend="chromadb", + storage_path="/tmp/does-not-matter", graph_enabled=True, + )) + session.commit() + finally: + session.close() + + monkeypatch.setattr(config_module, "get_kg_rag_config", lambda: {"enabled": True}) + + class _BoomStore: + def is_configured(self): + return True + + def is_available(self): + return True + + def delete_collection(self, collection_id): + raise RuntimeError("neo4j teardown boom") + + monkeypatch.setattr(gs_module, "get_graph_store", lambda: _BoomStore()) + + session2 = get_session_direct() + try: + cs.delete_collection(session2, cid) # must not raise + assert session2.get(Collection, cid) is None # row still deleted + finally: + session2.close() diff --git a/lamb-kb-server/tests/integration/test_ingestion_service_branches.py b/lamb-kb-server/tests/integration/test_ingestion_service_branches.py new file mode 100644 index 000000000..2e9c7b766 --- /dev/null +++ b/lamb-kb-server/tests/integration/test_ingestion_service_branches.py @@ -0,0 +1,276 @@ +"""Branch tests for ``services.ingestion_service`` worker logic. + +Drives ``execute_ingestion_job`` directly (invalid params, cooperative +cancellation, oversized-chunk re-split) and unit-tests the graph +index/delete helpers with fakes — the branches the happy-path pipeline +tests don't reach. +""" + +from __future__ import annotations + +import json +import tempfile +from types import SimpleNamespace +from uuid import uuid4 + +import pytest + +import services.ingestion_service as ing +from services.ingestion_service import ( + JobCancelledError, + _maybe_delete_graph_document, + _maybe_index_graph, + execute_ingestion_job, +) + + +def _mk_collection(session, **over): + from database.models import Collection + + cid = over.pop("id", f"col-ing-{uuid4().hex[:8]}") + fields = dict( + id=cid, + organization_id="org-ing", + name=f"ing-{uuid4().hex[:6]}", + chunking_strategy="simple", + chunking_params=json.dumps({"chunk_size": 200, "chunk_overlap": 20}), + embedding_vendor="fake", + embedding_model="fake-model", + vector_db_backend="chromadb", + backend_collection_id=cid, + storage_path=tempfile.mkdtemp(prefix="ing-"), + ) + fields.update(over) # caller overrides win, no kwarg collision + col = Collection(**fields) + session.add(col) + session.commit() + session.refresh(col) + return col + + +def _mk_job(session, collection_id, docs, status="processing"): + from database.models import IngestionJob + + jid = f"job-ing-{uuid4().hex[:8]}" + job = IngestionJob( + id=jid, + collection_id=collection_id, + organization_id="org-ing", + documents_json=json.dumps(docs), + documents_total=len(docs), + status=status, + ) + session.add(job) + session.commit() + session.refresh(job) + return job + + +_DOC = {"source_item_id": "s1", "title": "T", "text": "LAMB ingestion text body.", + "permalinks": {}, "pages": [], "extra_metadata": {}} + + +# --------------------------------------------------------------------------- +# execute_ingestion_job +# --------------------------------------------------------------------------- + + +def test_execute_invalid_chunking_params_raises(): + from database.connection import get_session_direct + + session = get_session_direct() + try: + col = _mk_collection( + session, chunking_params=json.dumps({"bogus_param": 1}) + ) + job = _mk_job(session, col.id, [_DOC]) + with pytest.raises(RuntimeError): + execute_ingestion_job(session, job, col, {"api_key": "k"}) + finally: + session.close() + + +def test_execute_cancelled_job_raises(): + from database.connection import get_session_direct + + session = get_session_direct() + try: + col = _mk_collection(session) + job = _mk_job(session, col.id, [_DOC], status="cancelled") + with pytest.raises(JobCancelledError): + execute_ingestion_job(session, job, col, {"api_key": "k"}) + finally: + session.close() + + +def test_execute_oversized_chunks_resplit(monkeypatch): + from database.connection import get_session_direct + from plugins.vector_db.chromadb_backend import ChromaDBBackend + from tests._fakes import FakeEmbedding + + # MAX between the two chunk sizes so one chunk is re-split (else branch) + # and one is kept as-is (if branch) — covers both sides of the guard loop. + monkeypatch.setattr(ing, "_MAX_EMBED_CHARS", 50) + monkeypatch.setattr(ing, "_RESPLIT_CHUNK_SIZE", 40) + monkeypatch.setattr(ing, "_RESPLIT_OVERLAP", 5) + + storage = tempfile.mkdtemp(prefix="ing-resplit-") + session = get_session_direct() + try: + # chunk_size 100 over "a*100 b*20" -> a 100-char chunk (>50, re-split) + # and a 20-char chunk (<=50, kept). + col = _mk_collection( + session, storage_path=storage, + chunking_params=json.dumps({"chunk_size": 100, "chunk_overlap": 0}), + ) + # The backend collection must exist for add_chunks. + ChromaDBBackend().create_collection( + collection_id=col.backend_collection_id, storage_path=storage, + embedding_function=FakeEmbedding(), + ) + mixed_text = "a" * 100 + " " + "b" * 20 + job = _mk_job(session, col.id, [dict(_DOC, text=mixed_text)]) + execute_ingestion_job(session, job, col, {"api_key": "k"}) + session.refresh(job) + assert job.documents_processed == 1 + assert job.chunks_created >= 1 + finally: + session.close() + + +# --------------------------------------------------------------------------- +# _maybe_index_graph +# --------------------------------------------------------------------------- + + +class _FakeItem: + def __init__(self, text, metadata): + self.text = text + self.metadata = metadata + + +class _FakeBackend: + def __init__(self, *, rows=None, raises=False): + self._rows = rows + self._raises = raises + + def get_chunks_by_source(self, **kw): + if self._raises: + raise RuntimeError("backend boom") + return self._rows or [] + + +def _coll(**over): + base = dict( + graph_enabled=True, vector_db_backend="chromadb", + backend_collection_id="bc", id="c1", storage_path="/tmp/x", + organization_id="org", + ) + base.update(over) + return SimpleNamespace(**base) + + +def test_maybe_index_graph_not_enabled(): + # graph_enabled False -> immediate return (no config read). + _maybe_index_graph( + collection=_coll(graph_enabled=False), docs_list=[_DOC], + backend=_FakeBackend(), embedding_function=None, + ) # must not raise + + +def test_maybe_index_graph_kg_disabled(monkeypatch): + import config as config_module + monkeypatch.setattr(config_module, "get_kg_rag_config", lambda: {"enabled": False}) + _maybe_index_graph( + collection=_coll(), docs_list=[_DOC], + backend=_FakeBackend(), embedding_function=None, + ) + + +def test_maybe_index_graph_backend_error_and_empty(monkeypatch): + import config as config_module + monkeypatch.setattr( + config_module, "get_kg_rag_config", + lambda: {"enabled": True, "index_on_ingest": True}, + ) + # Backend raises for the source -> warning + continue (no crash). + _maybe_index_graph( + collection=_coll(), docs_list=[_DOC], + backend=_FakeBackend(raises=True), embedding_function=None, + ) + # Backend returns no rows -> skip branch. + _maybe_index_graph( + collection=_coll(), docs_list=[_DOC], + backend=_FakeBackend(rows=[]), embedding_function=None, + ) + + +def test_maybe_index_graph_indexes_and_logs_error(monkeypatch): + import config as config_module + import services.graph_indexing as gi + + monkeypatch.setattr( + config_module, "get_kg_rag_config", + lambda: {"enabled": True, "index_on_ingest": True}, + ) + calls = {} + + def fake_index(**kw): + calls.update(kw) + return {"error": "neo4j_unavailable"} # exercises the error-log branch + + monkeypatch.setattr(gi, "index_chunks_for_collection", fake_index) + rows = [_FakeItem("chunk text", {"chunk_id": "ck1"})] + _maybe_index_graph( + collection=_coll(), docs_list=[_DOC], + backend=_FakeBackend(rows=rows), embedding_function=None, + openai_api_key="sk", + ) + assert calls["ids"] == ["ck1"] + assert calls["filename"] == "s1" + + +# --------------------------------------------------------------------------- +# _maybe_delete_graph_document +# --------------------------------------------------------------------------- + + +def test_maybe_delete_graph_document_not_enabled(): + _maybe_delete_graph_document(_coll(graph_enabled=False), "s1") + + +def test_maybe_delete_graph_document_kg_disabled(monkeypatch): + import config as config_module + monkeypatch.setattr(config_module, "get_kg_rag_config", lambda: {"enabled": False}) + _maybe_delete_graph_document(_coll(), "s1") + + +def test_maybe_delete_graph_document_calls_store(monkeypatch): + import config as config_module + import services.graph_store as gs_module + + monkeypatch.setattr(config_module, "get_kg_rag_config", lambda: {"enabled": True}) + deleted = {} + + class _Store: + def delete_document(self, cid, org, sid): + deleted["args"] = (cid, org, sid) + + monkeypatch.setattr(gs_module, "get_graph_store", lambda: _Store()) + _maybe_delete_graph_document(_coll(id="c9", organization_id="org9"), "s9") + assert deleted["args"] == ("c9", "org9", "s9") + + +def test_maybe_delete_graph_document_store_error(monkeypatch): + import config as config_module + import services.graph_store as gs_module + + monkeypatch.setattr(config_module, "get_kg_rag_config", lambda: {"enabled": True}) + + class _Store: + def delete_document(self, *a): + raise RuntimeError("neo4j down") + + monkeypatch.setattr(gs_module, "get_graph_store", lambda: _Store()) + # Exception is swallowed (logged) — must not raise. + _maybe_delete_graph_document(_coll(), "s1") diff --git a/lamb-kb-server/tests/integration/test_jobs_cancel.py b/lamb-kb-server/tests/integration/test_jobs_cancel.py new file mode 100644 index 000000000..b48d7d1b1 --- /dev/null +++ b/lamb-kb-server/tests/integration/test_jobs_cancel.py @@ -0,0 +1,63 @@ +"""Integration tests for ``POST /jobs/{id}/cancel``. + +Seeds an ``IngestionJob`` row directly so the cancel transitions can be +exercised deterministically (without racing a live worker). +""" + +from __future__ import annotations + +import pytest +from httpx import AsyncClient + +from tests._helpers import AUTH_HEADERS + + +def _seed_job(status: str) -> str: + from database.connection import get_session_direct + from database.models import IngestionJob + + session = get_session_direct() + try: + job_id = f"job-cancel-{status}" + # Remove any prior row from a previous run. + existing = session.query(IngestionJob).filter(IngestionJob.id == job_id).first() + if existing: + session.delete(existing) + session.commit() + session.add(IngestionJob( + id=job_id, + collection_id="col-x", + organization_id="org-x", + documents_json="[]", + documents_total=1, + status=status, + )) + session.commit() + return job_id + finally: + session.close() + + +@pytest.mark.asyncio +async def test_cancel_nonexistent_job_404(client_no_worker: AsyncClient) -> None: + resp = await client_no_worker.post( + "/jobs/does-not-exist/cancel", headers=AUTH_HEADERS + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_cancel_pending_job(client_no_worker: AsyncClient) -> None: + job_id = _seed_job("pending") + resp = await client_no_worker.post(f"/jobs/{job_id}/cancel", headers=AUTH_HEADERS) + assert resp.status_code == 200 + assert resp.json()["status"] == "cancelled" + + +@pytest.mark.asyncio +async def test_cancel_terminal_job_is_noop(client_no_worker: AsyncClient) -> None: + job_id = _seed_job("completed") + resp = await client_no_worker.post(f"/jobs/{job_id}/cancel", headers=AUTH_HEADERS) + assert resp.status_code == 200 + # Terminal job stays unchanged (no transition to cancelled). + assert resp.json()["status"] == "completed" diff --git a/lamb-kb-server/tests/integration/test_query_system_branches.py b/lamb-kb-server/tests/integration/test_query_system_branches.py new file mode 100644 index 000000000..72d0958cf --- /dev/null +++ b/lamb-kb-server/tests/integration/test_query_system_branches.py @@ -0,0 +1,65 @@ +"""Branch tests for the query + system routers. + +* ``query_collection`` surfacing ``question_entities`` at the top level when + a KG-RAG trace is attached (router line, no live graph needed). +* ``GET /llm-vendors`` returning the registered extraction vendors. +""" + +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from httpx import AsyncClient + +from database.connection import get_session +from dependencies import verify_token +from plugins.base import QueryResult +from tests._helpers import AUTH_HEADERS + + +def test_query_surfaces_question_entities(monkeypatch): + import routers.query as query_router + from services import query_service + + # Result carries a KG-RAG trace -> the router lifts question_entities up. + def fake_query(db, collection_id, body): + return [QueryResult( + text="hit", score=0.9, + metadata={"kg_rag": {"question_entities": ["alpha", "beta"]}}, + )] + + monkeypatch.setattr(query_service, "query_collection", fake_query) + + app = FastAPI() + app.include_router(query_router.router) + app.dependency_overrides[verify_token] = lambda: "tok" + app.dependency_overrides[get_session] = lambda: object() + tc = TestClient(app) + + resp = tc.post("/collections/c1/query", json={"query_text": "what is alpha?"}) + assert resp.status_code == 200 + assert resp.json()["entities"] == ["alpha", "beta"] + + +def test_query_no_trace_entities_none(monkeypatch): + import routers.query as query_router + from services import query_service + + monkeypatch.setattr( + query_service, "query_collection", + lambda db, cid, body: [QueryResult(text="hit", score=0.9, metadata={})], + ) + app = FastAPI() + app.include_router(query_router.router) + app.dependency_overrides[verify_token] = lambda: "tok" + app.dependency_overrides[get_session] = lambda: object() + resp = TestClient(app).post("/collections/c1/query", json={"query_text": "q"}) + assert resp.json()["entities"] is None + + +@pytest.mark.asyncio +async def test_list_llm_vendors_endpoint(client: AsyncClient): + resp = await client.get("/llm-vendors", headers=AUTH_HEADERS) + assert resp.status_code == 200 + assert "vendors" in resp.json() diff --git a/lamb-kb-server/tests/integration/test_retry.py b/lamb-kb-server/tests/integration/test_retry.py new file mode 100644 index 000000000..7fc53f50c --- /dev/null +++ b/lamb-kb-server/tests/integration/test_retry.py @@ -0,0 +1,158 @@ +"""Tests for failed-job retry with bounded in-memory credential retention. + +Covers the worker-level credential cache (retain across attempts, discard on +terminal/exhausted states, TTL purge) and the HTTP retry endpoints. +""" + +from datetime import UTC, datetime, timedelta +from uuid import uuid4 + +import pytest +from config import MAX_JOB_ATTEMPTS +from database.connection import get_session_direct +from database.models import IngestionJob +from httpx import AsyncClient +from tasks import worker + +from tests.conftest import AUTH_HEADERS + + +def _insert_job(status: str, attempts: int) -> str: + job_id = uuid4().hex + db = get_session_direct() + try: + db.add( + IngestionJob( + id=job_id, + collection_id="some-collection", + organization_id=f"org-{uuid4().hex[:8]}", + documents_json="[]", + status=status, + documents_total=0, + documents_processed=0, + chunks_created=0, + attempts=attempts, + ) + ) + db.commit() + finally: + db.close() + return job_id + + +# --------------------------------------------------------------------------- +# Worker-level credential cache +# --------------------------------------------------------------------------- + + +def test_store_and_retry_available() -> None: + job_id = uuid4().hex + worker.store_credentials(job_id, {"api_key": "k", "api_endpoint": ""}) + assert worker.retry_available(job_id) is True + worker._discard_credentials(job_id) + assert worker.retry_available(job_id) is False + + +def test_purge_expired_credentials_drops_old_entries() -> None: + job_id = uuid4().hex + worker.store_credentials(job_id, {"api_key": "k"}) + # Backdate the stored timestamp beyond the retention window. + worker._job_cred_stored_at[job_id] = datetime.now(UTC) - timedelta( + minutes=worker.RETRY_CACHE_TTL_MINUTES + 1 + ) + worker._purge_expired_credentials() + assert worker.retry_available(job_id) is False + + +def test_purge_keeps_fresh_entries() -> None: + job_id = uuid4().hex + worker.store_credentials(job_id, {"api_key": "k"}) + worker._purge_expired_credentials() + assert worker.retry_available(job_id) is True + worker._discard_credentials(job_id) + + +# --------------------------------------------------------------------------- +# HTTP retry endpoints +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_retry_failed_job_with_creds_requeues(client: AsyncClient) -> None: + job_id = _insert_job("failed", attempts=1) + worker.store_credentials(job_id, {"api_key": "k", "api_endpoint": ""}) + try: + resp = await client.post(f"/jobs/{job_id}/retry", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["status"] == "pending" + assert body["error_message"] is None + finally: + worker._discard_credentials(job_id) + + +@pytest.mark.asyncio +async def test_retry_without_creds_returns_410(client: AsyncClient) -> None: + job_id = _insert_job("failed", attempts=1) # no store_credentials + resp = await client.post(f"/jobs/{job_id}/retry", headers=AUTH_HEADERS) + assert resp.status_code == 410 + assert "re-submit" in resp.json()["detail"].lower() + + +@pytest.mark.asyncio +async def test_retry_non_failed_job_returns_409(client: AsyncClient) -> None: + job_id = _insert_job("completed", attempts=1) + worker.store_credentials(job_id, {"api_key": "k"}) + try: + resp = await client.post(f"/jobs/{job_id}/retry", headers=AUTH_HEADERS) + assert resp.status_code == 409 + finally: + worker._discard_credentials(job_id) + + +@pytest.mark.asyncio +async def test_retry_exhausted_attempts_returns_409(client: AsyncClient) -> None: + job_id = _insert_job("failed", attempts=MAX_JOB_ATTEMPTS) + worker.store_credentials(job_id, {"api_key": "k"}) + try: + resp = await client.post(f"/jobs/{job_id}/retry", headers=AUTH_HEADERS) + assert resp.status_code == 409 + assert "exhausted" in resp.json()["detail"].lower() + finally: + worker._discard_credentials(job_id) + + +@pytest.mark.asyncio +async def test_retry_unknown_job_returns_404(client: AsyncClient) -> None: + resp = await client.post( + "/jobs/00000000-0000-0000-0000-000000000000/retry", headers=AUTH_HEADERS + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_retry_available_endpoint(client: AsyncClient) -> None: + job_id = _insert_job("failed", attempts=1) + # Without creds → not available. + resp = await client.get(f"/jobs/{job_id}/retry-available", headers=AUTH_HEADERS) + assert resp.status_code == 200 + assert resp.json() == { + "available": False, + "attempts": 1, + "max_attempts": MAX_JOB_ATTEMPTS, + } + # With creds → available. + worker.store_credentials(job_id, {"api_key": "k"}) + try: + resp = await client.get( + f"/jobs/{job_id}/retry-available", headers=AUTH_HEADERS + ) + assert resp.json()["available"] is True + finally: + worker._discard_credentials(job_id) + + +@pytest.mark.asyncio +async def test_retry_requires_auth(client: AsyncClient) -> None: + resp = await client.post("/jobs/whatever/retry") + assert resp.status_code in (401, 403) diff --git a/lamb-kb-server/tests/integration/test_worker_cancellation.py b/lamb-kb-server/tests/integration/test_worker_cancellation.py new file mode 100644 index 000000000..ddb5a7200 --- /dev/null +++ b/lamb-kb-server/tests/integration/test_worker_cancellation.py @@ -0,0 +1,89 @@ +"""Worker cancellation-branch tests for ``tasks.worker._process_job_sync``. + +Covers the two cancellation exits: cancelled-before-pickup and the +cooperative ``JobCancelledError`` raised mid-ingestion. +""" + +from __future__ import annotations + +import json +from uuid import uuid4 + +import pytest + +import services.ingestion_service as ing +from database.connection import get_session_direct +from database.models import Collection, IngestionJob +from tasks.worker import _process_job_sync + + +def _seed_collection(session) -> str: + cid = f"col-wk-{uuid4().hex[:8]}" + session.add(Collection( + id=cid, organization_id="org-wk", name=f"wk-{uuid4().hex[:6]}", + chunking_strategy="simple", + chunking_params=json.dumps({"chunk_size": 200, "chunk_overlap": 20}), + embedding_vendor="fake", embedding_model="fake-model", + vector_db_backend="chromadb", backend_collection_id=cid, + storage_path="/tmp/wk-does-not-matter", + )) + session.commit() + return cid + + +def _seed_job(session, collection_id, status) -> str: + jid = f"job-wk-{uuid4().hex[:8]}" + session.add(IngestionJob( + id=jid, collection_id=collection_id, organization_id="org-wk", + documents_json=json.dumps([{"source_item_id": "s1", "title": "T", + "text": "body", "permalinks": {}, "pages": [], + "extra_metadata": {}}]), + documents_total=1, status=status, + )) + session.commit() + return jid + + +def test_process_job_cancelled_before_pickup(): + session = get_session_direct() + try: + cid = _seed_collection(session) + jid = _seed_job(session, cid, status="cancelled") + finally: + session.close() + + # Should return early without flipping back to processing. + _process_job_sync(jid) + + check = get_session_direct() + try: + job = check.query(IngestionJob).filter(IngestionJob.id == jid).first() + assert job.status == "cancelled" + finally: + check.close() + + +def test_process_job_cooperative_cancellation(monkeypatch): + session = get_session_direct() + try: + cid = _seed_collection(session) + jid = _seed_job(session, cid, status="pending") + finally: + session.close() + + # Force the ingestion to raise the cooperative-cancel signal once running. + def _raise_cancelled(db, job, collection, credentials): + raise ing.JobCancelledError("cancelled mid-flight") + + monkeypatch.setattr(ing, "execute_ingestion_job", _raise_cancelled) + + _process_job_sync(jid) + + # The worker leaves the row alone on cooperative cancel (does not mark + # completed/failed). Status stays whatever the canceller left it. + check = get_session_direct() + try: + job = check.query(IngestionJob).filter(IngestionJob.id == jid).first() + assert job.status != "completed" + finally: + check.close() diff --git a/lamb-kb-server/tests/unit/_neo4j_fakes.py b/lamb-kb-server/tests/unit/_neo4j_fakes.py new file mode 100644 index 000000000..4ead3d98a --- /dev/null +++ b/lamb-kb-server/tests/unit/_neo4j_fakes.py @@ -0,0 +1,142 @@ +"""In-memory fakes standing in for the neo4j driver/session/result. + +``graph_store.GraphStore`` only ever touches Neo4j through three shapes: + +* ``driver.session()`` used as a context manager, +* ``session.run(query, **params)`` returning a result that exposes + ``.data()`` / ``.single()`` / iteration, +* ``session.execute_write(fn, *args, **kwargs)`` running a unit-of-work + callback with a transaction object that itself exposes ``.run(...)``. + +These fakes reproduce exactly that surface so the store can be exercised +without a live Neo4j. Responses are programmed as an ordered queue; each +``run`` call pops the next one. A response may be a :class:`FakeResult`, a +list of row dicts (wrapped automatically), or a callable +``(query, params) -> FakeResult`` for query-dependent behaviour. +""" + +from __future__ import annotations + +from typing import Any, Callable, Dict, List, Optional, Union + +Response = Union["FakeResult", List[Dict[str, Any]], Callable[[str, dict], Any], None] + + +class FakeResult: + def __init__(self, rows: Optional[List[Dict[str, Any]]] = None, single=...): + self._rows = rows or [] + # `single` sentinel: default to first row (or None) when not provided. + if single is ...: + self._single = self._rows[0] if self._rows else None + else: + self._single = single + + def data(self) -> List[Dict[str, Any]]: + return list(self._rows) + + def single(self): + return self._single + + def __iter__(self): + return iter(self._rows) + + +def _coerce(resp: Response, query: str, params: dict) -> FakeResult: + if callable(resp) and not isinstance(resp, FakeResult): + resp = resp(query, params) + if resp is None: + return FakeResult() + if isinstance(resp, FakeResult): + return resp + if isinstance(resp, list): + return FakeResult(rows=resp) + raise TypeError(f"Unsupported fake response: {resp!r}") + + +class FakeSession: + """Context-manager session that records runs and replays a response queue. + + Also serves as its own transaction object for ``execute_write`` — the + unit-of-work callback receives this session and calls ``.run`` on it, + drawing from the same queue. + """ + + def __init__(self, responses: Optional[List[Response]] = None): + self.responses: List[Response] = list(responses or []) + self.runs: List[Dict[str, Any]] = [] + + # context-manager protocol (``with driver.session() as session:``) + def __enter__(self) -> "FakeSession": + return self + + def __exit__(self, *exc) -> bool: + return False + + def run(self, query: str, **params) -> FakeResult: + self.runs.append({"query": query, "params": params}) + resp = self.responses.pop(0) if self.responses else None + return _coerce(resp, query, params) + + def execute_write(self, fn: Callable, *args, **kwargs): + return fn(self, *args, **kwargs) + + def execute_read(self, fn: Callable, *args, **kwargs): + return fn(self, *args, **kwargs) + + +class FakeDriver: + def __init__( + self, + session: Optional[FakeSession] = None, + *, + connectivity_error: Optional[Exception] = None, + ): + self._session = session or FakeSession() + self.connectivity_error = connectivity_error + self.closed = False + + def session(self) -> FakeSession: + return self._session + + def verify_connectivity(self) -> None: + if self.connectivity_error is not None: + raise self.connectivity_error + + def close(self) -> None: + self.closed = True + + +def make_store( + *, + responses: Optional[List[Response]] = None, + session: Optional[FakeSession] = None, + driver: Optional[FakeDriver] = None, + schema_ready: bool = True, + config: Optional[Dict[str, Any]] = None, +): + """Build a ``GraphStore`` wired to a fake driver, bypassing real Neo4j. + + Returns ``(store, session)`` so tests can assert on ``session.runs``. + """ + from services.graph_store import GraphStore + + sess = session or FakeSession(responses) + drv = driver or FakeDriver(sess) + cfg = { + "enabled": False, # keep __init__ from creating a real driver + "neo4j_uri": "", + "neo4j_user": "neo4j", + "neo4j_password": "", + } + if config: + cfg.update(config) + store = GraphStore(kg_config=cfg) + store.driver = drv + store.enabled = True + # Populate connection fields so ``is_configured()`` is True without ever + # having created a real driver in __init__. + store.uri = "bolt://fake:7687" + store.user = "neo4j" + store.password = "fake-password" + store._schema_ready = schema_ready + return store, sess diff --git a/lamb-kb-server/tests/unit/conftest.py b/lamb-kb-server/tests/unit/conftest.py index e98051784..41f01a1bc 100644 --- a/lamb-kb-server/tests/unit/conftest.py +++ b/lamb-kb-server/tests/unit/conftest.py @@ -2,6 +2,7 @@ from __future__ import annotations +import importlib import shutil import tempfile from collections.abc import Iterator @@ -11,6 +12,21 @@ from tests._fakes import FakeEmbedding +@pytest.fixture() +def reload_config() -> Iterator[None]: + """Reload the ``config`` module after env-var mutations. + + Lives at the tier root so any unit test that needs env-driven config + re-evaluation can use it without re-defining the fixture locally. + Mirror of the same-named fixture in tests/unit/test_config.py — keep + them in sync. + """ + import config # noqa: PLC0415 + + yield + importlib.reload(config) + + @pytest.fixture def tmp_storage() -> Iterator[str]: """Per-test temp dir for vector DB persistence.""" diff --git a/lamb-kb-server/tests/unit/test_benchmark_service.py b/lamb-kb-server/tests/unit/test_benchmark_service.py new file mode 100644 index 000000000..04386e1c4 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_benchmark_service.py @@ -0,0 +1,246 @@ +"""Unit tests for ``services.benchmark.BenchmarkService``. + +Covers dataset listing/lookup/aliasing, the per-question scoring math +(precision/recall/MRR, retrieved-file extraction and normalization), the +trace extraction helper, and the full ``run`` / ``run_all`` orchestration +with a fake ``query_with_plugin``. +""" + +from __future__ import annotations + +import pytest +from fastapi import HTTPException + +import services.query_service as query_service_module +from schemas.benchmark import ( + BenchmarkQuestion, + BenchmarkQueryScore, + BenchmarkRunRequest, +) +from services.benchmark import _DATASETS, BenchmarkService + + +# --------------------------------------------------------------------------- +# datasets +# --------------------------------------------------------------------------- + + +def test_list_datasets_covers_all_registered(): + summaries = BenchmarkService.list_datasets() + assert {s.id for s in summaries} == set(_DATASETS) + assert all(s.question_count > 0 for s in summaries) + + +def test_get_dataset_returns_questions(): + ds = BenchmarkService.get_dataset("educational") + assert ds.id == "educational" + assert len(ds.questions) == ds.question_count + + +@pytest.mark.parametrize( + "alias,expected", + [ + ("default", "educational"), + ("SAMPLE", "educational"), + ("multi_hop", "paper"), + ("no_connections", "control"), + ("educational", "educational"), + ], +) +def test_canonical_dataset_id_aliases(alias, expected): + assert BenchmarkService._canonical_dataset_id(alias) == expected + + +def test_canonical_dataset_id_unknown_raises(): + with pytest.raises(HTTPException) as exc: + BenchmarkService._canonical_dataset_id("does-not-exist") + assert exc.value.status_code == 400 + + +def test_canonical_dataset_id_default_when_empty(): + assert BenchmarkService._canonical_dataset_id("") == "educational" + + +# --------------------------------------------------------------------------- +# filename + trace helpers +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "value,expected", + [ + (None, ""), + ("", ""), + (" ", ""), + ("/docs/Sub/Report.MD", "report.md"), + ("C:\\win\\Doc.PDF", "doc.pdf"), + ], +) +def test_normalize_filename(value, expected): + assert BenchmarkService._normalize_filename(value) == expected + + +def test_result_filename_key_precedence(): + assert BenchmarkService._result_filename({"filename": "a.md"}) == "a.md" + assert BenchmarkService._result_filename({"source": "b.md"}) == "b.md" + assert BenchmarkService._result_filename({"file_url": "/x/c.md"}) == "c.md" + assert BenchmarkService._result_filename({"unrelated": "x"}) == "" + + +def test_retrieved_files_dedups_and_caps(): + results = [ + {"metadata": {"filename": "a.md"}}, + {"metadata": {"filename": "a.md"}}, # dup + {"metadata": {}}, # no filename -> skipped + {"metadata": {"filename": "b.md"}}, + {"metadata": {"filename": "c.md"}}, + ] + out = BenchmarkService._retrieved_files(results, top_k=2) + assert out == ["a.md", "b.md"] # capped at top_k + + +def test_trace_from_results(): + assert BenchmarkService._trace_from_results( + [{"metadata": {"kg_rag": {"mode": "kg_rag"}}}] + ) == {"mode": "kg_rag"} + assert BenchmarkService._trace_from_results([{"metadata": {}}]) == {} + + +# --------------------------------------------------------------------------- +# scoring +# --------------------------------------------------------------------------- + + +def _question(relevant): + return BenchmarkQuestion( + id="q", question="q?", relevant_files=relevant, kind="single-hop" + ) + + +def test_score_response_no_relevant_files(): + score = BenchmarkService._score_response( + question=_question([]), + response={"results": [{"metadata": {"filename": "a.md"}}]}, + top_k=5, vector_ms=1.0, graph_ms=0.0, total_ms=1.0, + ) + # No relevant set -> precision/recall/mrr stay at defaults (0). + assert score.precision_at_k == 0.0 + assert score.retrieved_files == ["a.md"] + + +def test_score_response_with_hits(): + score = BenchmarkService._score_response( + question=_question(["a.md", "missing.md"]), + response={ + "results": [ + {"metadata": {"filename": "x.md"}}, # miss at rank 1 + {"metadata": {"filename": "a.md"}}, # hit at rank 2 + ] + }, + top_k=2, vector_ms=1.0, graph_ms=2.0, total_ms=3.0, + ) + assert score.precision_at_k == 0.5 # 1 hit / top_k(2) + assert score.recall_at_k == 0.5 # 1 of 2 relevant + assert score.mrr == 0.5 # first hit at rank 2 + + +def test_aggregate_empty_and_nonempty(): + assert BenchmarkService._aggregate([]).precision_at_k == 0.0 + rows = [ + BenchmarkQueryScore(precision_at_k=1.0, recall_at_k=1.0, mrr=1.0, + vector_ms=2.0, graph_ms=0.0, total_ms=2.0), + BenchmarkQueryScore(precision_at_k=0.0, recall_at_k=0.0, mrr=0.0, + vector_ms=4.0, graph_ms=0.0, total_ms=4.0), + ] + agg = BenchmarkService._aggregate(rows) + assert agg.precision_at_k == 0.5 + assert agg.avg_total_ms == 3.0 + + +# --------------------------------------------------------------------------- +# run / run_all orchestration +# --------------------------------------------------------------------------- + + +@pytest.fixture +def fake_query(monkeypatch): + """Stub query_with_plugin: baseline misses, kg_rag hits the relevant file.""" + + def _qwp(*, db, collection_id, query_text, plugin_name, plugin_params=None, + embedding_credentials=None): + if plugin_name == "kg_rag_query": + return { + "results": [ + {"metadata": {"filename": "rag_basics.md", + "kg_rag": {"graph_latency_ms": 5.0, + "vector_latency_ms": 3.0}}}, + ], + "timing": {"total_ms": 8.0}, + } + return { + "results": [{"metadata": {"filename": "unrelated.md"}}], + "timing": {"total_ms": 2.0}, + } + + monkeypatch.setattr(query_service_module, "query_with_plugin", _qwp) + return _qwp + + +def test_run_no_questions_raises(monkeypatch, fake_query): + # A custom dataset with zero questions -> 400. We force questions empty by + # passing an explicit empty list AND an unknown dataset is avoided by using + # a request whose dataset resolves but questions override to empty. + req = BenchmarkRunRequest(dataset_id="educational", questions=[]) + # questions=[] is falsy -> falls back to dataset.questions (non-empty), so + # to hit the guard we monkeypatch get_dataset to yield no questions. + monkeypatch.setattr( + BenchmarkService, "get_dataset", + classmethod(lambda cls, ds: type("D", (), { + "questions": [], "recommended_top_k": 5, "recommended_graph_depth": 2, + "id": "educational", "name": "n", "expected_behavior": "x", + })()), + ) + with pytest.raises(HTTPException) as exc: + BenchmarkService.run(db=None, collection_id="c1", request=req) + assert exc.value.status_code == 400 + + +def test_run_full_with_custom_question(fake_query): + req = BenchmarkRunRequest( + dataset_id="educational", + top_k=3, + graph_depth=2, + questions=[ + BenchmarkQuestion( + id="q1", question="what is rag?", + relevant_files=["rag_basics.md"], kind="single-hop", + ) + ], + rrf_k=15, # extra field forwarded to kg params + ) + resp = BenchmarkService.run(db=None, collection_id="c1", request=req) + assert resp.collection_id == "c1" + assert resp.dataset_id == "educational" + assert len(resp.results) == 1 + # KG hit the relevant file; baseline missed -> positive recall delta. + assert resp.kg_rag.recall_at_k == 1.0 + assert resp.baseline.recall_at_k == 0.0 + assert resp.comparison.delta_recall_at_k == 1.0 + + +def test_run_all_dedups_aliases(fake_query): + resp = BenchmarkService.run_all( + db=None, + collection_id="c1", + dataset_ids=["educational", "default", "sample"], # all alias to educational + ) + # All three collapse to one canonical run. + assert len(resp.runs) == 1 + assert resp.summary["datasets"] == 1 + assert "avg_delta_recall_at_k" in resp.summary + + +def test_run_all_empty_dataset_ids_uses_all(fake_query): + resp = BenchmarkService.run_all(db=None, collection_id="c1", dataset_ids=[]) + # Falls back to every registered dataset. + assert len(resp.runs) == len(_DATASETS) diff --git a/lamb-kb-server/tests/unit/test_branch_edge_cases.py b/lamb-kb-server/tests/unit/test_branch_edge_cases.py new file mode 100644 index 000000000..f99890195 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_branch_edge_cases.py @@ -0,0 +1,197 @@ +"""Edge-case coverage for small branch-added lines across the new KB server: +config int parsing, the auth TypeError guard, chunking-param numeric +validation, vector-DB base defaults, embedding-plugin error fallbacks, and +the plugin-discovery import-failure path. +""" + +from __future__ import annotations + +import sys +import types + +import httpx +import pytest +from fastapi import HTTPException + +from plugins.base import ( + Chunk, + EmbeddingFunction, + PluginParameter, + QueryResult, + VectorDBBackend, +) + + +# --------------------------------------------------------------------------- +# config._env_int +# --------------------------------------------------------------------------- + + +def test_env_int_invalid_returns_default(monkeypatch): + import config + + monkeypatch.setenv("KBTEST_INT", "not-an-int") + assert config._env_int("KBTEST_INT", 7, 1, 100) == 7 + + +def test_env_int_clamps(monkeypatch): + import config + + monkeypatch.setenv("KBTEST_INT", "999") + assert config._env_int("KBTEST_INT", 7, 1, 50) == 50 + + +# --------------------------------------------------------------------------- +# dependencies.verify_token TypeError guard +# --------------------------------------------------------------------------- + + +def test_verify_token_non_ascii_raises_401(): + import asyncio + + from dependencies import verify_token + + creds = types.SimpleNamespace(credentials="tök") # non-ASCII -> hmac TypeError + with pytest.raises(HTTPException) as exc: + asyncio.run(verify_token(creds)) + assert exc.value.status_code == 401 + + +# --------------------------------------------------------------------------- +# chunking _common.validate_chunking_params numeric branches +# --------------------------------------------------------------------------- + + +class _FakeStrategy: + name = "fake-strategy" + + def get_parameters(self): + return [PluginParameter(name="size", type="int", min_value=10, max_value=100)] + + +def test_validate_chunking_params_numeric_branches(): + from plugins.chunking._common import validate_chunking_params + + s = _FakeStrategy() + with pytest.raises(ValueError, match="must be numeric"): + validate_chunking_params(s, {"size": "x"}) + with pytest.raises(ValueError, match="below the declared"): + validate_chunking_params(s, {"size": 5}) + with pytest.raises(ValueError, match="above the declared"): + validate_chunking_params(s, {"size": 500}) + # In-range passes. + validate_chunking_params(s, {"size": 50}) + + +# --------------------------------------------------------------------------- +# VectorDBBackend / LLMExtractionFunction default no-op implementations +# --------------------------------------------------------------------------- + + +class _BareBackend(VectorDBBackend): + name = "bare" + + def create_collection(self, **kw): + return "id" + + def delete_collection(self, **kw): + return None + + def add_chunks(self, **kw): + return 0 + + def delete_by_source(self, **kw): + return 0 + + def query(self, **kw): + return [] + + +def test_vector_db_base_defaults_return_empty(): + be = _BareBackend() + assert be.get_chunks_by_id( + collection_id="c", storage_path="/x", chunk_ids=["a"], embedding_function=None + ) == [] + assert be.get_chunks_by_source( + collection_id="c", storage_path="/x", source_item_id="s", + embedding_function=None, + ) == [] + assert be.get_parameters() == [] + + +def test_llm_extraction_get_parameters_default(): + from plugins.llm_extraction.openai import OpenAIExtraction + + # The instance method is the base default (-> []); class_parameters is separate. + assert OpenAIExtraction(model="gpt-4o-mini").get_parameters() == [] + + +# --------------------------------------------------------------------------- +# embedding plugin error fallbacks +# --------------------------------------------------------------------------- + + +def test_ollama_embedding_attribute_error_fallback(monkeypatch): + fake = types.ModuleType("ollama") + + class FakeClient: + def __init__(self, **kw): + pass + + def embed(self, model, input): + return object() # no `.embeddings` attribute -> AttributeError + + def embeddings(self, model, prompt): + return {"embedding": [0.1, 0.2]} + + fake.Client = FakeClient + monkeypatch.setitem(sys.modules, "ollama", fake) + + from plugins.embedding.ollama import OllamaEmbedding + + out = OllamaEmbedding(api_endpoint="http://h:1234/api/embeddings")(["t1", "t2"]) + assert out == [[0.1, 0.2], [0.1, 0.2]] + + +def test_openai_embedding_connection_error_wrapped(monkeypatch): + import openai + + class FakeOpenAI: + def __init__(self, **kw): + self.embeddings = self + + def create(self, model, input): + raise openai.APIConnectionError(request=httpx.Request("POST", "http://x")) + + monkeypatch.setattr(openai, "OpenAI", FakeOpenAI) + + from plugins.embedding.openai import OpenAIEmbedding + + with pytest.raises(RuntimeError, match="cannot connect"): + OpenAIEmbedding(model="text-embedding-3-small", api_key="k")(["t"]) + + +# --------------------------------------------------------------------------- +# main._discover_plugins import-failure path +# --------------------------------------------------------------------------- + + +def test_discover_plugins_bad_package_is_swallowed(): + import main + + # A non-existent package -> import_module raises -> logged + return (no raise). + main._discover_plugins("nonexistent_pkg_xyz_123") + + +# --------------------------------------------------------------------------- +# database._run_lightweight_migrations non-SQLite early return +# --------------------------------------------------------------------------- + + +def test_lightweight_migrations_skips_non_sqlite(): + from database.connection import _run_lightweight_migrations + + # The helper only inspects ``str(engine.url)``; a non-SQLite URL hits the + # early return. A stub avoids needing a real Postgres driver installed. + stub_engine = types.SimpleNamespace(url="postgresql://user:pass@localhost:5432/db") + _run_lightweight_migrations(stub_engine) # must return without connecting diff --git a/lamb-kb-server/tests/unit/test_concept_extraction.py b/lamb-kb-server/tests/unit/test_concept_extraction.py new file mode 100644 index 000000000..4f44520f5 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_concept_extraction.py @@ -0,0 +1,373 @@ +"""Unit tests for ``services.concept_extraction``. + +Covers the pure validation/normalization helpers, the ``ConceptExtractor`` +backend-resolution logic, the sequential and threaded extraction paths, and +the JSON payload parsing (entity validation, relationship pruning, dropping +of entities that participate in no relationship). + +A ``FakeBackend`` (a real ``LLMExtractionFunction`` subclass) stands in for +the vendor SDK, so no network or API key is touched. +""" + +from __future__ import annotations + +import json +from typing import Any, Dict, List, Optional + +import pytest + +from plugins.base import LLMExtractionFunction +from services.concept_extraction import ( + ConceptExtractor, + ExtractedEntity, + GraphExtraction, + TextChunk, + _clean_text, + _confidence, + _is_valid_entity_name, + normalize_concept, + normalize_relation, +) + + +class FakeBackend(LLMExtractionFunction): + """Returns a payload chosen by matching a substring of the user text.""" + + name = "fake" + + def __init__(self, *, by_text: Optional[Dict[str, dict]] = None, default=None): + super().__init__(model="fake-model") + self._by_text = by_text or {} + self._default = default if default is not None else { + "entities": [], + "relationships": [], + } + self.calls: List[dict] = [] + + def chat_json(self, *, system: str, user: str, fallback_model=None): + self.calls.append( + {"system": system, "user": user, "fallback_model": fallback_model} + ) + decoded = json.loads(user) + text = decoded.get("text", "") + for needle, payload in self._by_text.items(): + if needle in text: + return payload + return self._default + + +# --------------------------------------------------------------------------- +# pure helpers +# --------------------------------------------------------------------------- + + +def test_clean_text_collapses_whitespace_and_truncates(): + assert _clean_text(" a\n\t b c ") == "a b c" + assert _clean_text(None) == "" + assert _clean_text("x" * 1000, limit=10) == "x" * 10 + + +@pytest.mark.parametrize( + "value,expected", + [ + ("0.5", 0.5), + (1.5, 1.0), # clamped up + (-2, 0.0), # clamped down + ("not-a-number", 1.0), # default + (None, 1.0), + ], +) +def test_confidence_clamps_and_defaults(value, expected): + assert _confidence(value) == expected + + +@pytest.mark.parametrize( + "name", + [ + "Machine Learning", + "GPT-4", + "Neo4j", + ], +) +def test_is_valid_entity_name_accepts(name): + assert _is_valid_entity_name(name) is True + + +@pytest.mark.parametrize( + "name", + [ + "a", # too short after normalization + "x" * 200, # too long + "!!!", # no alnum + "3.14 pies", # leading digit + "report.pdf", # file-extension suffix + "readme md", # ends with " md" + "a. b. c", # more than one period + "one two three four five six seven eight nine", # > 8 words + ], +) +def test_is_valid_entity_name_rejects(name): + assert _is_valid_entity_name(name) is False + + +def test_normalize_helpers_edge_cases(): + assert normalize_concept("Café Society!") == "cafe society" + assert normalize_relation("Depends On") == "depends_on" + assert normalize_relation("---") == "related_to" + assert len(normalize_relation("x" * 200)) == 64 + + +def test_graph_extraction_all_concepts(): + ext = GraphExtraction( + entities={"a": ExtractedEntity(name="a", display_name="A")} + ) + assert ext.all_concepts == {"a"} + + +# --------------------------------------------------------------------------- +# ConceptExtractor.__init__ backend resolution +# --------------------------------------------------------------------------- + + +def test_init_with_explicit_backend(): + be = FakeBackend() + ex = ConceptExtractor(kg_config={"extraction_max_workers": 4}, backend=be) + assert ex.backend is be + # workers clamped to [1, 16]. + assert ex.max_workers == 4 + + +def test_init_clamps_workers_to_16(): + ex = ConceptExtractor( + kg_config={"extraction_max_workers": 999}, backend=FakeBackend() + ) + assert ex.max_workers == 16 + + +def test_init_unregistered_vendor_sets_backend_none(caplog): + ex = ConceptExtractor( + kg_config={"extraction_max_workers": 1}, + vendor="totally-unknown-vendor", + ) + assert ex.backend is None + + +def test_init_builds_registered_vendor(): + # openai is a registered extraction vendor; building it must succeed even + # without a key (the key only matters at call time). + ex = ConceptExtractor( + kg_config={"extraction_max_workers": 1, "openai_api_key": ""}, + vendor="openai", + model="gpt-4o-mini", + ) + assert ex.backend is not None + assert ex.backend.name == "openai" + + +# --------------------------------------------------------------------------- +# extract_for_chunks +# --------------------------------------------------------------------------- + + +def _chunk(cid, text, parent=None, meta=None): + return TextChunk( + chunk_id=cid, + text=text, + parent_text=parent or "", + metadata=meta or {}, + ) + + +def test_extract_for_chunks_empty_returns_empty(): + ex = ConceptExtractor(backend=FakeBackend()) + out = ex.extract_for_chunks([]) + assert out.entities == {} + assert out.concepts_by_chunk == {} + + +def test_extract_for_chunks_backend_none_returns_skeleton(): + ex = ConceptExtractor(backend=FakeBackend()) + ex.backend = None + out = ex.extract_for_chunks([_chunk("c1", "hello")]) + assert out.concepts_by_chunk == {"c1": []} + assert out.entities == {} + + +def test_extract_for_chunks_sequential_single_group(): + payload = { + "entities": [ + {"name": "Machine Learning", "type": "concept", "confidence": 0.9}, + {"name": "Neural Network", "type": "concept"}, + ], + "relationships": [ + { + "source": "Machine Learning", + "target": "Neural Network", + "relation": "uses", + } + ], + } + be = FakeBackend(by_text={"deep learning": payload}) + ex = ConceptExtractor(kg_config={"extraction_max_workers": 1}, backend=be) + out = ex.extract_for_chunks([_chunk("c1", "deep learning text")]) + assert set(out.entities) == {"machine learning", "neural network"} + assert out.concepts_by_chunk["c1"] == ["machine learning", "neural network"] + assert len(out.relationships) == 1 + assert out.relationships[0].chunk_id == "c1" + + +def test_extract_for_chunks_threaded_multiple_groups(): + payload_a = { + "entities": [ + {"name": "Alpha One", "type": "concept"}, + {"name": "Alpha Two", "type": "concept"}, + ], + "relationships": [ + {"source": "Alpha One", "target": "Alpha Two", "relation": "links"} + ], + } + payload_b = { + "entities": [ + {"name": "Beta One", "type": "concept"}, + {"name": "Beta Two", "type": "concept"}, + ], + "relationships": [ + {"source": "Beta One", "target": "Beta Two", "relation": "links"} + ], + } + be = FakeBackend(by_text={"alpha-doc": payload_a, "beta-doc": payload_b}) + ex = ConceptExtractor(kg_config={"extraction_max_workers": 4}, backend=be) + chunks = [ + _chunk("a", "alpha-doc body", parent="alpha-doc body"), + _chunk("b", "beta-doc body", parent="beta-doc body"), + ] + out = ex.extract_for_chunks(chunks) + assert "alpha one" in out.entities and "beta one" in out.entities + assert out.concepts_by_chunk["a"] == ["alpha one", "alpha two"] + assert out.concepts_by_chunk["b"] == ["beta one", "beta two"] + # Two parent groups -> two backend calls. + assert len(be.calls) == 2 + + +def test_extract_groups_chunks_by_parent_text(): + payload = { + "entities": [ + {"name": "Shared Topic", "type": "concept"}, + {"name": "Other Topic", "type": "concept"}, + ], + "relationships": [ + {"source": "Shared Topic", "target": "Other Topic", "relation": "rel"} + ], + } + be = FakeBackend(default=payload) + ex = ConceptExtractor(kg_config={"extraction_max_workers": 1}, backend=be) + # Two chunks sharing one parent -> a single backend call, both chunks get + # the same concept list. + chunks = [ + _chunk("c1", "part one", parent="the parent"), + _chunk("c2", "part two", parent="the parent"), + ] + out = ex.extract_for_chunks(chunks) + assert len(be.calls) == 1 + assert out.concepts_by_chunk["c1"] == out.concepts_by_chunk["c2"] + + +# --------------------------------------------------------------------------- +# _extract_parent_text fallback-model wiring +# --------------------------------------------------------------------------- + + +def test_extract_parent_text_passes_fallback_when_models_differ(): + be = FakeBackend(default={"entities": [], "relationships": []}) + ex = ConceptExtractor( + kg_config={ + "extraction_max_workers": 1, + "chat_model": "gpt-4o-mini", + "extraction_model": "gpt-4o", + }, + backend=be, + ) + ex.extract_for_chunks([_chunk("c1", "anything")]) + # model != chat_model -> chat_model is offered as the fallback. + assert be.calls[0]["fallback_model"] == "gpt-4o-mini" + + +def test_extract_parent_text_no_fallback_when_models_equal(): + be = FakeBackend(default={"entities": [], "relationships": []}) + ex = ConceptExtractor( + kg_config={ + "extraction_max_workers": 1, + "chat_model": "gpt-4o-mini", + "extraction_model": "gpt-4o-mini", + }, + backend=be, + ) + ex.extract_for_chunks([_chunk("c1", "anything")]) + assert be.calls[0]["fallback_model"] is None + + +def test_extract_parent_text_backend_none_returns_empty_lists(): + # Direct call covering the defensive guard inside _extract_parent_text. + ex = ConceptExtractor(backend=FakeBackend()) + ex.backend = None + out = ex._extract_parent_text("some text", ["label"]) + assert out == {"entities": [], "relationships": []} + + +def test_extract_parent_text_non_dict_payload_is_safe(): + be = FakeBackend(default=["not", "a", "dict"]) # backend returns a list + ex = ConceptExtractor(kg_config={"extraction_max_workers": 1}, backend=be) + out = ex.extract_for_chunks([_chunk("c1", "x")]) + assert out.entities == {} + + +# --------------------------------------------------------------------------- +# _parse_payload edge cases +# --------------------------------------------------------------------------- + + +def test_parse_payload_filters_invalid_and_prunes_unconnected(): + payload = { + "entities": [ + {"name": "Valid Concept", "type": "concept"}, + {"name": "Connected A", "type": "concept"}, + {"name": "Connected B", "type": "concept"}, + {"name": "report.pdf"}, # invalid name -> dropped + "not-a-dict", # non-dict -> skipped + ], + "relationships": [ + # connects A<->B; "Valid Concept" has no relationship -> pruned. + {"source": "Connected A", "target": "Connected B", "relation": "uses"}, + {"source": "X", "target": "Connected A", "relation": "bad"}, # X unknown + {"source": "Connected A", "target": "Connected A", "relation": "self"}, + "non-dict-rel", + ], + } + ex = ConceptExtractor(backend=FakeBackend()) + out = ex._parse_payload(payload, "chunk-1") + # Only the two related concepts survive (pruning of unconnected entities). + assert set(out.entities) == {"connected a", "connected b"} + assert len(out.relationships) == 1 + assert out.relationships[0].relation == "uses" + + +def test_parse_payload_non_list_entities_and_relationships(): + ex = ConceptExtractor(backend=FakeBackend()) + out = ex._parse_payload( + {"entities": "nope", "relationships": "nope"}, "chunk-1" + ) + assert out.entities == {} + assert out.relationships == [] + assert out.concepts_by_chunk == {"chunk-1": []} + + +def test_parse_payload_keeps_entities_when_no_relationships(): + # With no relationships there is nothing to prune against, so valid + # entities are retained. + payload = { + "entities": [{"name": "Solo Concept", "type": "concept"}], + "relationships": [], + } + ex = ConceptExtractor(backend=FakeBackend()) + out = ex._parse_payload(payload, "chunk-1") + assert set(out.entities) == {"solo concept"} diff --git a/lamb-kb-server/tests/unit/test_graph_indexing.py b/lamb-kb-server/tests/unit/test_graph_indexing.py new file mode 100644 index 000000000..71aa686d2 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_graph_indexing.py @@ -0,0 +1,235 @@ +"""Unit tests for ``services.graph_indexing``. + +Drives ``index_chunks_for_collection`` through every return branch with a +fake ``ConceptExtractor`` and a fake graph store, and tests the ``_build_chunks`` +filter helper directly. +""" + +from __future__ import annotations + +import pytest + +import config as config_module +import services.graph_indexing as gi +import services.graph_store as gs_module +from services.concept_extraction import ( + ExtractedEntity, + ExtractedRelationship, + GraphExtraction, +) + + +class _Collection: + def __init__(self, **over): + self.id = "col-1" + self.name = "Coll" + self.organization_id = "org-1" + self.extraction_vendor = None + self.extraction_model = None + self.extraction_endpoint = None + self.__dict__.update(over) + + +class _FakeExtractor: + def __init__(self, *args, extraction=None, raises=False, **kwargs): + self._extraction = extraction or GraphExtraction() + self._raises = raises + _FakeExtractor.last_kwargs = kwargs + + def extract_for_chunks(self, chunks): + if self._raises: + raise RuntimeError("extraction boom") + return self._extraction + + +class _FakeGraphStore: + def __init__(self, *, configured=True, available=True, ingest_raises=False): + self._configured = configured + self._available = available + self._ingest_raises = ingest_raises + self.ingest_calls = [] + + def is_configured(self): + return self._configured + + def is_available(self): + return self._available + + def ingest_chunks(self, **kwargs): + self.ingest_calls.append(kwargs) + if self._ingest_raises: + raise RuntimeError("ingest boom") + return 5 + + +# --------------------------------------------------------------------------- +# _build_chunks +# --------------------------------------------------------------------------- + + +def test_build_chunks_filters_and_parent_text(): + chunks = gi._build_chunks( + ids=["c1", "", "c3", "c4"], + texts=["text one", "skipped", " ", "text four"], + metadatas=[{"parent_text": "PARENT"}, {}, {}], # c4 has no metadata entry + ) + # "" id dropped, " " whitespace text dropped. + ids = [c.chunk_id for c in chunks] + assert ids == ["c1", "c4"] + # parent_text taken from metadata when present. + assert chunks[0].parent_text == "PARENT" + # c4 has no metadata -> parent_text falls back to its own text. + assert chunks[1].parent_text == "text four" + + +def test_build_chunks_drops_ids_beyond_texts(): + chunks = gi._build_chunks(ids=["c1", "c2"], texts=["only one"], metadatas=[{}]) + assert [c.chunk_id for c in chunks] == ["c1"] + + +# --------------------------------------------------------------------------- +# index_chunks_for_collection +# --------------------------------------------------------------------------- + + +def _enable_kg(monkeypatch, **extra): + cfg = {"enabled": True, "openai_api_key": "sk-env"} + cfg.update(extra) + monkeypatch.setattr(config_module, "get_kg_rag_config", lambda: cfg) + + +def test_index_disabled(monkeypatch): + monkeypatch.setattr(config_module, "get_kg_rag_config", lambda: {"enabled": False}) + out = gi.index_chunks_for_collection( + collection=_Collection(), ids=["c1"], texts=["t"], metadatas=[{}] + ) + assert out == {"indexed": False, "chunks": 0, "reason": "kg_rag_disabled"} + + +def test_index_no_chunks(monkeypatch): + _enable_kg(monkeypatch) + out = gi.index_chunks_for_collection( + collection=_Collection(), ids=[], texts=[], metadatas=[] + ) + assert out["reason"] == "no_chunks" + + +def test_index_openai_requires_key(monkeypatch): + # enabled, openai vendor, but no key anywhere. + monkeypatch.setattr( + config_module, "get_kg_rag_config", + lambda: {"enabled": True, "openai_api_key": ""}, + ) + out = gi.index_chunks_for_collection( + collection=_Collection(), ids=["c1"], texts=["text"], metadatas=[{}], + openai_api_key="", + ) + assert out["reason"] == "no_openai_api_key" + assert "OpenAI API key" in out["error"] + + +def test_index_extraction_failure(monkeypatch): + _enable_kg(monkeypatch) + monkeypatch.setattr( + gi, "ConceptExtractor", + lambda *a, **k: _FakeExtractor(raises=True), + ) + out = gi.index_chunks_for_collection( + collection=_Collection(), ids=["c1"], texts=["text"], metadatas=[{}], + openai_api_key="sk-x", + ) + assert "concept_extraction_failed" in out["error"] + + +def test_index_neo4j_unavailable(monkeypatch): + _enable_kg(monkeypatch) + monkeypatch.setattr( + gi, "ConceptExtractor", lambda *a, **k: _FakeExtractor() + ) + monkeypatch.setattr( + gs_module, "get_graph_store", + lambda: _FakeGraphStore(available=False), + ) + out = gi.index_chunks_for_collection( + collection=_Collection(), ids=["c1"], texts=["text"], metadatas=[{}], + openai_api_key="sk-x", + ) + assert out["error"] == "neo4j_unavailable" + + +def test_index_graph_write_failure(monkeypatch): + _enable_kg(monkeypatch) + monkeypatch.setattr( + gi, "ConceptExtractor", lambda *a, **k: _FakeExtractor() + ) + monkeypatch.setattr( + gs_module, "get_graph_store", + lambda: _FakeGraphStore(ingest_raises=True), + ) + out = gi.index_chunks_for_collection( + collection=_Collection(), ids=["c1"], texts=["text"], metadatas=[{}], + openai_api_key="sk-x", + ) + assert "graph_write_failed" in out["error"] + assert "extraction_ms" in out + + +def test_index_success_with_filename_resolution(monkeypatch): + _enable_kg(monkeypatch) + extraction = GraphExtraction( + concepts_by_chunk={"c1": ["alpha"]}, + entities={"alpha": ExtractedEntity(name="alpha", display_name="Alpha")}, + relationships=[ + ExtractedRelationship(source="alpha", target="beta", relation="uses") + ], + ) + monkeypatch.setattr( + gi, "ConceptExtractor", lambda *a, **k: _FakeExtractor(extraction=extraction) + ) + store = _FakeGraphStore() + monkeypatch.setattr(gs_module, "get_graph_store", lambda: store) + out = gi.index_chunks_for_collection( + collection=_Collection(), + ids=["c1"], + texts=["body"], + metadatas=[{"source_label": "Doc Title"}], # filename resolved from here + openai_api_key="sk-x", + ) + assert out["indexed"] is True + assert out["chunks"] == 1 + assert out["entities"] == 1 + assert out["relationships"] == 1 + assert "extraction_ms" in out and "graph_ms" in out + # Filename was resolved from chunk metadata. + assert store.ingest_calls[0]["filename"] == "Doc Title" + + +def test_index_success_default_filename(monkeypatch): + _enable_kg(monkeypatch) + monkeypatch.setattr( + gi, "ConceptExtractor", lambda *a, **k: _FakeExtractor() + ) + store = _FakeGraphStore() + monkeypatch.setattr(gs_module, "get_graph_store", lambda: store) + out = gi.index_chunks_for_collection( + collection=_Collection(), ids=["c1"], texts=["body"], metadatas=[{}], + openai_api_key="sk-x", + ) + assert out["indexed"] is True + # No filename anywhere -> default sentinel. + assert store.ingest_calls[0]["filename"] == "collection_chunks" + + +def test_index_ollama_vendor_no_key_required(monkeypatch): + _enable_kg(monkeypatch, openai_api_key="") + monkeypatch.setattr( + gi, "ConceptExtractor", lambda *a, **k: _FakeExtractor() + ) + store = _FakeGraphStore() + monkeypatch.setattr(gs_module, "get_graph_store", lambda: store) + out = gi.index_chunks_for_collection( + collection=_Collection(extraction_vendor="ollama"), + ids=["c1"], texts=["body"], metadatas=[{}], + ) + # Ollama doesn't require an OpenAI key -> proceeds to indexing. + assert out["indexed"] is True diff --git a/lamb-kb-server/tests/unit/test_graph_store_core.py b/lamb-kb-server/tests/unit/test_graph_store_core.py new file mode 100644 index 000000000..47d79eb1b --- /dev/null +++ b/lamb-kb-server/tests/unit/test_graph_store_core.py @@ -0,0 +1,355 @@ +"""Unit tests for ``services.graph_store.GraphStore`` — core read/guard paths. + +Covers construction/configuration guards, schema bootstrap, collection and +document deletion, change listing with every filter branch, and the static +relationship-matching helpers. Heavier write/transaction paths (revert, +rename, merge, edit, curation, ingest) are covered in companion modules. +""" + +from __future__ import annotations + +import json + +import pytest + +from services import graph_store as gs_module +from services.graph_store import GraphStore, get_graph_store, utc_now +from tests.unit._neo4j_fakes import FakeDriver, FakeResult, FakeSession, make_store + + +# --------------------------------------------------------------------------- +# module-level helpers + construction +# --------------------------------------------------------------------------- + + +def test_utc_now_is_iso8601_utc(): + value = utc_now() + assert value.endswith("+00:00") + + +def test_get_graph_store_is_singleton(monkeypatch): + monkeypatch.setattr(gs_module, "_GRAPH_STORE", None) + a = get_graph_store() + b = get_graph_store() + assert a is b + + +def test_init_not_configured_leaves_driver_none(): + store = GraphStore(kg_config={"enabled": False}) + assert store.driver is None + assert store.is_configured() is False + + +def test_init_configured_creates_driver(monkeypatch): + created = {} + + class _FakeGraphDatabase: + @staticmethod + def driver(uri, auth): + created["uri"] = uri + created["auth"] = auth + return "the-driver" + + monkeypatch.setattr(gs_module, "GraphDatabase", _FakeGraphDatabase) + store = GraphStore( + kg_config={ + "enabled": True, + "neo4j_uri": "bolt://localhost:7687", + "neo4j_user": "neo4j", + "neo4j_password": "secret", + } + ) + assert store.driver == "the-driver" + assert created["auth"] == ("neo4j", "secret") + assert store.is_configured() is True + + +def test_init_driver_creation_failure_is_swallowed(monkeypatch): + class _BoomGraphDatabase: + @staticmethod + def driver(uri, auth): + raise RuntimeError("cannot connect") + + monkeypatch.setattr(gs_module, "GraphDatabase", _BoomGraphDatabase) + store = GraphStore( + kg_config={ + "enabled": True, + "neo4j_uri": "bolt://x", + "neo4j_user": "u", + "neo4j_password": "p", + } + ) + assert store.driver is None + + +def test_is_configured_false_when_graphdatabase_missing(monkeypatch): + monkeypatch.setattr(gs_module, "GraphDatabase", None) + store = GraphStore( + kg_config={ + "enabled": True, + "neo4j_uri": "bolt://x", + "neo4j_user": "u", + "neo4j_password": "p", + } + ) + assert store.is_configured() is False + + +# --------------------------------------------------------------------------- +# close / is_available +# --------------------------------------------------------------------------- + + +def test_close_calls_driver_close(): + store, _ = make_store() + store.close() + assert store.driver.closed is True + + +def test_close_noop_when_no_driver(): + store = GraphStore(kg_config={"enabled": False}) + store.close() # must not raise + + +def test_is_available_false_without_driver(): + store = GraphStore(kg_config={"enabled": False}) + assert store.is_available() is False + + +def test_is_available_true_when_connectivity_ok(): + store, _ = make_store() + assert store.is_available() is True + + +def test_is_available_false_on_connectivity_error(): + sess = FakeSession() + driver = FakeDriver(sess, connectivity_error=RuntimeError("down")) + store, _ = make_store(driver=driver) + assert store.is_available() is False + + +# --------------------------------------------------------------------------- +# ensure_schema +# --------------------------------------------------------------------------- + + +def test_ensure_schema_short_circuits_when_ready(): + store, sess = make_store(schema_ready=True) + assert store.ensure_schema() is True + assert sess.runs == [] # no statements executed + + +def test_ensure_schema_returns_false_when_unavailable(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, _ = make_store(driver=driver, schema_ready=False) + assert store.ensure_schema() is False + + +def test_ensure_schema_runs_all_statements(): + store, sess = make_store(schema_ready=False) + assert store.ensure_schema() is True + # 8 schema statements (constraints + indexes). + assert len(sess.runs) == 8 + assert store._schema_ready is True + + +def test_ensure_schema_handles_run_failure(): + class _BoomSession(FakeSession): + def run(self, query, **params): + raise RuntimeError("schema error") + + store, _ = make_store(session=_BoomSession(), schema_ready=False) + assert store.ensure_schema() is False + + +# --------------------------------------------------------------------------- +# delete_collection / delete_document +# --------------------------------------------------------------------------- + + +def test_delete_collection_runs_three_statements(): + store, sess = make_store() + store.delete_collection("col-1") + assert len(sess.runs) == 3 + assert sess.runs[0]["params"]["collection_id"] == "col-1" + + +def test_delete_collection_aborts_if_schema_not_ready(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, sess = make_store(driver=driver, schema_ready=False) + store.delete_collection("col-1") + assert sess.runs == [] + + +def test_delete_document_runs_three_statements(): + store, sess = make_store() + store.delete_document("col-1", "org-1", "doc.pdf") + assert len(sess.runs) == 3 + assert sess.runs[1]["params"]["filename"] == "doc.pdf" + assert sess.runs[2]["params"]["org_id"] == "org-1" + + +def test_delete_document_aborts_when_not_configured(): + # is_configured() is False (enabled stays False, no uri/password). + store = GraphStore(kg_config={"enabled": False}) + # Should return before touching any driver (driver is None). + store.delete_document("c", "o", "f") # must not raise + + +# --------------------------------------------------------------------------- +# list_changes filtering +# --------------------------------------------------------------------------- + + +def _event_row(**over): + row = { + "event_id": "e1", + "collection_id": "c1", + "org_id": "o1", + "operation": "add_document", + "actor": "user", + "timestamp": "2026-06-20T00:00:00Z", + "filename": "doc.pdf", + "concepts": ["alpha", "beta"], + "payload_json": None, + "document_id": "d1", + "file_id": 1, + } + row.update(over) + return row + + +def test_list_changes_returns_false_when_schema_unavailable(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, _ = make_store(driver=driver, schema_ready=False) + assert store.list_changes("c1", "o1") == [] + + +def test_list_changes_plain_no_filters(): + rows = [_event_row(), _event_row(event_id="e2")] + store, sess = make_store(responses=[FakeResult(rows=rows)]) + out = store.list_changes("c1", "o1", limit=5) + assert len(out) == 2 + assert sess.runs[0]["params"]["fetch_limit"] == 100 # clamped lower bound + + +def test_list_changes_limit_clamped_to_200(): + store, sess = make_store(responses=[FakeResult(rows=[])]) + store.list_changes("c1", "o1", limit=9999) + # fetch_limit = min(max(limit*10, 100), 1000) + assert sess.runs[0]["params"]["fetch_limit"] == 1000 + + +def test_list_changes_concept_filter_drops_relationship_ops(): + rows = [ + _event_row(operation="add_document"), + _event_row(event_id="e2", operation="manual_edit_relationship"), + _event_row(event_id="e3", operation="manual_curate_relationship"), + ] + store, _ = make_store(responses=[FakeResult(rows=rows)]) + out = store.list_changes("c1", "o1", concept="Alpha") + ops = {r["operation"] for r in out} + assert ops == {"add_document"} + + +def test_list_changes_relationship_filter_matches_payload(): + payload = json.dumps( + {"source": "alpha", "target": "beta", "relation": "uses"} + ) + rows = [ + _event_row(event_id="match", payload_json=payload), + _event_row(event_id="nomatch", payload_json=json.dumps( + {"source": "gamma", "target": "delta", "relation": "uses"} + )), + ] + store, _ = make_store(responses=[FakeResult(rows=rows)]) + out = store.list_changes( + "c1", + "o1", + relationship_source="Alpha", + relationship_target="Beta", + relationship_relation="uses", + ) + assert [r["event_id"] for r in out] == ["match"] + + +def test_list_changes_truncates_to_limit(): + rows = [_event_row(event_id=f"e{i}") for i in range(10)] + store, _ = make_store(responses=[FakeResult(rows=rows)]) + out = store.list_changes("c1", "o1", limit=3) + assert len(out) == 3 + + +# --------------------------------------------------------------------------- +# _event_payload / _event_matches_relationship +# --------------------------------------------------------------------------- + + +def test_event_payload_parses_and_guards(): + assert GraphStore._event_payload({"payload_json": '{"a": 1}'}) == {"a": 1} + assert GraphStore._event_payload({"payload_json": None}) == {} + assert GraphStore._event_payload({"payload_json": "not-json"}) == {} + # Valid JSON that isn't an object -> {}. + assert GraphStore._event_payload({"payload_json": "[1,2]"}) == {} + + +def test_event_matches_relationship_via_relationship_details(): + row = { + "payload_json": json.dumps( + {"relationship_details": [{"source": "alpha", "target": "beta", + "relation": "uses"}]} + ) + } + assert GraphStore._event_matches_relationship(row, "alpha", "beta", "uses") + assert not GraphStore._event_matches_relationship(row, "alpha", "zeta", None) + + +def test_event_matches_relationship_via_removed_relationships(): + row = { + "payload_json": json.dumps( + {"removed_relationships": [{"source": "alpha", "target": "beta", + "relation": "old"}]} + ) + } + assert GraphStore._event_matches_relationship(row, "alpha", None, None) + + +def test_event_matches_relationship_alternate_relation_keys(): + row = { + "payload_json": json.dumps( + {"source": "alpha", "target": "beta", "relation": "uses", + "new_relation": "depends_on"} + ) + } + # Matches against the new_relation value too. + assert GraphStore._event_matches_relationship(row, "alpha", "beta", "depends_on") + + +def test_event_matches_relationship_no_payload_match(): + row = {"payload_json": json.dumps({"unrelated": True})} + assert GraphStore._event_matches_relationship(row, "alpha", None, None) is False + + +# --------------------------------------------------------------------------- +# get_change +# --------------------------------------------------------------------------- + + +def test_get_change_returns_dict(): + row = _event_row() + row["chunk_ids"] = ["chunk-1", "chunk-2"] + store, _ = make_store(responses=[FakeResult(single=row)]) + out = store.get_change("c1", "o1", "e1") + assert out["event_id"] == "e1" + assert out["chunk_ids"] == ["chunk-1", "chunk-2"] + + +def test_get_change_returns_none_when_missing(): + store, _ = make_store(responses=[FakeResult(single=None)]) + assert store.get_change("c1", "o1", "missing") is None + + +def test_get_change_returns_none_when_schema_unavailable(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, _ = make_store(driver=driver, schema_ready=False) + assert store.get_change("c1", "o1", "e1") is None diff --git a/lamb-kb-server/tests/unit/test_graph_store_edit_curation.py b/lamb-kb-server/tests/unit/test_graph_store_edit_curation.py new file mode 100644 index 000000000..f719e7a85 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_graph_store_edit_curation.py @@ -0,0 +1,301 @@ +"""Unit tests for ``GraphStore`` relationship-edit and concept-curation paths, +plus the pure change-detection helpers they rely on. +""" + +from __future__ import annotations + +from services.graph_store import GraphStore +from tests.unit._neo4j_fakes import FakeDriver, FakeResult, FakeSession, make_store + + +# --------------------------------------------------------------------------- +# pure change-detection helpers +# --------------------------------------------------------------------------- + + +def test_optional_text_changed(): + assert GraphStore._optional_text_changed("a", None) is False # None = no request + assert GraphStore._optional_text_changed("a", "a") is False + assert GraphStore._optional_text_changed("a", "b") is True + assert GraphStore._optional_text_changed(None, "") is False + + +def test_optional_tags_changed(): + assert GraphStore._optional_tags_changed(["a"], None) is False + assert GraphStore._optional_tags_changed(["a"], ["a"]) is False + assert GraphStore._optional_tags_changed(["a"], ["a", "b"]) is True + assert GraphStore._optional_tags_changed(None, []) is False + + +def test_optional_weight_changed(): + assert GraphStore._optional_weight_changed(1.0, None) is False + assert GraphStore._optional_weight_changed(1.0, 1.0) is False + assert GraphStore._optional_weight_changed(1.0, 2.0) is True + # None current defaults to 1.0. + assert GraphStore._optional_weight_changed(None, 1.0) is False + # Non-numeric falls back to equality comparison. + assert GraphStore._optional_weight_changed("x", "x") is False + + +def test_optional_state_changed(): + assert GraphStore._optional_state_changed("verified", None) is False + assert GraphStore._optional_state_changed(None, "unverified") is False # default + assert GraphStore._optional_state_changed("unverified", "verified") is True + + +def test_relationship_update_has_changes(): + base = { + "old_weight": 1.0, + "old_description": "d", + "old_evidence": "e", + "old_notes": "n", + "old_tags": ["t"], + "old_verification_state": "verified", + } + # Relation change alone -> True. + assert GraphStore._relationship_update_has_changes( + base, current_relation="uses", target_relation="depends_on", + weight=None, description=None, evidence=None, notes=None, tags=None, + verification_state=None, + ) + # Nothing changes -> False. + assert not GraphStore._relationship_update_has_changes( + base, current_relation="uses", target_relation="uses", + weight=None, description=None, evidence=None, notes=None, tags=None, + verification_state=None, + ) + + +def test_concept_update_has_changes(): + row = {"old_notes": "n", "old_tags": ["t"], "old_verification_state": "verified"} + assert GraphStore._concept_update_has_changes( + row, notes="new", tags=None, verification_state=None + ) + assert not GraphStore._concept_update_has_changes( + row, notes=None, tags=None, verification_state=None + ) + + +# --------------------------------------------------------------------------- +# edit_relationship +# --------------------------------------------------------------------------- + + +def _rel_row(**over): + row = { + "old_weight": 1.0, + "old_description": "d", + "old_evidence": "e", + "old_chunk_id": "ch1", + "old_notes": "n", + "old_tags": ["t"], + "old_verification_state": "verified", + } + row.update(over) + return row + + +def test_edit_relationship_schema_unavailable(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, _ = make_store(driver=driver, schema_ready=False) + out = store.edit_relationship( + "c1", "o1", source_name="a", target_name="b", relation="uses" + ) + assert out == {"ok": False, "reason": "neo4j_not_available"} + + +def test_edit_relationship_invalid_identity(): + store, sess = make_store(responses=[]) + out = store.edit_relationship( + "c1", "o1", source_name="!!!", target_name="b", relation="uses" + ) + assert out["reason"] == "invalid_relationship_identity" + assert sess.runs == [] + + +def test_edit_relationship_not_found(): + store, _ = make_store(responses=[FakeResult(single=None)]) + out = store.edit_relationship( + "c1", "o1", source_name="a", target_name="b", relation="uses" + ) + assert out["reason"] == "relationship_not_found" + + +def test_edit_relationship_rejected_expunges(): + responses = [ + FakeResult(single=_rel_row()), # rel lookup + FakeResult(single={"event_id": "ev"}), # audit event + FakeResult(), # delete relationship + ] + store, sess = make_store(responses=responses) + out = store.edit_relationship( + "c1", "o1", source_name="Alpha", target_name="Beta", relation="uses", + verification_state="rejected", + ) + assert out["ok"] is True + assert out["operation"] == "manual_expunge_relationship" + assert out["details"]["expunged"] is True + assert len(sess.runs) == 3 + + +def test_edit_relationship_no_change(): + store, sess = make_store(responses=[FakeResult(single=_rel_row())]) + out = store.edit_relationship( + "c1", "o1", source_name="Alpha", target_name="Beta", relation="uses", + ) + assert out["reason"] == "no_change" + assert out["details"]["changed"] is False + assert len(sess.runs) == 1 + + +def test_edit_relationship_same_relation_update(): + responses = [ + FakeResult(single=_rel_row()), # rel lookup + FakeResult(), # SET update + FakeResult(single={"event_id": "ev2"}), # audit event + ] + store, sess = make_store(responses=responses) + out = store.edit_relationship( + "c1", "o1", source_name="Alpha", target_name="Beta", relation="uses", + weight=5.0, + ) + assert out["ok"] is True + assert out["operation"] == "manual_edit_relationship" + assert out["event_id"] == "ev2" + assert out["details"]["new_relation"] == "uses" + assert len(sess.runs) == 3 + + +def test_edit_relationship_changed_relation_recreates(): + responses = [ + FakeResult(single=_rel_row()), # rel lookup + FakeResult(), # MERGE new + DELETE old + FakeResult(single={"event_id": "ev3"}), # audit event + ] + store, sess = make_store(responses=responses) + out = store.edit_relationship( + "c1", "o1", source_name="Alpha", target_name="Beta", relation="uses", + new_relation="depends on", + ) + assert out["ok"] is True + assert out["details"]["old_relation"] == "uses" + assert out["details"]["new_relation"] == "depends_on" + + +# --------------------------------------------------------------------------- +# update_concept_curation +# --------------------------------------------------------------------------- + + +def _concept_row(**over): + row = { + "name": "alpha", + "old_notes": "n", + "old_tags": ["t"], + "old_verification_state": "verified", + } + row.update(over) + return row + + +def test_curation_schema_unavailable(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, _ = make_store(driver=driver, schema_ready=False) + out = store.update_concept_curation("c1", "o1", "Alpha") + assert out == {"ok": False, "reason": "neo4j_not_available"} + + +def test_curation_invalid_name(): + store, sess = make_store(responses=[]) + out = store.update_concept_curation("c1", "o1", "!!!") + assert out["reason"] == "invalid_concept_name" + assert sess.runs == [] + + +def test_curation_concept_not_found(): + store, _ = make_store(responses=[FakeResult(single=None)]) + out = store.update_concept_curation("c1", "o1", "Alpha", notes="x") + assert out["reason"] == "concept_not_found" + + +def test_curation_no_change(): + # notes/tags/state all None -> no change. + store, sess = make_store(responses=[FakeResult(single=_concept_row())]) + out = store.update_concept_curation("c1", "o1", "Alpha") + assert out["reason"] == "no_change" + assert len(sess.runs) == 1 + + +def test_curation_notes_only_update(): + responses = [ + FakeResult(single=_concept_row(old_notes="old")), # concept lookup + FakeResult(), # SET notes/tags + FakeResult(single={"event_id": "ce"}), # audit event + ] + store, sess = make_store(responses=responses) + out = store.update_concept_curation("c1", "o1", "Alpha", notes="new notes") + assert out["ok"] is True + assert out["operation"] == "manual_curate_concept" + assert len(sess.runs) == 3 + + +def test_curation_verified_promotes_org_node(): + responses = [ + FakeResult(single=_concept_row(old_verification_state="verified")), # lookup + FakeResult(single={"vs": None}), # per-collection mentions vs + FakeResult(), # SET mentions vs + FakeResult(), # promote org-level concept to verified + FakeResult(), # SET concept notes/tags + FakeResult(single={"event_id": "ce2"}), # audit event + ] + store, sess = make_store(responses=responses) + out = store.update_concept_curation( + "c1", "o1", "Alpha", verification_state="verified" + ) + assert out["ok"] is True + assert out["operation"] == "manual_curate_concept" + assert len(sess.runs) == 6 + + +def test_curation_unverified_does_not_promote_org_node(): + # verification_state is set but not "verified" -> mentions vs is cleared and + # the org-level promote is skipped. + responses = [ + FakeResult(single=_concept_row(old_verification_state="verified")), # lookup + FakeResult(single={"vs": "verified"}), # per-collection mentions vs + FakeResult(), # SET mentions vs (cleared to null) + FakeResult(), # SET concept notes/tags + FakeResult(single={"event_id": "ce4"}), # audit event + ] + store, sess = make_store(responses=responses) + out = store.update_concept_curation( + "c1", "o1", "Alpha", verification_state="unverified" + ) + assert out["ok"] is True + assert out["operation"] == "manual_curate_concept" + # No org-promote query (that path only runs for "verified"). + assert len(sess.runs) == 5 + + +def test_curation_rejected_expunges_concept(): + responses = [ + FakeResult(single=_concept_row()), # concept lookup + FakeResult(single={"vs": "verified"}), # per-collection mentions vs + FakeResult(single={"chunk_ids": ["ch1"]}), # chunk ids + FakeResult(rows=[{"source": "alpha", "target": "beta", + "relation": "uses"}]), # relationships + FakeResult(single={"event_id": "ce3"}), # audit event + FakeResult(), # delete mentions + FakeResult(), # delete relationships + FakeResult(), # orphan delete + ] + store, sess = make_store(responses=responses) + out = store.update_concept_curation( + "c1", "o1", "Alpha", verification_state="rejected" + ) + assert out["ok"] is True + assert out["operation"] == "manual_expunge_concept" + assert out["details"]["expunged"] is True + assert out["details"]["removed_chunk_mentions"] == 1 + assert out["details"]["removed_relationships"] == 1 + assert len(sess.runs) == 8 diff --git a/lamb-kb-server/tests/unit/test_graph_store_ingest_expand.py b/lamb-kb-server/tests/unit/test_graph_store_ingest_expand.py new file mode 100644 index 000000000..057c4adc8 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_graph_store_ingest_expand.py @@ -0,0 +1,167 @@ +"""Unit tests for ``GraphStore.ingest_chunks`` / ``_ingest_tx`` and the +``expand_from_concept_names`` graph-traversal seeding path. +""" + +from __future__ import annotations + +from services.concept_extraction import ( + ExtractedEntity, + ExtractedRelationship, + TextChunk, +) +from tests.unit._neo4j_fakes import FakeDriver, FakeResult, FakeSession, make_store + + +# --------------------------------------------------------------------------- +# ingest_chunks +# --------------------------------------------------------------------------- + + +def _chunk(cid="ch1", text="body", parent="parent", meta=None): + return TextChunk(chunk_id=cid, text=text, parent_text=parent, metadata=meta or {}) + + +def test_ingest_empty_chunks_returns_zero(): + store, sess = make_store() + n = store.ingest_chunks( + collection={"id": "c1"}, + file_id=1, + filename="f.pdf", + chunks=[], + concepts_by_chunk={}, + entities={}, + relationships=[], + ) + assert n == 0 + assert sess.runs == [] + + +def test_ingest_schema_unavailable_returns_zero(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, _ = make_store(driver=driver, schema_ready=False) + n = store.ingest_chunks( + collection={"id": "c1"}, + file_id=1, + filename="f.pdf", + chunks=[_chunk()], + concepts_by_chunk={}, + entities={}, + relationships=[], + ) + assert n == 0 + + +def test_ingest_full_path_counts_and_writes(): + chunk = _chunk( + meta={ + "permalink_original": "http://x/orig", + "permalink_full_markdown": "http://x/md", + "permalink_page": "http://x/p1", + "section_title": "Intro", + "source_label": "Chunk 1", + "filename": "f.pdf", + } + ) + entities = {"alpha": ExtractedEntity(name="alpha", display_name="Alpha")} + relationships = [ + ExtractedRelationship(source="alpha", target="beta", relation="uses") + ] + store, sess = make_store( + config={"name": "Coll"}, + ) + n = store.ingest_chunks( + collection={ + "id": "c1", + "organization_id": "o1", + "name": "Coll", + "description": "Desc", + }, + file_id=7, + filename="f.pdf", + chunks=[chunk], + concepts_by_chunk={"ch1": ["alpha"]}, + entities=entities, + relationships=relationships, + ) + # 1 chunk + entity_map(alpha + beta from relationship) + 1 relationship + 1 event. + assert n == 1 + 2 + 1 + 1 + # _ingest_tx issued multiple writes: org/coll/doc + 2 entities + chunk + + # mention + relationship + event. + assert len(sess.runs) == 1 + 2 + 1 + 1 + 1 + 1 + + +def test_ingest_uses_owner_fallback_for_org_id(): + store, sess = make_store() + n = store.ingest_chunks( + collection={"id": "c1", "owner": "owner-org"}, + file_id=None, # -> int(0) + filename="f.pdf", + chunks=[_chunk()], + concepts_by_chunk={}, + entities={}, + relationships=[], + ) + # 1 chunk + 0 entities + 0 rels + 1 = 2. + assert n == 2 + # org/coll/doc write carries the owner fallback org id. + assert sess.runs[0]["params"]["org_id"] == "owner-org" + + +# --------------------------------------------------------------------------- +# expand_from_concept_names +# --------------------------------------------------------------------------- + + +def test_expand_empty_names_returns_empty(): + store, sess = make_store() + out = store.expand_from_concept_names("c1", "o1", ["", " "], depth=2, limit=10) + assert out["entry_concepts"] == [] + assert out["expanded_chunk_ids"] == [] + assert "graph_latency_ms" in out + assert sess.runs == [] # returns before any query + + +def test_expand_schema_unavailable_warns(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, _ = make_store(driver=driver, schema_ready=False) + out = store.expand_from_concept_names("c1", "o1", ["Alpha"], depth=2, limit=10) + assert out["entry_concepts"] == [] + assert any("KG expansion skipped" in c.get("warning", "") + for c in out["latest_changes"]) + + +def test_expand_no_entry_concepts(): + store, sess = make_store(responses=[FakeResult(rows=[])]) # matched empty + out = store.expand_from_concept_names("c1", "o1", ["Alpha"], depth=2, limit=10) + assert out["entry_concepts"] == [] + assert len(sess.runs) == 1 + + +def test_expand_full_path_collects_chunks(): + responses = [ + FakeResult(rows=[{"name": "alpha", "mentions": 2}]), # matched + FakeResult(rows=[{"chunk_id": "ch1", "mentions": 5}]), # direct + FakeResult(rows=[{"chunk_id": "ch2"}, {"chunk_id": "ch1"}]), # related (dedup) + ] + store, sess = make_store(responses=responses) + out = store.expand_from_concept_names( + "c1", "o1", ["Alpha"], depth=2, limit=10 + ) + assert out["entry_concepts"] == ["alpha"] + # ch1 from direct, ch2 from related; ch1 deduped on second appearance. + assert out["expanded_chunk_ids"] == ["ch1", "ch2"] + assert len(sess.runs) == 3 + + +def test_expand_clamps_depth_and_limit(): + responses = [ + FakeResult(rows=[{"name": "alpha", "mentions": 1}]), + FakeResult(rows=[]), + FakeResult(rows=[]), + ] + store, sess = make_store(responses=responses) + # Negative limit exercises the max(1, ...) lower clamp; depth=99 clamps to 4 + # (interpolated into the related-query string). + store.expand_from_concept_names("c1", "o1", ["Alpha"], depth=99, limit=-5) + assert sess.runs[1]["params"]["limit"] == 1 + assert "RELATES_TO*1..4" in sess.runs[2]["query"] diff --git a/lamb-kb-server/tests/unit/test_graph_store_rename_merge.py b/lamb-kb-server/tests/unit/test_graph_store_rename_merge.py new file mode 100644 index 000000000..7723012cc --- /dev/null +++ b/lamb-kb-server/tests/unit/test_graph_store_rename_merge.py @@ -0,0 +1,143 @@ +"""Unit tests for ``GraphStore`` concept rename/merge curation operations. + +These exercise ``rename_concept`` / ``merge_concepts`` and, through them, the +shared ``_move_concept_in_collection_tx`` (7 Cypher writes) and +``_manual_change_event_tx`` audit-event helper. +""" + +from __future__ import annotations + +from tests.unit._neo4j_fakes import FakeDriver, FakeResult, FakeSession, make_store + + +def _move_responses(*, source_found=True, counts=None, event_id="evt-1"): + """Response queue for one successful _move + audit event. + + _move issues: source lookup, MERGE target, removed_between, mentions, + outgoing, incoming, deleted_source. Then the caller records an event. + """ + counts = counts or {} + if not source_found: + return [FakeResult(single=None)] + return [ + FakeResult(single={ + "name": "alpha", "old_notes": None, "old_tags": None, + "old_verification_state": None, + }), + FakeResult(), # MERGE target + FakeResult(single={"count": counts.get("removed_between", 0)}), + FakeResult(single={"count": counts.get("mentions", 3)}), + FakeResult(single={"count": counts.get("outgoing", 1)}), + FakeResult(single={"count": counts.get("incoming", 2)}), + FakeResult(single={"count": counts.get("deleted_source", 1)}), + ] + ([FakeResult(single={"event_id": event_id})] if event_id else []) + + +# --------------------------------------------------------------------------- +# rename_concept +# --------------------------------------------------------------------------- + + +def test_rename_schema_unavailable(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, _ = make_store(driver=driver, schema_ready=False) + out = store.rename_concept("c1", "o1", "Old", "New") + assert out == {"ok": False, "reason": "neo4j_not_available"} + + +def test_rename_invalid_name(): + store, sess = make_store(responses=[]) + out = store.rename_concept("c1", "o1", "!!!", "New") + assert out["reason"] == "invalid_concept_name" + assert sess.runs == [] + + +def test_rename_equal_names(): + store, _ = make_store(responses=[]) + out = store.rename_concept("c1", "o1", "Same Name", "same name") + assert out["reason"] == "concept_names_are_equal" + + +def test_rename_source_not_found(): + store, sess = make_store(responses=_move_responses(source_found=False)) + out = store.rename_concept("c1", "o1", "Old Name", "New Name") + assert out["ok"] is False + assert out["reason"] == "source_concept_not_found" + assert len(sess.runs) == 1 + + +def test_rename_success(): + store, sess = make_store( + responses=_move_responses(counts={"mentions": 5}, event_id="evt-9") + ) + out = store.rename_concept("c1", "o1", "Old Name", "New Name", reason="cleanup") + assert out["ok"] is True + assert out["operation"] == "manual_rename_concept" + assert out["event_id"] == "evt-9" + assert out["details"]["mentions"] == 5 + # 7 move writes + 1 audit event. + assert len(sess.runs) == 8 + + +# --------------------------------------------------------------------------- +# merge_concepts +# --------------------------------------------------------------------------- + + +def test_merge_schema_unavailable(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, _ = make_store(driver=driver, schema_ready=False) + out = store.merge_concepts("c1", "o1", ["a"], "b") + assert out == {"ok": False, "reason": "neo4j_not_available"} + + +def test_merge_invalid_request_no_sources(): + store, sess = make_store(responses=[]) + # Only source equals target after normalization -> nothing to merge. + out = store.merge_concepts("c1", "o1", ["Target"], "Target") + assert out["reason"] == "invalid_merge_request" + assert sess.runs == [] + + +def test_merge_invalid_request_empty_target(): + store, _ = make_store(responses=[]) + out = store.merge_concepts("c1", "o1", ["a"], "!!!") + assert out["reason"] == "invalid_merge_request" + + +def test_merge_all_sources_missing(): + store, sess = make_store(responses=_move_responses(source_found=False)) + out = store.merge_concepts("c1", "o1", ["Ghost"], "Target") + assert out["ok"] is False + assert out["reason"] == "source_concepts_not_found" + assert out["missing"] == ["ghost"] + + +def test_merge_success_one_source(): + store, sess = make_store( + responses=_move_responses(counts={"mentions": 4}, event_id="merge-evt") + ) + out = store.merge_concepts("c1", "o1", ["Source One"], "Target", reason="dedupe") + assert out["ok"] is True + assert out["operation"] == "manual_merge_concepts" + assert out["event_id"] == "merge-evt" + assert out["details"]["target"] == "target" + assert len(out["details"]["moved"]) == 1 + # 7 move writes + 1 audit event. + assert len(sess.runs) == 8 + + +def test_merge_partial_missing_source_still_succeeds(): + # First source moves; second source is missing -> still ok overall. + responses = ( + _move_responses(event_id=None) # first source: 7 writes, no event yet + + [FakeResult(single=None)] # second source: not found (1 write) + + [FakeResult(single={"event_id": "merge-evt"})] # audit event + ) + store, sess = make_store(responses=responses) + out = store.merge_concepts( + "c1", "o1", ["First Source", "Second Source"], "Target" + ) + assert out["ok"] is True + assert out["details"]["missing"] == ["second source"] + assert len(out["details"]["moved"]) == 1 diff --git a/lamb-kb-server/tests/unit/test_graph_store_restore.py b/lamb-kb-server/tests/unit/test_graph_store_restore.py new file mode 100644 index 000000000..5c95f9f5f --- /dev/null +++ b/lamb-kb-server/tests/unit/test_graph_store_restore.py @@ -0,0 +1,211 @@ +"""Unit tests for ``GraphStore`` expunge-restore reverts and remaining guards. + +When ``revert_change`` targets a ``manual_expunge_relationship`` or +``manual_expunge_concept`` event, it delegates into the restore unit-of-work +helpers (``_restore_relationship_payload_tx``, ``_restore_concept_node_tx``, +``_restore_expunged_*_tx``, ``_create_revert_event_tx``). These tests drive +those paths and also close a few residual guard branches. +""" + +from __future__ import annotations + +import json + +from services.graph_store import GraphStore +from tests.unit._neo4j_fakes import FakeDriver, FakeResult, FakeSession, make_store + + +# --------------------------------------------------------------------------- +# _relationship_restore_weight (pure) +# --------------------------------------------------------------------------- + + +def test_relationship_restore_weight(): + assert GraphStore._relationship_restore_weight(2.5) == 2.5 + assert GraphStore._relationship_restore_weight(None) == 1.0 + assert GraphStore._relationship_restore_weight("not-a-number") == 1.0 + + +# --------------------------------------------------------------------------- +# revert of manual_expunge_relationship +# --------------------------------------------------------------------------- + + +def test_revert_expunged_relationship_success(): + payload = { + "source": "alpha", + "target": "beta", + "relation": "uses", + "old_weight": 2.0, + "old_description": "d", + "old_evidence": "e", + "old_chunk_id": "ch1", + "old_notes": "n", + "old_tags": ["t"], + } + event_row = { + "operation": "manual_expunge_relationship", + "filename": "f.pdf", + "concepts": ["alpha", "beta"], + "payload_json": json.dumps(payload), + "document_id": None, + } + responses = [ + FakeResult(single=event_row), # event lookup + FakeResult(), # restore source Concept + FakeResult(), # restore target Concept + FakeResult(), # MERGE relationship + FakeResult(single={"revert_event_id": "rev-1"}), # audit revert event + ] + store, sess = make_store(responses=responses) + out = store.revert_change("c1", "o1", "e1", reason="undo") + assert out["reverted"] is True + assert out["operation"] == "manual_expunge_relationship" + assert out["revert_event_id"] == "rev-1" + assert out["chunk_ids"] == [] + assert len(sess.runs) == 5 + + +def test_revert_expunged_relationship_invalid_payload(): + # Empty source -> restore returns None -> invalid_expunge_payload. + payload = {"source": "", "target": "beta", "relation": "uses"} + event_row = { + "operation": "manual_expunge_relationship", + "filename": "f.pdf", + "concepts": [], + "payload_json": json.dumps(payload), + "document_id": None, + } + responses = [ + FakeResult(single=event_row), # event lookup + FakeResult(), # restore target Concept (source skipped) + ] + store, _ = make_store(responses=responses) + out = store.revert_change("c1", "o1", "e1") + assert out["reverted"] is False + assert out["reason"] == "invalid_expunge_payload" + + +# --------------------------------------------------------------------------- +# revert of manual_expunge_concept +# --------------------------------------------------------------------------- + + +def test_revert_expunged_concept_success(): + payload = { + "concept": "alpha", + "old_notes": "n", + "old_tags": ["t"], + "removed_chunk_mentions": ["ch1", "ch2"], + "removed_relationships": [ + {"source": "alpha", "target": "beta", "relation": "uses"}, + "not-a-dict", # skipped + ], + } + event_row = { + "operation": "manual_expunge_concept", + "filename": "f.pdf", + "concepts": ["alpha"], + "payload_json": json.dumps(payload), + "document_id": None, + } + responses = [ + FakeResult(single=event_row), # event lookup + FakeResult(), # restore concept node + FakeResult(), # restore chunk mention ch1 + FakeResult(), # restore chunk mention ch2 + FakeResult(), # restore rel: source node + FakeResult(), # restore rel: target node + FakeResult(), # restore rel: MERGE relationship + FakeResult(single={"revert_event_id": "rev-2"}), # audit revert event + ] + store, sess = make_store(responses=responses) + out = store.revert_change("c1", "o1", "e1") + assert out["reverted"] is True + assert out["operation"] == "manual_expunge_concept" + assert out["chunk_ids"] == ["ch1", "ch2"] + assert out["revert_event_id"] == "rev-2" + assert len(sess.runs) == 8 + + +def test_revert_expunged_concept_skips_unrestorable_relationship(): + # A removed relationship whose source is empty cannot be restored and is + # skipped (the ``if restored:`` false branch). + payload = { + "concept": "alpha", + "removed_chunk_mentions": [], + "removed_relationships": [ + {"source": "", "target": "beta", "relation": "uses"}, # unrestorable + ], + } + event_row = { + "operation": "manual_expunge_concept", + "filename": "f.pdf", + "concepts": ["alpha"], + "payload_json": json.dumps(payload), + "document_id": None, + } + responses = [ + FakeResult(single=event_row), # event lookup + FakeResult(), # restore concept node + FakeResult(), # restore rel: target node (source skipped) + FakeResult(single={"revert_event_id": "rev-3"}), # audit event + ] + store, _ = make_store(responses=responses) + out = store.revert_change("c1", "o1", "e1") + assert out["reverted"] is True + assert out["chunk_ids"] == [] + + +def test_revert_expunged_concept_invalid_payload(): + # No concept name available -> restore returns "" -> invalid_expunge_payload. + payload = {"concept": ""} + event_row = { + "operation": "manual_expunge_concept", + "filename": "f.pdf", + "concepts": [""], + "payload_json": json.dumps(payload), + "document_id": None, + } + store, _ = make_store(responses=[FakeResult(single=event_row)]) + out = store.revert_change("c1", "o1", "e1") + assert out["reverted"] is False + assert out["reason"] == "invalid_expunge_payload" + + +# --------------------------------------------------------------------------- +# residual guard branches +# --------------------------------------------------------------------------- + + +def test_delete_document_aborts_when_schema_not_ready(): + # is_configured() is True (make_store sets uri/password) but the schema + # bootstrap fails -> early return at the ensure_schema guard. + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, sess = make_store(driver=driver, schema_ready=False) + store.delete_document("c1", "o1", "doc.pdf") + assert sess.runs == [] + + +def test_snapshot_without_chunks_skips_document_missing_id(): + concepts = [{ + "name": "alpha", "display_name": "Alpha", "entity_type": "concept", + "description": "", "confidence": 0.9, "notes": "", "tags": [], + "chunk_count": 1, + }] + coll_vs = [{"name": "alpha", "vs": "verified"}] + documents = [ + {"document_id": None, "filename": "ghost", "chunk_count": 0, + "concepts": ["alpha"]}, # skipped in document-mention loop + ] + store, _ = make_store( + responses=[ + FakeResult(rows=concepts), + FakeResult(rows=coll_vs), + FakeResult(rows=documents), + FakeResult(rows=[]), # relationships + ] + ) + out = store.get_collection_graph("c1", "o1", include_chunks=False) + # No document-mention edges were produced (the only doc had no id). + assert all(not e["id"].startswith("document-mention:") for e in out["edges"]) diff --git a/lamb-kb-server/tests/unit/test_graph_store_revert.py b/lamb-kb-server/tests/unit/test_graph_store_revert.py new file mode 100644 index 000000000..ac8bd507c --- /dev/null +++ b/lamb-kb-server/tests/unit/test_graph_store_revert.py @@ -0,0 +1,139 @@ +"""Unit tests for ``GraphStore.revert_change`` and its automatic-ingestion +unit-of-work ``_revert_change_tx``. + +The transaction issues a deterministic sequence of ``tx.run`` calls; we +program the fake session's response queue to match and assert each guard +branch plus the full revert path (weight decrement, document/chunk deletion, +orphan cleanup, and the audit ``revert_change`` event). +""" + +from __future__ import annotations + +import json + +from tests.unit._neo4j_fakes import FakeDriver, FakeResult, FakeSession, make_store + + +def test_revert_returns_unavailable_when_schema_down(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, _ = make_store(driver=driver, schema_ready=False) + out = store.revert_change("c1", "o1", "e1") + assert out == {"reverted": False, "reason": "neo4j_not_available"} + + +def test_revert_change_not_found(): + store, _ = make_store(responses=[FakeResult(single=None)]) + out = store.revert_change("c1", "o1", "missing") + assert out["reverted"] is False + assert out["reason"] == "change_not_found" + + +def test_revert_unsupported_operation(): + event_row = { + "operation": "manual_rename_concept", + "filename": "", + "concepts": [], + "payload_json": "{}", + "document_id": None, + } + store, _ = make_store(responses=[FakeResult(single=event_row)]) + out = store.revert_change("c1", "o1", "e1") + assert out["reason"] == "unsupported_operation" + assert out["operation"] == "manual_rename_concept" + + +def test_revert_automatic_ingestion_without_document(): + event_row = { + "operation": "automatic_ingestion", + "filename": "doc.pdf", + "concepts": ["alpha"], + "payload_json": "{}", + "document_id": None, + } + store, _ = make_store(responses=[FakeResult(single=event_row)]) + out = store.revert_change("c1", "o1", "e1") + assert out["reason"] == "change_has_no_document" + + +def test_revert_automatic_ingestion_full_path(): + payload = json.dumps( + { + "relationship_details": [ + {"source": "alpha", "target": "beta", "relation": "uses", + "confidence": 1.0}, + "not-a-dict", # skipped by the loop + ] + } + ) + event_row = { + "operation": "automatic_ingestion", + "filename": "doc.pdf", + "concepts": ["alpha", "beta"], + "payload_json": payload, + "document_id": "d1", + } + responses = [ + FakeResult(single=event_row), # event lookup + FakeResult(single={"chunk_ids": ["chunk-1", "chunk-2"]}), # chunk ids + FakeResult(), # relationship decrement + FakeResult(), # delete doc + chunks + FakeResult(), # orphan concept cleanup + FakeResult(single={"revert_event_id": "rev-1"}), # audit event + ] + store, sess = make_store(responses=responses) + out = store.revert_change("c1", "o1", "e1", actor="me", reason="oops") + assert out["reverted"] is True + assert out["operation"] == "automatic_ingestion" + assert out["document_id"] == "d1" + assert out["chunk_ids"] == ["chunk-1", "chunk-2"] + assert out["revert_event_id"] == "rev-1" + # event + chunk + 1 relationship + delete + orphan + audit = 6 runs. + assert len(sess.runs) == 6 + + +def test_revert_full_path_uses_relationships_key_fallback(): + # No "relationship_details"; the tx falls back to "relationships". + payload = json.dumps( + {"relationships": [{"source": "a", "target": "b", "relation": "rel"}]} + ) + event_row = { + "operation": "automatic_ingestion", + "filename": "doc.pdf", + "concepts": [], + "payload_json": payload, + "document_id": "d1", + } + responses = [ + FakeResult(single=event_row), + FakeResult(single=None), # chunk row None -> chunk_ids == [] + FakeResult(), # relationship decrement + FakeResult(), # delete + FakeResult(), # orphan + FakeResult(single=None), # audit event row None -> revert_event_id None + ] + store, _ = make_store(responses=responses) + out = store.revert_change("c1", "o1", "e1") + assert out["reverted"] is True + assert out["chunk_ids"] == [] + assert out["revert_event_id"] is None + + +def test_revert_bad_payload_json_is_tolerated(): + event_row = { + "operation": "automatic_ingestion", + "filename": "doc.pdf", + "concepts": [], + "payload_json": "not-json{", + "document_id": "d1", + } + responses = [ + FakeResult(single=event_row), + FakeResult(single={"chunk_ids": []}), + FakeResult(), # delete (no relationships to decrement) + FakeResult(), # orphan + FakeResult(single={"revert_event_id": "rev-9"}), + ] + store, _ = make_store(responses=responses) + out = store.revert_change("c1", "o1", "e1") + assert out["reverted"] is True + assert out["revert_event_id"] == "rev-9" diff --git a/lamb-kb-server/tests/unit/test_graph_store_snapshot.py b/lamb-kb-server/tests/unit/test_graph_store_snapshot.py new file mode 100644 index 000000000..cc8e4ea99 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_graph_store_snapshot.py @@ -0,0 +1,173 @@ +"""Unit tests for ``GraphStore.get_collection_graph`` (the snapshot builder). + +The method issues up to five sequential Cypher reads: + 1. concepts, 2. per-collection verification states, 3. documents, + 4. relationships, 5. chunks (only when ``include_chunks``). +We program the fake session's response queue in that order and assert the +node/edge assembly, counts, filter echo, and the early-return branches. +""" + +from __future__ import annotations + +from services.graph_store import GraphStore +from tests.unit._neo4j_fakes import FakeDriver, FakeResult, FakeSession, make_store + + +def _concept(name, display=None, **extra): + row = { + "name": name, + "display_name": display or name.title(), + "entity_type": "concept", + "description": f"about {name}", + "confidence": 0.9, + "notes": "", + "tags": [], + "chunk_count": 1, + } + row.update(extra) + return row + + +def test_snapshot_returns_empty_when_schema_unavailable(): + driver = FakeDriver(FakeSession(), connectivity_error=RuntimeError("x")) + store, _ = make_store(driver=driver, schema_ready=False) + out = store.get_collection_graph("c1", "o1", concept="Alpha", limit=10) + assert out["nodes"] == [] and out["edges"] == [] + assert out["counts"] == {"concepts": 0, "documents": 0, "chunks": 0, "edges": 0} + # Filters echo the request even on the empty path. + assert out["filters"]["concept"] == "Alpha" + assert out["filters"]["limit"] == 10 + + +def test_snapshot_empty_when_no_concepts_match(): + store, sess = make_store(responses=[FakeResult(rows=[])]) # concept query empty + out = store.get_collection_graph("c1", "o1") + assert out["counts"]["concepts"] == 0 + # Only the concept query ran; it short-circuits before verification query. + assert len(sess.runs) == 1 + + +def test_snapshot_full_path_with_chunks(): + concepts = [_concept("alpha"), _concept("beta")] + coll_vs = [{"name": "alpha", "vs": "verified"}, {"name": "beta", "vs": None}] + documents = [ + { + "document_id": "d1", + "filename": "doc.pdf", + "file_id": 7, + "chunk_count": 2, + "concepts": ["alpha", "beta"], + }, + # row with missing document_id -> skipped during node assembly. + {"document_id": None, "filename": "ghost", "chunk_count": 0, "concepts": []}, + ] + relationships = [ + { + "source": "alpha", + "target": "beta", + "type": "RELATES_TO", + "relation": "uses", + "weight": 2.0, + "description": "d", + "evidence": "e", + "chunk_id": "chunk-1", + "notes": "", + "tags": [], + "verification_state": "verified", + } + ] + chunks = [ + { + "chunk_id": "chunk-1", + "source_label": "Chunk 1", + "filename": "doc.pdf", + "document_id": "d1", + "text_preview": "preview", + "permalink_original": "", + "permalink_full_markdown": "", + "permalink_page": "", + "concepts": ["alpha"], + }, + # missing chunk_id -> skipped. + {"chunk_id": None, "concepts": []}, + ] + store, sess = make_store( + responses=[ + FakeResult(rows=concepts), + FakeResult(rows=coll_vs), + FakeResult(rows=documents), + FakeResult(rows=relationships), + FakeResult(rows=chunks), + ] + ) + out = store.get_collection_graph("c1", "o1", include_chunks=True) + + node_ids = {n["id"] for n in out["nodes"]} + assert "concept:alpha" in node_ids + assert "concept:beta" in node_ids + assert "document:d1" in node_ids + assert "chunk:chunk-1" in node_ids + + edge_ids = {e["id"] for e in out["edges"]} + assert "relationship:alpha:uses:beta" in edge_ids + assert "mention:chunk-1:alpha" in edge_ids + assert "contains:d1:chunk-1" in edge_ids + + # verification_state for alpha comes from the per-collection map. + alpha = next(n for n in out["nodes"] if n["id"] == "concept:alpha") + assert alpha["data"]["verification_state"] == "verified" + beta = next(n for n in out["nodes"] if n["id"] == "concept:beta") + assert beta["data"]["verification_state"] == "unverified" + + assert out["counts"]["concepts"] == 2 + assert out["counts"]["documents"] == 2 # both rows counted (len) + assert out["counts"]["chunks"] == 2 + assert len(sess.runs) == 5 + + +def test_snapshot_rejected_concept_is_filtered_out(): + concepts = [_concept("alpha"), _concept("beta")] + coll_vs = [{"name": "alpha", "vs": "verified"}, {"name": "beta", "vs": "rejected"}] + store, _ = make_store( + responses=[ + FakeResult(rows=concepts), + FakeResult(rows=coll_vs), + FakeResult(rows=[]), # documents + FakeResult(rows=[]), # relationships + FakeResult(rows=[]), # chunks + ] + ) + out = store.get_collection_graph("c1", "o1", include_chunks=True) + node_ids = {n["id"] for n in out["nodes"]} + assert "concept:alpha" in node_ids + assert "concept:beta" not in node_ids # rejected in this collection + assert out["counts"]["concepts"] == 1 + + +def test_snapshot_without_chunks_builds_document_mention_edges(): + concepts = [_concept("alpha")] + coll_vs = [{"name": "alpha", "vs": "verified"}] + documents = [ + { + "document_id": "d1", + "filename": "doc.pdf", + "file_id": 1, + "chunk_count": 1, + "concepts": ["alpha"], + } + ] + store, sess = make_store( + responses=[ + FakeResult(rows=concepts), + FakeResult(rows=coll_vs), + FakeResult(rows=documents), + FakeResult(rows=[]), # relationships + ] + ) + out = store.get_collection_graph("c1", "o1", include_chunks=False) + # No chunk query issued. + assert len(sess.runs) == 4 + edge_ids = {e["id"] for e in out["edges"]} + assert "document-mention:d1:alpha" in edge_ids + assert out["filters"]["include_chunks"] is False + assert out["counts"]["chunks"] == 0 diff --git a/lamb-kb-server/tests/unit/test_kg_rag.py b/lamb-kb-server/tests/unit/test_kg_rag.py new file mode 100644 index 000000000..3bd9dfa4f --- /dev/null +++ b/lamb-kb-server/tests/unit/test_kg_rag.py @@ -0,0 +1,309 @@ +"""Focused unit tests for the KG-RAG migration. + +Covered: + +* ``config.get_kg_rag_config`` reads env vars and applies sensible defaults + (disabled by default, sane bounds on ``graph_depth`` / ``limit_factor`` / + ``extraction_max_workers``). +* The KG-RAG query plugin returns the baseline unchanged with explicit + warnings when the feature is disabled, when the collection has not opted + in, and when there are no seed chunks. +* The concept-extraction module's pure helpers (``normalize_concept`` / + ``normalize_relation``) are deterministic and robust to Unicode noise. + +These tests intentionally avoid Neo4j / OpenAI — the heavy paths are +exercised by integration / e2e suites that are gated on the optional +``kg-rag`` extra and Docker services. +""" + +from __future__ import annotations + +import importlib + +import pytest + + +# --------------------------------------------------------------------------- +# config +# --------------------------------------------------------------------------- + + +def test_kg_rag_config_disabled_by_default(reload_config, monkeypatch): + for var in ( + "KG_RAG_ENABLED", + "KG_RAG_INDEX_ON_INGEST", + "KG_RAG_OPENAI_API_KEY", + "KG_RAG_NEO4J_URI", + "OPENAI_API_KEY", + ): + monkeypatch.delenv(var, raising=False) + import config + + importlib.reload(config) + cfg = config.get_kg_rag_config() + assert cfg["enabled"] is False + assert cfg["graph_depth"] == 2 + assert cfg["limit_factor"] == 4 + assert cfg["extraction_max_workers"] == 4 + + +def test_kg_rag_config_clamps_graph_depth(reload_config, monkeypatch): + monkeypatch.setenv("KG_RAG_ENABLED", "true") + monkeypatch.setenv("KG_RAG_GRAPH_DEPTH", "99") + monkeypatch.setenv("KG_RAG_LIMIT_FACTOR", "999") + monkeypatch.setenv("KG_RAG_EXTRACTION_MAX_WORKERS", "999") + import config + + importlib.reload(config) + cfg = config.get_kg_rag_config() + assert cfg["enabled"] is True + assert cfg["graph_depth"] == 4 + assert cfg["limit_factor"] == 20 + assert cfg["extraction_max_workers"] == 16 + + +def test_kg_rag_config_falls_back_to_openai_api_key(reload_config, monkeypatch): + monkeypatch.delenv("KG_RAG_OPENAI_API_KEY", raising=False) + monkeypatch.setenv("OPENAI_API_KEY", "sk-test-fallback") + import config + + importlib.reload(config) + cfg = config.get_kg_rag_config() + assert cfg["openai_api_key"] == "sk-test-fallback" + + +# --------------------------------------------------------------------------- +# concept extraction (pure helpers) +# --------------------------------------------------------------------------- + + +def test_normalize_concept_strips_accents_and_punct(): + from services.concept_extraction import normalize_concept + + assert normalize_concept("Café Society!") == "cafe society" + assert normalize_concept(" multiple spaces ") == "multiple spaces" + assert normalize_concept("") == "" + + +def test_normalize_relation_keeps_short_relation_keys(): + from services.concept_extraction import normalize_relation + + assert normalize_relation("Depends On") == "depends_on" + assert normalize_relation(" improves -- ") == "improves" + # Fallback for empty / unusable input. + assert normalize_relation("---") == "related_to" + + +# --------------------------------------------------------------------------- +# KG-RAG query plugin: graceful degradation paths +# --------------------------------------------------------------------------- + + +class _StubCollection: + """Minimal stand-in for the ORM Collection row.""" + + def __init__(self, *, graph_enabled: bool = True): + self.id = "stub-collection" + self.organization_id = "stub-org" + self.graph_enabled = graph_enabled + self.backend_collection_id = "stub-backend" + self.storage_path = "/tmp/does-not-matter" + + +def test_plugin_returns_baseline_when_kg_rag_disabled(monkeypatch): + monkeypatch.delenv("KG_RAG_ENABLED", raising=False) + + from plugins.kg_rag_query import KGRAGQueryPlugin + + plugin = KGRAGQueryPlugin() + baseline = [{"similarity": 0.9, "data": "hello", "metadata": {"document_id": "c1"}}] + out = plugin.augment( + db=None, + collection=_StubCollection(), + backend=None, + embedding_function=None, + query_text="anything", + baseline_results=baseline, + params={}, + ) + # Same payload, with a kg_rag trace attached that explains the no-op. + assert len(out) == 1 + trace = out[0]["metadata"].get("kg_rag") + assert trace is not None + assert trace["enabled"] is False + assert "KG-RAG is disabled" in " ".join(trace["warnings"]) + + +def test_plugin_returns_baseline_when_collection_not_opted_in(monkeypatch): + monkeypatch.setenv("KG_RAG_ENABLED", "true") + + from plugins.kg_rag_query import KGRAGQueryPlugin + + plugin = KGRAGQueryPlugin() + out = plugin.augment( + db=None, + collection=_StubCollection(graph_enabled=False), + backend=None, + embedding_function=None, + query_text="anything", + baseline_results=[ + {"similarity": 0.9, "data": "x", "metadata": {"document_id": "c1"}} + ], + params={}, + ) + trace = out[0]["metadata"]["kg_rag"] + assert any("graph_enabled=false" in w for w in trace["warnings"]) + + +def test_plugin_returns_baseline_when_no_seed_ids(monkeypatch): + monkeypatch.setenv("KG_RAG_ENABLED", "true") + + from plugins.kg_rag_query import KGRAGQueryPlugin + + plugin = KGRAGQueryPlugin() + out = plugin.augment( + db=None, + collection=_StubCollection(), + backend=None, + embedding_function=None, + query_text="anything", + # Baseline result has no chunk-like id, so seed extraction is empty. + baseline_results=[ + {"similarity": 0.9, "data": "x", "metadata": {}} + ], + params={}, + ) + trace = out[0]["metadata"]["kg_rag"] + assert any("No vector seed chunks" in w for w in trace["warnings"]) + + +# --------------------------------------------------------------------------- +# Benchmark route body shape — regression test for the FastAPI Body(embed=True) +# bug. The LAMB proxy sends a flat JSON body; the route must accept it without +# wrapping fields under ``request``. +# --------------------------------------------------------------------------- + + +def test_benchmark_run_request_validates_flat_proxy_body(): + """The exact JSON shape sent by ``KnowledgeStoreClient.run_benchmark`` + must validate as a ``BenchmarkRunRequest`` — including the embedded + ``embedding_credentials`` sub-object.""" + from schemas.benchmark import BenchmarkRunRequest + + proxy_body = { + "dataset_id": "educational", + "top_k": 5, + "graph_depth": 2, + "threshold": 0.0, + "embedding_credentials": { + "api_key": "sk-test", + "api_endpoint": "", + }, + } + parsed = BenchmarkRunRequest.model_validate(proxy_body) + assert parsed.dataset_id == "educational" + assert parsed.top_k == 5 + assert parsed.embedding_credentials.api_key == "sk-test" + + +def test_benchmark_run_request_credentials_optional(): + """Direct API callers may omit ``embedding_credentials`` entirely.""" + from schemas.benchmark import BenchmarkRunRequest + + parsed = BenchmarkRunRequest.model_validate( + {"dataset_id": "educational", "top_k": 5} + ) + # Default factory produces an empty-string credentials object. + assert parsed.embedding_credentials.api_key == "" + assert parsed.embedding_credentials.api_endpoint == "" + + +# --------------------------------------------------------------------------- +# Schema migration: graph_enabled column auto-added on init_db +# --------------------------------------------------------------------------- + + +def test_init_db_adds_graph_enabled_to_legacy_collections_table( + tmp_path, monkeypatch +): + """A DB that was created before this branch (no graph_enabled column) + should get the column added by init_db without losing existing rows. + + The bug we're guarding: ``Base.metadata.create_all`` is a no-op on an + existing table, so without _run_lightweight_migrations any query + against ``collections`` would raise ``no such column``. + + We call the lightweight-migrations helper directly here rather than + spinning up a second init_db — that keeps this test from clobbering + the session-wide ``_engine`` / ``_SessionLocal`` that the rest of the + suite depends on. + """ + import sqlite3 + + from sqlalchemy import create_engine + + db_path = tmp_path / "legacy.db" + + # Hand-roll the legacy schema (pre-branch) and seed a row. + conn = sqlite3.connect(str(db_path)) + conn.executescript( + """ + CREATE TABLE collections ( + id TEXT PRIMARY KEY, + organization_id TEXT NOT NULL, + name TEXT NOT NULL, + description TEXT, + chunking_strategy TEXT NOT NULL, + chunking_params TEXT, + embedding_vendor TEXT NOT NULL, + embedding_model TEXT NOT NULL, + embedding_endpoint TEXT, + vector_db_backend TEXT NOT NULL, + backend_collection_id TEXT, + storage_path TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'ready', + error_message TEXT, + document_count INTEGER NOT NULL DEFAULT 0, + chunk_count INTEGER NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP + ); + INSERT INTO collections ( + id, organization_id, name, chunking_strategy, embedding_vendor, + embedding_model, vector_db_backend, storage_path + ) VALUES ( + 'legacy-1', 'org-1', 'legacy', 'simple', 'fake', + 'fake-model', 'chromadb', '/tmp/legacy' + ); + """ + ) + conn.commit() + conn.close() + + legacy_engine = create_engine(f"sqlite:///{db_path}") + from database.connection import _run_lightweight_migrations + + _run_lightweight_migrations(legacy_engine) + legacy_engine.dispose() + + # After the migration runs, the legacy row should still be present AND + # the new column must exist with the documented default. + conn = sqlite3.connect(str(db_path)) + cur = conn.execute("PRAGMA table_info(collections)") + columns = {row[1] for row in cur.fetchall()} + assert "graph_enabled" in columns + + row = conn.execute( + "SELECT graph_enabled FROM collections WHERE id = 'legacy-1'" + ).fetchone() + assert row == (0,) # NOT NULL default 0 + + # The migration must be idempotent — running it again is a no-op. + legacy_engine = create_engine(f"sqlite:///{db_path}") + _run_lightweight_migrations(legacy_engine) + legacy_engine.dispose() + cur = conn.execute("PRAGMA table_info(collections)") + assert ( + sum(1 for row in cur.fetchall() if row[1] == "graph_enabled") == 1 + ) + conn.close() diff --git a/lamb-kb-server/tests/unit/test_kg_rag_query_plugin.py b/lamb-kb-server/tests/unit/test_kg_rag_query_plugin.py new file mode 100644 index 000000000..762e45504 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_kg_rag_query_plugin.py @@ -0,0 +1,396 @@ +"""Unit tests for ``plugins.kg_rag_query.KGRAGQueryPlugin`` active paths. + +The graceful-degradation no-ops (disabled / not-opted-in / no-seed) are +covered in ``test_kg_rag.py``. Here we drive the *active* augmentation path +with a fake graph store + fake vector backend, plus the static helpers +(``_as_bool``, ``_result_chunk_id``, ``_rrf_merge``, ``_attach_trace``, +``_fetch_expanded_results``) and the LLM-based question-entity extractor. +""" + +from __future__ import annotations + +import types + +import pytest + +import config as config_module +import services.graph_store as gs_module +from plugins.kg_rag_query import KGRAGQueryPlugin + + +class _Collection: + def __init__(self, **over): + self.id = "col-1" + self.organization_id = "org-1" + self.graph_enabled = True + self.backend_collection_id = "backend-col" + self.storage_path = "/tmp/store" + self.extraction_vendor = None + self.extraction_model = None + self.extraction_endpoint = None + self.__dict__.update(over) + + +class _FakeItem: + def __init__(self, text, score, metadata): + self.text = text + self.score = score + self.metadata = metadata + + +class _FakeBackend: + def __init__(self, *, items=None, raises=False): + self._items = items or [] + self._raises = raises + self.calls = [] + + def get_chunks_by_id(self, **kwargs): + self.calls.append(kwargs) + if self._raises: + raise RuntimeError("backend boom") + return self._items + + +class _FakeGraphStore: + def __init__(self, *, configured=True, expansion=None, raises=False): + self._configured = configured + self._expansion = expansion or {} + self._raises = raises + + def is_configured(self): + return self._configured + + def expand_from_concept_names(self, **kwargs): + if self._raises: + raise RuntimeError("expand boom") + return self._expansion + + +# --------------------------------------------------------------------------- +# static helpers +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "value,expected", + [ + (True, True), + (False, False), + (None, False), + ("yes", True), + ("ON", True), + ("enabled", True), + ("0", False), + ("nope", False), + ], +) +def test_as_bool(value, expected): + assert KGRAGQueryPlugin._as_bool(value) is expected + + +def test_result_chunk_id_precedence(): + assert KGRAGQueryPlugin._result_chunk_id( + {"metadata": {"document_id": "d", "chunk_id": "c"}} + ) == "d" + assert KGRAGQueryPlugin._result_chunk_id( + {"metadata": {"child_chunk_id": "cc", "chunk_id": "c"}} + ) == "cc" + assert KGRAGQueryPlugin._result_chunk_id({"metadata": {"chunk_id": "c"}}) == "c" + assert KGRAGQueryPlugin._result_chunk_id({"metadata": {}}) is None + assert KGRAGQueryPlugin._result_chunk_id({}) is None + + +def test_attach_trace_toggle(): + results = [{"data": "x", "metadata": {"a": 1}}] + trace = {"mode": "kg_rag"} + # include_trace False -> unchanged objects. + assert KGRAGQueryPlugin._attach_trace(results, trace, False) is results + traced = KGRAGQueryPlugin._attach_trace(results, trace, True) + assert traced[0]["metadata"]["kg_rag"] == trace + assert traced[0]["metadata"]["a"] == 1 + + +def test_rrf_merge_fuses_and_boosts_overlap(): + plugin = KGRAGQueryPlugin() + baseline = [ + {"similarity": 0.9, "data": "A", "metadata": {"chunk_id": "a"}}, + {"similarity": 0.8, "data": "B", "metadata": {"chunk_id": "b"}}, + ] + expanded = [ + {"similarity": 0.7, "data": "A", "metadata": {"chunk_id": "a"}}, # overlap + {"similarity": 0.6, "data": "C", "metadata": {"chunk_id": "c"}}, + ] + merged = plugin._rrf_merge( + baseline_results=baseline, expanded_results=expanded, top_k=5 + ) + by_id = {KGRAGQueryPlugin._result_chunk_id(m): m for m in merged} + # "a" appears in both lists -> fusion_lists has both. + assert by_id["a"]["fusion_lists"] == ["baseline", "graph"] + assert "rrf_score" in by_id["a"] + # Overlapping "a" should outrank single-list entries. + assert merged[0]["metadata"]["chunk_id"] == "a" + + +def test_fetch_expanded_results_empty_ids(): + plugin = KGRAGQueryPlugin() + out = plugin._fetch_expanded_results( + backend=_FakeBackend(), collection=_Collection(), + embedding_function=None, expanded_ids=[], return_parent_context=True, + ) + assert out == [] + + +def test_fetch_expanded_results_backend_error_returns_empty(): + plugin = KGRAGQueryPlugin() + out = plugin._fetch_expanded_results( + backend=_FakeBackend(raises=True), collection=_Collection(), + embedding_function=None, expanded_ids=["c1"], return_parent_context=True, + ) + assert out == [] + + +def test_fetch_expanded_results_with_parent_context(): + items = [ + _FakeItem("child text", 0.5, {"parent_text": "PARENT", "chunk_id": "c1"}), + _FakeItem("only child", None, {"chunk_id": "c2"}), + ] + plugin = KGRAGQueryPlugin() + out = plugin._fetch_expanded_results( + backend=_FakeBackend(items=items), collection=_Collection(), + embedding_function=None, expanded_ids=["c1", "c2"], + return_parent_context=True, + ) + # First item uses parent_text as data; origin tag added; parent_text popped. + assert out[0]["data"] == "PARENT" + assert out[0]["metadata"]["kg_rag_origin"] == "graph_expansion" + assert "parent_text" not in out[0]["metadata"] + # Second item had no score -> default 0.72; data falls back to text. + assert out[1]["similarity"] == 0.72 + assert out[1]["data"] == "only child" + + +def test_fetch_expanded_results_without_parent_context(): + items = [_FakeItem("child", 0.5, {"parent_text": "PARENT", "chunk_id": "c1"})] + plugin = KGRAGQueryPlugin() + out = plugin._fetch_expanded_results( + backend=_FakeBackend(items=items), collection=_Collection(), + embedding_function=None, expanded_ids=["c1"], + return_parent_context=False, + ) + # data uses the child text (parent context not requested). + assert out[0]["data"] == "child" + + +# --------------------------------------------------------------------------- +# augment active path +# --------------------------------------------------------------------------- + + +def _baseline(): + return [{"similarity": 0.9, "data": "seed", "metadata": {"document_id": "seed-1"}}] + + +def test_augment_neo4j_not_configured(monkeypatch): + monkeypatch.setattr( + config_module, "get_kg_rag_config", lambda: {"enabled": True, "graph_depth": 2} + ) + monkeypatch.setattr( + gs_module, "get_graph_store", lambda: _FakeGraphStore(configured=False) + ) + out = KGRAGQueryPlugin().augment( + db=None, collection=_Collection(), backend=_FakeBackend(), + embedding_function=None, query_text="q", baseline_results=_baseline(), + params={}, + ) + trace = out[0]["metadata"]["kg_rag"] + assert any("Neo4j is not configured" in w for w in trace["warnings"]) + + +def test_augment_full_path_merges_expansion(monkeypatch): + monkeypatch.setattr( + config_module, "get_kg_rag_config", lambda: {"enabled": True, "graph_depth": 2} + ) + expansion = { + "entry_concepts": ["alpha"], + "expanded_chunk_ids": ["exp-1"], + "graph_latency_ms": 12.0, + } + monkeypatch.setattr( + gs_module, "get_graph_store", lambda: _FakeGraphStore(expansion=expansion) + ) + monkeypatch.setattr( + KGRAGQueryPlugin, "_extract_question_entities", + classmethod(lambda cls, q, c: ["alpha"]), + ) + backend = _FakeBackend( + items=[_FakeItem("expanded text", 0.6, {"chunk_id": "exp-1"})] + ) + out = KGRAGQueryPlugin().augment( + db=None, collection=_Collection(), backend=backend, + embedding_function=None, query_text="what is alpha?", + baseline_results=_baseline(), params={"top_k": 5}, + ) + trace = out[0]["metadata"]["kg_rag"] + assert trace["graph_expanded"] is True + assert trace["entry_concepts"] == ["alpha"] + assert trace["expanded_chunk_ids"] == ["exp-1"] + # Both seed and expanded chunks present in the merged output. + ids = {KGRAGQueryPlugin._result_chunk_id(r) for r in out} + assert "seed-1" in ids and "exp-1" in ids + + +def test_augment_expansion_failure_adds_warning(monkeypatch): + monkeypatch.setattr( + config_module, "get_kg_rag_config", lambda: {"enabled": True, "graph_depth": 2} + ) + monkeypatch.setattr( + gs_module, "get_graph_store", lambda: _FakeGraphStore(raises=True) + ) + monkeypatch.setattr( + KGRAGQueryPlugin, "_extract_question_entities", + classmethod(lambda cls, q, c: ["alpha"]), + ) + out = KGRAGQueryPlugin().augment( + db=None, collection=_Collection(), backend=_FakeBackend(), + embedding_function=None, query_text="q", baseline_results=_baseline(), + params={}, + ) + trace = out[0]["metadata"]["kg_rag"] + assert any("Graph expansion failed" in w for w in trace["warnings"]) + + +def test_augment_no_expanded_chunks_warns(monkeypatch): + monkeypatch.setattr( + config_module, "get_kg_rag_config", lambda: {"enabled": True, "graph_depth": 2} + ) + monkeypatch.setattr( + gs_module, "get_graph_store", + lambda: _FakeGraphStore(expansion={"entry_concepts": [], "expanded_chunk_ids": []}), + ) + monkeypatch.setattr( + KGRAGQueryPlugin, "_extract_question_entities", + classmethod(lambda cls, q, c: []), + ) + out = KGRAGQueryPlugin().augment( + db=None, collection=_Collection(), backend=_FakeBackend(), + embedding_function=None, query_text="q", baseline_results=_baseline(), + params={}, + ) + trace = out[0]["metadata"]["kg_rag"] + assert any("Graph returned no additional chunks" in w for w in trace["warnings"]) + + +# --------------------------------------------------------------------------- +# _extract_question_entities (LLM-based) +# --------------------------------------------------------------------------- + + +class _FakeExtractionBackend: + def __init__(self, payload): + self._payload = payload + + def chat_json(self, *, system, user, fallback_model=None): + return self._payload + + +@pytest.fixture(autouse=True) +def _clear_question_cache(): + KGRAGQueryPlugin._question_cache.clear() + yield + KGRAGQueryPlugin._question_cache.clear() + + +def test_extract_entities_empty_question(): + assert KGRAGQueryPlugin._extract_question_entities(" ", _Collection()) == [] + + +def test_extract_entities_openai_no_key(monkeypatch): + monkeypatch.setattr( + config_module, "get_kg_rag_config", + lambda: {"openai_api_key": "", "chat_model": "gpt-4o-mini"}, + ) + assert KGRAGQueryPlugin._extract_question_entities("q", _Collection()) == [] + + +def test_extract_entities_success_and_cache(monkeypatch): + monkeypatch.setattr( + config_module, "get_kg_rag_config", + lambda: {"openai_api_key": "sk-x", "chat_model": "gpt-4o-mini"}, + ) + import plugins.base as base_module + + built = {"count": 0} + + def fake_build(vendor, **kwargs): + built["count"] += 1 + return _FakeExtractionBackend({"entities": ["alpha", " beta ", ""]}) + + monkeypatch.setattr(base_module.LLMExtractionRegistry, "build", fake_build) + out = KGRAGQueryPlugin._extract_question_entities("what is alpha?", _Collection()) + assert out == ["alpha", "beta"] + # Second identical call is served from cache (build not called again). + out2 = KGRAGQueryPlugin._extract_question_entities("what is alpha?", _Collection()) + assert out2 == ["alpha", "beta"] + assert built["count"] == 1 + + +def test_extract_entities_vendor_unavailable(monkeypatch): + monkeypatch.setattr( + config_module, "get_kg_rag_config", + lambda: {"openai_api_key": "sk-x", "chat_model": "gpt-4o-mini"}, + ) + import plugins.base as base_module + + def fake_build(vendor, **kwargs): + raise ValueError("not registered") + + monkeypatch.setattr(base_module.LLMExtractionRegistry, "build", fake_build) + assert KGRAGQueryPlugin._extract_question_entities("q", _Collection()) == [] + + +def test_extract_entities_non_dict_payload(monkeypatch): + monkeypatch.setattr( + config_module, "get_kg_rag_config", + lambda: {"openai_api_key": "sk-x", "chat_model": "gpt-4o-mini"}, + ) + import plugins.base as base_module + + monkeypatch.setattr( + base_module.LLMExtractionRegistry, "build", + lambda vendor, **kwargs: _FakeExtractionBackend(["not", "a", "dict"]), + ) + assert KGRAGQueryPlugin._extract_question_entities("q", _Collection()) == [] + + +def test_extract_entities_custom_vendor_skips_key_check(monkeypatch): + # A non-openai vendor doesn't require an OpenAI key. + monkeypatch.setattr( + config_module, "get_kg_rag_config", + lambda: {"openai_api_key": "", "chat_model": "gpt-4o-mini"}, + ) + import plugins.base as base_module + + monkeypatch.setattr( + base_module.LLMExtractionRegistry, "build", + lambda vendor, **kwargs: _FakeExtractionBackend({"entities": ["x"]}), + ) + coll = _Collection(extraction_vendor="ollama", extraction_model="llama3.1:8b") + assert KGRAGQueryPlugin._extract_question_entities("q", coll) == ["x"] + + +def test_extract_entities_cache_eviction(monkeypatch): + monkeypatch.setattr( + config_module, "get_kg_rag_config", + lambda: {"openai_api_key": "sk-x", "chat_model": "gpt-4o-mini"}, + ) + import plugins.base as base_module + + monkeypatch.setattr( + base_module.LLMExtractionRegistry, "build", + lambda vendor, **kwargs: _FakeExtractionBackend({"entities": ["e"]}), + ) + # Fill the cache beyond its bound to exercise FIFO eviction. + monkeypatch.setattr(KGRAGQueryPlugin, "_QUESTION_CACHE_SIZE", 3) + for i in range(5): + KGRAGQueryPlugin._extract_question_entities(f"question {i}", _Collection()) + assert len(KGRAGQueryPlugin._question_cache) <= 3 diff --git a/lamb-kb-server/tests/unit/test_llm_extraction.py b/lamb-kb-server/tests/unit/test_llm_extraction.py new file mode 100644 index 000000000..c31d33c76 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_llm_extraction.py @@ -0,0 +1,369 @@ +"""Unit tests for the KG-RAG LLM-extraction backends. + +Both ``OpenAIExtraction`` and ``OllamaExtraction`` lazily import their +vendor SDK inside ``chat_json``. We inject fakes via ``sys.modules`` so the +tests never touch the network, and exercise every branch: + +* missing SDK (ImportError) +* missing API key (OpenAI only) +* happy path (valid JSON object) +* custom endpoint / auth header forwarding +* primary-model failure with fallback success / fallback failure / no fallback +* invalid JSON and non-dict JSON bodies + +``class_parameters`` is asserted to keep the plugin-UI schema honest. +""" + +from __future__ import annotations + +import sys +import types +from typing import Any + +import pytest + + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +def _make_openai_module(*, behaviors): + """Return a fake ``openai`` module whose ``OpenAI`` client replays + ``behaviors`` (a list of either a content string or an Exception) for + successive ``chat.completions.create`` calls.""" + calls: list[dict[str, Any]] = [] + + class _Message: + def __init__(self, content): + self.content = content + + class _Choice: + def __init__(self, content): + self.message = _Message(content) + + class _Response: + def __init__(self, content): + self.choices = [_Choice(content)] + + class _Completions: + def create(self, **kwargs): + calls.append(kwargs) + behavior = behaviors.pop(0) + if isinstance(behavior, Exception): + raise behavior + return _Response(behavior) + + class _Chat: + def __init__(self): + self.completions = _Completions() + + class FakeOpenAI: + last_init_kwargs: dict[str, Any] = {} + + def __init__(self, **kwargs): + FakeOpenAI.last_init_kwargs = kwargs + self.chat = _Chat() + + mod = types.ModuleType("openai") + mod.OpenAI = FakeOpenAI + mod._calls = calls + mod._FakeOpenAI = FakeOpenAI + return mod + + +def _make_ollama_module(*, behaviors): + """Fake ``ollama`` module; ``Client.chat`` replays ``behaviors`` (each a + dict response or an Exception).""" + calls: list[dict[str, Any]] = [] + + class FakeClient: + last_init_kwargs: dict[str, Any] = {} + + def __init__(self, **kwargs): + FakeClient.last_init_kwargs = kwargs + + def chat(self, **kwargs): + calls.append(kwargs) + behavior = behaviors.pop(0) + if isinstance(behavior, Exception): + raise behavior + return behavior + + mod = types.ModuleType("ollama") + mod.Client = FakeClient + mod._calls = calls + mod._FakeClient = FakeClient + return mod + + +@pytest.fixture +def patch_openai(monkeypatch): + def _install(behaviors): + mod = _make_openai_module(behaviors=list(behaviors)) + monkeypatch.setitem(sys.modules, "openai", mod) + return mod + + return _install + + +@pytest.fixture +def patch_ollama(monkeypatch): + def _install(behaviors): + mod = _make_ollama_module(behaviors=list(behaviors)) + monkeypatch.setitem(sys.modules, "ollama", mod) + return mod + + return _install + + +# --------------------------------------------------------------------------- +# OpenAIExtraction +# --------------------------------------------------------------------------- + + +def test_openai_missing_sdk_returns_empty(monkeypatch): + from plugins.llm_extraction.openai import OpenAIExtraction + + # Make `from openai import OpenAI` raise ImportError. + monkeypatch.setitem(sys.modules, "openai", None) + out = OpenAIExtraction(api_key="sk-x").chat_json(system="s", user="u") + assert out == {} + + +def test_openai_no_api_key_returns_empty(monkeypatch, patch_openai): + monkeypatch.delenv("KG_RAG_OPENAI_API_KEY", raising=False) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + patch_openai(['{"x": 1}']) + + from plugins.llm_extraction.openai import OpenAIExtraction + + out = OpenAIExtraction(api_key="").chat_json(system="s", user="u") + assert out == {} + + +def test_openai_uses_env_key_fallback(monkeypatch, patch_openai): + monkeypatch.delenv("KG_RAG_OPENAI_API_KEY", raising=False) + monkeypatch.setenv("OPENAI_API_KEY", "sk-env") + mod = patch_openai(['{"entities": []}']) + + from plugins.llm_extraction.openai import OpenAIExtraction + + out = OpenAIExtraction(api_key="").chat_json(system="s", user="u") + assert out == {"entities": []} + assert mod._FakeOpenAI.last_init_kwargs["api_key"] == "sk-env" + + +def test_openai_happy_path_and_endpoint(monkeypatch, patch_openai): + mod = patch_openai(['{"entities": ["a"], "relationships": []}']) + + from plugins.llm_extraction.openai import OpenAIExtraction + + out = OpenAIExtraction( + model="gpt-4o", api_key="sk-x", api_endpoint="https://proxy.example/v1/" + ).chat_json(system="s", user="u") + assert out == {"entities": ["a"], "relationships": []} + # base_url is forwarded with the trailing slash stripped. + assert mod._FakeOpenAI.last_init_kwargs["base_url"] == "https://proxy.example/v1" + # The requested model is used. + assert mod._calls[0]["model"] == "gpt-4o" + assert mod._calls[0]["response_format"] == {"type": "json_object"} + + +def test_openai_fallback_model_success(patch_openai): + mod = patch_openai([RuntimeError("json mode unsupported"), '{"ok": true}']) + + from plugins.llm_extraction.openai import OpenAIExtraction + + out = OpenAIExtraction(model="gpt-5-nano", api_key="sk-x").chat_json( + system="s", user="u", fallback_model="gpt-4o-mini" + ) + assert out == {"ok": True} + assert mod._calls[0]["model"] == "gpt-5-nano" + assert mod._calls[1]["model"] == "gpt-4o-mini" + + +def test_openai_fallback_model_also_fails(patch_openai): + patch_openai([RuntimeError("boom"), RuntimeError("boom2")]) + + from plugins.llm_extraction.openai import OpenAIExtraction + + out = OpenAIExtraction(model="m1", api_key="sk-x").chat_json( + system="s", user="u", fallback_model="m2" + ) + assert out == {} + + +def test_openai_failure_no_fallback(patch_openai): + patch_openai([RuntimeError("boom")]) + + from plugins.llm_extraction.openai import OpenAIExtraction + + out = OpenAIExtraction(model="m1", api_key="sk-x").chat_json(system="s", user="u") + assert out == {} + + +def test_openai_failure_fallback_equals_primary(patch_openai): + # fallback_model == primary -> treated as "no usable fallback". + patch_openai([RuntimeError("boom")]) + + from plugins.llm_extraction.openai import OpenAIExtraction + + out = OpenAIExtraction(model="m1", api_key="sk-x").chat_json( + system="s", user="u", fallback_model="m1" + ) + assert out == {} + + +def test_openai_invalid_json_returns_empty(patch_openai): + patch_openai(["this is not json"]) + + from plugins.llm_extraction.openai import OpenAIExtraction + + assert OpenAIExtraction(api_key="sk-x").chat_json(system="s", user="u") == {} + + +def test_openai_non_dict_json_returns_empty(patch_openai): + patch_openai(["[1, 2, 3]"]) + + from plugins.llm_extraction.openai import OpenAIExtraction + + assert OpenAIExtraction(api_key="sk-x").chat_json(system="s", user="u") == {} + + +def test_openai_null_content_defaults_to_empty_object(patch_openai): + # message.content is None -> falls back to "{}". + patch_openai([None]) + + from plugins.llm_extraction.openai import OpenAIExtraction + + assert OpenAIExtraction(api_key="sk-x").chat_json(system="s", user="u") == {} + + +def test_openai_class_parameters(): + from plugins.llm_extraction.openai import OpenAIExtraction + + params = OpenAIExtraction.class_parameters() + names = [p.name for p in params] + assert "model" in names and "api_endpoint" in names + + +# --------------------------------------------------------------------------- +# OllamaExtraction +# --------------------------------------------------------------------------- + + +def test_ollama_missing_sdk_returns_empty(monkeypatch): + monkeypatch.setitem(sys.modules, "ollama", None) + + from plugins.llm_extraction.ollama import OllamaExtraction + + assert OllamaExtraction().chat_json(system="s", user="u") == {} + + +def test_ollama_happy_path_default_host(monkeypatch, patch_ollama): + monkeypatch.delenv("OLLAMA_HOST", raising=False) + mod = patch_ollama([{"message": {"content": '{"entities": ["x"]}'}}]) + + from plugins.llm_extraction.ollama import OllamaExtraction + + out = OllamaExtraction(model="qwen2.5:7b").chat_json(system="s", user="u") + assert out == {"entities": ["x"]} + assert mod._FakeClient.last_init_kwargs["host"] == "http://localhost:11434" + # No auth header without an api_key. + assert mod._FakeClient.last_init_kwargs["headers"] == {} + assert mod._calls[0]["format"] == "json" + + +def test_ollama_env_host_and_auth_header(monkeypatch, patch_ollama): + monkeypatch.setenv("OLLAMA_HOST", "http://ollama-host:11434") + mod = patch_ollama([{"message": {"content": "{}"}}]) + + from plugins.llm_extraction.ollama import OllamaExtraction + + OllamaExtraction(api_key="secret-token").chat_json(system="s", user="u") + # api_endpoint empty -> env host wins. + assert mod._FakeClient.last_init_kwargs["host"] == "http://ollama-host:11434" + assert ( + mod._FakeClient.last_init_kwargs["headers"]["Authorization"] + == "Bearer secret-token" + ) + + +def test_ollama_explicit_endpoint_overrides_env(monkeypatch, patch_ollama): + monkeypatch.setenv("OLLAMA_HOST", "http://ignored:11434") + mod = patch_ollama([{"message": {"content": "{}"}}]) + + from plugins.llm_extraction.ollama import OllamaExtraction + + OllamaExtraction(api_endpoint="http://explicit:9999").chat_json( + system="s", user="u" + ) + assert mod._FakeClient.last_init_kwargs["host"] == "http://explicit:9999" + + +def test_ollama_fallback_success(patch_ollama): + mod = patch_ollama( + [RuntimeError("model missing"), {"message": {"content": '{"ok": 1}'}}] + ) + + from plugins.llm_extraction.ollama import OllamaExtraction + + out = OllamaExtraction(model="llama3.1:70b").chat_json( + system="s", user="u", fallback_model="llama3.2:3b" + ) + assert out == {"ok": 1} + assert mod._calls[1]["model"] == "llama3.2:3b" + + +def test_ollama_fallback_also_fails(patch_ollama): + patch_ollama([RuntimeError("a"), RuntimeError("b")]) + + from plugins.llm_extraction.ollama import OllamaExtraction + + assert ( + OllamaExtraction(model="m1").chat_json( + system="s", user="u", fallback_model="m2" + ) + == {} + ) + + +def test_ollama_failure_no_fallback(patch_ollama): + patch_ollama([RuntimeError("boom")]) + + from plugins.llm_extraction.ollama import OllamaExtraction + + assert OllamaExtraction(model="m1").chat_json(system="s", user="u") == {} + + +def test_ollama_invalid_json(patch_ollama): + patch_ollama([{"message": {"content": "not json"}}]) + + from plugins.llm_extraction.ollama import OllamaExtraction + + assert OllamaExtraction().chat_json(system="s", user="u") == {} + + +def test_ollama_non_dict_json(patch_ollama): + patch_ollama([{"message": {"content": "[1,2]"}}]) + + from plugins.llm_extraction.ollama import OllamaExtraction + + assert OllamaExtraction().chat_json(system="s", user="u") == {} + + +def test_ollama_missing_message_defaults(patch_ollama): + # response without "message" -> .get(...) default "{}" -> empty dict. + patch_ollama([{}]) + + from plugins.llm_extraction.ollama import OllamaExtraction + + assert OllamaExtraction().chat_json(system="s", user="u") == {} + + +def test_ollama_class_parameters(): + from plugins.llm_extraction.ollama import OllamaExtraction + + names = [p.name for p in OllamaExtraction.class_parameters()] + assert "model" in names and "api_endpoint" in names diff --git a/lamb-kb-server/tests/unit/test_migrations.py b/lamb-kb-server/tests/unit/test_migrations.py new file mode 100644 index 000000000..1a2be6c54 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_migrations.py @@ -0,0 +1,56 @@ +"""Alembic migration round-trip test. + +Runs the real migration chain (``upgrade head`` then ``downgrade base``) +against an isolated temp database via the Alembic CLI, so a broken or +non-reversible revision fails the suite. +""" + +from __future__ import annotations + +import os +import sqlite3 +import subprocess +import sys +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parents[2] / "backend" + + +def _alembic(args: list[str], data_dir: str) -> subprocess.CompletedProcess: + env = {**os.environ, "DATA_DIR": data_dir, "LAMB_API_TOKEN": "test"} + return subprocess.run( + [sys.executable, "-m", "alembic", "-c", "alembic.ini", *args], + cwd=_BACKEND, + env=env, + capture_output=True, + text=True, + ) + + +def _tables(db_path: Path) -> set[str]: + con = sqlite3.connect(db_path) + try: + rows = con.execute( + "SELECT name FROM sqlite_master WHERE type='table'" + ).fetchall() + finally: + con.close() + return {r[0] for r in rows} + + +def test_upgrade_then_downgrade_roundtrip(tmp_path): + data_dir = str(tmp_path) + db_path = tmp_path / "kb-server.db" + + up = _alembic(["upgrade", "head"], data_dir) + assert up.returncode == 0, up.stderr + + tables = _tables(db_path) + assert {"collections", "ingestion_jobs", "alembic_version"} <= tables + + down = _alembic(["downgrade", "base"], data_dir) + assert down.returncode == 0, down.stderr + + tables = _tables(db_path) + assert "collections" not in tables + assert "ingestion_jobs" not in tables diff --git a/lamb-kb-server/tests/unit/test_query_service.py b/lamb-kb-server/tests/unit/test_query_service.py new file mode 100644 index 000000000..5d3b29e5d --- /dev/null +++ b/lamb-kb-server/tests/unit/test_query_service.py @@ -0,0 +1,234 @@ +"""Unit tests for ``services.query_service``. + +Exercises both entry points (``query_collection`` and ``query_with_plugin``) +with fake registries, a fake vector backend, and a stub DB session — including +the KG-RAG auto-routing branch, its graceful degradation on failure, the +threshold filter, and the 404 / 503 error paths. +""" + +from __future__ import annotations + +import pytest +from fastapi import HTTPException + +import config as config_module +import plugins.kg_rag_query as kg_module +import services.query_service as qs +from plugins.base import QueryResult +from schemas.query import QueryRequest + + +class _FakeQuery: + def __init__(self, result): + self._result = result + + def filter(self, *a, **k): + return self + + def first(self): + return self._result + + +class _FakeDB: + def __init__(self, collection): + self._c = collection + + def query(self, _model): + return _FakeQuery(self._c) + + +class _Collection: + def __init__(self, **over): + self.id = "col-1" + self.embedding_vendor = "openai" + self.embedding_model = "text-embedding-3-small" + self.embedding_endpoint = "" + self.vector_db_backend = "chromadb" + self.backend_collection_id = "backend-col" + self.storage_path = "/tmp/store" + self.graph_enabled = False + self.__dict__.update(over) + + +class _FakeBackend: + def __init__(self, results): + self._results = results + + def query(self, **kwargs): + return self._results + + +@pytest.fixture +def patch_registries(monkeypatch): + """Patch EmbeddingRegistry.build + VectorDBRegistry.get. Returns a setter + for the backend the registry hands back (None -> 503).""" + monkeypatch.setattr( + qs.EmbeddingRegistry, "build", + classmethod(lambda cls, vendor, **kwargs: object()), + ) + + state = {"backend": None} + + monkeypatch.setattr( + qs.VectorDBRegistry, "get", + classmethod(lambda cls, name: state["backend"]), + ) + + def _set_backend(backend): + state["backend"] = backend + + return _set_backend + + +def _req(top_k=5): + return QueryRequest.model_validate( + { + "query_text": "what is alpha?", + "top_k": top_k, + "embedding_credentials": {"api_key": "sk-x", "api_endpoint": ""}, + } + ) + + +# --------------------------------------------------------------------------- +# query_collection +# --------------------------------------------------------------------------- + + +def test_query_collection_not_found(patch_registries): + with pytest.raises(HTTPException) as exc: + qs.query_collection(_FakeDB(None), "missing", _req()) + assert exc.value.status_code == 404 + + +def test_query_collection_backend_unavailable(patch_registries): + patch_registries(None) # registry returns no backend + with pytest.raises(HTTPException) as exc: + qs.query_collection(_FakeDB(_Collection()), "col-1", _req()) + assert exc.value.status_code == 503 + + +def test_query_collection_plain_vector(patch_registries): + patch_registries(_FakeBackend([QueryResult(text="hit", score=0.9, metadata={})])) + out = qs.query_collection(_FakeDB(_Collection()), "col-1", _req()) + assert len(out) == 1 + assert out[0].text == "hit" + + +def test_query_collection_graph_enabled_but_kg_disabled( + patch_registries, monkeypatch +): + monkeypatch.setattr(config_module, "KG_RAG_ENABLED", False) + patch_registries(_FakeBackend([QueryResult(text="hit", score=0.9, metadata={})])) + out = qs.query_collection( + _FakeDB(_Collection(graph_enabled=True)), "col-1", _req() + ) + # No augmentation applied -> baseline result preserved. + assert out[0].text == "hit" + + +def test_query_collection_kg_rag_augments(patch_registries, monkeypatch): + monkeypatch.setattr(config_module, "KG_RAG_ENABLED", True) + patch_registries(_FakeBackend([QueryResult(text="hit", score=0.9, metadata={})])) + + def fake_augment(self, **kwargs): + return [ + {"similarity": 0.95, "data": "augmented", "metadata": {"kg_rag": {}}}, + ] + + monkeypatch.setattr(kg_module.KGRAGQueryPlugin, "augment", fake_augment) + out = qs.query_collection( + _FakeDB(_Collection(graph_enabled=True)), "col-1", _req() + ) + assert out[0].text == "augmented" + assert out[0].score == 0.95 + + +def test_query_collection_kg_rag_failure_degrades(patch_registries, monkeypatch): + monkeypatch.setattr(config_module, "KG_RAG_ENABLED", True) + patch_registries(_FakeBackend([QueryResult(text="baseline", score=0.8, metadata={})])) + + def boom(self, **kwargs): + raise RuntimeError("augment failed") + + monkeypatch.setattr(kg_module.KGRAGQueryPlugin, "augment", boom) + out = qs.query_collection( + _FakeDB(_Collection(graph_enabled=True)), "col-1", _req() + ) + # Falls back to the vector baseline. + assert out[0].text == "baseline" + + +# --------------------------------------------------------------------------- +# query_with_plugin +# --------------------------------------------------------------------------- + + +def test_query_with_plugin_not_found(patch_registries): + with pytest.raises(HTTPException) as exc: + qs.query_with_plugin(db=_FakeDB(None), collection_id="x", query_text="q") + assert exc.value.status_code == 404 + + +def test_query_with_plugin_backend_unavailable(patch_registries): + patch_registries(None) + with pytest.raises(HTTPException) as exc: + qs.query_with_plugin( + db=_FakeDB(_Collection()), collection_id="col-1", query_text="q" + ) + assert exc.value.status_code == 503 + + +def test_query_with_plugin_simple_applies_threshold(patch_registries): + patch_registries( + _FakeBackend( + [ + QueryResult(text="keep", score=0.9, metadata={}), + QueryResult(text="drop", score=0.1, metadata={}), + ] + ) + ) + out = qs.query_with_plugin( + db=_FakeDB(_Collection()), + collection_id="col-1", + query_text="q", + plugin_params={"top_k": 5, "threshold": 0.5}, + ) + texts = [r["data"] for r in out["results"]] + assert texts == ["keep"] + assert out["query"] == "q" + assert out["top_k"] == 5 + assert "total_ms" in out["timing"] + + +def test_query_with_plugin_kg_rag_augments(patch_registries, monkeypatch): + patch_registries(_FakeBackend([QueryResult(text="seed", score=0.9, metadata={})])) + + def fake_augment(self, **kwargs): + return [{"similarity": 0.9, "data": "graph-added", "metadata": {}}] + + monkeypatch.setattr(kg_module.KGRAGQueryPlugin, "augment", fake_augment) + out = qs.query_with_plugin( + db=_FakeDB(_Collection()), + collection_id="col-1", + query_text="q", + plugin_name="kg_rag_query", + ) + assert out["results"][0]["data"] == "graph-added" + + +def test_query_with_plugin_kg_rag_failure_degrades(patch_registries, monkeypatch): + patch_registries(_FakeBackend([QueryResult(text="seed", score=0.9, metadata={})])) + + def boom(self, **kwargs): + raise RuntimeError("augment failed") + + monkeypatch.setattr(kg_module.KGRAGQueryPlugin, "augment", boom) + out = qs.query_with_plugin( + db=_FakeDB(_Collection()), + collection_id="col-1", + query_text="q", + plugin_name="kg_rag_query", + ) + # Degrades to the formatted baseline. + assert out["results"][0]["data"] == "seed" diff --git a/lamb-kb-server/tests/unit/test_routers_benchmarks.py b/lamb-kb-server/tests/unit/test_routers_benchmarks.py new file mode 100644 index 000000000..f7ba57cac --- /dev/null +++ b/lamb-kb-server/tests/unit/test_routers_benchmarks.py @@ -0,0 +1,108 @@ +"""Router tests for ``routers.benchmarks`` via FastAPI TestClient. + +Auth and DB session dependencies are overridden; the BenchmarkService layer +is monkeypatched so these tests stay at the routing/serialization layer. +""" + +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from database.connection import get_session +from dependencies import verify_token +from routers import benchmarks as benchmarks_router +from schemas.benchmark import ( + BenchmarkComparison, + BenchmarkDataset, + BenchmarkDatasetSummary, + BenchmarkMetrics, + BenchmarkRunAllResponse, + BenchmarkRunResponse, +) +from services.benchmark import BenchmarkService + + +@pytest.fixture +def client(): + app = FastAPI() + app.include_router(benchmarks_router.router) + app.dependency_overrides[verify_token] = lambda: "token" + app.dependency_overrides[get_session] = lambda: object() + return TestClient(app) + + +def _summary(ds_id="educational"): + return BenchmarkDatasetSummary( + id=ds_id, name="N", description="D", expected_behavior="x", + recommended_top_k=5, recommended_graph_depth=2, question_count=3, + ) + + +def test_list_datasets(client, monkeypatch): + monkeypatch.setattr(BenchmarkService, "list_datasets", classmethod( + lambda cls: [_summary()])) + resp = client.get("/benchmarks/datasets") + assert resp.status_code == 200 + assert resp.json()[0]["id"] == "educational" + + +def test_get_dataset(client, monkeypatch): + monkeypatch.setattr(BenchmarkService, "get_dataset", classmethod( + lambda cls, ds: BenchmarkDataset( + id=ds, name="N", description="D", expected_behavior="x", + recommended_top_k=5, recommended_graph_depth=2, question_count=0, + questions=[], + ))) + resp = client.get("/benchmarks/datasets/educational") + assert resp.status_code == 200 + assert resp.json()["id"] == "educational" + + +def _run_response(): + metrics = BenchmarkMetrics() + comparison = BenchmarkComparison( + delta_precision_at_k=0.0, delta_recall_at_k=0.0, delta_mrr=0.0, + graph_overhead_ms=0.0, expected_behavior="x", + ) + return BenchmarkRunResponse( + collection_id="c1", dataset_id="educational", dataset_name="N", + top_k=5, graph_depth=2, baseline=metrics, kg_rag=metrics, + comparison=comparison, results=[], + ) + + +def test_run_collection_benchmark(client, monkeypatch): + captured = {} + + def fake_run(cls, *, db, collection_id, request, embedding_credentials): + captured["creds"] = embedding_credentials + captured["collection_id"] = collection_id + return _run_response() + + monkeypatch.setattr(BenchmarkService, "run", classmethod(fake_run)) + resp = client.post( + "/benchmarks/collections/c1/run", + json={"dataset_id": "educational", + "embedding_credentials": {"api_key": "sk-x", "api_endpoint": ""}}, + ) + assert resp.status_code == 200 + assert resp.json()["collection_id"] == "c1" + assert captured["creds"] == {"api_key": "sk-x", "api_endpoint": ""} + + +def test_run_all_collection_benchmarks(client, monkeypatch): + def fake_run_all(cls, *, db, collection_id, dataset_ids, threshold, + embedding_credentials): + return BenchmarkRunAllResponse( + collection_id=collection_id, runs=[], summary={"datasets": 0}) + + monkeypatch.setattr(BenchmarkService, "run_all", classmethod(fake_run_all)) + resp = client.post( + "/benchmarks/collections/c1/run-all", + json={"dataset_ids": ["educational"], "threshold": 0.0, + "embedding_credentials": {"api_key": "sk-x", "api_endpoint": ""}}, + ) + assert resp.status_code == 200 + assert resp.json()["collection_id"] == "c1" diff --git a/lamb-kb-server/tests/unit/test_routers_graph.py b/lamb-kb-server/tests/unit/test_routers_graph.py new file mode 100644 index 000000000..545f61c1e --- /dev/null +++ b/lamb-kb-server/tests/unit/test_routers_graph.py @@ -0,0 +1,466 @@ +"""Router tests for ``routers.graph`` via FastAPI TestClient. + +Auth + DB-session dependencies are overridden; ``get_graph_store`` and the +service layer are monkeypatched so the tests exercise the routing, +validation, status-code mapping, and response shaping in isolation. +""" + +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import config as config_module +import routers.graph as graph_router +from database.connection import get_session +from dependencies import verify_token + + +class _Collection: + def __init__(self, **over): + self.id = "col-1" + self.organization_id = "org-1" + self.vector_db_backend = "chromadb" + self.embedding_vendor = "fake" + self.embedding_model = "fake-model" + self.embedding_endpoint = "" + self.backend_collection_id = "backend-col" + self.storage_path = "/tmp/store" + self.graph_enabled = False + self.__dict__.update(over) + + +class _FakeQuery: + def __init__(self, result): + self._result = result + + def filter(self, *a, **k): + return self + + def first(self): + return self._result + + +class _FakeDB: + def __init__(self, collection): + self._c = collection + self.committed = False + + def query(self, _model): + return _FakeQuery(self._c) + + def commit(self): + self.committed = True + + +class _FakeGraphStore: + def __init__(self, *, configured=True, available=True, **methods): + self._configured = configured + self._available = available + self._methods = methods + + def is_configured(self): + return self._configured + + def is_available(self): + return self._available + + def __getattr__(self, name): + # Return a stub that yields the preconfigured value for any graph op. + result = self.__dict__.get("_methods", {}).get(name, {}) + return lambda **kwargs: result + + +def _make_client(collection, store, monkeypatch, *, kg_enabled=True): + monkeypatch.setattr(graph_router, "get_graph_store", lambda: store) + monkeypatch.setattr( + config_module, "get_kg_rag_config", + lambda: {"enabled": kg_enabled, "index_on_ingest": True}, + ) + app = FastAPI() + app.include_router(graph_router.router) + app.dependency_overrides[verify_token] = lambda: "token" + db = _FakeDB(collection) + app.dependency_overrides[get_session] = lambda: db + client = TestClient(app) + client._db = db # expose for assertions + return client + + +# --------------------------------------------------------------------------- +# status +# --------------------------------------------------------------------------- + + +def test_status(monkeypatch): + client = _make_client(_Collection(), _FakeGraphStore(), monkeypatch) + resp = client.get("/graph/status") + assert resp.status_code == 200 + body = resp.json() + assert body["enabled"] is True + assert body["neo4j_configured"] is True + assert body["neo4j_available"] is True + + +def test_status_unavailable_when_connectivity_raises(monkeypatch): + class _Boom(_FakeGraphStore): + def is_available(self): + raise RuntimeError("down") + + client = _make_client(_Collection(), _Boom(), monkeypatch) + resp = client.get("/graph/status") + assert resp.json()["neo4j_available"] is False + + +# --------------------------------------------------------------------------- +# snapshot / changes / get_change +# --------------------------------------------------------------------------- + + +def test_snapshot_404_when_collection_missing(monkeypatch): + client = _make_client(None, _FakeGraphStore(), monkeypatch) + assert client.get("/graph/collections/col-1/snapshot").status_code == 404 + + +def test_snapshot_503_when_not_configured(monkeypatch): + client = _make_client( + _Collection(), _FakeGraphStore(configured=False), monkeypatch + ) + assert client.get("/graph/collections/col-1/snapshot").status_code == 503 + + +def test_snapshot_503_when_unavailable(monkeypatch): + client = _make_client( + _Collection(), _FakeGraphStore(available=False), monkeypatch + ) + assert client.get("/graph/collections/col-1/snapshot").status_code == 503 + + +def test_snapshot_success(monkeypatch): + snapshot = { + "collection_id": "col-1", "nodes": [], "edges": [], + "filters": {}, "counts": {}, + } + store = _FakeGraphStore(get_collection_graph=snapshot) + client = _make_client(_Collection(), store, monkeypatch) + resp = client.get("/graph/collections/col-1/snapshot?concept=alpha&limit=10") + assert resp.status_code == 200 + assert resp.json()["collection_id"] == "col-1" + + +def test_list_changes_success(monkeypatch): + changes = [{ + "event_id": "e1", "collection_id": "col-1", "org_id": "org-1", + "operation": "add_document", "concepts": [], + }] + store = _FakeGraphStore(list_changes=changes) + client = _make_client(_Collection(), store, monkeypatch) + resp = client.get("/graph/collections/col-1/changes?limit=5") + assert resp.status_code == 200 + assert resp.json()[0]["event_id"] == "e1" + + +def test_concept_changes_and_document_changes(monkeypatch): + store = _FakeGraphStore(list_changes=[]) + client = _make_client(_Collection(), store, monkeypatch) + assert client.get( + "/graph/collections/col-1/concepts/alpha/changes").status_code == 200 + assert client.get( + "/graph/collections/col-1/documents/doc1/changes").status_code == 200 + + +def test_get_change_success_and_404(monkeypatch): + change = { + "event_id": "e1", "collection_id": "col-1", "org_id": "org-1", + "operation": "add_document", "concepts": [], "chunk_ids": [], + } + client = _make_client( + _Collection(), _FakeGraphStore(get_change=change), monkeypatch + ) + assert client.get("/graph/collections/col-1/changes/e1").status_code == 200 + + client2 = _make_client( + _Collection(), _FakeGraphStore(get_change=None), monkeypatch + ) + assert client2.get("/graph/collections/col-1/changes/missing").status_code == 404 + + +# --------------------------------------------------------------------------- +# revert +# --------------------------------------------------------------------------- + + +def test_revert_success(monkeypatch): + result = {"reverted": True, "event_id": "e1", "chunk_ids": []} + client = _make_client( + _Collection(), _FakeGraphStore(revert_change=result), monkeypatch + ) + resp = client.post( + "/graph/collections/col-1/changes/e1/revert", json={"actor": "me"} + ) + assert resp.status_code == 200 + assert resp.json()["reverted"] is True + + +def test_revert_not_found(monkeypatch): + result = {"reverted": False, "reason": "change_not_found"} + client = _make_client( + _Collection(), _FakeGraphStore(revert_change=result), monkeypatch + ) + resp = client.post( + "/graph/collections/col-1/changes/x/revert", json={} + ) + assert resp.status_code == 404 + + +def test_revert_unsupported_400(monkeypatch): + result = {"reverted": False, "reason": "unsupported_operation"} + client = _make_client( + _Collection(), _FakeGraphStore(revert_change=result), monkeypatch + ) + resp = client.post( + "/graph/collections/col-1/changes/x/revert", json={} + ) + assert resp.status_code == 400 + + +# --------------------------------------------------------------------------- +# curation operations (rename / merge / curate / edit / curate-rel) +# --------------------------------------------------------------------------- + + +def test_rename_concept_ok_and_400(monkeypatch): + ok = {"ok": True, "operation": "manual_rename_concept", "event_id": "e"} + client = _make_client( + _Collection(), _FakeGraphStore(rename_concept=ok), monkeypatch + ) + resp = client.patch( + "/graph/collections/col-1/concepts/alpha/rename", + json={"new_name": "beta"}, + ) + assert resp.status_code == 200 + + bad = {"ok": False, "reason": "invalid_concept_name"} + client2 = _make_client( + _Collection(), _FakeGraphStore(rename_concept=bad), monkeypatch + ) + resp2 = client2.patch( + "/graph/collections/col-1/concepts/alpha/rename", + json={"new_name": "beta"}, + ) + assert resp2.status_code == 400 + + +def test_merge_concepts_ok_and_400(monkeypatch): + ok = {"ok": True, "operation": "manual_merge_concepts", "event_id": "e"} + client = _make_client( + _Collection(), _FakeGraphStore(merge_concepts=ok), monkeypatch + ) + resp = client.post( + "/graph/collections/col-1/concepts/merge", + json={"source_names": ["a"], "target_name": "b"}, + ) + assert resp.status_code == 200 + + bad = {"ok": False, "reason": "invalid_merge_request"} + client2 = _make_client( + _Collection(), _FakeGraphStore(merge_concepts=bad), monkeypatch + ) + resp2 = client2.post( + "/graph/collections/col-1/concepts/merge", + json={"source_names": ["a"], "target_name": "b"}, + ) + assert resp2.status_code == 400 + + +def test_curate_concept_ok_and_400(monkeypatch): + ok = {"ok": True, "operation": "manual_curate_concept", "event_id": "e"} + client = _make_client( + _Collection(), _FakeGraphStore(update_concept_curation=ok), monkeypatch + ) + resp = client.patch( + "/graph/collections/col-1/concepts/alpha/curation", + json={"notes": "n", "verification_state": "verified"}, + ) + assert resp.status_code == 200 + + bad = {"ok": False, "reason": "concept_not_found"} + client2 = _make_client( + _Collection(), _FakeGraphStore(update_concept_curation=bad), monkeypatch + ) + resp2 = client2.patch( + "/graph/collections/col-1/concepts/alpha/curation", + json={"verification_state": "verified"}, + ) + assert resp2.status_code == 400 + + +def test_edit_relationship_ok_and_400(monkeypatch): + ok = {"ok": True, "operation": "manual_edit_relationship", "event_id": "e"} + client = _make_client( + _Collection(), _FakeGraphStore(edit_relationship=ok), monkeypatch + ) + resp = client.patch( + "/graph/collections/col-1/relationships", + json={"source_concept": "a", "target_concept": "b", "relation": "uses", + "new_relation": "depends_on"}, + ) + assert resp.status_code == 200 + + bad = {"ok": False, "reason": "relationship_not_found"} + client2 = _make_client( + _Collection(), _FakeGraphStore(edit_relationship=bad), monkeypatch + ) + resp2 = client2.patch( + "/graph/collections/col-1/relationships", + json={"source_concept": "a", "target_concept": "b", "relation": "uses"}, + ) + assert resp2.status_code == 400 + + +def test_curate_relationship_ok_and_400(monkeypatch): + ok = {"ok": True, "operation": "manual_curate_relationship", "event_id": "e"} + client = _make_client( + _Collection(), _FakeGraphStore(edit_relationship=ok), monkeypatch + ) + resp = client.patch( + "/graph/collections/col-1/relationships/curation", + json={"source_concept": "a", "target_concept": "b", "relation": "uses", + "verification_state": "verified"}, + ) + assert resp.status_code == 200 + + bad = {"ok": False, "reason": "relationship_not_found"} + client2 = _make_client( + _Collection(), _FakeGraphStore(edit_relationship=bad), monkeypatch + ) + resp2 = client2.patch( + "/graph/collections/col-1/relationships/curation", + json={"source_concept": "a", "target_concept": "b", "relation": "uses"}, + ) + assert resp2.status_code == 400 + + +def test_status_not_configured_skips_availability(monkeypatch): + client = _make_client( + _Collection(), _FakeGraphStore(configured=False), monkeypatch + ) + body = client.get("/graph/status").json() + assert body["neo4j_configured"] is False + assert body["neo4j_available"] is False + + +# --------------------------------------------------------------------------- +# migrate +# --------------------------------------------------------------------------- + + +def test_migrate_disabled_503(monkeypatch): + client = _make_client( + _Collection(), _FakeGraphStore(), monkeypatch, kg_enabled=False + ) + resp = client.post("/graph/collections/col-1/migrate") + assert resp.status_code == 503 + + +def test_migrate_collection_404(monkeypatch): + client = _make_client(None, _FakeGraphStore(), monkeypatch) + assert client.post("/graph/collections/col-1/migrate").status_code == 404 + + +def test_migrate_backend_unavailable_503(monkeypatch): + import plugins.base as base_module + + monkeypatch.setattr(base_module.VectorDBRegistry, "get", + classmethod(lambda cls, name: None)) + client = _make_client(_Collection(), _FakeGraphStore(), monkeypatch) + assert client.post("/graph/collections/col-1/migrate").status_code == 503 + + +class _FakeBackend: + def __init__(self, batches=None, raises=None): + self._batches = batches or [] + self._raises = raises + + def iter_all_chunks(self, **kwargs): + if self._raises: + raise self._raises + yield from self._batches + + +def _patch_migrate_backend(monkeypatch, backend): + import plugins.base as base_module + + monkeypatch.setattr(base_module.VectorDBRegistry, "get", + classmethod(lambda cls, name: backend)) + monkeypatch.setattr(base_module.EmbeddingRegistry, "build", + classmethod(lambda cls, vendor, **kw: object())) + + +def test_migrate_not_implemented_400(monkeypatch): + backend = _FakeBackend(raises=NotImplementedError()) + _patch_migrate_backend(monkeypatch, backend) + client = _make_client(_Collection(), _FakeGraphStore(), monkeypatch) + assert client.post("/graph/collections/col-1/migrate").status_code == 400 + + +def test_migrate_backend_collection_missing_404(monkeypatch): + backend = _FakeBackend(raises=RuntimeError("missing")) + _patch_migrate_backend(monkeypatch, backend) + client = _make_client(_Collection(), _FakeGraphStore(), monkeypatch) + assert client.post("/graph/collections/col-1/migrate").status_code == 404 + + +def test_migrate_no_chunks_sets_graph_enabled(monkeypatch): + backend = _FakeBackend(batches=[]) # no batches -> no ids + _patch_migrate_backend(monkeypatch, backend) + client = _make_client(_Collection(), _FakeGraphStore(), monkeypatch) + resp = client.post("/graph/collections/col-1/migrate") + assert resp.status_code == 200 + body = resp.json() + assert body["graph_enabled"] is True + assert body["indexed"] is False + assert client._db.committed is True + + +def test_migrate_success(monkeypatch): + backend = _FakeBackend(batches=[(["c1"], ["chunk text"], [{"filename": "f.md"}])]) + _patch_migrate_backend(monkeypatch, backend) + import services.graph_indexing as gi + + monkeypatch.setattr( + gi, "index_chunks_for_collection", + lambda **kwargs: {"indexed": True, "chunks": 1, "entities": 2}, + ) + client = _make_client(_Collection(), _FakeGraphStore(), monkeypatch) + resp = client.post( + "/graph/collections/col-1/migrate", + headers={"X-OpenAI-Api-Key": "sk-x"}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "success" + assert body["chunks_seen"] == 1 + assert body["indexed"] is True + + +def test_migrate_skips_empty_text_chunks(monkeypatch): + # First text is empty -> skipped; only c2 is collected. + backend = _FakeBackend( + batches=[(["c1", "c2"], ["", "real text"], [{}, {"filename": "f.md"}])] + ) + _patch_migrate_backend(monkeypatch, backend) + import services.graph_indexing as gi + + captured = {} + + def fake_index(**kwargs): + captured["ids"] = kwargs["ids"] + return {"indexed": True, "chunks": len(kwargs["ids"])} + + monkeypatch.setattr(gi, "index_chunks_for_collection", fake_index) + client = _make_client(_Collection(), _FakeGraphStore(), monkeypatch) + resp = client.post("/graph/collections/col-1/migrate") + assert resp.status_code == 200 + assert captured["ids"] == ["c2"] diff --git a/lamb-kb-server/tests/unit/test_schemas_content_validation.py b/lamb-kb-server/tests/unit/test_schemas_content_validation.py new file mode 100644 index 000000000..f38b4c8c3 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_schemas_content_validation.py @@ -0,0 +1,44 @@ +"""Unit tests for ``schemas.content`` validators (extra_metadata guards).""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from schemas.content import AddContentRequest, DocumentInputPayload + + +def _doc(**over): + base = dict(source_item_id="i1", title="T", text="body") + base.update(over) + return base + + +def test_extra_metadata_accepts_primitives(): + doc = DocumentInputPayload(**_doc(extra_metadata={ + "s": "x", "i": 1, "f": 1.5, "b": True, + })) + assert doc.extra_metadata["i"] == 1 + + +def test_extra_metadata_rejects_none_value(): + # pydantic's typed-dict value union rejects None before the custom + # validator's defensive check is reached. + with pytest.raises(ValidationError): + DocumentInputPayload(**_doc(extra_metadata={"k": None})) + + +def test_extra_metadata_rejects_non_primitive(): + with pytest.raises(ValidationError): + DocumentInputPayload(**_doc(extra_metadata={"k": ["a", "b"]})) + + +def test_add_content_request_requires_documents(): + # Field(min_length=1) rejects an empty documents list. + with pytest.raises(ValidationError): + AddContentRequest(documents=[]) + + +def test_add_content_request_valid(): + req = AddContentRequest(documents=[DocumentInputPayload(**_doc())]) + assert len(req.documents) == 1 diff --git a/lamb-kb-server/tests/unit/test_schemas_graph.py b/lamb-kb-server/tests/unit/test_schemas_graph.py new file mode 100644 index 000000000..b89e984e6 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_schemas_graph.py @@ -0,0 +1,213 @@ +"""Unit tests for ``schemas.graph`` — the KG-RAG traceability/curation models. + +These are pure Pydantic models, but they were previously never imported by +any test, leaving the whole module uncovered. We instantiate every model in +both its minimal (required-only) and fully-populated forms, and assert the +documented defaults so a future field rename or default change is caught. +""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from schemas.graph import ( + GraphChangeDetail, + GraphChangeEvent, + GraphConceptCurationRequest, + GraphConceptMergeRequest, + GraphConceptRenameRequest, + GraphEdge, + GraphManualOperationResponse, + GraphNode, + GraphRelationshipCurationRequest, + GraphRelationshipEditRequest, + GraphRevertRequest, + GraphRevertResponse, + GraphSnapshotResponse, +) + + +def test_graph_change_event_minimal_defaults(): + ev = GraphChangeEvent( + event_id="e1", + collection_id="c1", + org_id="org1", + operation="add_document", + ) + assert ev.actor is None + assert ev.concepts == [] + assert ev.payload_json is None + assert ev.file_id is None + + +def test_graph_change_event_full_payload(): + ev = GraphChangeEvent( + event_id="e1", + collection_id="c1", + org_id="org1", + operation="add_document", + actor="user@example.com", + timestamp="2026-06-20T00:00:00Z", + filename="doc.pdf", + concepts=["alpha", "beta"], + payload_json='{"k": "v"}', + document_id="d1", + file_id=42, + ) + assert ev.concepts == ["alpha", "beta"] + assert ev.file_id == 42 + + +def test_graph_change_event_requires_core_fields(): + with pytest.raises(ValidationError): + GraphChangeEvent(event_id="e1") # missing collection_id/org_id/operation + + +def test_graph_change_detail_extends_event_with_chunk_ids(): + detail = GraphChangeDetail( + event_id="e1", + collection_id="c1", + org_id="org1", + operation="add_document", + chunk_ids=["chunk-1", "chunk-2"], + ) + # Inherited fields plus the new list. + assert detail.operation == "add_document" + assert detail.chunk_ids == ["chunk-1", "chunk-2"] + # Default still applies for the chunk list. + assert GraphChangeDetail( + event_id="e2", collection_id="c1", org_id="org1", operation="op" + ).chunk_ids == [] + + +def test_graph_revert_request_defaults(): + req = GraphRevertRequest() + assert req.actor == "graph-traceability-api" + assert req.reason == "" + custom = GraphRevertRequest(actor="me", reason="mistake") + assert custom.actor == "me" + + +def test_graph_revert_response(): + resp = GraphRevertResponse(reverted=True) + assert resp.reverted is True + assert resp.chunk_ids == [] + full = GraphRevertResponse( + reverted=False, + reason="not found", + event_id="e1", + revert_event_id="r1", + operation="add_document", + document_id="d1", + chunk_ids=["c1"], + ) + assert full.reason == "not found" + assert full.chunk_ids == ["c1"] + + +def test_graph_node_and_edge(): + node = GraphNode(id="n1", type="concept", label="Alpha") + assert node.data == {} + edge = GraphEdge(id="e1", type="related_to", source="n1", target="n2") + assert edge.label is None + assert edge.weight is None + assert edge.data == {} + weighted = GraphEdge( + id="e2", + type="related_to", + source="n1", + target="n2", + label="relates", + weight=0.75, + data={"evidence": "x"}, + ) + assert weighted.weight == 0.75 + + +def test_graph_snapshot_response_defaults(): + snap = GraphSnapshotResponse(collection_id="c1") + assert snap.nodes == [] + assert snap.edges == [] + assert snap.filters == {} + assert snap.counts == {} + populated = GraphSnapshotResponse( + collection_id="c1", + nodes=[GraphNode(id="n1", type="concept", label="A")], + edges=[GraphEdge(id="e1", type="rel", source="n1", target="n1")], + filters={"verification_state": "verified"}, + counts={"nodes": 1, "edges": 1}, + ) + assert populated.counts["nodes"] == 1 + + +def test_graph_manual_operation_response(): + resp = GraphManualOperationResponse(ok=True) + assert resp.operation is None + assert resp.details == {} + resp2 = GraphManualOperationResponse( + ok=False, operation="rename", event_id="e1", reason="blocked", details={"x": 1} + ) + assert resp2.details == {"x": 1} + + +def test_concept_rename_and_merge_requests(): + rename = GraphConceptRenameRequest(new_name="Beta") + assert rename.actor == "graph-curation-api" + assert rename.reason == "" + + merge = GraphConceptMergeRequest( + source_names=["a", "b"], target_name="c" + ) + assert merge.source_names == ["a", "b"] + assert merge.actor == "graph-curation-api" + + with pytest.raises(ValidationError): + GraphConceptMergeRequest(target_name="c") # missing source_names + + +def test_concept_curation_request_defaults(): + req = GraphConceptCurationRequest() + assert req.notes is None + assert req.tags is None + assert req.verification_state is None + assert req.actor == "graph-curation-api" + full = GraphConceptCurationRequest( + notes="n", tags=["t1"], verification_state="verified", reason="cleanup" + ) + assert full.tags == ["t1"] + + +def test_relationship_edit_request(): + req = GraphRelationshipEditRequest( + source_concept="a", target_concept="b", relation="related_to" + ) + assert req.new_relation is None + assert req.weight is None + assert req.actor == "graph-curation-api" + full = GraphRelationshipEditRequest( + source_concept="a", + target_concept="b", + relation="related_to", + new_relation="depends_on", + weight=0.5, + description="d", + evidence="e", + notes="n", + tags=["t"], + verification_state="verified", + ) + assert full.new_relation == "depends_on" + assert full.weight == 0.5 + + with pytest.raises(ValidationError): + GraphRelationshipEditRequest(source_concept="a") # missing target/relation + + +def test_relationship_curation_request(): + req = GraphRelationshipCurationRequest( + source_concept="a", target_concept="b", relation="related_to" + ) + assert req.notes is None + assert req.actor == "graph-curation-api" + assert req.reason == "" diff --git a/lamb-kb-server/tests/unit/test_vector_db_chromadb_lookup.py b/lamb-kb-server/tests/unit/test_vector_db_chromadb_lookup.py new file mode 100644 index 000000000..4ec153496 --- /dev/null +++ b/lamb-kb-server/tests/unit/test_vector_db_chromadb_lookup.py @@ -0,0 +1,179 @@ +"""Functional tests for ``ChromaDBBackend`` lookup / query / iteration methods. + +Uses a real ``PersistentClient`` over a temp dir with the deterministic +``FakeEmbedding`` — covering ``query`` (incl. parent-text rewrite), +``get_chunks_by_id``, ``get_chunks_by_source``, ``delete_by_source``, +``iter_all_chunks``, and the embedding-adapter / add_chunks error paths. +""" + +from __future__ import annotations + +import pytest + +from plugins.base import Chunk +from plugins.vector_db.chromadb_backend import ChromaDBBackend + + +@pytest.fixture +def populated(tmp_storage, fake_embedding): + """A collection with three chunks (two sharing a source_item_id).""" + be = ChromaDBBackend() + cid = "kb_lookup_test" + be.create_collection( + collection_id=cid, storage_path=tmp_storage, embedding_function=fake_embedding + ) + chunks = [ + Chunk(text="alpha doc one", metadata={ + "source_item_id": "src-1", "parent_text": "PARENT ONE"}), + Chunk(text="alpha doc two", metadata={"source_item_id": "src-1"}), + Chunk(text="beta document", metadata={"source_item_id": "src-2"}), + ] + be.add_chunks( + collection_id=cid, storage_path=tmp_storage, chunks=chunks, + embedding_function=fake_embedding, + ) + return be, cid, tmp_storage, fake_embedding + + +def test_query_returns_results_with_chunk_id_and_parent_text(populated): + be, cid, path, ef = populated + results = be.query( + collection_id=cid, storage_path=path, query_text="alpha doc one", + top_k=3, embedding_function=ef, + ) + assert results + # The backend surfaces the internal chunk id on every result. + assert results[0].metadata.get("chunk_id") + # The chunk stored with parent_text returns the parent as its text. + parent_hit = [r for r in results if r.text == "PARENT ONE"] + assert parent_hit, "parent_text should be surfaced as the result text" + + +def test_get_chunks_by_source(populated): + be, cid, path, ef = populated + rows = be.get_chunks_by_source( + collection_id=cid, storage_path=path, source_item_id="src-1", + embedding_function=ef, + ) + assert len(rows) == 2 + assert all(r.score == 0.72 for r in rows) + assert all(r.metadata.get("chunk_id") for r in rows) + # Unknown source -> empty. + assert be.get_chunks_by_source( + collection_id=cid, storage_path=path, source_item_id="nope", + embedding_function=ef, + ) == [] + + +def test_get_chunks_by_id_roundtrip_and_parent(populated): + be, cid, path, ef = populated + # First grab real chunk ids via a source lookup. + src_rows = be.get_chunks_by_source( + collection_id=cid, storage_path=path, source_item_id="src-1", + embedding_function=ef, + ) + ids = [r.metadata["chunk_id"] for r in src_rows] + out = be.get_chunks_by_id( + collection_id=cid, storage_path=path, chunk_ids=ids, embedding_function=ef, + ) + assert len(out) == 2 + assert all(r.score == 0.72 for r in out) + # document_id is defaulted to the chunk id. + assert all(r.metadata.get("document_id") for r in out) + # The chunk that had parent_text returns it as text (popped from metadata). + assert any(r.text == "PARENT ONE" for r in out) + + +def test_get_chunks_by_id_empty_and_error(populated, tmp_storage, fake_embedding): + be, cid, path, ef = populated + assert be.get_chunks_by_id( + collection_id=cid, storage_path=path, chunk_ids=[], embedding_function=ef, + ) == [] + # Unknown collection -> the except branch returns []. + assert be.get_chunks_by_id( + collection_id="does-not-exist", storage_path=path, chunk_ids=["x"], + embedding_function=ef, + ) == [] + + +def test_get_chunks_by_source_error(populated): + be, _cid, path, ef = populated + assert be.get_chunks_by_source( + collection_id="missing", storage_path=path, source_item_id="src-1", + embedding_function=ef, + ) == [] + + +def test_delete_by_source_counts(populated): + be, cid, path, ef = populated + deleted = be.delete_by_source( + collection_id=cid, storage_path=path, source_item_id="src-1", + ) + assert deleted == 2 + # Deleting again finds nothing. + assert be.delete_by_source( + collection_id=cid, storage_path=path, source_item_id="src-1", + ) == 0 + + +def test_iter_all_chunks_paginates(populated): + be, cid, path, ef = populated + batches = list(be.iter_all_chunks( + collection_id=cid, storage_path=path, embedding_function=ef, batch_size=2, + )) + total = sum(len(ids) for ids, _texts, _metas in batches) + assert total == 3 + # batch_size=2 over 3 chunks -> at least two batches. + assert len(batches) >= 2 + + +def test_delete_absent_collection_when_dir_exists(tmp_storage, fake_embedding): + # Create one collection so the client + dir exist, then delete a DIFFERENT + # (non-existent) collection name -> the fast-path guard is bypassed and the + # NotFoundError "already absent" branch is exercised. + be = ChromaDBBackend() + be.create_collection( + collection_id="kb_present", storage_path=tmp_storage, + embedding_function=fake_embedding, + ) + # Must not raise. + be.delete_collection(collection_id="kb_absent", storage_path=tmp_storage) + + +def test_add_chunks_empty_is_noop(tmp_storage, fake_embedding): + be = ChromaDBBackend() + cid = "kb_empty" + be.create_collection( + collection_id=cid, storage_path=tmp_storage, embedding_function=fake_embedding + ) + assert be.add_chunks( + collection_id=cid, storage_path=tmp_storage, chunks=[], + embedding_function=fake_embedding, + ) == 0 + + +def test_add_chunks_embedding_error_wrapped(tmp_storage): + """An embedding function that raises surfaces as a clean RuntimeError.""" + from plugins.base import EmbeddingFunction + + class _BoomEmbedding(EmbeddingFunction): + name = "boom" + + def __init__(self): + super().__init__(model="boom") + + def __call__(self, input): + raise ValueError("embedding backend exploded") + + be = ChromaDBBackend() + cid = "kb_boom" + boom = _BoomEmbedding() + be.create_collection( + collection_id=cid, storage_path=tmp_storage, embedding_function=boom + ) + with pytest.raises(RuntimeError): + be.add_chunks( + collection_id=cid, storage_path=tmp_storage, + chunks=[Chunk(text="x", metadata={"source_item_id": "s"})], + embedding_function=boom, + ) diff --git a/library-manager/.gitignore b/library-manager/.gitignore index 0d58eae0b..feab5ac52 100644 --- a/library-manager/.gitignore +++ b/library-manager/.gitignore @@ -30,4 +30,5 @@ Thumbs.db .pytest_cache/ htmlcov/ .coverage -tests/.yt_cache/ +# Note: tests/.yt_cache/ cassettes ARE committed so the YouTube tests run +# offline on a fresh clone. Do not ignore them. diff --git a/library-manager/README.md b/library-manager/README.md index df1839c91..9d0d9ecf5 100644 --- a/library-manager/README.md +++ b/library-manager/README.md @@ -97,11 +97,24 @@ LAMB_API_TOKEN=your-token uvicorn main:app --host 0.0.0.0 --port 9091 --reload ### Running tests +The suite is organised into three tiers — **unit** (direct module calls), +**integration** (in-process ASGI against the real DB + worker), and **e2e** +(real HTTP to a `uvicorn` subprocess, plus a Docker smoke test). See +`tests/README.md` for the full tier contract. + ```bash -pytest tests/ -v +./scripts/run_tests.sh # all tiers + the 95% line+branch coverage gate +pytest tests/unit/ -q # a single tier +pytest -m "not slow" -q # skip subprocess/Docker tests +pytest tests/unit/ tests/integration/ --cov=backend --cov-branch \ + --cov-report=term-missing --cov-fail-under=95 ``` -Tests use an in-process ASGI client (httpx + FastAPI) with a temporary SQLite database. No external services are required — the YouTube test uses a live video with reliable subtitles, and the URL import test verifies the error path since Firecrawl is not available in the test environment. +No external services are required. External SDKs (markitdown, Firecrawl, +PyMuPDF, OpenAI) are mocked at their import boundary so the plugins' own logic +still runs; YouTube transcripts use an offline cassette cache in +`tests/.yt_cache/`. The e2e Docker test skips cleanly when no Docker daemon is +available. Combined line+branch coverage of `backend/` is enforced at ≥95%. ### Docker @@ -151,4 +164,12 @@ SQLite with WAL mode (`data/library-manager.db`). Tables: | `content_images` | Extracted images linked to content items | | `import_jobs` | Persistent job queue for async processing | -Tables are created automatically on first startup via SQLAlchemy `create_all`. +The schema is managed with Alembic. Migrations live in `backend/migrations/` and +are run automatically up to `head` at startup (`init_db` → `_run_migrations`). +Databases created before Alembic adoption (tables present, no `alembic_version`) +are stamped to the baseline revision before upgrading, so existing data is +preserved. To create a new migration after changing the models: + +```bash +cd backend && alembic -c alembic.ini revision --autogenerate -m "describe change" +``` diff --git a/library-manager/backend/alembic.ini b/library-manager/backend/alembic.ini new file mode 100644 index 000000000..aaedeb5ed --- /dev/null +++ b/library-manager/backend/alembic.ini @@ -0,0 +1,48 @@ +# Alembic configuration for the LAMB Library Manager. +# +# The database URL is NOT set here — env.py derives it from config.DB_PATH so +# the DATA_DIR env override (used by Docker and tests) stays authoritative. + +[alembic] +# Migration scripts live in backend/migrations (deliberately NOT named +# "alembic" so the directory cannot shadow the installed alembic package on +# sys.path). +script_location = migrations +prepend_sys_path = . + +# Use a stable, sortable file naming scheme. +file_template = %%(year)d%%(month).2d%%(day).2d_%%(rev)s_%%(slug)s + +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARNING +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARNING +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/library-manager/backend/database/connection.py b/library-manager/backend/database/connection.py index 444d6369d..cb0f616e0 100644 --- a/library-manager/backend/database/connection.py +++ b/library-manager/backend/database/connection.py @@ -1,20 +1,20 @@ """Database connection management. Provides a single engine and session factory for the Library Manager's -SQLite database. All tables are created on first call to ``init_db``. +SQLite database. The schema is brought to ``head`` with Alembic on the first +call to ``init_db`` (see ``_run_migrations``). """ import logging from collections.abc import Generator +from pathlib import Path from config import DB_PATH -from sqlalchemy import create_engine, event +from sqlalchemy import create_engine, event, inspect from sqlalchemy.engine import Engine from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.pool import NullPool -from database.models import Base - logger = logging.getLogger(__name__) _engine: Engine | None = None @@ -82,27 +82,39 @@ def init_db() -> None: event.listen(_engine, "connect", _enable_sqlite_wal) - Base.metadata.create_all(bind=_engine) - _apply_lightweight_migrations(_engine) + _run_migrations() _SessionLocal = sessionmaker(bind=_engine, expire_on_commit=False) logger.info("Database initialized at %s", DB_PATH) -def _apply_lightweight_migrations(engine: Engine) -> None: - """Idempotently add additive schema deltas not handled by ``create_all``. +def _run_migrations() -> None: + """Bring the database schema to ``head`` using Alembic. - ``Base.metadata.create_all`` creates missing tables but never adds new - columns to existing ones. Every new column we ship lives here, gated by - a ``PRAGMA table_info`` check so re-running is a no-op. This runs on - every startup, matching the project's existing on-boot schema setup. + Databases created by the historical ``create_all`` + ad-hoc-ALTER path + already have the tables but no ``alembic_version`` row. Those are stamped + to the baseline revision first so the baseline migration is not re-applied + on top of the existing schema; any later revisions then run normally. """ - with engine.begin() as conn: - cols = {row[1] for row in conn.exec_driver_sql("PRAGMA table_info(content_items)")} - if "folder_id" not in cols: - conn.exec_driver_sql("ALTER TABLE content_items ADD COLUMN folder_id TEXT") - logger.info("Schema migration: added content_items.folder_id") + from alembic import command # noqa: PLC0415 + from alembic.config import Config # noqa: PLC0415 + from alembic.script import ScriptDirectory # noqa: PLC0415 + + backend_dir = Path(__file__).resolve().parents[1] + cfg = Config(str(backend_dir / "alembic.ini")) + cfg.set_main_option("script_location", str(backend_dir / "migrations")) + + inspector = inspect(_engine) + has_schema = inspector.has_table("content_items") + has_version = inspector.has_table("alembic_version") + if has_schema and not has_version: + base_rev = ScriptDirectory.from_config(cfg).get_base() + command.stamp(cfg, base_rev) + logger.info("Stamped pre-Alembic database to baseline revision %s", base_rev) + + command.upgrade(cfg, "head") + logger.info("Database schema migrated to head") def get_session() -> Generator[Session, None, None]: diff --git a/library-manager/backend/dependencies.py b/library-manager/backend/dependencies.py index 3b530ccf6..674d475d8 100644 --- a/library-manager/backend/dependencies.py +++ b/library-manager/backend/dependencies.py @@ -21,6 +21,11 @@ async def verify_token( Uses ``hmac.compare_digest`` for timing-safe comparison to prevent timing attacks that could reveal the token character by character. + Operands are encoded to bytes first because ``hmac.compare_digest`` + raises ``TypeError`` on ``str`` inputs containing non-ASCII characters; + comparing bytes keeps the comparison timing-safe while ensuring any + invalid token (including non-ASCII) is rejected with 401 rather than + crashing with a 500. Args: credentials: The Authorization header parsed by HTTPBearer. @@ -31,7 +36,9 @@ async def verify_token( Raises: HTTPException: 401 if token is missing or invalid. """ - if not hmac.compare_digest(credentials.credentials, LAMB_API_TOKEN): + if not hmac.compare_digest( + credentials.credentials.encode("utf-8"), LAMB_API_TOKEN.encode("utf-8") + ): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid service token.", diff --git a/library-manager/backend/migrations/env.py b/library-manager/backend/migrations/env.py new file mode 100644 index 000000000..c181aa49a --- /dev/null +++ b/library-manager/backend/migrations/env.py @@ -0,0 +1,83 @@ +"""Alembic migration environment for the LAMB Library Manager. + +Derives the database URL from ``config.DB_PATH`` (so the ``DATA_DIR`` env +override stays authoritative), targets ``database.models.Base.metadata`` for +autogenerate, and enables SQLite batch mode + WAL/foreign-key pragmas so +migrations behave identically to the running application. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +from alembic import context +from sqlalchemy import create_engine, event, pool + +# Make the backend package importable when Alembic runs from the CLI. +_BACKEND_DIR = Path(__file__).resolve().parents[1] +if str(_BACKEND_DIR) not in sys.path: + sys.path.insert(0, str(_BACKEND_DIR)) + +from config import DB_PATH # noqa: E402 +from database.models import Base # noqa: E402 + +config = context.config + +# NB: deliberately not calling logging.config.fileConfig here — migrations run +# in-process at startup (see database.connection._run_migrations) and we must +# not reconfigure the application's logging. + +target_metadata = Base.metadata + +_DB_URL = f"sqlite:///{DB_PATH}" + + +def _enable_sqlite_pragmas(dbapi_conn, _record) -> None: # noqa: ANN001 + """Match the application's per-connection SQLite pragmas.""" + cursor = dbapi_conn.cursor() + cursor.execute("PRAGMA journal_mode=WAL") + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode (emit SQL without a DBAPI connection).""" + context.configure( + url=_DB_URL, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + render_as_batch=True, + ) + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations in 'online' mode against a live SQLite connection.""" + engine = create_engine( + _DB_URL, + poolclass=pool.NullPool, + connect_args={"check_same_thread": False}, + ) + event.listen(engine, "connect", _enable_sqlite_pragmas) + + with engine.connect() as connection: + context.configure( + connection=connection, + target_metadata=target_metadata, + # SQLite cannot ALTER most things in place; batch mode rebuilds + # tables via copy-and-swap so future migrations work. + render_as_batch=True, + ) + with context.begin_transaction(): + context.run_migrations() + + engine.dispose() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/library-manager/backend/migrations/script.py.mako b/library-manager/backend/migrations/script.py.mako new file mode 100644 index 000000000..5716ac760 --- /dev/null +++ b/library-manager/backend/migrations/script.py.mako @@ -0,0 +1,25 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} +""" +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: str | None = ${repr(down_revision)} +branch_labels: str | Sequence[str] | None = ${repr(branch_labels)} +depends_on: str | Sequence[str] | None = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/library-manager/backend/migrations/versions/20260624_77979a6f7fd8_baseline_schema.py b/library-manager/backend/migrations/versions/20260624_77979a6f7fd8_baseline_schema.py new file mode 100644 index 000000000..d4572a8ca --- /dev/null +++ b/library-manager/backend/migrations/versions/20260624_77979a6f7fd8_baseline_schema.py @@ -0,0 +1,158 @@ +"""baseline schema + +Revision ID: 77979a6f7fd8 +Revises: +Create Date: 2026-06-24 12:43:13.742764 +""" +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '77979a6f7fd8' +down_revision: str | None = None +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('import_jobs', + sa.Column('id', sa.String(), nullable=False), + sa.Column('content_item_id', sa.String(), nullable=False), + sa.Column('library_id', sa.String(), nullable=False), + sa.Column('organization_id', sa.String(), nullable=False), + sa.Column('source_type', sa.String(), nullable=False), + sa.Column('plugin_name', sa.String(), nullable=False), + sa.Column('plugin_params', sa.Text(), nullable=True), + sa.Column('source_path', sa.String(), nullable=True), + sa.Column('source_url', sa.String(), nullable=True), + sa.Column('title', sa.String(), nullable=False), + sa.Column('status', sa.String(), nullable=False), + sa.Column('error_message', sa.Text(), nullable=True), + sa.Column('attempts', sa.Integer(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.Column('updated_at', sa.DateTime(), nullable=False), + sa.Column('started_at', sa.DateTime(), nullable=True), + sa.Column('completed_at', sa.DateTime(), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + with op.batch_alter_table('import_jobs', schema=None) as batch_op: + batch_op.create_index('idx_import_jobs_item', ['content_item_id'], unique=False) + batch_op.create_index('idx_import_jobs_status', ['status'], unique=False) + batch_op.create_index('idx_import_jobs_status_created', ['status', 'created_at'], unique=False) + + op.create_table('organizations', + sa.Column('id', sa.String(), nullable=False), + sa.Column('name', sa.String(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint('id') + ) + op.create_table('libraries', + sa.Column('id', sa.String(), nullable=False), + sa.Column('organization_id', sa.String(), nullable=False), + sa.Column('name', sa.String(), nullable=False), + sa.Column('import_config', sa.Text(), nullable=True), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint(['organization_id'], ['organizations.id'], ondelete='CASCADE'), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('organization_id', 'name', name='uq_library_org_name') + ) + op.create_table('content_folders', + sa.Column('id', sa.String(), nullable=False), + sa.Column('library_id', sa.String(), nullable=False), + sa.Column('parent_folder_id', sa.String(), nullable=True), + sa.Column('name', sa.String(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.Column('updated_at', sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint(['library_id'], ['libraries.id'], ondelete='CASCADE'), + sa.ForeignKeyConstraint(['parent_folder_id'], ['content_folders.id'], ondelete='CASCADE'), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('library_id', 'parent_folder_id', 'name', name='uq_folder_sibling_name') + ) + with op.batch_alter_table('content_folders', schema=None) as batch_op: + batch_op.create_index('idx_content_folders_library', ['library_id'], unique=False) + batch_op.create_index('idx_content_folders_parent', ['parent_folder_id'], unique=False) + + op.create_table('content_items', + sa.Column('id', sa.String(), nullable=False), + sa.Column('library_id', sa.String(), nullable=False), + sa.Column('organization_id', sa.String(), nullable=False), + sa.Column('folder_id', sa.String(), nullable=True), + sa.Column('title', sa.String(), nullable=False), + sa.Column('source_type', sa.String(), nullable=False), + sa.Column('original_filename', sa.String(), nullable=True), + sa.Column('content_type', sa.String(), nullable=True), + sa.Column('file_size', sa.Integer(), nullable=True), + sa.Column('source_url', sa.String(), nullable=True), + sa.Column('base_path', sa.String(), nullable=False), + sa.Column('original_path', sa.String(), nullable=True), + sa.Column('full_markdown_path', sa.String(), nullable=True), + sa.Column('page_count', sa.Integer(), nullable=False), + sa.Column('image_count', sa.Integer(), nullable=False), + sa.Column('source_ref', sa.Text(), nullable=True), + sa.Column('permalink_base', sa.String(), nullable=False), + sa.Column('metadata', sa.Text(), nullable=True), + sa.Column('import_plugin', sa.String(), nullable=False), + sa.Column('import_params', sa.Text(), nullable=True), + sa.Column('status', sa.String(), nullable=False), + sa.Column('error_message', sa.Text(), nullable=True), + sa.Column('processing_stats', sa.Text(), nullable=True), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.Column('updated_at', sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint(['folder_id'], ['content_folders.id'], ondelete='SET NULL'), + sa.ForeignKeyConstraint(['library_id'], ['libraries.id'], ondelete='CASCADE'), + sa.ForeignKeyConstraint(['organization_id'], ['organizations.id'], ondelete='CASCADE'), + sa.PrimaryKeyConstraint('id') + ) + with op.batch_alter_table('content_items', schema=None) as batch_op: + batch_op.create_index('idx_content_items_folder', ['folder_id'], unique=False) + batch_op.create_index('idx_content_items_library', ['library_id'], unique=False) + batch_op.create_index('idx_content_items_org', ['organization_id'], unique=False) + batch_op.create_index('idx_content_items_status', ['status'], unique=False) + + op.create_table('content_images', + sa.Column('id', sa.String(), nullable=False), + sa.Column('content_item_id', sa.String(), nullable=False), + sa.Column('image_path', sa.String(), nullable=False), + sa.Column('llm_description', sa.Text(), nullable=True), + sa.Column('page_number', sa.Integer(), nullable=True), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint(['content_item_id'], ['content_items.id'], ondelete='CASCADE'), + sa.PrimaryKeyConstraint('id') + ) + with op.batch_alter_table('content_images', schema=None) as batch_op: + batch_op.create_index('idx_content_images_item', ['content_item_id'], unique=False) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('content_images', schema=None) as batch_op: + batch_op.drop_index('idx_content_images_item') + + op.drop_table('content_images') + with op.batch_alter_table('content_items', schema=None) as batch_op: + batch_op.drop_index('idx_content_items_status') + batch_op.drop_index('idx_content_items_org') + batch_op.drop_index('idx_content_items_library') + batch_op.drop_index('idx_content_items_folder') + + op.drop_table('content_items') + with op.batch_alter_table('content_folders', schema=None) as batch_op: + batch_op.drop_index('idx_content_folders_parent') + batch_op.drop_index('idx_content_folders_library') + + op.drop_table('content_folders') + op.drop_table('libraries') + op.drop_table('organizations') + with op.batch_alter_table('import_jobs', schema=None) as batch_op: + batch_op.drop_index('idx_import_jobs_status_created') + batch_op.drop_index('idx_import_jobs_status') + batch_op.drop_index('idx_import_jobs_item') + + op.drop_table('import_jobs') + # ### end Alembic commands ### diff --git a/library-manager/backend/plugins/youtube_transcript_import.py b/library-manager/backend/plugins/youtube_transcript_import.py index 018b35c89..14ca223b6 100644 --- a/library-manager/backend/plugins/youtube_transcript_import.py +++ b/library-manager/backend/plugins/youtube_transcript_import.py @@ -311,7 +311,7 @@ def _download_and_parse( # file in the temp dir. try: files = os.listdir(tmp) - except OSError: + except OSError: # pragma: no cover - defensive (just-written tmp dir) return [] sub_files = [f for f in files if f.endswith((".srt", ".vtt"))] if not sub_files: @@ -323,7 +323,7 @@ def _download_and_parse( try: with open(path, encoding="utf-8") as fh: content = fh.read() - except OSError: + except OSError: # pragma: no cover - defensive (just-written subtitle file) return [] # _parse_srt_content already handles SRT and VTT alike (VTT is a diff --git a/library-manager/backend/routers/capabilities.py b/library-manager/backend/routers/capabilities.py index 06872b49e..a59b38d7b 100644 --- a/library-manager/backend/routers/capabilities.py +++ b/library-manager/backend/routers/capabilities.py @@ -220,6 +220,6 @@ async def get_item_capability_content( if payload.mime == "application/json": return JSONResponse(content=payload.body) - if isinstance(payload.body, bytes): + if isinstance(payload.body, bytes): # pragma: no cover - no built-in handler emits bytes return Response(content=payload.body, media_type=payload.mime) return Response(content=str(payload.body), media_type=payload.mime) diff --git a/library-manager/backend/services/folder_service.py b/library-manager/backend/services/folder_service.py index b7801cc0c..d8c139144 100644 --- a/library-manager/backend/services/folder_service.py +++ b/library-manager/backend/services/folder_service.py @@ -272,10 +272,15 @@ def delete_folder(db: Session, folder_id: str) -> str: .all() ) for sub in subfolders: - sub.parent_folder_id = new_parent_id + # Compute the deduped name against the destination parent BEFORE + # reassigning parent_folder_id. _next_available_name issues queries + # that trigger an autoflush; if the row were already reparented with + # its still-colliding name, that flush would violate + # uq_folder_sibling_name when the parent is non-NULL. sub.name = _next_available_name( db, folder.library_id, new_parent_id, sub.name, exclude_id=sub.id ) + sub.parent_folder_id = new_parent_id db.delete(folder) db.commit() @@ -385,14 +390,13 @@ def _is_descendant(db: Session, ancestor_id: str, candidate_id: str) -> bool: cursor: str | None = candidate_id seen: set[str] = set() while cursor is not None: - if cursor in seen: - # Defensive: data corruption shouldn't loop us forever + if cursor in seen: # pragma: no cover - defensive cycle guard (corrupt data) return False seen.add(cursor) if cursor == ancestor_id: return True parent = get_folder(db, cursor) - if parent is None: + if parent is None: # pragma: no cover - defensive (orphaned parent ref) return False cursor = parent.parent_folder_id return False diff --git a/library-manager/pyproject.toml b/library-manager/pyproject.toml index 4fcb68490..7448d67f5 100644 --- a/library-manager/pyproject.toml +++ b/library-manager/pyproject.toml @@ -11,6 +11,8 @@ dependencies = [ # Database "sqlalchemy>=2.0.30", + # Schema migrations (run at startup via init_db). + "alembic>=1.13", # Markdown processing "markdown2>=2.4.0", @@ -46,6 +48,7 @@ all = [ dev = [ "pytest>=8.0", "pytest-asyncio>=0.23", + "pytest-cov>=5.0", "httpx>=0.27.0", "ruff>=0.4.0", ] @@ -68,6 +71,38 @@ select = ["E", "F", "W", "I", "UP", "B", "SIM"] # B905: zip strict — prefer concise zip() over strict=False noise everywhere. ignore = ["B008", "B904", "B905"] +[tool.ruff.lint.per-file-ignores] +# Autogenerated Alembic migrations: long DDL lines are unavoidable and the +# Sequence import is only used in some revisions. +"backend/migrations/versions/*" = ["E501", "F401", "I001"] + [tool.pytest.ini_options] asyncio_mode = "auto" testpaths = ["tests"] +markers = [ + "unit: tier 1 — direct module calls, no FastAPI app, no HTTP", + "integration: tier 2 — in-process ASGI, real DB + real worker", + "e2e: tier 3 — real HTTP to a uvicorn subprocess (+ optional Docker)", + "slow: takes more than ~5s (e.g. subprocess boot, Docker build)", + "needs_docker: requires a reachable Docker daemon; skips otherwise", +] +filterwarnings = [ + "ignore::DeprecationWarning", + "ignore::PendingDeprecationWarning", +] + +[tool.coverage.run] +source = ["backend"] +branch = true +omit = [ + "backend/migrations/*", + "*/tests/*", +] + +[tool.coverage.report] +exclude_lines = [ + "pragma: no cover", + "if TYPE_CHECKING:", + "raise NotImplementedError", + "if __name__ == .__main__.:", +] diff --git a/library-manager/scripts/run_tests.sh b/library-manager/scripts/run_tests.sh new file mode 100755 index 000000000..d656dac69 --- /dev/null +++ b/library-manager/scripts/run_tests.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# +# Library Manager test runner. +# +# Runs the three tiers separately (each must pass on its own), then a combined +# coverage-gated run over the in-process tiers (unit + integration). The e2e +# tier is a behavioural signal, not a coverage number: its work happens in a +# uvicorn subprocess / Docker container that the parent pytest-cov cannot +# instrument. +# +# Usage: +# ./scripts/run_tests.sh # all tiers + 95% combined gate +# ./scripts/run_tests.sh fast # skip slow tests (no subprocess/Docker) +# +set -euo pipefail + +cd "$(dirname "$0")/.." + +PYTEST="${PYTEST:-python -m pytest}" +COV_MIN="${COV_MIN:-95}" + +if [[ "${1:-}" == "fast" ]]; then + echo "== Fast run (skipping slow / e2e) ==" + $PYTEST tests/unit/ tests/integration/ -m "not slow" -q \ + --cov=backend --cov-branch --cov-report=term-missing + exit 0 +fi + +echo "== Tier 1: unit ==" +$PYTEST tests/unit/ -q + +echo "== Tier 2: integration ==" +$PYTEST tests/integration/ -q + +echo "== Tier 3: e2e ==" +$PYTEST tests/e2e/ -q + +echo "== Combined coverage gate (unit + integration, fail-under ${COV_MIN}%) ==" +$PYTEST tests/unit/ tests/integration/ \ + --cov=backend --cov-branch \ + --cov-report=term-missing \ + --cov-report=html:htmlcov \ + --cov-fail-under="${COV_MIN}" \ + -q --no-header + +echo "All tiers passed and coverage gate (>=${COV_MIN}%) satisfied." diff --git a/library-manager/tests/.yt_cache/abc123_en.json b/library-manager/tests/.yt_cache/abc123_en.json new file mode 100644 index 000000000..18406e630 --- /dev/null +++ b/library-manager/tests/.yt_cache/abc123_en.json @@ -0,0 +1 @@ +{"pieces": [{"text": "Hello", "start": 0, "duration": 1}], "source": "manual"} \ No newline at end of file diff --git a/library-manager/tests/.yt_cache/dQw4w9WgXcQ_en.json b/library-manager/tests/.yt_cache/dQw4w9WgXcQ_en.json new file mode 100644 index 000000000..0c9d97cb5 --- /dev/null +++ b/library-manager/tests/.yt_cache/dQw4w9WgXcQ_en.json @@ -0,0 +1 @@ +{"pieces": [{"text": "\u266a We're no strangers to love \u266a", "start": 18.64, "duration": 3.2399999999999984}, {"text": "\u266a You know the rules and so do I \u266a", "start": 22.64, "duration": 4.32}, {"text": "\u266a A full commitment's what I'm thinking of \u266a", "start": 27.04, "duration": 4.0}, {"text": "\u266a You wouldn't get this from any other guy \u266a", "start": 31.12, "duration": 3.9599999999999973}, {"text": "\u266a I just wanna tell you how I'm feeling \u266a", "start": 35.16, "duration": 4.3600000000000065}, {"text": "\u266a Gotta make you understand \u266a", "start": 40.52, "duration": 2.3999999999999986}, {"text": "\u266a Never gonna give you up \u266a", "start": 43.0, "duration": 2.1199999999999974}, {"text": "\u266a Never gonna let you down \u266a", "start": 45.2, "duration": 1.8799999999999955}, {"text": "\u266a Never gonna run around and desert you \u266a", "start": 47.32, "duration": 3.799999999999997}, {"text": "\u266a Never gonna make you cry \u266a", "start": 51.48, "duration": 2.0}, {"text": "\u266a Never gonna say goodbye \u266a", "start": 53.6, "duration": 1.9200000000000017}, {"text": "\u266a Never gonna tell a lie and hurt you \u266a", "start": 55.72, "duration": 3.6400000000000006}, {"text": "\u266a We've known each other for so long \u266a", "start": 60.8, "duration": 4.0}, {"text": "\u266a Your heart's been aching but you're too shy to say it \u266a", "start": 64.88, "duration": 4.160000000000011}, {"text": "\u266a Inside we both know what's been going \u266a", "start": 69.12, "duration": 3.839999999999989}, {"text": "\u266a We know the game and we're gonna play it \u266a", "start": 73.36, "duration": 3.8400000000000034}, {"text": "\u266a And if you ask me how I'm feeling \u266a", "start": 77.4, "duration": 4.640000000000001}, {"text": "\u266a Don't tell me you're too blind to see \u266a", "start": 82.4, "duration": 2.839999999999989}, {"text": "\u266a Never gonna give you up \u266a", "start": 85.32, "duration": 1.960000000000008}, {"text": "\u266a Never gonna let you down \u266a", "start": 87.36, "duration": 1.9599999999999937}, {"text": "\u266a Never gonna run around and desert you \u266a", "start": 89.44, "duration": 4.280000000000001}, {"text": "\u266a Never gonna make you cry \u266a", "start": 93.8, "duration": 1.7999999999999972}, {"text": "\u266a Never gonna say goodbye \u266a", "start": 95.76, "duration": 2.239999999999995}, {"text": "\u266a Never gonna tell a lie and hurt you \u266a", "start": 98.08, "duration": 3.960000000000008}, {"text": "\u266a Never gonna give you up \u266a", "start": 102.2, "duration": 1.9200000000000017}, {"text": "\u266a Never gonna let you down \u266a", "start": 104.28, "duration": 2.0799999999999983}, {"text": "\u266a Never gonna run around and desert you \u266a", "start": 106.48, "duration": 3.5999999999999943}, {"text": "\u266a Never gonna make you cry \u266a", "start": 110.76, "duration": 1.9599999999999937}, {"text": "\u266a Never gonna say goodbye \u266a", "start": 112.8, "duration": 1.8800000000000097}, {"text": "\u266a Never gonna tell a lie and hurt you \u266a", "start": 114.96, "duration": 3.8000000000000114}, {"text": "\u266a \u266a", "start": 119.84, "duration": 3.1199999999999903}, {"text": "\u266a \u266a", "start": 123.72, "duration": 3.6400000000000006}, {"text": "\u266a Never gonna give, never gonna give \u266a", "start": 128.48, "duration": 1.6400000000000148}, {"text": "\u266a \u266a", "start": 130.24, "duration": 1.3199999999999932}, {"text": "\u266a Never gonna give, never gonna give \u266a", "start": 132.48, "duration": 1.7600000000000193}, {"text": "\u266a \u266a", "start": 134.36, "duration": 1.5599999999999739}, {"text": "\u266a We've known each other for so long \u266a", "start": 136.76, "duration": 4.320000000000022}, {"text": "\u266a Your heart's been aching but you're too shy to say it \u266a", "start": 141.2, "duration": 4.0}, {"text": "\u266a Inside we both know what's been going \u266a", "start": 145.28, "duration": 3.8400000000000034}, {"text": "\u266a We know the game and we're gonna play it \u266a", "start": 149.52, "duration": 3.6799999999999784}, {"text": "\u266a I just wanna tell you how I'm feeling \u266a", "start": 153.36, "duration": 4.679999999999978}, {"text": "\u266a Gotta make you understand \u266a", "start": 158.64, "duration": 2.680000000000007}, {"text": "\u266a Never gonna give you up \u266a", "start": 161.4, "duration": 1.960000000000008}, {"text": "\u266a Never gonna let you down \u266a", "start": 163.44, "duration": 2.1999999999999886}, {"text": "\u266a Never gonna run around and desert you \u266a", "start": 165.72, "duration": 4.0}, {"text": "\u266a Never gonna make you cry \u266a", "start": 169.8, "duration": 1.839999999999975}, {"text": "\u266a Never gonna say goodbye \u266a", "start": 171.8, "duration": 2.1599999999999966}, {"text": "\u266a Never gonna tell a lie and hurt you \u266a", "start": 174.04, "duration": 3.5600000000000023}, {"text": "\u266a Never gonna give you up \u266a", "start": 178.2, "duration": 2.0400000000000205}, {"text": "\u266a Never gonna let you down \u266a", "start": 180.32, "duration": 2.1200000000000045}, {"text": "\u266a Never gonna run around and desert you \u266a", "start": 182.52, "duration": 4.119999999999976}, {"text": "\u266a Never gonna make you cry \u266a", "start": 186.72, "duration": 2.0}, {"text": "\u266a Never gonna say goodbye \u266a", "start": 188.84, "duration": 1.8799999999999955}, {"text": "\u266a Never gonna tell a lie and hurt you \u266a", "start": 190.84, "duration": 4.359999999999985}, {"text": "\u266a Never gonna give you up \u266a", "start": 195.28, "duration": 1.8400000000000034}, {"text": "\u266a Never gonna let you down \u266a", "start": 197.2, "duration": 2.0400000000000205}, {"text": "\u266a Never gonna run around and desert you \u266a", "start": 199.4, "duration": 3.719999999999999}, {"text": "\u266a Never gonna make you cry \u266a", "start": 203.36, "duration": 2.2399999999999807}, {"text": "\u266a Never gonna say goodbye \u266a", "start": 205.68, "duration": 2.1599999999999966}, {"text": "\u266a Never gonna tell a lie and hurt you \u266a", "start": 207.92, "duration": 3.4000000000000057}], "source": "manual"} \ No newline at end of file diff --git a/library-manager/tests/README.md b/library-manager/tests/README.md new file mode 100644 index 000000000..d5616b2ea --- /dev/null +++ b/library-manager/tests/README.md @@ -0,0 +1,73 @@ +# Library Manager — Test Suite + +A three-tier suite (**unit / integration / e2e**) with an enforced +**95% line + branch** coverage gate. The goal: when this suite is green, a +change to the Library Manager is known-good. + +## The tier contract + +| Tier | Mechanism | Proves | Speed | +|------|-----------|--------|-------| +| **unit** | Direct module calls. No FastAPI app, no HTTP. | Pure logic: config, auth dependency, DB engine setup, ORM constraints, schemas, every import plugin's internals, the plugin/capability registries, and the service layer. | fast | +| **integration** | In-process ASGI (`httpx.ASGITransport`) against the **real** SQLite DB and the **real** background worker. | Routers, request validation, the async import pipeline end-to-end, worker concurrency/timeout/recovery, lifespan, route ordering. | medium | +| **e2e** | **Real HTTP** to a `uvicorn` subprocess on loopback (+ one Docker smoke test of the shipped image). | The real process boots and serves; crash recovery across restart; the single-instance file lock; graceful shutdown; multi-tenant isolation; the full HTTP error-code matrix. | slow | + +### The one rule: mock only at the boundary + +Tests run against real systems — real SQLite, the real worker, the real +filesystem, the real plugin code. The **only** things mocked are the +third-party SDKs the plugins call out to (`markitdown`, `firecrawl`, PyMuPDF +`fitz`, `openai`), and they are patched at their exact import boundary +(`tests/_fakes.py`) so the plugin's own logic — SSRF guards, page splitting, +image extraction, error humanisation — still executes. YouTube transcripts +use the offline cassette cache in `tests/.yt_cache/`. Nothing in the unit or +integration tiers touches the network. + +## Layout + +``` +tests/ + conftest.py # env + sys.path + session DB init + auto-tag by directory + _helpers.py # AUTH_HEADERS, poll_until_ready (async+sync), payload factories + _fakes.py # boundary doubles: markitdown / firecrawl / fitz / openai + youtube_cache.py # offline YouTube transcript cache installer + .yt_cache/ # cached transcripts + unit/ # tier 1 (+ conftest: tmp_storage, db_session, registry reset) + integration/ # tier 2 (+ conftest: client, client_no_worker, library, org_id) + e2e/ # tier 3 (+ conftest, _server.py, _docker.py) +``` + +Tier markers are applied **automatically** by directory +(`pytest_collection_modifyitems` in the root `conftest.py`) — no per-test +decoration. `pytest -m unit` just works. + +## Running + +```bash +# everything, with the 95% combined gate +./scripts/run_tests.sh + +# a single tier +pytest tests/unit/ -q +pytest tests/integration/ -q +pytest tests/e2e/ -q # Docker smoke test skips cleanly if no daemon + +# by marker +pytest -m unit -q +pytest -m "not slow" -q # skip subprocess/Docker tiers + +# coverage only (in-process tiers) +pytest tests/unit/ tests/integration/ --cov=backend --cov-branch \ + --cov-report=term-missing --cov-fail-under=95 +``` + +## Coverage scope + +The 95% line+branch gate is computed over `unit/` + `integration/` — the +tiers that exercise `backend/` in-process. The e2e tier provides behavioural +assurance (a real process really boots, recovers, and locks), not a coverage +number: a subprocess/container can't be instrumented by the parent +`pytest-cov`. A small, audited set of genuinely unreachable defensive lines +(optional-dependency `ImportError` guards, `if __name__ == "__main__"`) carry +`# pragma: no cover`. `backend/migrations/` is excluded — schema DDL is +covered by the migration round-trip test, not line coverage. diff --git a/library-manager/tests/_fakes.py b/library-manager/tests/_fakes.py new file mode 100644 index 000000000..e50d90b34 --- /dev/null +++ b/library-manager/tests/_fakes.py @@ -0,0 +1,239 @@ +"""Boundary doubles for the external libraries the import plugins call. + +The rule (borrowed from the Knowledge Store suite): **mock only at the +boundary**. We replace the third-party SDK/client at the exact point each +plugin imports it, so the plugin's own code — SSRF guards, page splitting, +image extraction, metadata building, error humanisation — still runs for +real. Nothing here touches the network. + +Each plugin lazy-imports its dependency *inside* the function that uses it +(e.g. ``from markitdown import MarkItDown`` inside ``import_content``), so +patching the attribute on the dependency's module (``markitdown.MarkItDown``) +is picked up at call time. That is what every ``patch_*`` context manager +below does. + +Usage:: + + from _fakes import patch_markitdown + with patch_markitdown(text="# Title\\n\\npage one\\n\\n---\\n\\npage two"): + result = MarkItDownImportPlugin().import_content(path) +""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any +from unittest import mock + +# --------------------------------------------------------------------------- +# markitdown — used by markitdown_import, markitdown_plus_import, and the +# url_import direct-fetch fallback. Call shape: +# md = MarkItDown(); result = md.convert(src); result.text_content +# --------------------------------------------------------------------------- + + +class FakeConversionResult: + """Stand-in for ``markitdown.MarkItDown().convert(...)``'s return value.""" + + def __init__(self, text_content: str) -> None: + self.text_content = text_content + + +def make_fake_markitdown(text: str = "converted text", raises: Exception | None = None): + """Build a fake ``MarkItDown`` class returning ``text`` from ``convert``. + + If ``raises`` is given, ``convert`` raises it (to exercise error paths). + """ + + class _FakeMarkItDown: + def __init__(self, *args: Any, **kwargs: Any) -> None: + self._args = args + self._kwargs = kwargs + + def convert(self, src: Any, *args: Any, **kwargs: Any) -> FakeConversionResult: + if raises is not None: + raise raises + return FakeConversionResult(text) + + return _FakeMarkItDown + + +@contextmanager +def patch_markitdown(text: str = "converted text", raises: Exception | None = None): + """Patch ``markitdown.MarkItDown`` with a fake for the duration of the block.""" + with mock.patch("markitdown.MarkItDown", make_fake_markitdown(text=text, raises=raises)): + yield + + +# --------------------------------------------------------------------------- +# firecrawl — used by url_import's deep-crawl path. Call shape: +# app = FirecrawlApp(api_key=..., api_url=...) +# result = app.scrape(url, formats=["markdown"], timeout=...) +# result.markdown / result.metadata +# --------------------------------------------------------------------------- + + +class FakeFirecrawlDoc: + """Stand-in for a Firecrawl scrape result document.""" + + def __init__(self, markdown: str = "# Scraped\n\nbody", metadata: Any | None = None) -> None: + self.markdown = markdown + self.metadata = metadata if metadata is not None else {"title": "Scraped", "sourceURL": ""} + + +def make_fake_firecrawl( + markdown: str = "# Scraped\n\nbody", + metadata: Any | None = None, + raises: Exception | None = None, +): + """Build a fake ``FirecrawlApp`` class whose ``scrape`` returns one doc.""" + + class _FakeFirecrawlApp: + def __init__(self, *args: Any, **kwargs: Any) -> None: + self._args = args + self._kwargs = kwargs + + def scrape(self, url: str, *args: Any, **kwargs: Any) -> FakeFirecrawlDoc: + if raises is not None: + raise raises + return FakeFirecrawlDoc(markdown=markdown, metadata=metadata) + + return _FakeFirecrawlApp + + +@contextmanager +def patch_firecrawl( + markdown: str = "# Scraped\n\nbody", + metadata: Any | None = None, + raises: Exception | None = None, +): + """Patch ``firecrawl.FirecrawlApp`` with a fake for the duration of the block.""" + with mock.patch( + "firecrawl.FirecrawlApp", + make_fake_firecrawl(markdown=markdown, metadata=metadata, raises=raises), + ): + yield + + +# --------------------------------------------------------------------------- +# PyMuPDF (fitz) — used by markitdown_plus_import._extract_images. Call shape: +# doc = fitz.open(path); len(doc); page = doc[i] +# page.get_images(full=True) -> [(xref, ...), ...] +# doc.extract_image(xref) -> {"image": bytes, "ext": "png"} +# page.get_pixmap(dpi=144).tobytes("png") -> bytes +# doc.close() +# --------------------------------------------------------------------------- + + +class _FakePixmap: + def __init__(self, data: bytes) -> None: + self._data = data + + def tobytes(self, *_args: Any, **_kwargs: Any) -> bytes: + return self._data + + +class _FakeFitzPage: + def __init__(self, image_xrefs: list[int], page_png: bytes) -> None: + self._image_xrefs = image_xrefs + self._page_png = page_png + + def get_images(self, full: bool = False) -> list[tuple]: + # Real PyMuPDF returns tuples whose first element is the xref. + return [(xref, 0, 0, 0, 0, "", "", "") for xref in self._image_xrefs] + + def get_pixmap(self, *_args: Any, **_kwargs: Any) -> _FakePixmap: + return _FakePixmap(self._page_png) + + +class FakeFitzDoc: + """Stand-in for ``fitz.open(...)`` covering the calls _extract_images makes.""" + + def __init__( + self, + pages: int = 1, + images_per_page: int = 1, + image_bytes: bytes = b"\x89PNG-fake-image", + page_png: bytes = b"\x89PNG-fake-page", + extension: str = "png", + ) -> None: + self._image_bytes = image_bytes + self._extension = extension + # Give each image a unique xref so extract_image is exercised per image. + xref = 1 + self._pages: list[_FakeFitzPage] = [] + for _ in range(pages): + xrefs = list(range(xref, xref + images_per_page)) + xref += images_per_page + self._pages.append(_FakeFitzPage(xrefs, page_png)) + self.closed = False + + def __len__(self) -> int: + return len(self._pages) + + def __getitem__(self, idx: int) -> _FakeFitzPage: + return self._pages[idx] + + def extract_image(self, xref: int) -> dict[str, Any]: + return {"image": self._image_bytes, "ext": self._extension} + + def close(self) -> None: + self.closed = True + + +@contextmanager +def patch_fitz(doc: FakeFitzDoc | None = None, raises: Exception | None = None): + """Patch ``fitz.open`` to return a :class:`FakeFitzDoc` (or raise).""" + fake_doc = doc if doc is not None else FakeFitzDoc() + + def _fake_open(*_args: Any, **_kwargs: Any) -> FakeFitzDoc: + if raises is not None: + raise raises + return fake_doc + + with mock.patch("fitz.open", _fake_open): + yield fake_doc + + +# --------------------------------------------------------------------------- +# openai — used by markitdown_plus_import._describe_image (llm mode). Call shape: +# client = OpenAI(api_key=...) +# resp = client.chat.completions.create(model=..., messages=..., max_tokens=...) +# resp.choices[0].message.content ; resp.usage.total_tokens +# --------------------------------------------------------------------------- + + +def make_fake_openai(description: str = "A fake description.", raises: Exception | None = None): + """Build a fake ``OpenAI`` class whose vision call returns ``description``.""" + + class _Msg: + def __init__(self, content: str) -> None: + self.message = mock.Mock(content=content) + + class _Resp: + def __init__(self, content: str) -> None: + self.choices = [_Msg(content)] + self.usage = mock.Mock(total_tokens=42) + + class _Completions: + def create(self, *args: Any, **kwargs: Any) -> _Resp: + if raises is not None: + raise raises + return _Resp(description) + + class _Chat: + def __init__(self) -> None: + self.completions = _Completions() + + class _FakeOpenAI: + def __init__(self, *args: Any, **kwargs: Any) -> None: + self.chat = _Chat() + + return _FakeOpenAI + + +@contextmanager +def patch_openai_vision(description: str = "A fake description.", raises: Exception | None = None): + """Patch ``openai.OpenAI`` with a fake vision client for the block.""" + with mock.patch("openai.OpenAI", make_fake_openai(description=description, raises=raises)): + yield diff --git a/library-manager/tests/_helpers.py b/library-manager/tests/_helpers.py new file mode 100644 index 000000000..28954b5c2 --- /dev/null +++ b/library-manager/tests/_helpers.py @@ -0,0 +1,125 @@ +"""Shared, dependency-light test helpers. + +Consolidates the upload / poll / payload helpers that were copy-pasted across +the old flat test modules. Importable from any tier as ``from _helpers import +...`` (``tests/`` is on ``sys.path``). + +Nothing here imports the application at module scope, so this module is safe to +import from the root ``conftest.py`` before the app is loaded. +""" + +from __future__ import annotations + +import io +import time +import uuid +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: # pragma: no cover - typing only + import httpx + +# The single bearer token every tier authenticates with (matches the +# ``LAMB_API_TOKEN`` set in the root conftest before the app is imported). +AUTH_HEADERS: dict[str, str] = {"Authorization": "Bearer test-token"} + + +# --------------------------------------------------------------------------- +# Identifier factories — unique per call so tests isolate without DB resets. +# --------------------------------------------------------------------------- + + +def unique_id(prefix: str) -> str: + """Return a short unique identifier like ``lib-3f9a1c2b``.""" + return f"{prefix}-{uuid.uuid4().hex[:8]}" + + +def library_payload(**overrides: Any) -> dict[str, Any]: + """Build a valid ``POST /libraries`` body with a fresh id/org. + + Pass keyword overrides to change any field (e.g. ``organization_id``). + """ + lib_id = overrides.pop("id", None) or unique_id("lib") + payload = { + "id": lib_id, + "organization_id": "org-test", + "name": f"Test Library {lib_id[-8:]}", + } + payload.update(overrides) + return payload + + +def text_file(content: str = "# Hello\n\nWorld.", filename: str = "doc.md") -> dict[str, Any]: + """Return ``files=`` kwargs for a multipart upload of a small text file.""" + return {"file": (filename, io.BytesIO(content.encode("utf-8")), "text/markdown")} + + +def file_bytes( + data: bytes, filename: str, mime: str = "application/octet-stream" +) -> dict[str, Any]: + """Return ``files=`` kwargs for a multipart upload of raw bytes.""" + return {"file": (filename, io.BytesIO(data), mime)} + + +# --------------------------------------------------------------------------- +# Async polling (integration tier — in-process ASGI client) +# --------------------------------------------------------------------------- + + +async def poll_until_ready( + client: httpx.AsyncClient, + lib_id: str, + item_id: str, + *, + timeout: float = 30.0, + interval: float = 0.2, +) -> str: + """Poll an item's status until it leaves the in-flight states. + + Returns the terminal status (``"ready"`` or ``"failed"``). Raises + ``AssertionError`` with the last seen status if the item is still + ``pending``/``processing`` after ``timeout`` seconds. + """ + deadline = time.monotonic() + timeout + last = "unknown" + while time.monotonic() < deadline: + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/status", headers=AUTH_HEADERS + ) + if resp.status_code == 200: + last = resp.json()["status"] + if last in ("ready", "failed"): + return last + import asyncio # noqa: PLC0415 + + await asyncio.sleep(interval) + raise AssertionError( + f"item {item_id} did not reach a terminal status within {timeout}s (last={last})" + ) + + +# --------------------------------------------------------------------------- +# Sync polling (e2e tier — real HTTP to a subprocess) +# --------------------------------------------------------------------------- + + +def poll_until_ready_sync( + client: httpx.Client, + lib_id: str, + item_id: str, + *, + timeout: float = 30.0, + interval: float = 0.3, +) -> str: + """Synchronous twin of :func:`poll_until_ready` for the e2e tier.""" + deadline = time.monotonic() + timeout + last = "unknown" + while time.monotonic() < deadline: + resp = client.get(f"/libraries/{lib_id}/items/{item_id}/status", headers=AUTH_HEADERS) + if resp.status_code == 200: + last = resp.json()["status"] + if last in ("ready", "failed"): + return last + time.sleep(interval) + raise AssertionError( + f"item {item_id} did not reach a terminal status within {timeout}s (last={last})" + ) diff --git a/library-manager/tests/conftest.py b/library-manager/tests/conftest.py index 301098d87..3b4923ecc 100644 --- a/library-manager/tests/conftest.py +++ b/library-manager/tests/conftest.py @@ -1,7 +1,23 @@ -"""Shared fixtures for Library Manager tests. - -Provides an in-process async test client (httpx + FastAPI TestClient pattern) -with a fresh temporary database for each test session. +"""Root test configuration for the Library Manager suite. + +This module runs **before the application is imported** so it can fix the +environment the app reads at import time (``config.py`` freezes +``DATA_DIR``, ``LAMB_API_TOKEN`` etc. into module-level constants the moment +it is imported). + +Responsibilities that belong to *every* tier live here: + +* Set a deterministic, env-independent environment (token, temp data dir). +* Put ``backend/`` and ``tests/`` on ``sys.path`` (no editable install needed). +* Initialise the database + discover plugins once per session. +* Install the offline YouTube transcript cache. +* **Auto-tag** every collected test with its tier marker (``unit`` / + ``integration`` / ``e2e``) based on the directory it lives in, so + ``pytest -m unit`` works with zero per-test decoration. + +Tier-specific fixtures (the ASGI ``client``, the uvicorn subprocess, …) live +in the per-tier ``conftest.py`` files under ``unit/``, ``integration/`` and +``e2e/``. """ import os @@ -10,81 +26,77 @@ import tempfile import pytest -from httpx import ASGITransport, AsyncClient -# Set environment BEFORE importing the app — config reads at import time. +# --- Environment must be set BEFORE importing the app (config reads at import) --- _TEST_DIR = tempfile.mkdtemp(prefix="lm-test-") os.environ["LAMB_API_TOKEN"] = "test-token" os.environ["DATA_DIR"] = _TEST_DIR -os.environ["LOG_LEVEL"] = "WARNING" +os.environ.setdefault("LOG_LEVEL", "WARNING") -# Add backend and tests directories to sys.path. +# --- Make backend/ and tests/ importable as top-level packages --- _TESTS_DIR = os.path.dirname(__file__) -_BACKEND_DIR = os.path.join(os.path.dirname(__file__), "..", "backend") -if _TESTS_DIR not in sys.path: - sys.path.insert(0, _TESTS_DIR) -if _BACKEND_DIR not in sys.path: - sys.path.insert(0, _BACKEND_DIR) +_BACKEND_DIR = os.path.join(_TESTS_DIR, "..", "backend") +for _p in (_TESTS_DIR, _BACKEND_DIR): + if _p not in sys.path: + sys.path.insert(0, _p) +# Imported after sys.path is configured. ``_helpers`` is dependency-light +# (no app imports at module scope); ``init_db`` is used by the session fixture. from database.connection import init_db # noqa: E402 -from main import app # noqa: E402 - -AUTH_HEADERS = {"Authorization": "Bearer test-token"} @pytest.fixture(scope="session", autouse=True) -def _setup_db(): - """Initialize the database and discover plugins for the test session.""" +def _setup_session(): + """Initialise the database and discover plugins once for the whole run. + + Mirrors the production startup sequence (``ensure_directories`` → + ``init_db`` → ``_discover_plugins``) plus the offline YouTube cache so + no test ever hits the network for a transcript. Tears down the temp data + directory at the end of the session. + """ from config import ensure_directories # noqa: PLC0415 from main import _discover_plugins # noqa: PLC0415 + ensure_directories() init_db() _discover_plugins() from youtube_cache import install_youtube_cache # noqa: PLC0415 + install_youtube_cache() yield - shutil.rmtree(_TEST_DIR, ignore_errors=True) - - -@pytest.fixture -async def client(): - """Provide an async httpx client with the background worker running. - Starts the import worker so async jobs are processed, and stops it - after the test. + # Quiesce any lingering import-worker threads before removing the data + # directory. ``stop_worker`` shuts the executor down with ``wait=False`` + # (correct for production fast-shutdown), so a thread from the last + # worker-backed test can still be mid-commit at session end and would + # otherwise log a spurious "unable to open database file" after the dir + # is gone. Draining here keeps the run's output clean. + try: + from tasks import worker # noqa: PLC0415 + + worker._running = False + if worker._executor is not None: + worker._executor.shutdown(wait=True, cancel_futures=True) + except Exception: + pass - Yields: - An ``AsyncClient`` that sends requests in-process (no network). - """ - from tasks.worker import start_worker, stop_worker # noqa: PLC0415 - - await start_worker() - transport = ASGITransport(app=app) - async with AsyncClient(transport=transport, base_url="http://test") as ac: - yield ac - await stop_worker() + shutil.rmtree(_TEST_DIR, ignore_errors=True) -@pytest.fixture -async def library(client: AsyncClient) -> dict: - """Create a test library and return its metadata. +def pytest_collection_modifyitems(config, items): + """Auto-tag each test with its tier marker based on its directory. - Yields: - Dict with ``id``, ``organization_id``, ``name``. + A test under ``tests/unit/`` gets ``@pytest.mark.unit`` for free, and so + on. This keeps marker assignment impossible to forget and lets + ``pytest -m integration`` select a tier without any decorator drift. """ - import uuid # noqa: PLC0415 - - lib_id = f"lib-{uuid.uuid4().hex[:8]}" - resp = await client.post( - "/libraries", - headers=AUTH_HEADERS, - json={ - "id": lib_id, - "organization_id": "org-test", - "name": f"Test Library {lib_id[-8:]}", - }, - ) - assert resp.status_code == 201 - return resp.json() + for item in items: + path = str(item.fspath).replace(os.sep, "/") + if "/tests/unit/" in path: + item.add_marker(pytest.mark.unit) + elif "/tests/integration/" in path: + item.add_marker(pytest.mark.integration) + elif "/tests/e2e/" in path: + item.add_marker(pytest.mark.e2e) diff --git a/library-manager/tests/e2e/__init__.py b/library-manager/tests/e2e/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/library-manager/tests/e2e/_docker.py b/library-manager/tests/e2e/_docker.py new file mode 100644 index 000000000..1bfb34fd4 --- /dev/null +++ b/library-manager/tests/e2e/_docker.py @@ -0,0 +1,115 @@ +"""Docker helpers for the e2e smoke test of the real shipping image. + +The Library Manager has no external container dependency, so the bulk of the +e2e tier runs as a plain subprocess (see ``_server.py``). This module exists +only for the one test that validates the actual deployment artifact: build +``library-manager/Dockerfile``, run it, and confirm it boots, serves, and +enforces the single-instance file lock inside the container. + +Everything here degrades gracefully: if the Docker daemon is unavailable the +test skips with an actionable message rather than failing. +""" + +from __future__ import annotations + +import socket +import subprocess +import time +from pathlib import Path + +import httpx + +# library-manager/ — the Docker build context (contains Dockerfile + backend/). +PROJECT_DIR = Path(__file__).resolve().parents[2] +IMAGE_TAG = "lamb-library-manager:e2e-test" + + +def docker_available() -> bool: + """True if a usable Docker daemon is reachable.""" + try: + r = subprocess.run( + ["docker", "info"], + capture_output=True, + timeout=15, + ) + return r.returncode == 0 + except (FileNotFoundError, subprocess.SubprocessError): + return False + + +def free_port() -> int: + """Return an unused loopback TCP port for the host-side mapping.""" + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def build_image(timeout: float = 600.0) -> None: + """Build the real image from the project Dockerfile.""" + subprocess.run( + ["docker", "build", "-t", IMAGE_TAG, "."], + cwd=str(PROJECT_DIR), + check=True, + timeout=timeout, + ) + + +class Container: + """A running test container, with host-port mapping and exec helpers.""" + + def __init__(self, name: str, host_port: int) -> None: + self.name = name + self.host_port = host_port + self.base_url = f"http://127.0.0.1:{host_port}" + + @classmethod + def run(cls, name: str, host_port: int | None = None) -> Container: + """Start a detached container mapping host_port → 9091.""" + port = host_port or free_port() + subprocess.run( + [ + "docker", "run", "-d", "--rm", + "--name", name, + "-e", "LAMB_API_TOKEN=test-token", + "-e", "LOG_LEVEL=INFO", + "-p", f"{port}:9091", + IMAGE_TAG, + ], + check=True, + capture_output=True, + timeout=60, + ) + return cls(name, port) + + def wait_healthy(self, timeout: float = 60.0) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + r = httpx.get(f"{self.base_url}/health", timeout=2.0) + if r.status_code == 200: + return + except httpx.HTTPError: + pass + time.sleep(0.5) + raise RuntimeError(f"container {self.name} not healthy within {timeout}s:\n{self.logs()}") + + def exec(self, *cmd: str) -> subprocess.CompletedProcess: + """Run a command inside the container.""" + return subprocess.run( + ["docker", "exec", self.name, *cmd], + capture_output=True, + text=True, + timeout=30, + ) + + def logs(self) -> str: + try: + r = subprocess.run( + ["docker", "logs", self.name], capture_output=True, text=True, timeout=15 + ) + return r.stdout + r.stderr + except subprocess.SubprocessError: # pragma: no cover - defensive + return "(no logs)" + + def stop(self) -> None: + subprocess.run(["docker", "rm", "-f", self.name], capture_output=True, timeout=30) diff --git a/library-manager/tests/e2e/_server.py b/library-manager/tests/e2e/_server.py new file mode 100644 index 000000000..a03f65905 --- /dev/null +++ b/library-manager/tests/e2e/_server.py @@ -0,0 +1,178 @@ +"""Subprocess control for the e2e tier. + +Spawns the real application as a ``uvicorn`` subprocess listening on a +loopback TCP port — exactly how it runs in production, minus the container. +Gives tests fine-grained process control (graceful ``terminate``, hard +``kill``, restart against the same data directory) that a container makes +clumsy, which is why the behaviour-heavy e2e tests use this rather than +Docker. + +``stdout``/``stderr`` are redirected to a log file inside the data directory +(never a PIPE) to avoid the classic fill-the-pipe-buffer deadlock when the +parent isn't draining it. +""" + +from __future__ import annotations + +import os +import signal +import socket +import sqlite3 +import subprocess +import sys +import tempfile +import time +from pathlib import Path + +import httpx + +# library-manager/backend — the working directory uvicorn runs from. +BACKEND_DIR = Path(__file__).resolve().parents[2] / "backend" + + +def free_port() -> int: + """Return an unused loopback TCP port.""" + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +class ServerProcess: + """A managed ``uvicorn main:app`` subprocess for one data directory. + + Parameters mirror the knobs the e2e tests need: + + * ``data_dir`` — reuse an existing dir to test restart/recovery; omit for + a fresh throwaway dir. + * ``port`` — omit to grab a free one. + * ``env`` — extra environment (e.g. ``LM_MAX_JOB_ATTEMPTS``). + """ + + def __init__( + self, + data_dir: str | os.PathLike | None = None, + port: int | None = None, + env: dict[str, str] | None = None, + ) -> None: + self.data_dir = Path(data_dir) if data_dir else Path(tempfile.mkdtemp(prefix="lm-e2e-")) + self.owns_data_dir = data_dir is None + self.port = port or free_port() + self.base_url = f"http://127.0.0.1:{self.port}" + self.env = env or {} + self.proc: subprocess.Popen | None = None + self.log_path = self.data_dir / "server.log" + + # -- lifecycle ---------------------------------------------------------- + + def start(self, wait: bool = True, timeout: float = 30.0) -> ServerProcess: + """Launch the subprocess; optionally block until ``/health`` is 200.""" + self.data_dir.mkdir(parents=True, exist_ok=True) + env = os.environ.copy() + env["LAMB_API_TOKEN"] = "test-token" + env["DATA_DIR"] = str(self.data_dir) + env.setdefault("LOG_LEVEL", "INFO") + env.update(self.env) + log = open(self.log_path, "ab") # noqa: SIM115 — closed in stop() + self._log_file = log + self.proc = subprocess.Popen( + [ + sys.executable, "-m", "uvicorn", "main:app", + "--host", "127.0.0.1", "--port", str(self.port), + ], + cwd=str(BACKEND_DIR), + env=env, + stdout=log, + stderr=subprocess.STDOUT, + ) + if wait: + self.wait_healthy(timeout=timeout) + return self + + def wait_healthy(self, timeout: float = 30.0) -> None: + """Block until the server answers ``/health`` 200, or raise.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if self.proc is not None and self.proc.poll() is not None: + raise RuntimeError( + f"server exited early (code={self.proc.returncode}):\n{self.read_log()}" + ) + try: + r = httpx.get(f"{self.base_url}/health", timeout=1.0) + if r.status_code == 200: + return + except httpx.HTTPError: + pass + time.sleep(0.2) + raise RuntimeError(f"server not healthy within {timeout}s:\n{self.read_log()}") + + def wait_exit(self, timeout: float = 15.0) -> int | None: + """Block until the process exits; return its code (or None on timeout).""" + if self.proc is None: + return None + try: + return self.proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + return None + + def kill(self) -> None: + """Hard kill (SIGKILL) — simulates an uncontrolled crash.""" + if self.proc is not None and self.proc.poll() is None: + self.proc.send_signal(signal.SIGKILL) + self.proc.wait(timeout=10) + self._close_log() + + def terminate(self) -> int | None: + """Graceful shutdown (SIGTERM) — exercises the lifespan shutdown path.""" + if self.proc is None: + return None + if self.proc.poll() is None: + self.proc.send_signal(signal.SIGTERM) + code = self.wait_exit(timeout=15.0) + if code is None: # pragma: no cover - defensive + self.proc.send_signal(signal.SIGKILL) + self.proc.wait(timeout=10) + self._close_log() + return code + + def stop(self) -> None: + """Best-effort teardown: terminate the process, free the data dir.""" + if self.proc is not None and self.proc.poll() is None: + self.proc.send_signal(signal.SIGTERM) + if self.wait_exit(timeout=10.0) is None: # pragma: no cover - defensive + self.proc.send_signal(signal.SIGKILL) + self.proc.wait(timeout=10) + self._close_log() + if self.owns_data_dir: + import shutil # noqa: PLC0415 + + shutil.rmtree(self.data_dir, ignore_errors=True) + + # -- helpers ------------------------------------------------------------ + + def client(self, **kwargs) -> httpx.Client: + """A real HTTP client bound to this server's base URL with auth.""" + from _helpers import AUTH_HEADERS # noqa: PLC0415 + + headers = {**AUTH_HEADERS, **kwargs.pop("headers", {})} + return httpx.Client(base_url=self.base_url, headers=headers, timeout=30.0, **kwargs) + + @property + def db_path(self) -> Path: + """Path to this server's SQLite database file.""" + return self.data_dir / "library-manager.db" + + def sqlite(self) -> sqlite3.Connection: + """Open the server's SQLite DB directly (only while the server is stopped).""" + return sqlite3.connect(str(self.db_path)) + + def read_log(self) -> str: + """Return the captured server log (for failure diagnostics).""" + try: + return self.log_path.read_text(errors="replace") + except OSError: # pragma: no cover - defensive + return "(no log)" + + def _close_log(self) -> None: + log = getattr(self, "_log_file", None) + if log is not None and not log.closed: + log.close() diff --git a/library-manager/tests/e2e/conftest.py b/library-manager/tests/e2e/conftest.py new file mode 100644 index 000000000..2843f1aa3 --- /dev/null +++ b/library-manager/tests/e2e/conftest.py @@ -0,0 +1,34 @@ +"""E2E-tier fixtures: real HTTP to a uvicorn subprocess (+ optional Docker). + +The shared ``server`` fixture covers the read-mostly e2e tests (smoke, auth +boundary, multitenancy, error matrix, real imports). Tests that need to +control the process lifecycle directly (crash recovery, single-instance lock, +graceful shutdown) construct :class:`ServerProcess` themselves. +""" + +from __future__ import annotations + +from collections.abc import Iterator + +import httpx +import pytest + +from ._server import ServerProcess + + +@pytest.fixture(scope="session") +def server() -> Iterator[ServerProcess]: + """A single uvicorn subprocess shared across read-mostly e2e tests.""" + proc = ServerProcess() + proc.start() + try: + yield proc + finally: + proc.stop() + + +@pytest.fixture +def http(server: ServerProcess) -> Iterator[httpx.Client]: + """A real HTTP client bound to the shared server, with auth headers.""" + with server.client() as client: + yield client diff --git a/library-manager/tests/e2e/test_auth_boundary.py b/library-manager/tests/e2e/test_auth_boundary.py new file mode 100644 index 000000000..8b25b4345 --- /dev/null +++ b/library-manager/tests/e2e/test_auth_boundary.py @@ -0,0 +1,45 @@ +"""E2E auth boundary: bad credentials always yield 401, never 500. + +Hits the real socket with no header, the wrong token, and deliberately +malformed / non-Latin-1 / garbage tokens. The service must reject them +cleanly (401) and must never crash (500) on a token it cannot parse. +""" + +from __future__ import annotations + +import httpx +import pytest + +pytestmark = pytest.mark.slow + + +def test_missing_authorization_header(server) -> None: + """No Authorization header → 401 on a protected route.""" + with httpx.Client(base_url=server.base_url, timeout=10.0) as c: + resp = c.get("/plugins") + assert resp.status_code == 401, resp.text + + +def test_wrong_bearer_token(server) -> None: + """A syntactically valid but incorrect bearer token → 401.""" + with server.client(headers={"Authorization": "Bearer wrong-token"}) as c: + resp = c.get("/plugins") + assert resp.status_code == 401, resp.text + + +def test_non_latin1_token_is_401_not_500(server) -> None: + """A non-Latin-1 token must be rejected (401), never crash the server (500).""" + # httpx encodes headers as latin-1; pass raw bytes to smuggle non-ASCII + # through so the server receives a token it cannot compare normally. + headers = {"Authorization": "Bearer tökén-ñøn-latin".encode()} + with httpx.Client(base_url=server.base_url, timeout=10.0) as c: + resp = c.get("/plugins", headers=headers) + assert resp.status_code == 401, resp.text + + +def test_garbage_authorization_value_is_401_not_500(server) -> None: + """A malformed Authorization value (not a Bearer scheme) → 401, not 500.""" + with httpx.Client(base_url=server.base_url, timeout=10.0) as c: + resp = c.get("/plugins", headers={"Authorization": "@@@garbage@@@"}) + assert resp.status_code == 401, resp.text + assert resp.status_code != 500 diff --git a/library-manager/tests/e2e/test_crash_recovery.py b/library-manager/tests/e2e/test_crash_recovery.py new file mode 100644 index 000000000..2e8421171 --- /dev/null +++ b/library-manager/tests/e2e/test_crash_recovery.py @@ -0,0 +1,132 @@ +"""E2E crash recovery: stale 'processing' jobs are recovered on restart. + +Simulates an uncontrolled crash (SIGKILL) of an in-flight import by hand-editing +the SQLite state so a job is left ``status='processing'`` (exactly the post-crash +DB shape of a job that was mid-flight). A fresh process on the same data +directory runs ``recover_stale_jobs`` at startup: + +* attempts < LM_MAX_JOB_ATTEMPTS → reset to 'pending', reprocessed → 'ready'. +* attempts >= LM_MAX_JOB_ATTEMPTS → marked 'failed' (no retry). + +Note: a *successful* import deletes the temp upload file, so to test the +recover-to-ready path we queue the import, kill before processing completes, +and make sure the source temp file is present (recreating it if the worker had +already consumed it) — mirroring a crash that happened while the job was still +in flight and its source still on disk. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from _helpers import library_payload, poll_until_ready_sync, text_file + +from ._server import ServerProcess + +pytestmark = pytest.mark.slow + +_CONTENT = "# Recover\n\nbody text to recover after a crash" + + +def _queue_import(server: ServerProcess) -> tuple[str, str, str]: + """Create a library + queue a simple import; return (lib_id, item_id, job_id).""" + with server.client() as c: + lib = library_payload() + assert c.post("/libraries", json=lib).status_code == 201 + resp = c.post( + f"/libraries/{lib['id']}/import/file", + files=text_file(_CONTENT, "recover.md"), + data={"plugin_name": "simple_import", "title": "Recover Doc"}, + ) + assert resp.status_code == 202, resp.text + body = resp.json() + return lib["id"], body["item_id"], body["job_id"] + + +def _job_source_path(server: ServerProcess, job_id: str) -> str: + """Read the temp source_path recorded for a job (server must be stopped).""" + conn = server.sqlite() + try: + row = conn.execute( + "SELECT source_path FROM import_jobs WHERE id=?", (job_id,) + ).fetchone() + finally: + conn.close() + assert row and row[0], "job has no source_path" + return row[0] + + +def _set_stale_state(server: ServerProcess, job_id: str, item_id: str, attempts: int) -> None: + """Rewrite the DB to the post-crash shape: job 'processing', item 'processing'.""" + conn = server.sqlite() + try: + conn.execute( + "UPDATE import_jobs SET status='processing', attempts=? WHERE id=?", + (attempts, job_id), + ) + conn.execute( + "UPDATE content_items SET status='processing' WHERE id=?", + (item_id,), + ) + conn.commit() + finally: + conn.close() + + +def test_stale_job_recovered_to_ready() -> None: + """A killed in-flight job (attempts None: + """A killed job at/over LM_MAX_JOB_ATTEMPTS is marked failed, not retried.""" + server = ServerProcess(env={"LM_MAX_JOB_ATTEMPTS": "3"}) + server.start() + restarted = None + try: + lib_id, item_id, job_id = _queue_import(server) + + server.kill() + _set_stale_state(server, job_id, item_id, attempts=3) + + restarted = ServerProcess( + data_dir=server.data_dir, env={"LM_MAX_JOB_ATTEMPTS": "3"} + ) + restarted.start() + with restarted.client() as c: + status = poll_until_ready_sync(c, lib_id, item_id, timeout=30) + assert status == "failed", status + finally: + if restarted is not None: + restarted.stop() + server.stop() diff --git a/library-manager/tests/e2e/test_docker_image.py b/library-manager/tests/e2e/test_docker_image.py new file mode 100644 index 000000000..3156451fb --- /dev/null +++ b/library-manager/tests/e2e/test_docker_image.py @@ -0,0 +1,84 @@ +"""E2E Docker smoke: the real shipping image boots, serves, and runs non-root. + +Builds ``library-manager/Dockerfile`` once, runs the container, and validates +the actual deployment artifact end-to-end: health, plugin listing, non-root +user, and one real import over the mapped host port. Skips cleanly when no +Docker daemon is available. +""" + +from __future__ import annotations + +from collections.abc import Iterator + +import httpx +import pytest +from _helpers import AUTH_HEADERS, library_payload, poll_until_ready_sync, text_file + +from ._docker import Container, build_image, docker_available + +pytestmark = [pytest.mark.slow, pytest.mark.needs_docker] + +if not docker_available(): + pytest.skip("Docker daemon not available", allow_module_level=True) + +_EXPECTED_PLUGINS = { + "simple_import", + "markitdown_import", + "markitdown_plus_import", + "url_import", + "youtube_transcript_import", +} + + +@pytest.fixture(scope="module") +def container() -> Iterator[Container]: + """Build the real image and run it once for the module.""" + build_image(timeout=900.0) + c = Container.run("lamb-library-manager-e2e") + try: + c.wait_healthy(timeout=90.0) + yield c + finally: + c.stop() + + +def test_health_over_mapped_port(container: Container) -> None: + """GET /health on the mapped host port → 200 with status ok.""" + r = httpx.get(f"{container.base_url}/health", timeout=10.0) + assert r.status_code == 200, r.text + assert r.json()["status"] == "ok", r.text + + +def test_plugins_listed_in_container(container: Container) -> None: + """The containerized service lists all five import plugins.""" + r = httpx.get( + f"{container.base_url}/plugins", headers=AUTH_HEADERS, timeout=10.0 + ) + assert r.status_code == 200, r.text + names = {p["name"] for p in r.json()["plugins"]} + assert names >= _EXPECTED_PLUGINS, names + + +def test_runs_as_non_root(container: Container) -> None: + """The container process runs as the unprivileged 'appuser', not root.""" + result = container.exec("whoami") + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "appuser", result.stdout + + +def test_real_import_in_container(container: Container) -> None: + """A real simple_import inside the container reaches 'ready' over HTTP.""" + with httpx.Client( + base_url=container.base_url, headers=AUTH_HEADERS, timeout=30.0 + ) as c: + lib = library_payload() + assert c.post("/libraries", json=lib).status_code == 201 + resp = c.post( + f"/libraries/{lib['id']}/import/file", + files=text_file("# Docker\n\ncontainer import body", "docker.md"), + data={"plugin_name": "simple_import", "title": "Docker Doc"}, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + status = poll_until_ready_sync(c, lib["id"], item_id, timeout=40) + assert status == "ready", status diff --git a/library-manager/tests/e2e/test_error_matrix.py b/library-manager/tests/e2e/test_error_matrix.py new file mode 100644 index 000000000..51ace8dc0 --- /dev/null +++ b/library-manager/tests/e2e/test_error_matrix.py @@ -0,0 +1,104 @@ +"""E2E HTTP error-code matrix over the real socket. + +One assertion per status code the service can return, each exercising the real +code path: 400, 401, 404, 409, 413, 422. Proves the service maps failures to +the documented codes rather than leaking 500s. +""" + +from __future__ import annotations + +import httpx +import pytest +from _helpers import library_payload, text_file + +from ._server import ServerProcess + +pytestmark = pytest.mark.slow + + +def test_401_no_auth(server) -> None: + """No bearer token on a protected route → 401.""" + with httpx.Client(base_url=server.base_url, timeout=10.0) as c: + resp = c.get("/plugins") + assert resp.status_code == 401, resp.text + + +def test_404_missing_library(http: httpx.Client) -> None: + """GET on a non-existent library → 404.""" + resp = http.get("/libraries/does-not-exist") + assert resp.status_code == 404, resp.text + + +def test_404_missing_item(http: httpx.Client) -> None: + """GET an item under a real library that has no such item → 404.""" + lib = library_payload() + assert http.post("/libraries", json=lib).status_code == 201 + resp = http.get(f"/libraries/{lib['id']}/items/no-such-item") + assert resp.status_code == 404, resp.text + + +def test_409_duplicate_library_id(http: httpx.Client) -> None: + """Creating a library whose id already exists → 409 conflict.""" + lib = library_payload() + assert http.post("/libraries", json=lib).status_code == 201 + # Same id, different name → primary-key conflict. + dup = library_payload(id=lib["id"], name="Different Name") + resp = http.post("/libraries", json=dup) + assert resp.status_code == 409, resp.text + + +def test_422_invalid_query_param(http: httpx.Client) -> None: + """limit=0 violates the ge=1 constraint → 422 validation error.""" + resp = http.get("/libraries", params={"organization_id": "org-x", "limit": 0}) + assert resp.status_code == 422, resp.text + + +def test_422_missing_required_body_field(http: httpx.Client) -> None: + """POST /libraries without the required 'id' field → 422.""" + resp = http.post("/libraries", json={"organization_id": "org-x", "name": "No Id"}) + assert resp.status_code == 422, resp.text + + +def test_400_empty_file(http: httpx.Client) -> None: + """Uploading a 0-byte file → 400 (empty file rejected before queueing).""" + lib = library_payload() + assert http.post("/libraries", json=lib).status_code == 201 + files = {"file": ("empty.md", b"", "text/markdown")} + resp = http.post( + f"/libraries/{lib['id']}/import/file", + files=files, + data={"plugin_name": "simple_import", "title": "Empty"}, + ) + assert resp.status_code == 400, resp.text + + +def test_400_unknown_plugin(http: httpx.Client) -> None: + """Importing with an unregistered plugin name → 400.""" + lib = library_payload() + assert http.post("/libraries", json=lib).status_code == 201 + resp = http.post( + f"/libraries/{lib['id']}/import/file", + files=text_file("# X\n\nbody", "x.md"), + data={"plugin_name": "nonexistent_plugin", "title": "X"}, + ) + assert resp.status_code == 400, resp.text + + +def test_413_oversize_upload() -> None: + """An upload above MAX_UPLOAD_SIZE_BYTES → 413 (dedicated 1 KB-cap server).""" + server = ServerProcess(env={"MAX_UPLOAD_SIZE_BYTES": "1024"}) + server.start() + try: + with server.client() as c: + lib = library_payload() + assert c.post("/libraries", json=lib).status_code == 201 + big = b"x" * 4096 # 4 KB > 1 KB cap + files = {"file": ("big.md", big, "text/markdown")} + resp = c.post( + f"/libraries/{lib['id']}/import/file", + files=files, + data={"plugin_name": "simple_import", "title": "Big"}, + ) + assert resp.status_code == 413, resp.text + finally: + server.stop() diff --git a/library-manager/tests/e2e/test_graceful_shutdown.py b/library-manager/tests/e2e/test_graceful_shutdown.py new file mode 100644 index 000000000..fa71102b9 --- /dev/null +++ b/library-manager/tests/e2e/test_graceful_shutdown.py @@ -0,0 +1,39 @@ +"""E2E graceful shutdown: SIGTERM runs the lifespan shutdown and exits 0. + +A SIGTERM to uvicorn triggers the FastAPI ``lifespan`` shutdown (which calls +``stop_worker``) and the process exits cleanly with code 0. +""" + +from __future__ import annotations + +import pytest + +from ._server import ServerProcess + +pytestmark = pytest.mark.slow + + +def test_sigterm_exits_cleanly() -> None: + """SIGTERM runs the lifespan shutdown to completion and the process exits. + + Uvicorn (launched via ``python -m uvicorn``) runs the FastAPI lifespan + shutdown — ``stop_worker`` — to completion on SIGTERM, then exits. The exit + code is either 0 or the signal code (-SIGTERM); both are clean given the + lifespan shutdown completed (verified in the log). + """ + import signal # noqa: PLC0415 + + server = ServerProcess() + server.start() + try: + server.wait_healthy(timeout=5) + code = server.terminate() + # Process is no longer running. + assert server.proc.poll() is not None, "process should no longer be running" + # Clean shutdown: exit 0 or terminated by the SIGTERM we sent. + assert code in (0, -signal.SIGTERM), f"unexpected exit code {code}\n{server.read_log()}" + # The lifespan shutdown ran to completion. + log = server.read_log() + assert "Application shutdown complete" in log, log + finally: + server.stop() diff --git a/library-manager/tests/e2e/test_multitenancy.py b/library-manager/tests/e2e/test_multitenancy.py new file mode 100644 index 000000000..bba746d62 --- /dev/null +++ b/library-manager/tests/e2e/test_multitenancy.py @@ -0,0 +1,50 @@ +"""E2E multi-tenancy: the org filter isolates libraries between organizations. + +``GET /libraries?organization_id=`` is the tenant boundary — LAMB always +scopes by org. Libraries created under org A must never appear in org B's list +and vice versa. +""" + +from __future__ import annotations + +import httpx +import pytest +from _helpers import library_payload, unique_id + +pytestmark = pytest.mark.slow + + +def test_org_filter_isolates_libraries(http: httpx.Client) -> None: + """Libraries are listed only under their own organization, never another's.""" + org_a = unique_id("orgA") + org_b = unique_id("orgB") + + lib_a = library_payload(organization_id=org_a) + lib_b = library_payload(organization_id=org_b) + assert http.post("/libraries", json=lib_a).status_code == 201 + assert http.post("/libraries", json=lib_b).status_code == 201 + + resp_a = http.get("/libraries", params={"organization_id": org_a}) + assert resp_a.status_code == 200, resp_a.text + ids_a = {lib["id"] for lib in resp_a.json()["libraries"]} + assert lib_a["id"] in ids_a + assert lib_b["id"] not in ids_a + + resp_b = http.get("/libraries", params={"organization_id": org_b}) + assert resp_b.status_code == 200, resp_b.text + ids_b = {lib["id"] for lib in resp_b.json()["libraries"]} + assert lib_b["id"] in ids_b + assert lib_a["id"] not in ids_b + + +def test_org_a_library_absent_from_org_b(http: httpx.Client) -> None: + """A library created under org A is not visible in org B's list at all.""" + org_a = unique_id("orgA") + org_b = unique_id("orgB") + lib_a = library_payload(organization_id=org_a) + assert http.post("/libraries", json=lib_a).status_code == 201 + + resp_b = http.get("/libraries", params={"organization_id": org_b}) + assert resp_b.status_code == 200, resp_b.text + body = resp_b.json() + assert all(lib["id"] != lib_a["id"] for lib in body["libraries"]), body diff --git a/library-manager/tests/e2e/test_real_imports.py b/library-manager/tests/e2e/test_real_imports.py new file mode 100644 index 000000000..51c135322 --- /dev/null +++ b/library-manager/tests/e2e/test_real_imports.py @@ -0,0 +1,112 @@ +"""E2E real local imports end-to-end over HTTP. + +Drives the only network-free plugins — ``simple_import`` (.txt/.md) and +``markitdown_import`` on a local .html file — through the real subprocess and +its real background worker. No mocking. Asserts each item reaches 'ready', +serves its converted content, and exposes metadata + source_ref. +""" + +from __future__ import annotations + +import httpx +import pytest +from _helpers import file_bytes, library_payload, poll_until_ready_sync, text_file + +pytestmark = pytest.mark.slow + + +def _import_to_ready( + http: httpx.Client, lib_id: str, *, files: dict, plugin: str, title: str +) -> str: + """Queue a file import, poll to a terminal state, and return the item_id.""" + resp = http.post( + f"/libraries/{lib_id}/import/file", + files=files, + data={"plugin_name": plugin, "title": title}, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + status = poll_until_ready_sync(http, lib_id, item_id, timeout=40) + assert status == "ready", f"expected ready, got {status}" + return item_id + + +def test_simple_import_txt(http: httpx.Client) -> None: + """simple_import of a .txt file reaches 'ready' and serves the verbatim text.""" + lib = library_payload() + assert http.post("/libraries", json=lib).status_code == 201 + + content = "Plain text body. Lorem ipsum dolor sit amet." + files = {"file": ("notes.txt", content.encode("utf-8"), "text/plain")} + item_id = _import_to_ready( + http, lib["id"], files=files, plugin="simple_import", title="Notes" + ) + + got = http.get(f"/libraries/{lib['id']}/items/{item_id}/content") + assert got.status_code == 200, got.text + assert "Lorem ipsum" in got.text + + +def test_simple_import_md(http: httpx.Client) -> None: + """simple_import of a .md file reaches 'ready' and serves the markdown.""" + lib = library_payload() + assert http.post("/libraries", json=lib).status_code == 201 + + content = "# Heading\n\nA markdown paragraph with **bold** text." + item_id = _import_to_ready( + http, + lib["id"], + files=text_file(content, "doc.md"), + plugin="simple_import", + title="Markdown Doc", + ) + + got = http.get(f"/libraries/{lib['id']}/items/{item_id}/content") + assert got.status_code == 200, got.text + assert "markdown paragraph" in got.text + + +def test_markitdown_import_html(http: httpx.Client) -> None: + """markitdown_import converts a local .html file to markdown → 'ready'.""" + lib = library_payload() + assert http.post("/libraries", json=lib).status_code == 201 + + html = ( + "Report" + "

    Quarterly Report

    " + "

    Revenue grew by forty two percent this quarter.

    " + "" + ) + item_id = _import_to_ready( + http, + lib["id"], + files=file_bytes(html.encode("utf-8"), "report.html", "text/html"), + plugin="markitdown_import", + title="HTML Report", + ) + + got = http.get(f"/libraries/{lib['id']}/items/{item_id}/content") + assert got.status_code == 200, got.text + assert "forty two percent" in got.text + + +def test_metadata_and_source_ref_endpoints(http: httpx.Client) -> None: + """A ready item exposes its metadata.json and source_ref.json over HTTP.""" + lib = library_payload() + assert http.post("/libraries", json=lib).status_code == 201 + + item_id = _import_to_ready( + http, + lib["id"], + files=text_file("# Meta\n\nbody", "meta.md"), + plugin="simple_import", + title="Meta Doc", + ) + + meta = http.get(f"/libraries/{lib['id']}/items/{item_id}/metadata") + assert meta.status_code == 200, meta.text + assert isinstance(meta.json(), dict) + + src = http.get(f"/libraries/{lib['id']}/items/{item_id}/source_ref") + assert src.status_code == 200, src.text + assert isinstance(src.json(), dict) diff --git a/library-manager/tests/e2e/test_server_smoke.py b/library-manager/tests/e2e/test_server_smoke.py new file mode 100644 index 000000000..c8375c853 --- /dev/null +++ b/library-manager/tests/e2e/test_server_smoke.py @@ -0,0 +1,71 @@ +"""E2E smoke: the real uvicorn process boots, serves, and imports a file. + +Drives the shared ``server`` subprocess over real loopback HTTP — no mocking. +Proves the process actually answers, lists its plugins, rejects anonymous +callers, and completes one real ``simple_import`` end-to-end. +""" + +from __future__ import annotations + +import httpx +import pytest +from _helpers import library_payload, poll_until_ready_sync, text_file + +pytestmark = pytest.mark.slow + +# The five import plugins the service registers at startup. +_EXPECTED_PLUGINS = { + "simple_import", + "markitdown_import", + "markitdown_plus_import", + "url_import", + "youtube_transcript_import", +} + + +def test_health_ok(http: httpx.Client) -> None: + """GET /health returns 200 with status 'ok' (db + worker both up).""" + resp = http.get("/health") + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["status"] == "ok", body + assert body["service"] == "library-manager", body + + +def test_plugins_lists_all_five(http: httpx.Client) -> None: + """GET /plugins lists the five registered import plugins.""" + resp = http.get("/plugins") + assert resp.status_code == 200, resp.text + names = {p["name"] for p in resp.json()["plugins"]} + assert names >= _EXPECTED_PLUGINS, names + + +def test_protected_route_rejects_anonymous(server) -> None: + """A request to a protected route with no auth is rejected with 401.""" + with httpx.Client(base_url=server.base_url, timeout=10.0) as anon: + resp = anon.get("/plugins") + assert resp.status_code == 401, resp.text + + +def test_real_simple_import_roundtrip(http: httpx.Client) -> None: + """A real simple_import of a .md file reaches 'ready' and serves its text.""" + lib = library_payload() + assert http.post("/libraries", json=lib).status_code == 201 + + content = "# Smoke Test\n\nThe quick brown fox jumps over the lazy dog." + resp = http.post( + f"/libraries/{lib['id']}/import/file", + files=text_file(content, "smoke.md"), + data={"plugin_name": "simple_import", "title": "Smoke Doc"}, + ) + assert resp.status_code == 202, resp.text + body = resp.json() + assert body["status"] == "processing", body + item_id = body["item_id"] + + status = poll_until_ready_sync(http, lib["id"], item_id, timeout=30) + assert status == "ready", status + + got = http.get(f"/libraries/{lib['id']}/items/{item_id}/content") + assert got.status_code == 200, got.text + assert "quick brown fox" in got.text diff --git a/library-manager/tests/e2e/test_single_instance_lock.py b/library-manager/tests/e2e/test_single_instance_lock.py new file mode 100644 index 000000000..f90f02342 --- /dev/null +++ b/library-manager/tests/e2e/test_single_instance_lock.py @@ -0,0 +1,51 @@ +"""E2E single-instance lock: two processes can't share a data directory. + +The first process acquires an exclusive ``fcntl`` flock on the data dir's +``.lock`` file at startup. A second process pointed at the same dir must fail +to start (its ``init_db`` raises RuntimeError → lifespan startup fails → the +process exits) while the first stays healthy. +""" + +from __future__ import annotations + +import time + +import pytest + +from ._server import ServerProcess + +pytestmark = pytest.mark.slow + + +def test_second_instance_cannot_acquire_lock() -> None: + """A second process on the same data dir fails to become healthy.""" + a = ServerProcess() + a.start() + b = None + try: + # B shares A's data dir → the flock is already held. + b = ServerProcess(data_dir=a.data_dir, port=None) + b.start(wait=False) + + # B must exit (non-zero) within a few seconds, or never go healthy. + deadline = time.monotonic() + 8.0 + b_exited = False + while time.monotonic() < deadline: + if b.proc is not None and b.proc.poll() is not None: + b_exited = True + break + time.sleep(0.2) + + assert b_exited, f"B should have exited but is still running:\n{b.read_log()}" + assert b.proc.returncode not in (0, None), b.read_log() + + # A is unaffected and still serving. + a.wait_healthy(timeout=5) + + # The failure must mention the single-instance lock. + log = b.read_log() + assert "Another Library Manager instance is using" in log, log + finally: + if b is not None: + b.stop() + a.stop() diff --git a/library-manager/tests/integration/__init__.py b/library-manager/tests/integration/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/library-manager/tests/integration/conftest.py b/library-manager/tests/integration/conftest.py new file mode 100644 index 000000000..18e45f5cd --- /dev/null +++ b/library-manager/tests/integration/conftest.py @@ -0,0 +1,68 @@ +"""Integration-tier fixtures: in-process ASGI, real DB + real worker. + +These tests drive the FastAPI app through ``httpx.ASGITransport`` (no socket, +no subprocess) against the real SQLite database and the real background import +worker. Isolation between tests comes from unique library/organization ids +rather than resetting the shared session database. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator + +import pytest +from _helpers import AUTH_HEADERS, library_payload +from httpx import ASGITransport, AsyncClient + + +@pytest.fixture +async def client() -> AsyncIterator[AsyncClient]: + """ASGI client with the background import worker running. + + Use for any test that imports content and waits for it to reach a + terminal status. + """ + from main import app # noqa: PLC0415 + from tasks import worker # noqa: PLC0415 + + await worker.start_worker() + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + yield ac + await worker.stop_worker() + # ``stop_worker`` shuts the executor down with ``wait=False`` (production + # fast-shutdown). Drain it here so no import thread outlives this test and + # later touches the shared session DB during teardown. + if worker._executor is not None: + worker._executor.shutdown(wait=True, cancel_futures=True) + + +@pytest.fixture +async def client_no_worker() -> AsyncIterator[AsyncClient]: + """ASGI client with NO worker running. + + Use when a test drives the worker loop manually (e.g. dispatching a + single job, asserting stale-job recovery) and does not want the polling + loop racing it. + """ + from main import app # noqa: PLC0415 + + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + yield ac + + +@pytest.fixture +async def library(client: AsyncClient) -> dict: + """Create a fresh library (unique id under ``org-test``) and return it.""" + resp = await client.post("/libraries", headers=AUTH_HEADERS, json=library_payload()) + assert resp.status_code == 201, resp.text + return resp.json() + + +@pytest.fixture +def org_id() -> str: + """A unique organization id, for multi-tenant isolation assertions.""" + from _helpers import unique_id # noqa: PLC0415 + + return unique_id("org") diff --git a/library-manager/tests/integration/test_auth.py b/library-manager/tests/integration/test_auth.py new file mode 100644 index 000000000..fa7e0315c --- /dev/null +++ b/library-manager/tests/integration/test_auth.py @@ -0,0 +1,82 @@ +"""Auth behavior across every router. + +The Library Manager has no user-level ACL: a single bearer token (the +``LAMB_API_TOKEN``) protects every endpoint via ``Depends(verify_token)``. +This module asserts the three outcomes that token check produces on a +representative endpoint of each router, plus the one unprotected endpoint +(``GET /health``). + +Source: ``backend/dependencies.py`` (``verify_token`` → +``HTTPBearer`` raises 401 "Not authenticated" when the header is absent, and +``verify_token`` raises 401 "Invalid service token." when the token is +wrong) and every router declaring ``dependencies=[Depends(verify_token)]``. +""" + +from __future__ import annotations + +import pytest +from _helpers import AUTH_HEADERS +from httpx import AsyncClient + +_WRONG_HEADERS = {"Authorization": "Bearer wrong-token"} + +# One representative protected GET per router. ``library`` is created by the +# fixture so the path resolves before auth would even matter (auth runs first +# regardless). Paths that need no existing resource use a literal id. +_PROTECTED_ENDPOINTS = [ + pytest.param("/libraries/{lib}", id="libraries"), + pytest.param("/libraries/{lib}/tree", id="folders"), + pytest.param("/libraries/{lib}/items", id="content"), + pytest.param("/libraries/{lib}/export", id="importing-and-content-export"), + pytest.param("/capabilities", id="capabilities"), + pytest.param("/plugins", id="system-plugins"), +] + + +@pytest.mark.parametrize("path", _PROTECTED_ENDPOINTS) +async def test_missing_header_returns_401( + client: AsyncClient, library: dict, path: str +) -> None: + """No Authorization header on a protected endpoint yields 401 (HTTPBearer).""" + resp = await client.get(path.format(lib=library["id"])) + assert resp.status_code == 401, resp.text + + +@pytest.mark.parametrize("path", _PROTECTED_ENDPOINTS) +async def test_wrong_token_returns_401( + client: AsyncClient, library: dict, path: str +) -> None: + """A bearer token that does not match returns 401 (verify_token).""" + resp = await client.get(path.format(lib=library["id"]), headers=_WRONG_HEADERS) + assert resp.status_code == 401, resp.text + + +@pytest.mark.parametrize("path", _PROTECTED_ENDPOINTS) +async def test_correct_token_returns_200( + client: AsyncClient, library: dict, path: str +) -> None: + """The correct bearer token passes auth and reaches the handler (200).""" + resp = await client.get(path.format(lib=library["id"]), headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + + +async def test_health_needs_no_auth(client: AsyncClient) -> None: + """``GET /health`` is the one endpoint reachable without any header.""" + resp = await client.get("/health") + assert resp.status_code == 200, resp.text + + +async def test_protected_write_endpoint_missing_header_401(client: AsyncClient) -> None: + """A protected POST (create library) also rejects a missing header with 401.""" + resp = await client.post("/libraries", json={"id": "x", "organization_id": "o", "name": "n"}) + assert resp.status_code == 401, resp.text + + +async def test_protected_write_endpoint_wrong_token_401(client: AsyncClient) -> None: + """A protected POST with a wrong token returns 401, before body validation.""" + resp = await client.post( + "/libraries", + headers=_WRONG_HEADERS, + json={"id": "x", "organization_id": "o", "name": "n"}, + ) + assert resp.status_code == 401, resp.text diff --git a/library-manager/tests/integration/test_capabilities.py b/library-manager/tests/integration/test_capabilities.py new file mode 100644 index 000000000..de93f6cb2 --- /dev/null +++ b/library-manager/tests/integration/test_capabilities.py @@ -0,0 +1,358 @@ +"""Capability handler endpoints (``backend/routers/capabilities.py``). + +Ports ``tests/test_capabilities.py`` (HTTP portions) into the integration +tier and strengthens it. Covers the registry list, per-item capability list +(including the legacy ``["text"]`` fallback), capability dispatch (text → +markdown, pages → JSON, images → JSON gallery), the 404 paths (unknown +capability, no handler, handler-unavailable) and the 500 path (handler +crash). Critically it asserts ROUTE ORDERING from ``backend/main.py``: + +* ``/content/{capability}`` (item_router) wins over the legacy literal + ``/content/pages`` and ``/content/images`` list routes — so those URLs + return the renderer JSON shape, not bare filename lists. +* ``/content/images/file/{filename}`` (raw_router) wins over the legacy + ``/content/images/{image_name}`` route and serves raw bytes. +* The legacy single-page ``/content/pages/{page}`` and image-bytes + ``/content/images/{name}`` routes still resolve (extra path segments). +""" + +from __future__ import annotations + +import io +import json +from unittest import mock + +from _fakes import FakeFitzDoc, patch_fitz, patch_markitdown +from _helpers import AUTH_HEADERS, poll_until_ready, text_file +from httpx import AsyncClient + + +async def _upload_simple(client: AsyncClient, lib_id: str, content: str, + *, title: str = "Doc", filename: str = "doc.md") -> str: + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=text_file(content, filename), + data={"plugin_name": "simple_import", "title": title}, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready", status + return item_id + + +async def _upload_pdf(client: AsyncClient, lib_id: str, *, title: str = "PDF") -> str: + """Build an item with pages + images via faked markitdown_plus.""" + files = {"file": ("doc.pdf", io.BytesIO(b"%PDF-fake"), "application/pdf")} + with patch_markitdown(text="p1\n\n---\n\np2"), patch_fitz( + FakeFitzDoc(pages=1, images_per_page=1) + ): + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=files, + data={ + "plugin_name": "markitdown_plus_import", + "title": title, + "plugin_params": '{"image_descriptions": "basic"}', + }, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready", status + return item_id + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + + +async def test_capabilities_registry_lists_handlers(client: AsyncClient) -> None: + """GET /capabilities lists every registered handler.""" + resp = await client.get("/capabilities", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + names = {row["capability"] for row in resp.json()["capabilities"]} + assert {"text", "pages", "images"}.issubset(names) + + +# --------------------------------------------------------------------------- +# Per-item capability list +# --------------------------------------------------------------------------- + + +async def test_item_capabilities_text_only(client: AsyncClient, library: dict) -> None: + """A markdown upload exposes only the text capability.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Hi\n\nbody") + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/capabilities", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["item_id"] == item_id + assert data["capabilities"] == ["text"] + + +async def test_item_capabilities_pdf_has_pages_images(client: AsyncClient, library: dict) -> None: + """A faked PDF exposes text, pages and images.""" + lib_id = library["id"] + item_id = await _upload_pdf(client, lib_id) + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/capabilities", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + assert set(resp.json()["capabilities"]) == {"text", "pages", "images"} + + +async def test_item_capabilities_legacy_fallback(client: AsyncClient, library: dict) -> None: + """An item whose metadata.json lacks 'capabilities' defaults to ['text'].""" + from services import content_service # noqa: PLC0415 + + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Legacy") + + detail = await client.get(f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS) + org_id = library["organization_id"] + base = content_service.get_item_base_path(org_id, lib_id, item_id) + meta_path = base / "metadata.json" + raw = json.loads(meta_path.read_text(encoding="utf-8")) + raw.pop("capabilities", None) + meta_path.write_text(json.dumps(raw), encoding="utf-8") + assert detail.status_code == 200, detail.text + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/capabilities", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + assert resp.json()["capabilities"] == ["text"] + + +async def test_item_capabilities_missing_item_404(client: AsyncClient, library: dict) -> None: + """Capabilities for an unknown item is 404.""" + resp = await client.get( + f"/libraries/{library['id']}/items/ghost/capabilities", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# Capability dispatch — text +# --------------------------------------------------------------------------- + + +async def test_content_text_returns_markdown(client: AsyncClient, library: dict) -> None: + """GET /content/text returns the full markdown body (text/markdown).""" + lib_id = library["id"] + body = "# Title\n\nthe quick brown fox" + item_id = await _upload_simple(client, lib_id, body) + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/text", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + assert "text/markdown" in resp.headers["content-type"] + assert resp.text == body + + +# --------------------------------------------------------------------------- +# Route ordering — the heart of this module +# --------------------------------------------------------------------------- + + +async def test_content_pages_returns_renderer_json(client: AsyncClient, library: dict) -> None: + """/content/pages resolves to the capability handler → renderer JSON list.""" + lib_id = library["id"] + item_id = await _upload_pdf(client, lib_id) + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/pages", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + assert "application/json" in resp.headers["content-type"] + body = resp.json() + # Renderer shape: a list of {page, markdown}, NOT the legacy {pages, count}. + assert isinstance(body, list) + assert [p["page"] for p in body] == [1, 2] + assert body[0]["markdown"] == "p1" + + +async def test_content_images_returns_renderer_json(client: AsyncClient, library: dict) -> None: + """/content/images resolves to the capability handler → renderer JSON gallery.""" + lib_id = library["id"] + item_id = await _upload_pdf(client, lib_id) + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/images", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + assert "application/json" in resp.headers["content-type"] + body = resp.json() + # Renderer shape: a list of {filename, mime, url}, NOT legacy {images, count}. + assert isinstance(body, list) + assert body + assert "filename" in body[0] + assert "/content/images/file/" in body[0]["url"] + + +async def test_image_file_raw_route_wins_and_serves_bytes( + client: AsyncClient, library: dict +) -> None: + """/content/images/file/{fn} (raw_router) beats the legacy route and serves bytes.""" + lib_id = library["id"] + item_id = await _upload_pdf(client, lib_id) + + gallery = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/images", headers=AUTH_HEADERS + ) + filename = gallery.json()[0]["filename"] + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/images/file/{filename}", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 200, resp.text + assert resp.content == b"\x89PNG-fake-image" + assert "image/png" in resp.headers["content-type"] + + +async def test_image_file_raw_route_missing_404(client: AsyncClient, library: dict) -> None: + """The raw image-file route 404s for a missing filename.""" + lib_id = library["id"] + item_id = await _upload_pdf(client, lib_id) + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/images/file/missing.png", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 404, resp.text + + +async def test_image_file_raw_route_missing_item_404(client: AsyncClient, library: dict) -> None: + """The raw image-file route 404s for an unknown item.""" + resp = await client.get( + f"/libraries/{library['id']}/items/ghost/content/images/file/x.png", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 404, resp.text + + +async def test_legacy_single_page_route_still_works(client: AsyncClient, library: dict) -> None: + """The legacy /content/pages/{page} route still serves a single page's markdown.""" + lib_id = library["id"] + item_id = await _upload_pdf(client, lib_id) + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/pages/page_001", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 200, resp.text + assert "text/markdown" in resp.headers["content-type"] + assert resp.text == "p1" + + +async def test_legacy_image_bytes_route_still_works(client: AsyncClient, library: dict) -> None: + """The legacy /content/images/{name} route still serves raw image bytes.""" + lib_id = library["id"] + item_id = await _upload_pdf(client, lib_id) + + gallery = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/images", headers=AUTH_HEADERS + ) + filename = gallery.json()[0]["filename"] + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/images/{filename}", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 200, resp.text + assert resp.content == b"\x89PNG-fake-image" + + +# --------------------------------------------------------------------------- +# Dispatch error paths +# --------------------------------------------------------------------------- + + +async def test_unknown_capability_404(client: AsyncClient, library: dict) -> None: + """A capability value not in the enum returns 404.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "x") + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/sparkles", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + assert "unknown capability" in resp.json()["detail"].lower() + + +async def test_capability_no_handler_404(client: AsyncClient, library: dict) -> None: + """A valid enum value with no registered handler (audio) returns 404.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "x") + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/audio", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +async def test_capability_dispatch_missing_item_404(client: AsyncClient, library: dict) -> None: + """Capability dispatch on an unknown item is 404.""" + resp = await client.get( + f"/libraries/{library['id']}/items/ghost/content/text", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +async def test_capability_unavailable_for_item_404(client: AsyncClient, library: dict) -> None: + """Asking for pages on a text-only item → HandlerUnavailable → 404.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "no pages here") + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/pages", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +async def test_capability_dispatch_item_dir_missing_404( + client: AsyncClient, library: dict +) -> None: + """When the on-disk item dir is gone but the DB row remains → 404.""" + import shutil # noqa: PLC0415 + + from services import content_service # noqa: PLC0415 + + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# gone on disk") + base = content_service.get_item_base_path(library["organization_id"], lib_id, item_id) + shutil.rmtree(base) + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/text", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + assert "missing on disk" in resp.json()["detail"].lower() + + +async def test_capability_handler_crash_500(client: AsyncClient, library: dict) -> None: + """An unexpected handler exception is surfaced as a 500.""" + from plugins.content_handlers.capability import Capability, CapabilityRegistry # noqa: PLC0415 + + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# crashy") + + # ``CapabilityRegistry.get`` returns a fresh instance per call, so patch + # the handler class's ``get`` method (every instance shares it). + handler_cls = CapabilityRegistry._handlers[Capability.TEXT] + with mock.patch.object(handler_cls, "get", side_effect=RuntimeError("boom")): + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/text", headers=AUTH_HEADERS + ) + assert resp.status_code == 500, resp.text + assert "failed" in resp.json()["detail"].lower() diff --git a/library-manager/tests/integration/test_content_serving.py b/library-manager/tests/integration/test_content_serving.py new file mode 100644 index 000000000..f1a603f8a --- /dev/null +++ b/library-manager/tests/integration/test_content_serving.py @@ -0,0 +1,562 @@ +"""Content retrieval and management routes (``backend/routers/content.py``). + +Ports ``tests/test_content_serving.py`` into the integration tier and +strengthens it: exercises listing (limit/offset/status/ids), full detail, +status, every content-serving format, the pages/images/original/metadata/ +source_ref endpoints, deletion, and the 404 paths for each. Items with real +pages + images are produced by a faked ``markitdown_plus_import`` against a +``.pdf`` upload (markitdown + fitz are mocked at the boundary, so the bytes +on disk are irrelevant). +""" + +from __future__ import annotations + +import io + +from _fakes import FakeFitzDoc, patch_fitz, patch_markitdown +from _helpers import AUTH_HEADERS, poll_until_ready, text_file +from httpx import AsyncClient + + +async def _upload_simple(client: AsyncClient, lib_id: str, content: str, + *, title: str = "Doc", filename: str = "doc.md") -> str: + """Upload via simple_import, poll to ready, return the item id.""" + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=text_file(content, filename), + data={"plugin_name": "simple_import", "title": title}, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready", status + return item_id + + +async def _upload_pdf_with_pages_and_images( + client: AsyncClient, lib_id: str, *, title: str = "PDF Doc" +) -> str: + """Upload a fake .pdf through markitdown_plus → 2 pages + 1 image.""" + files = {"file": ("doc.pdf", io.BytesIO(b"%PDF-fake"), "application/pdf")} + with patch_markitdown(text="p1\n\n---\n\np2"), patch_fitz( + FakeFitzDoc(pages=1, images_per_page=1) + ): + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=files, + data={ + "plugin_name": "markitdown_plus_import", + "title": title, + "plugin_params": '{"image_descriptions": "basic"}', + }, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready", status + return item_id + + +# --------------------------------------------------------------------------- +# Listing: limit, offset, status, ids +# --------------------------------------------------------------------------- + + +async def test_list_items_returns_uploaded(client: AsyncClient, library: dict) -> None: + """GET /items returns the items uploaded into the library.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Listed") + + resp = await client.get(f"/libraries/{lib_id}/items", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["total"] == 1 + assert data["items"][0]["id"] == item_id + + +async def test_list_items_missing_library_404(client: AsyncClient) -> None: + """GET /items on an unknown library returns 404.""" + resp = await client.get("/libraries/no-such-lib/items", headers=AUTH_HEADERS) + assert resp.status_code == 404, resp.text + + +async def test_list_items_limit_offset(client: AsyncClient, library: dict) -> None: + """limit + offset paginate; total stays the full count.""" + lib_id = library["id"] + await _upload_simple(client, lib_id, "# A", title="A") + await _upload_simple(client, lib_id, "# B", title="B") + + page1 = await client.get( + f"/libraries/{lib_id}/items", headers=AUTH_HEADERS, params={"limit": 1, "offset": 0} + ) + page2 = await client.get( + f"/libraries/{lib_id}/items", headers=AUTH_HEADERS, params={"limit": 1, "offset": 1} + ) + assert page1.status_code == 200, page1.text + assert page2.status_code == 200, page2.text + assert page1.json()["total"] == 2 + assert len(page1.json()["items"]) == 1 + assert len(page2.json()["items"]) == 1 + assert page1.json()["items"][0]["id"] != page2.json()["items"][0]["id"] + + +async def test_list_items_limit_out_of_range_422(client: AsyncClient, library: dict) -> None: + """limit below 1 or above 500 is rejected by query validation.""" + lib_id = library["id"] + too_small = await client.get( + f"/libraries/{lib_id}/items", headers=AUTH_HEADERS, params={"limit": 0} + ) + too_big = await client.get( + f"/libraries/{lib_id}/items", headers=AUTH_HEADERS, params={"limit": 501} + ) + assert too_small.status_code == 422, too_small.text + assert too_big.status_code == 422, too_big.text + + +async def test_list_items_filter_by_status(client: AsyncClient, library: dict) -> None: + """status filter returns only items in that status.""" + lib_id = library["id"] + await _upload_simple(client, lib_id, "# Ready") + + resp = await client.get( + f"/libraries/{lib_id}/items", headers=AUTH_HEADERS, params={"status": "ready"} + ) + assert resp.status_code == 200, resp.text + assert all(i["status"] == "ready" for i in resp.json()["items"]) + + none = await client.get( + f"/libraries/{lib_id}/items", headers=AUTH_HEADERS, params={"status": "failed"} + ) + assert none.status_code == 200, none.text + assert none.json()["total"] == 0 + + +async def test_list_items_filter_by_ids_csv(client: AsyncClient, library: dict) -> None: + """ids CSV filter returns only the requested items (whitespace tolerant).""" + lib_id = library["id"] + item1 = await _upload_simple(client, lib_id, "# 1", title="1") + item2 = await _upload_simple(client, lib_id, "# 2", title="2") + await _upload_simple(client, lib_id, "# 3", title="3") + + resp = await client.get( + f"/libraries/{lib_id}/items", + headers=AUTH_HEADERS, + params={"ids": f" {item1} , {item2} , "}, + ) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["total"] == 2 + assert {i["id"] for i in data["items"]} == {item1, item2} + + +# --------------------------------------------------------------------------- +# Item detail and status +# --------------------------------------------------------------------------- + + +async def test_get_item_detail(client: AsyncClient, library: dict) -> None: + """GET /items/{id} returns the full detail dict with permalink_base.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Detail") + + resp = await client.get(f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + detail = resp.json() + assert detail["id"] == item_id + assert "metadata" in detail + assert "import_params" in detail + assert detail["permalink_base"] + + +async def test_get_item_wrong_library_404(client: AsyncClient, library: dict) -> None: + """An item id valid in one library 404s when queried under another.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# X") + + other = await client.post("/libraries", headers=AUTH_HEADERS, json={ + "id": "lib-other-cs", "organization_id": "org-test", "name": "Other CS Lib", + }) + assert other.status_code in (201, 409), other.text + resp = await client.get( + f"/libraries/lib-other-cs/items/{item_id}", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +async def test_get_item_missing_404(client: AsyncClient, library: dict) -> None: + """GET /items/{id} for an unknown item returns 404.""" + resp = await client.get( + f"/libraries/{library['id']}/items/no-such-item", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +async def test_get_item_status(client: AsyncClient, library: dict) -> None: + """GET /items/{id}/status returns terminal status and stats.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Status") + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/status", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["item_id"] == item_id + assert data["status"] == "ready" + assert data["error_message"] is None + + +async def test_get_status_missing_404(client: AsyncClient, library: dict) -> None: + """Status of an unknown item is 404.""" + resp = await client.get( + f"/libraries/{library['id']}/items/nope/status", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# Full content: markdown / text / html +# --------------------------------------------------------------------------- + + +async def test_content_default_markdown(client: AsyncClient, library: dict) -> None: + """GET /content defaults to text/markdown and returns the body.""" + lib_id = library["id"] + body = "# Title\n\nbody text" + item_id = await _upload_simple(client, lib_id, body) + + resp = await client.get(f"/libraries/{lib_id}/items/{item_id}/content", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + assert "text/markdown" in resp.headers["content-type"] + assert resp.text == body + + +async def test_content_format_text(client: AsyncClient, library: dict) -> None: + """?format=text returns text/plain.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Plain\n\npara") + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content", + headers=AUTH_HEADERS, + params={"format": "text"}, + ) + assert resp.status_code == 200, resp.text + assert "text/plain" in resp.headers["content-type"] + + +async def test_content_format_html(client: AsyncClient, library: dict) -> None: + """?format=html renders markdown to HTML.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Heading\n\nparagraph") + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content", + headers=AUTH_HEADERS, + params={"format": "html"}, + ) + assert resp.status_code == 200, resp.text + assert "text/html" in resp.headers["content-type"] + assert " None: + """GET /content for an unknown item is 404.""" + resp = await client.get( + f"/libraries/{library['id']}/items/missing/content", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# Pages (item built via faked markitdown_plus) +# --------------------------------------------------------------------------- + + +async def test_list_pages_and_get_page(client: AsyncClient, library: dict) -> None: + """A 2-page PDF lists its pages and serves each by name.""" + lib_id = library["id"] + item_id = await _upload_pdf_with_pages_and_images(client, lib_id) + + listing = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/pages", headers=AUTH_HEADERS + ) + assert listing.status_code == 200, listing.text + # The capability item_router wins on /content/pages — renderer JSON shape. + body = listing.json() + assert isinstance(body, list) + assert [p["page"] for p in body] == [1, 2] + + page = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/pages/page_001", + headers=AUTH_HEADERS, + ) + assert page.status_code == 200, page.text + assert page.text == "p1" + + +async def test_get_page_html_format(client: AsyncClient, library: dict) -> None: + """A page can be rendered as HTML via ?format=html.""" + lib_id = library["id"] + item_id = await _upload_pdf_with_pages_and_images(client, lib_id) + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/pages/page_002.md", + headers=AUTH_HEADERS, + params={"format": "html"}, + ) + assert resp.status_code == 200, resp.text + assert "text/html" in resp.headers["content-type"] + + +async def test_get_page_missing_404(client: AsyncClient, library: dict) -> None: + """A nonexistent page name returns 404.""" + lib_id = library["id"] + item_id = await _upload_pdf_with_pages_and_images(client, lib_id) + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/pages/page_999", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 404, resp.text + + +async def test_list_pages_missing_item_404(client: AsyncClient, library: dict) -> None: + """Listing pages for an unknown item is 404.""" + resp = await client.get( + f"/libraries/{library['id']}/items/ghost/content/pages", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# Images +# --------------------------------------------------------------------------- + + +async def test_list_images_and_get_image_bytes(client: AsyncClient, library: dict) -> None: + """An item with extracted images lists them and serves the raw bytes.""" + lib_id = library["id"] + item_id = await _upload_pdf_with_pages_and_images(client, lib_id) + + listing = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/images", headers=AUTH_HEADERS + ) + assert listing.status_code == 200, listing.text + # Capability renderer JSON shape (item_router wins on /content/images too). + gallery = listing.json() + assert isinstance(gallery, list) + assert gallery + filename = gallery[0]["filename"] + + # Legacy image-bytes route /content/images/{name} still works. + img = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/images/{filename}", + headers=AUTH_HEADERS, + ) + assert img.status_code == 200, img.text + assert img.content == b"\x89PNG-fake-image" + + +async def test_get_image_missing_404(client: AsyncClient, library: dict) -> None: + """A nonexistent image returns 404.""" + lib_id = library["id"] + item_id = await _upload_pdf_with_pages_and_images(client, lib_id) + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/images/missing.png", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 404, resp.text + + +async def test_list_images_missing_item_404(client: AsyncClient, library: dict) -> None: + """Listing images for an unknown item is 404.""" + resp = await client.get( + f"/libraries/{library['id']}/items/ghost/content/images", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# Original file +# --------------------------------------------------------------------------- + + +async def test_get_original_file(client: AsyncClient, library: dict) -> None: + """The original upload is servable by filename.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Orig\n\nbody", filename="orig.md") + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/original/orig.md", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + assert "Orig" in resp.text + + +async def test_get_original_missing_404(client: AsyncClient, library: dict) -> None: + """A nonexistent original filename returns 404.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# X", filename="x.md") + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/original/nope.pdf", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +async def test_get_original_missing_item_404(client: AsyncClient, library: dict) -> None: + """Original for an unknown item is 404.""" + resp = await client.get( + f"/libraries/{library['id']}/items/ghost/original/x.md", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# metadata / source_ref +# --------------------------------------------------------------------------- + + +async def test_get_metadata_has_permalinks(client: AsyncClient, library: dict) -> None: + """metadata.json carries permalinks and item_id.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Meta\n\nbody") + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/metadata", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + meta = resp.json() + assert meta["item_id"] == item_id + assert meta["permalinks"]["full_markdown"].endswith("/content/full.md") + + +async def test_get_metadata_missing_item_404(client: AsyncClient, library: dict) -> None: + """Metadata for an unknown item is 404.""" + resp = await client.get( + f"/libraries/{library['id']}/items/ghost/metadata", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +async def test_get_source_ref(client: AsyncClient, library: dict) -> None: + """source_ref.json describes the file import source.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Src") + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/source_ref", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + src = resp.json() + assert src["type"] == "file" + assert "original_filename" in src + + +async def test_get_source_ref_missing_item_404(client: AsyncClient, library: dict) -> None: + """source_ref for an unknown item is 404.""" + resp = await client.get( + f"/libraries/{library['id']}/items/ghost/source_ref", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# Deletion +# --------------------------------------------------------------------------- + + +async def test_delete_item_removes_content(client: AsyncClient, library: dict) -> None: + """Deleting an item makes it (and its content) inaccessible.""" + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Doomed") + + resp = await client.delete(f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + assert item_id in resp.json()["message"] + + gone = await client.get(f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS) + assert gone.status_code == 404, gone.text + content = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content", headers=AUTH_HEADERS + ) + assert content.status_code == 404, content.text + + +async def test_delete_missing_item_404(client: AsyncClient, library: dict) -> None: + """Deleting an unknown item returns 404.""" + resp = await client.delete( + f"/libraries/{library['id']}/items/ghost", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# On-disk-missing paths (DB row exists, files gone) +# --------------------------------------------------------------------------- + + +async def test_content_missing_on_disk_404(client: AsyncClient, library: dict) -> None: + """When full.md is gone but the DB row remains, /content is 404.""" + from services import content_service # noqa: PLC0415 + + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Disk") + base = content_service.get_item_base_path(library["organization_id"], lib_id, item_id) + (base / "content" / "full.md").unlink() + + resp = await client.get(f"/libraries/{lib_id}/items/{item_id}/content", headers=AUTH_HEADERS) + assert resp.status_code == 404, resp.text + assert "on disk" in resp.json()["detail"].lower() + + +async def test_metadata_missing_on_disk_404(client: AsyncClient, library: dict) -> None: + """When metadata.json is gone but the DB row remains, /metadata is 404.""" + from services import content_service # noqa: PLC0415 + + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Meta gone") + base = content_service.get_item_base_path(library["organization_id"], lib_id, item_id) + (base / "metadata.json").unlink() + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/metadata", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +async def test_source_ref_missing_on_disk_404(client: AsyncClient, library: dict) -> None: + """When source_ref.json is gone but the DB row remains, /source_ref is 404.""" + from services import content_service # noqa: PLC0415 + + lib_id = library["id"] + item_id = await _upload_simple(client, lib_id, "# Src gone") + base = content_service.get_item_base_path(library["organization_id"], lib_id, item_id) + (base / "source_ref.json").unlink() + + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/source_ref", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +async def test_get_page_missing_item_404(client: AsyncClient, library: dict) -> None: + """Single-page route 404s for an unknown item.""" + resp = await client.get( + f"/libraries/{library['id']}/items/ghost/content/pages/page_001", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 404, resp.text + + +async def test_get_image_bytes_missing_item_404(client: AsyncClient, library: dict) -> None: + """Legacy image-bytes route 404s for an unknown item.""" + resp = await client.get( + f"/libraries/{library['id']}/items/ghost/content/images/x.png", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 404, resp.text diff --git a/library-manager/tests/integration/test_edge_cases.py b/library-manager/tests/integration/test_edge_cases.py new file mode 100644 index 000000000..7e021ce17 --- /dev/null +++ b/library-manager/tests/integration/test_edge_cases.py @@ -0,0 +1,240 @@ +"""Import + content-serving edge cases and error paths. + +Ports the import/content-serving portions of ``tests/test_edge_cases.py`` +and strengthens them. Covers: the upload size limit (413), 0-byte rejection +(400), Unicode + spaced filenames, path traversal on the page / image / +original routes, malformed JSON in plugin_params / api_keys (400), wrong +file extension for a plugin (400), a nonexistent plugin (400), and a plugin +whose supported_source_types excludes the requested source (400). + +Sources: ``backend/routers/importing.py``, ``backend/routers/content.py``. +""" + +from __future__ import annotations + +import io + +import pytest +from _helpers import AUTH_HEADERS, poll_until_ready, text_file +from httpx import AsyncClient + + +async def _upload_ready(client: AsyncClient, lib_id: str, content: str, + *, title: str = "Doc", filename: str = "doc.md") -> str: + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=text_file(content, filename), + data={"plugin_name": "simple_import", "title": title}, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready", status + return item_id + + +# --------------------------------------------------------------------------- +# Upload size limit +# --------------------------------------------------------------------------- + + +async def test_upload_exceeds_size_limit_413( + client: AsyncClient, library: dict, monkeypatch: pytest.MonkeyPatch +) -> None: + """A file larger than MAX_UPLOAD_SIZE_BYTES returns 413.""" + import config # noqa: PLC0415 + import routers.importing as importing_mod # noqa: PLC0415 + + # importing.py did ``from config import MAX_UPLOAD_SIZE_BYTES`` — patch + # both the source and the module-level name it bound. + monkeypatch.setattr(config, "MAX_UPLOAD_SIZE_BYTES", 100) + monkeypatch.setattr(importing_mod, "MAX_UPLOAD_SIZE_BYTES", 100) + + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files={"file": ("big.md", io.BytesIO(b"x" * 200), "text/markdown")}, + data={"plugin_name": "simple_import", "title": "Too Big"}, + ) + assert resp.status_code == 413, resp.text + + +# --------------------------------------------------------------------------- +# 0-byte rejection +# --------------------------------------------------------------------------- + + +async def test_empty_file_rejected_400(client: AsyncClient, library: dict) -> None: + """A 0-byte upload is refused with 400 and creates no item.""" + lib_id = library["id"] + before = await client.get(f"/libraries/{lib_id}/items", headers=AUTH_HEADERS) + count_before = before.json()["total"] + + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files={"file": ("empty.md", io.BytesIO(b""), "text/markdown")}, + data={"plugin_name": "simple_import", "title": "Empty"}, + ) + assert resp.status_code == 400, resp.text + detail = resp.json()["detail"].lower() + assert "empty" in detail or "0 bytes" in detail + + after = await client.get(f"/libraries/{lib_id}/items", headers=AUTH_HEADERS) + assert after.json()["total"] == count_before + + +# --------------------------------------------------------------------------- +# Unicode + spaced filenames +# --------------------------------------------------------------------------- + + +async def test_unicode_filename_accepted(client: AsyncClient, library: dict) -> None: + """A file with a Unicode name imports and serves its content.""" + lib_id = library["id"] + content = "# Documento académico\n\nacentos: é è ê ë." + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files={"file": ("documento_académico.md", io.BytesIO(content.encode()), "text/markdown")}, + data={"plugin_name": "simple_import", "title": "Académico"}, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + assert await poll_until_ready(client, lib_id, item_id) == "ready" + + served = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content", headers=AUTH_HEADERS + ) + assert "académico" in served.text + + +async def test_spaced_filename_accepted(client: AsyncClient, library: dict) -> None: + """A file with spaces in its name imports successfully.""" + lib_id = library["id"] + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files={"file": ("my document file.md", io.BytesIO(b"# Spaced"), "text/markdown")}, + data={"plugin_name": "simple_import", "title": "Spaced"}, + ) + assert resp.status_code == 202, resp.text + assert await poll_until_ready(client, lib_id, resp.json()["item_id"]) == "ready" + + +# --------------------------------------------------------------------------- +# Path traversal +# --------------------------------------------------------------------------- + + +async def test_path_traversal_page_blocked(client: AsyncClient, library: dict) -> None: + """A ../ page name never escapes the item dir (404/400).""" + lib_id = library["id"] + item_id = await _upload_ready(client, lib_id, "# Safe") + + for attempt in ("../../metadata.json", "..%2F..%2Fetc%2Fpasswd"): + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/pages/{attempt}", + headers=AUTH_HEADERS, + ) + assert resp.status_code in (400, 404), resp.text + assert "passwd" not in resp.text.lower() or resp.status_code != 200 + + +async def test_path_traversal_image_blocked(client: AsyncClient, library: dict) -> None: + """A ../ image name never escapes the item dir (404/400).""" + lib_id = library["id"] + item_id = await _upload_ready(client, lib_id, "# Safe") + + for attempt in ("../../../metadata.json", "..%2F..%2F..%2Fetc%2Fpasswd"): + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/images/{attempt}", + headers=AUTH_HEADERS, + ) + assert resp.status_code in (400, 404), resp.text + + +async def test_path_traversal_original_blocked(client: AsyncClient, library: dict) -> None: + """A ../ original filename never escapes the item dir (404/400).""" + lib_id = library["id"] + item_id = await _upload_ready(client, lib_id, "# Safe") + + for attempt in ("../../metadata.json", "..%2F..%2Fetc%2Fpasswd"): + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/original/{attempt}", + headers=AUTH_HEADERS, + ) + assert resp.status_code in (400, 404), resp.text + + +# --------------------------------------------------------------------------- +# Malformed JSON in form fields +# --------------------------------------------------------------------------- + + +async def test_malformed_plugin_params_json_400(client: AsyncClient, library: dict) -> None: + """Invalid JSON in plugin_params returns 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files={"file": ("t.md", io.BytesIO(b"# T"), "text/markdown")}, + data={"plugin_name": "simple_import", "title": "T", "plugin_params": "{not json"}, + ) + assert resp.status_code == 400, resp.text + assert "plugin_params" in resp.json()["detail"].lower() + + +async def test_malformed_api_keys_json_400(client: AsyncClient, library: dict) -> None: + """Invalid JSON in api_keys returns 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files={"file": ("t.md", io.BytesIO(b"# T"), "text/markdown")}, + data={"plugin_name": "simple_import", "title": "T", "api_keys": "not json"}, + ) + assert resp.status_code == 400, resp.text + assert "api_keys" in resp.json()["detail"].lower() + + +# --------------------------------------------------------------------------- +# Plugin selection errors +# --------------------------------------------------------------------------- + + +async def test_nonexistent_plugin_400(client: AsyncClient, library: dict) -> None: + """An unknown plugin name returns 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files={"file": ("t.md", io.BytesIO(b"# T"), "text/markdown")}, + data={"plugin_name": "nope_plugin", "title": "Bad"}, + ) + assert resp.status_code == 400, resp.text + assert "not found" in resp.json()["detail"].lower() + + +async def test_wrong_extension_for_plugin_400(client: AsyncClient, library: dict) -> None: + """An extension the plugin doesn't accept is rejected with 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files={"file": ("fake.pdf", io.BytesIO(b"not a pdf"), "application/pdf")}, + data={"plugin_name": "simple_import", "title": "Wrong Ext"}, + ) + assert resp.status_code == 400, resp.text + assert "does not support" in resp.json()["detail"].lower() + + +async def test_file_plugin_does_not_support_file_source_400( + client: AsyncClient, library: dict +) -> None: + """A URL-only plugin (url_import) rejects file uploads with 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files={"file": ("t.md", io.BytesIO(b"# T"), "text/markdown")}, + data={"plugin_name": "url_import", "title": "Bad"}, + ) + assert resp.status_code == 400, resp.text + assert "does not support file" in resp.json()["detail"].lower() diff --git a/library-manager/tests/integration/test_export_import.py b/library-manager/tests/integration/test_export_import.py new file mode 100644 index 000000000..e7f053c06 --- /dev/null +++ b/library-manager/tests/integration/test_export_import.py @@ -0,0 +1,264 @@ +"""Library export / import via ZIP (``backend/routers/content.py`` + +``backend/services/export_service.py``). + +Ports ``tests/test_export_import.py`` and strengthens it: validates the +manifest (format_version 1.0, type, content/{item}/ layout), that only +``ready`` items are exported, a full export→import round-trip with NEW ids +and permalinks, and the import error paths (invalid zip, missing manifest, +non-.zip filename, bad org id, empty file, oversize). +""" + +from __future__ import annotations + +import io +import json +import zipfile + +import pytest +from _helpers import AUTH_HEADERS, poll_until_ready, text_file +from httpx import AsyncClient + + +async def _upload_ready(client: AsyncClient, lib_id: str, content: str, + *, title: str = "Doc", filename: str = "doc.md") -> str: + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=text_file(content, filename), + data={"plugin_name": "simple_import", "title": title}, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready", status + return item_id + + +# --------------------------------------------------------------------------- +# Export +# --------------------------------------------------------------------------- + + +async def test_export_produces_valid_zip(client: AsyncClient, library: dict) -> None: + """Export streams a ZIP with a 1.0 manifest and content/{item}/ files.""" + lib_id = library["id"] + await _upload_ready(client, lib_id, "# Exported\n\nbody", title="Export Doc") + + resp = await client.get(f"/libraries/{lib_id}/export", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + assert "application/zip" in resp.headers.get("content-type", "") + assert "attachment" in resp.headers.get("content-disposition", "") + + zf = zipfile.ZipFile(io.BytesIO(resp.content)) + names = zf.namelist() + assert "manifest.json" in names + manifest = json.loads(zf.read("manifest.json")) + assert manifest["format_version"] == "1.0" + assert manifest["type"] == "library_export" + assert len(manifest["items"]) == 1 + assert manifest["items"][0]["title"] == "Export Doc" + + item_id = manifest["items"][0]["id"] + content_files = [n for n in names if n.startswith(f"content/{item_id}/")] + assert any(n.endswith("metadata.json") for n in content_files) + assert any(n.endswith("full.md") for n in content_files) + + +async def test_export_excludes_non_ready_items(client: AsyncClient, library: dict) -> None: + """Only ready items appear in the export manifest.""" + lib_id = library["id"] + ready_id = await _upload_ready(client, lib_id, "# Ready", title="Ready") + + # Force a failed item by faking simple_import to blow up. simple_import + # reads the file directly, so an undecodable byte sequence with a forced + # plugin error isn't available — instead inject a failed row directly. + from database.connection import get_session_direct # noqa: PLC0415 + from database.models import ContentItem # noqa: PLC0415 + + db = get_session_direct() + failed_id = "item-failed-export" + try: + db.add(ContentItem( + id=failed_id, library_id=lib_id, organization_id="org-test", + title="Failed One", source_type="file", import_plugin="simple_import", + base_path="/nonexistent/failed", permalink_base="/docs/x/y/z", + status="failed", error_message="boom", + )) + db.commit() + + resp = await client.get(f"/libraries/{lib_id}/export", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + manifest = json.loads(zipfile.ZipFile(io.BytesIO(resp.content)).read("manifest.json")) + ids = {i["id"] for i in manifest["items"]} + assert ready_id in ids + assert failed_id not in ids + finally: + db.query(ContentItem).filter(ContentItem.id == failed_id).delete() + db.commit() + db.close() + + +async def test_export_missing_library_404(client: AsyncClient) -> None: + """Export of an unknown library is 404.""" + resp = await client.get("/libraries/no-such-lib/export", headers=AUTH_HEADERS) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# Round-trip import +# --------------------------------------------------------------------------- + + +async def test_import_roundtrip_new_ids_and_permalinks( + client: AsyncClient, library: dict +) -> None: + """Exported ZIP re-imports into a NEW library with new ids + permalinks.""" + lib_id = library["id"] + item_id = await _upload_ready( + client, lib_id, "# Roundtrip\n\nsurvives", title="Roundtrip Doc" + ) + + exported = await client.get(f"/libraries/{lib_id}/export", headers=AUTH_HEADERS) + zip_data = exported.content + + resp = await client.post( + "/libraries/import", + headers=AUTH_HEADERS, + params={"organization_id": "org-imported"}, + files={"file": ("export.zip", io.BytesIO(zip_data), "application/zip")}, + ) + assert resp.status_code == 201, resp.text + result = resp.json() + assert result["item_count"] == 1 + new_lib_id = result["library_id"] + assert new_lib_id != lib_id + + items = await client.get(f"/libraries/{new_lib_id}/items", headers=AUTH_HEADERS) + assert items.status_code == 200, items.text + new_items = items.json()["items"] + assert len(new_items) == 1 + new_item_id = new_items[0]["id"] + assert new_item_id != item_id + assert new_items[0]["status"] == "ready" + + content = await client.get( + f"/libraries/{new_lib_id}/items/{new_item_id}/content", headers=AUTH_HEADERS + ) + assert content.status_code == 200, content.text + assert "Roundtrip" in content.text + + meta = await client.get( + f"/libraries/{new_lib_id}/items/{new_item_id}/metadata", headers=AUTH_HEADERS + ) + assert meta.status_code == 200, meta.text + full = meta.json()["permalinks"]["full_markdown"] + assert new_lib_id in full + assert new_item_id in full + + +# --------------------------------------------------------------------------- +# Import error paths +# --------------------------------------------------------------------------- + + +async def test_import_invalid_zip_400(client: AsyncClient) -> None: + """A .zip-named file that isn't a real ZIP returns 400.""" + resp = await client.post( + "/libraries/import", + headers=AUTH_HEADERS, + params={"organization_id": "org-bad"}, + files={"file": ("not-a-zip.zip", io.BytesIO(b"not a zip"), "application/zip")}, + ) + assert resp.status_code == 400, resp.text + + +async def test_import_zip_missing_manifest_400(client: AsyncClient) -> None: + """A valid ZIP without manifest.json returns 400.""" + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr("random.txt", "hello") + buf.seek(0) + resp = await client.post( + "/libraries/import", + headers=AUTH_HEADERS, + params={"organization_id": "org-nomanifest"}, + files={"file": ("x.zip", buf, "application/zip")}, + ) + assert resp.status_code == 400, resp.text + + +async def test_import_non_zip_filename_400(client: AsyncClient) -> None: + """A file not ending in .zip is rejected with 400.""" + resp = await client.post( + "/libraries/import", + headers=AUTH_HEADERS, + params={"organization_id": "org-test"}, + files={"file": ("file.txt", io.BytesIO(b"data"), "text/plain")}, + ) + assert resp.status_code == 400, resp.text + + +async def test_import_bad_org_id_400(client: AsyncClient) -> None: + """An org id that violates the [a-zA-Z0-9_-]+ regex returns 400.""" + resp = await client.post( + "/libraries/import", + headers=AUTH_HEADERS, + params={"organization_id": "bad org/../id"}, + files={"file": ("x.zip", io.BytesIO(b"PK"), "application/zip")}, + ) + assert resp.status_code == 400, resp.text + + +async def test_import_missing_org_id_422(client: AsyncClient) -> None: + """The organization_id query param is required.""" + resp = await client.post( + "/libraries/import", + headers=AUTH_HEADERS, + files={"file": ("x.zip", io.BytesIO(b"PK"), "application/zip")}, + ) + assert resp.status_code == 422, resp.text + + +async def test_import_empty_file_400(client: AsyncClient) -> None: + """A 0-byte upload is rejected with 400.""" + resp = await client.post( + "/libraries/import", + headers=AUTH_HEADERS, + params={"organization_id": "org-empty"}, + files={"file": ("x.zip", io.BytesIO(b""), "application/zip")}, + ) + assert resp.status_code == 400, resp.text + + +async def test_import_duplicate_name_409(client: AsyncClient, library: dict) -> None: + """Importing the same library name twice into one org returns 409.""" + lib_id = library["id"] + await _upload_ready(client, lib_id, "# Dup", title="Dup Doc") + zip_data = (await client.get(f"/libraries/{lib_id}/export", headers=AUTH_HEADERS)).content + + org = "org-dup-import" + first = await client.post( + "/libraries/import", headers=AUTH_HEADERS, params={"organization_id": org}, + files={"file": ("e.zip", io.BytesIO(zip_data), "application/zip")}, + ) + assert first.status_code == 201, first.text + second = await client.post( + "/libraries/import", headers=AUTH_HEADERS, params={"organization_id": org}, + files={"file": ("e.zip", io.BytesIO(zip_data), "application/zip")}, + ) + assert second.status_code == 409, second.text + + +async def test_import_oversize_413(client: AsyncClient, monkeypatch: pytest.MonkeyPatch) -> None: + """A ZIP exceeding MAX_ZIP_IMPORT_SIZE_BYTES returns 413.""" + import config # noqa: PLC0415 + + # The route imports the name locally from config at call time. + monkeypatch.setattr(config, "MAX_ZIP_IMPORT_SIZE_BYTES", 10) + resp = await client.post( + "/libraries/import", + headers=AUTH_HEADERS, + params={"organization_id": "org-big"}, + files={"file": ("x.zip", io.BytesIO(b"x" * 100), "application/zip")}, + ) + assert resp.status_code == 413, resp.text diff --git a/library-manager/tests/integration/test_folders.py b/library-manager/tests/integration/test_folders.py new file mode 100644 index 000000000..08fb438bf --- /dev/null +++ b/library-manager/tests/integration/test_folders.py @@ -0,0 +1,396 @@ +"""Folder + tree endpoints. + +Ports ``tests/test_folders.py`` and strengthens the status-code assertions +to the exact codes the ``FolderError`` subclasses map to (see +``backend/services/folder_service.py`` and ``_raise_for_folder_error`` in +``backend/routers/folders.py``): + +* ``FolderNotFoundError`` -> 404 +* ``FolderConflictError`` -> 409 +* ``FolderCycleError`` -> 400 +* ``FolderValidationError`` -> 400 + +Source: ``backend/routers/folders.py``, ``backend/schemas/folders.py``. +""" + +from __future__ import annotations + +import pytest +from _helpers import AUTH_HEADERS, library_payload, unique_id +from httpx import AsyncClient + + +async def _create_folder( + client: AsyncClient, lib_id: str, name: str, parent_id: str | None = None +) -> dict: + resp = await client.post( + f"/libraries/{lib_id}/folders", + headers=AUTH_HEADERS, + json={"name": name, "parent_folder_id": parent_id}, + ) + assert resp.status_code == 201, resp.text + return resp.json() + + +async def _create_library(client: AsyncClient) -> dict: + resp = await client.post("/libraries", headers=AUTH_HEADERS, json=library_payload()) + assert resp.status_code == 201, resp.text + return resp.json() + + +# --------------------------------------------------------------------------- +# Tree +# --------------------------------------------------------------------------- + + +async def test_empty_library_tree(client: AsyncClient, library: dict) -> None: + """A fresh library's tree has its id and two empty lists.""" + resp = await client.get(f"/libraries/{library['id']}/tree", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + tree = resp.json() + assert tree["library_id"] == library["id"] + assert tree["folders"] == [] + assert tree["items"] == [] + + +async def test_tree_returns_flat_lists(client: AsyncClient, library: dict) -> None: + """The tree returns every folder as a flat list the frontend nests.""" + a = await _create_folder(client, library["id"], "A") + b = await _create_folder(client, library["id"], "B", a["id"]) + c = await _create_folder(client, library["id"], "C") + resp = await client.get(f"/libraries/{library['id']}/tree", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + ids = {f["id"] for f in resp.json()["folders"]} + assert ids == {a["id"], b["id"], c["id"]} + + +async def test_tree_on_missing_library_404(client: AsyncClient) -> None: + """The tree endpoint 404s for an unknown library.""" + resp = await client.get(f"/libraries/{unique_id('missing')}/tree", headers=AUTH_HEADERS) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# Create +# --------------------------------------------------------------------------- + + +async def test_create_folder_at_root(client: AsyncClient, library: dict) -> None: + """A root folder has a null parent and a server-assigned id.""" + folder = await _create_folder(client, library["id"], "Q1 Research") + assert folder["name"] == "Q1 Research" + assert folder["parent_folder_id"] is None + assert folder["id"] + + +async def test_create_folder_nested(client: AsyncClient, library: dict) -> None: + """A nested folder records its parent id.""" + parent = await _create_folder(client, library["id"], "Papers") + child = await _create_folder(client, library["id"], "Biology", parent["id"]) + assert child["parent_folder_id"] == parent["id"] + + +async def test_duplicate_sibling_name_409(client: AsyncClient, library: dict) -> None: + """A second sibling with the same name returns 409 (FolderConflictError).""" + await _create_folder(client, library["id"], "Drafts") + resp = await client.post( + f"/libraries/{library['id']}/folders", headers=AUTH_HEADERS, json={"name": "Drafts"} + ) + assert resp.status_code == 409, resp.text + + +async def test_same_name_different_parents_allowed(client: AsyncClient, library: dict) -> None: + """Uniqueness is scoped per parent: same name under two parents is fine.""" + a = await _create_folder(client, library["id"], "A") + b = await _create_folder(client, library["id"], "B") + await _create_folder(client, library["id"], "Drafts", a["id"]) + await _create_folder(client, library["id"], "Drafts", b["id"]) + + +async def test_create_under_missing_parent_404(client: AsyncClient, library: dict) -> None: + """A non-existent parent folder id returns 404 (FolderNotFoundError).""" + resp = await client.post( + f"/libraries/{library['id']}/folders", + headers=AUTH_HEADERS, + json={"name": "Orphan", "parent_folder_id": unique_id("nope")}, + ) + assert resp.status_code == 404, resp.text + + +async def test_create_under_missing_library_404(client: AsyncClient) -> None: + """Creating a folder in an unknown library returns 404.""" + resp = await client.post( + f"/libraries/{unique_id('missing')}/folders", + headers=AUTH_HEADERS, + json={"name": "X"}, + ) + assert resp.status_code == 404, resp.text + + +async def test_cross_library_parent_400(client: AsyncClient) -> None: + """A parent from another library returns 400 (FolderValidationError).""" + lib_a = await _create_library(client) + lib_b = await _create_library(client) + folder_in_b = await _create_folder(client, lib_b["id"], "Foo") + resp = await client.post( + f"/libraries/{lib_a['id']}/folders", + headers=AUTH_HEADERS, + json={"name": "Bar", "parent_folder_id": folder_in_b["id"]}, + ) + assert resp.status_code == 400, resp.text + + +# --------------------------------------------------------------------------- +# Rename +# --------------------------------------------------------------------------- + + +async def test_rename_folder(client: AsyncClient, library: dict) -> None: + """Renaming a folder returns 200 with the new name.""" + folder = await _create_folder(client, library["id"], "Old") + resp = await client.put( + f"/libraries/{library['id']}/folders/{folder['id']}", + headers=AUTH_HEADERS, + json={"name": "New"}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["name"] == "New" + + +async def test_rename_collision_409(client: AsyncClient, library: dict) -> None: + """Renaming onto an existing sibling name returns 409.""" + await _create_folder(client, library["id"], "Existing") + other = await _create_folder(client, library["id"], "Other") + resp = await client.put( + f"/libraries/{library['id']}/folders/{other['id']}", + headers=AUTH_HEADERS, + json={"name": "Existing"}, + ) + assert resp.status_code == 409, resp.text + + +async def test_rename_missing_folder_404(client: AsyncClient, library: dict) -> None: + """Renaming an unknown folder returns 404 (router pre-check).""" + resp = await client.put( + f"/libraries/{library['id']}/folders/{unique_id('nope')}", + headers=AUTH_HEADERS, + json={"name": "Whatever"}, + ) + assert resp.status_code == 404, resp.text + + +async def test_rename_folder_from_other_library_404(client: AsyncClient) -> None: + """A folder id that belongs to another library is treated as not found.""" + lib_a = await _create_library(client) + lib_b = await _create_library(client) + folder_in_b = await _create_folder(client, lib_b["id"], "Foo") + resp = await client.put( + f"/libraries/{lib_a['id']}/folders/{folder_in_b['id']}", + headers=AUTH_HEADERS, + json={"name": "Renamed"}, + ) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# Move +# --------------------------------------------------------------------------- + + +async def test_move_folder(client: AsyncClient, library: dict) -> None: + """Re-parenting a folder returns 200 with the new parent id.""" + a = await _create_folder(client, library["id"], "A") + b = await _create_folder(client, library["id"], "B") + resp = await client.put( + f"/libraries/{library['id']}/folders/{b['id']}/move", + headers=AUTH_HEADERS, + json={"parent_folder_id": a["id"]}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["parent_folder_id"] == a["id"] + + +async def test_move_folder_to_root(client: AsyncClient, library: dict) -> None: + """Moving with parent_folder_id null re-homes the folder at the root.""" + parent = await _create_folder(client, library["id"], "P") + child = await _create_folder(client, library["id"], "C", parent["id"]) + resp = await client.put( + f"/libraries/{library['id']}/folders/{child['id']}/move", + headers=AUTH_HEADERS, + json={"parent_folder_id": None}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["parent_folder_id"] is None + + +async def test_move_into_self_400(client: AsyncClient, library: dict) -> None: + """Moving a folder into itself returns 400 (FolderCycleError).""" + folder = await _create_folder(client, library["id"], "Self") + resp = await client.put( + f"/libraries/{library['id']}/folders/{folder['id']}/move", + headers=AUTH_HEADERS, + json={"parent_folder_id": folder["id"]}, + ) + assert resp.status_code == 400, resp.text + + +async def test_move_into_descendant_400(client: AsyncClient, library: dict) -> None: + """Moving a folder under its own descendant returns 400 (FolderCycleError).""" + a = await _create_folder(client, library["id"], "A") + b = await _create_folder(client, library["id"], "B", a["id"]) + c = await _create_folder(client, library["id"], "C", b["id"]) + resp = await client.put( + f"/libraries/{library['id']}/folders/{a['id']}/move", + headers=AUTH_HEADERS, + json={"parent_folder_id": c["id"]}, + ) + assert resp.status_code == 400, resp.text + + +async def test_move_to_missing_destination_404(client: AsyncClient, library: dict) -> None: + """Moving under a non-existent destination returns 404.""" + folder = await _create_folder(client, library["id"], "F") + resp = await client.put( + f"/libraries/{library['id']}/folders/{folder['id']}/move", + headers=AUTH_HEADERS, + json={"parent_folder_id": unique_id("nope")}, + ) + assert resp.status_code == 404, resp.text + + +async def test_move_missing_folder_404(client: AsyncClient, library: dict) -> None: + """Moving an unknown folder returns 404 (router pre-check).""" + resp = await client.put( + f"/libraries/{library['id']}/folders/{unique_id('nope')}/move", + headers=AUTH_HEADERS, + json={"parent_folder_id": None}, + ) + assert resp.status_code == 404, resp.text + + +async def test_move_collision_409(client: AsyncClient, library: dict) -> None: + """Moving into a parent that already has a same-named child returns 409.""" + parent = await _create_folder(client, library["id"], "Parent") + await _create_folder(client, library["id"], "Dup", parent["id"]) + loose = await _create_folder(client, library["id"], "Dup") + resp = await client.put( + f"/libraries/{library['id']}/folders/{loose['id']}/move", + headers=AUTH_HEADERS, + json={"parent_folder_id": parent["id"]}, + ) + assert resp.status_code == 409, resp.text + + +async def test_move_across_library_400(client: AsyncClient) -> None: + """Moving a folder under a destination in another library returns 400. + + The router's pre-check resolves the folder within ``lib_a`` first, so the + moved folder and the destination must both be addressable from lib_a's + path. We therefore exercise the service-level cross-library guard by + routing through the folder's own library but targeting a foreign parent. + """ + lib_a = await _create_library(client) + lib_b = await _create_library(client) + folder_a = await _create_folder(client, lib_a["id"], "Mover") + dest_b = await _create_folder(client, lib_b["id"], "Dest") + resp = await client.put( + f"/libraries/{lib_a['id']}/folders/{folder_a['id']}/move", + headers=AUTH_HEADERS, + json={"parent_folder_id": dest_b["id"]}, + ) + assert resp.status_code == 400, resp.text + + +# --------------------------------------------------------------------------- +# Delete (reparents children + items up) +# --------------------------------------------------------------------------- + + +async def test_delete_folder_reparents_subfolders(client: AsyncClient, library: dict) -> None: + """Deleting a middle folder re-homes its subfolders under its parent.""" + a = await _create_folder(client, library["id"], "A") + b = await _create_folder(client, library["id"], "B", a["id"]) + c = await _create_folder(client, library["id"], "C", b["id"]) + resp = await client.delete( + f"/libraries/{library['id']}/folders/{b['id']}", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + assert resp.json()["items_reparented_to"] == a["id"] + + tree = (await client.get(f"/libraries/{library['id']}/tree", headers=AUTH_HEADERS)).json() + folders_by_id = {f["id"]: f for f in tree["folders"]} + assert b["id"] not in folders_by_id + assert folders_by_id[c["id"]]["parent_folder_id"] == a["id"] + + +async def test_delete_collision_renames_subfolder(client: AsyncClient, library: dict) -> None: + """A reparented subfolder colliding with a sibling gets a numeric suffix.""" + a = await _create_folder(client, library["id"], "A") + await _create_folder(client, library["id"], "Drafts") # root-level Drafts + await _create_folder(client, library["id"], "Drafts", a["id"]) # A/Drafts + resp = await client.delete( + f"/libraries/{library['id']}/folders/{a['id']}", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + + tree = (await client.get(f"/libraries/{library['id']}/tree", headers=AUTH_HEADERS)).json() + root_names = sorted(f["name"] for f in tree["folders"] if f["parent_folder_id"] is None) + assert "Drafts" in root_names + assert any(n.startswith("Drafts (") for n in root_names) + + +async def test_delete_missing_folder_404(client: AsyncClient, library: dict) -> None: + """Deleting an unknown folder returns 404.""" + resp = await client.delete( + f"/libraries/{library['id']}/folders/{unique_id('nope')}", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +async def test_delete_folder_from_other_library_404(client: AsyncClient) -> None: + """Deleting a folder via the wrong library's path returns 404.""" + lib_a = await _create_library(client) + lib_b = await _create_library(client) + folder_in_b = await _create_folder(client, lib_b["id"], "Foo") + resp = await client.delete( + f"/libraries/{lib_a['id']}/folders/{folder_in_b['id']}", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# Name validation (schema) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "bad_name", + [ + "", + " ", + "x" * 200, + "with/slash", + "with\\backslash", + "with\x00null", + "with\nnewline", + ], +) +async def test_invalid_folder_names_rejected( + client: AsyncClient, library: dict, bad_name: str +) -> None: + """Invalid names fail validation: 422 (min/max length) or 400 (validator).""" + resp = await client.post( + f"/libraries/{library['id']}/folders", headers=AUTH_HEADERS, json={"name": bad_name} + ) + assert resp.status_code in (400, 422), resp.text + + +async def test_folder_name_trimmed(client: AsyncClient, library: dict) -> None: + """A surrounding-whitespace name is trimmed by the validator before storing.""" + resp = await client.post( + f"/libraries/{library['id']}/folders", + headers=AUTH_HEADERS, + json={"name": " Padded "}, + ) + assert resp.status_code == 201, resp.text + assert resp.json()["name"] == "Padded" diff --git a/library-manager/tests/integration/test_import_pipeline.py b/library-manager/tests/integration/test_import_pipeline.py new file mode 100644 index 000000000..93a60c037 --- /dev/null +++ b/library-manager/tests/integration/test_import_pipeline.py @@ -0,0 +1,679 @@ +"""Full HTTP import → ready pipeline for every import plugin. + +Each test drives a real import through the ASGI client and the real +background worker, mocking only the external boundary (``markitdown``, +``firecrawl``, ``fitz``, ``openai``) so the plugin, service, worker, and +content-writing code all run for real. The YouTube path uses the offline +transcript cache installed by the root conftest. + +Source under test: ``routers/importing.py``, ``services/import_service.py``, +``tasks/worker.py`` and the five plugins. +""" + +from __future__ import annotations + +import pytest +from _fakes import ( + FakeFitzDoc, + patch_firecrawl, + patch_fitz, + patch_markitdown, + patch_openai_vision, +) +from _helpers import ( + AUTH_HEADERS, + file_bytes, + library_payload, + poll_until_ready, + text_file, +) +from database.connection import get_session_direct +from database.models import ContentImage, ImportJob +from httpx import AsyncClient +from tasks import worker + +# 202 response shape from importing.py: {"item_id", "job_id", "status": "processing"}. + + +async def _assert_accepted(resp) -> tuple[str, str]: + """Assert a 202 import response and return (item_id, job_id).""" + assert resp.status_code == 202, resp.text + body = resp.json() + assert body["status"] == "processing", body + return body["item_id"], body["job_id"] + + +# --------------------------------------------------------------------------- +# simple_import +# --------------------------------------------------------------------------- + + +async def test_simple_import_text_file_ready(client: AsyncClient, library: dict) -> None: + """A plain text file imports to ready with the content preserved verbatim.""" + lib_id = library["id"] + content = "# Hello Pipeline\n\nThe quick brown fox." + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=text_file(content, "simple.md"), + data={"plugin_name": "simple_import", "title": "Simple Doc"}, + ) + item_id, _ = await _assert_accepted(resp) + + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready" + + body = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content", headers=AUTH_HEADERS + ) + assert body.status_code == 200, body.text + assert "The quick brown fox." in body.text + + detail = await client.get( + f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS + ) + assert detail.status_code == 200, detail.text + assert detail.json()["source_type"] == "file" + + +# --------------------------------------------------------------------------- +# markitdown_import +# --------------------------------------------------------------------------- + + +async def test_markitdown_import_pdf_with_pages(client: AsyncClient, library: dict) -> None: + """A .pdf upload through markitdown_import yields a ready item with pages.""" + lib_id = library["id"] + with patch_markitdown(text="page one body\n\n---\n\npage two body"): + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=file_bytes(b"%PDF-fake", "doc.pdf", "application/pdf"), + data={"plugin_name": "markitdown_import", "title": "MD Doc"}, + ) + item_id, _ = await _assert_accepted(resp) + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready" + + detail = await client.get( + f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS + ) + assert detail.json()["page_count"] == 2, detail.text + + +async def test_markitdown_import_docx_ready(client: AsyncClient, library: dict) -> None: + """A .docx upload through markitdown_import reaches ready with content.""" + lib_id = library["id"] + with patch_markitdown(text="just one page, no breaks"): + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=file_bytes(b"PK-fake-docx", "doc.docx", "application/octet-stream"), + data={"plugin_name": "markitdown_import", "title": "Docx Doc"}, + ) + item_id, _ = await _assert_accepted(resp) + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready" + + body = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content", headers=AUTH_HEADERS + ) + assert "just one page" in body.text + + +# --------------------------------------------------------------------------- +# markitdown_plus_import +# --------------------------------------------------------------------------- + + +async def test_markitdown_plus_basic_pages_and_images( + client: AsyncClient, library: dict +) -> None: + """markitdown_plus on a PDF extracts pages + images with basic descriptions.""" + lib_id = library["id"] + fake_doc = FakeFitzDoc(pages=1, images_per_page=1) + with patch_markitdown(text="alpha\n\n---\n\nbeta"), patch_fitz(fake_doc): + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=file_bytes(b"%PDF-fake", "rich.pdf", "application/pdf"), + data={"plugin_name": "markitdown_plus_import", "title": "Rich Doc"}, + ) + item_id, _ = await _assert_accepted(resp) + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready" + + detail = await client.get( + f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS + ) + data = detail.json() + assert data["page_count"] == 2, data + # 1 embedded image + 1 rasterized page render = 2 images. + assert data["image_count"] >= 1, data + assert data["metadata"]["image_descriptions_mode"] == "basic", data + + # Basic mode records filename-style descriptions on each image row. + session = get_session_direct() + try: + descs = [ + img.llm_description + for img in session.query(ContentImage) + .filter(ContentImage.content_item_id == item_id) + .all() + ] + finally: + session.close() + assert descs, "no image rows recorded" + assert all(d and d.startswith("Image: ") for d in descs), descs + + +async def test_markitdown_plus_llm_descriptions(client: AsyncClient, library: dict) -> None: + """markitdown_plus with image_descriptions=llm records LLM descriptions.""" + lib_id = library["id"] + fake_doc = FakeFitzDoc(pages=1, images_per_page=1) + with ( + patch_markitdown(text="alpha"), + patch_fitz(fake_doc), + patch_openai_vision(description="A vivid llm description."), + ): + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=file_bytes(b"%PDF-fake", "llm.pdf", "application/pdf"), + data={ + "plugin_name": "markitdown_plus_import", + "title": "LLM Doc", + "plugin_params": '{"image_descriptions": "llm"}', + "api_keys": '{"openai_vision": "sk-x"}', + }, + ) + item_id, _ = await _assert_accepted(resp) + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready" + + detail = await client.get( + f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS + ) + data = detail.json() + assert data["metadata"]["image_descriptions_mode"] == "llm", data + + # LLM descriptions are written to the content_images rows verbatim. + session = get_session_direct() + try: + descs = [ + img.llm_description + for img in session.query(ContentImage) + .filter(ContentImage.content_item_id == item_id) + .all() + ] + finally: + session.close() + assert descs, "no image rows recorded" + assert any(d == "A vivid llm description." for d in descs), descs + + +# --------------------------------------------------------------------------- +# url_import +# --------------------------------------------------------------------------- + + +async def test_url_import_markitdown_fallback(client: AsyncClient, library: dict) -> None: + """URL import with no firecrawl key falls back to markitdown direct fetch.""" + lib_id = library["id"] + with patch_markitdown(text="# Page\n\nfetched body"): + resp = await client.post( + f"/libraries/{lib_id}/import/url", + headers=AUTH_HEADERS, + json={ + "url": "https://example.com/article", + "plugin_name": "url_import", + "title": "URL Doc", + }, + ) + item_id, _ = await _assert_accepted(resp) + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready" + + detail = await client.get( + f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS + ) + assert detail.json()["metadata"]["fetch_method"] == "markitdown", detail.text + + +async def test_url_import_firecrawl(client: AsyncClient, library: dict) -> None: + """URL import with a firecrawl_key uses the firecrawl fetch path.""" + lib_id = library["id"] + with patch_firecrawl(markdown="# Scraped\n\ndeep crawl body"): + resp = await client.post( + f"/libraries/{lib_id}/import/url", + headers=AUTH_HEADERS, + json={ + "url": "https://example.com/deep", + "plugin_name": "url_import", + "title": "Firecrawl Doc", + "api_keys": {"firecrawl_key": "fc-x"}, + }, + ) + item_id, _ = await _assert_accepted(resp) + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready" + + detail = await client.get( + f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS + ) + assert detail.json()["metadata"]["fetch_method"] == "firecrawl", detail.text + + +# --------------------------------------------------------------------------- +# youtube +# --------------------------------------------------------------------------- + + +async def test_youtube_import_via_offline_cache(client: AsyncClient, library: dict) -> None: + """A YouTube import completes offline via the cached transcript fixture.""" + lib_id = library["id"] + resp = await client.post( + f"/libraries/{lib_id}/import/youtube", + headers=AUTH_HEADERS, + json={ + "video_url": "https://www.youtube.com/watch?v=dQw4w9WgXcQ", + "language": "en", + "plugin_name": "youtube_transcript_import", + "title": "YouTube Doc", + }, + ) + item_id, _ = await _assert_accepted(resp) + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready" + + sref = await client.get( + f"/libraries/{lib_id}/items/{item_id}/source_ref", headers=AUTH_HEADERS + ) + assert sref.status_code == 200, sref.text + data = sref.json() + assert data["type"] == "youtube", data + assert data["video_id"] == "dQw4w9WgXcQ", data + + +# --------------------------------------------------------------------------- +# API-key-in-memory-only guarantee +# --------------------------------------------------------------------------- + + +async def test_api_keys_never_persisted_to_db( + client_no_worker: AsyncClient, library: dict +) -> None: + """API keys sent with an import never land in any import_jobs column. + + Drives the queue without the worker so the keys remain in the in-memory + store; asserts they are absent from every column of the DB row, present + in ``worker._job_api_keys`` while queued, and popped exactly once. + """ + lib_id = library["id"] + secret = "sk-super-secret-value-123" + resp = await client_no_worker.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=file_bytes(b"%PDF-fake", "secret.pdf", "application/pdf"), + data={ + "plugin_name": "markitdown_plus_import", + "title": "Secret Doc", + "api_keys": f'{{"openai_vision": "{secret}"}}', + }, + ) + item_id, job_id = await _assert_accepted(resp) + + # The key lives only in the in-memory store, not in any DB column. + assert worker._job_api_keys.get(job_id) == {"openai_vision": secret} + + session = get_session_direct() + try: + row = session.query(ImportJob).filter(ImportJob.id == job_id).first() + assert row is not None + for column in ImportJob.__table__.columns: + value = getattr(row, column.name) + assert secret not in str(value), ( + f"secret leaked into import_jobs.{column.name}" + ) + finally: + session.close() + + # The worker pops the keys exactly once when it starts processing. + popped = worker._job_api_keys.pop(job_id, None) + assert popped == {"openai_vision": secret} + assert job_id not in worker._job_api_keys + + +# --------------------------------------------------------------------------- +# Router validation / error paths +# --------------------------------------------------------------------------- + + +async def test_file_import_unknown_library_404(client: AsyncClient) -> None: + """File import into a missing library returns 404.""" + resp = await client.post( + "/libraries/nope-lib/import/file", + headers=AUTH_HEADERS, + files=text_file("x", "a.md"), + data={"plugin_name": "simple_import", "title": "X"}, + ) + assert resp.status_code == 404, resp.text + + +async def test_file_import_unknown_plugin_400(client: AsyncClient, library: dict) -> None: + """File import with an unregistered plugin returns 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files=text_file("x", "a.md"), + data={"plugin_name": "does_not_exist", "title": "X"}, + ) + assert resp.status_code == 400, resp.text + + +async def test_file_import_plugin_not_file_source_400( + client: AsyncClient, library: dict +) -> None: + """File import using a URL-only plugin is rejected with 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files=text_file("x", "a.md"), + data={"plugin_name": "url_import", "title": "X"}, + ) + assert resp.status_code == 400, resp.text + + +async def test_file_import_unsupported_extension_400( + client: AsyncClient, library: dict +) -> None: + """An extension the plugin doesn't accept is rejected with 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files=file_bytes(b"data", "a.xyz", "application/octet-stream"), + data={"plugin_name": "simple_import", "title": "X"}, + ) + assert resp.status_code == 400, resp.text + + +async def test_file_import_bad_plugin_params_json_400( + client: AsyncClient, library: dict +) -> None: + """Malformed plugin_params JSON yields 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files=text_file("x", "a.md"), + data={"plugin_name": "simple_import", "title": "X", "plugin_params": "{not json"}, + ) + assert resp.status_code == 400, resp.text + + +async def test_file_import_bad_api_keys_json_400( + client: AsyncClient, library: dict +) -> None: + """Malformed api_keys JSON yields 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files=text_file("x", "a.md"), + data={"plugin_name": "simple_import", "title": "X", "api_keys": "{bad"}, + ) + assert resp.status_code == 400, resp.text + + +async def test_file_import_empty_file_400(client: AsyncClient, library: dict) -> None: + """A 0-byte upload is refused before any DB row is created.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files=file_bytes(b"", "empty.md", "text/markdown"), + data={"plugin_name": "simple_import", "title": "Empty"}, + ) + assert resp.status_code == 400, resp.text + assert "empty" in resp.json()["detail"].lower() + + +async def test_file_import_invalid_folder_400(client: AsyncClient, library: dict) -> None: + """Importing into a non-existent folder surfaces a 400 from the service.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files=text_file("x", "a.md"), + data={ + "plugin_name": "simple_import", + "title": "X", + "folder_id": "no-such-folder", + }, + ) + assert resp.status_code == 400, resp.text + + +async def test_file_import_into_valid_folder(client: AsyncClient, library: dict) -> None: + """Importing into a folder that belongs to the library succeeds and is filed there.""" + lib_id = library["id"] + folder = await client.post( + f"/libraries/{lib_id}/folders", headers=AUTH_HEADERS, json={"name": "inbox"} + ) + assert folder.status_code == 201, folder.text + folder_id = folder.json()["id"] + + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=text_file("foldered body", "f.md"), + data={"plugin_name": "simple_import", "title": "Filed", "folder_id": folder_id}, + ) + item_id, _ = await _assert_accepted(resp) + assert await poll_until_ready(client, lib_id, item_id) == "ready" + + detail = await client.get( + f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS + ) + assert detail.json()["folder_id"] == folder_id, detail.text + + +async def test_file_import_queue_failure_cleans_temp( + client: AsyncClient, library: dict, monkeypatch +) -> None: + """A non-ValueError during queueing propagates and the temp file is cleaned. + + The endpoint's ``except Exception`` branch unlinks the temp upload and + re-raises; ASGITransport surfaces the original error to the caller. + """ + import routers.importing as importing_mod # noqa: PLC0415 + + upload_dir = importing_mod._UPLOAD_DIR + + def boom(*args, **kwargs): + raise RuntimeError("queue exploded") + + monkeypatch.setattr(importing_mod.import_service, "queue_file_import", boom) + before = set(upload_dir.iterdir()) if upload_dir.exists() else set() + + with pytest.raises(RuntimeError, match="queue exploded"): + await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files=text_file("x", "a.md"), + data={"plugin_name": "simple_import", "title": "X"}, + ) + + after = set(upload_dir.iterdir()) if upload_dir.exists() else set() + assert after == before, "temp upload should have been cleaned up on failure" + + +async def test_file_import_too_large_413( + client: AsyncClient, library: dict, monkeypatch +) -> None: + """An upload exceeding the size cap is rejected with 413.""" + import routers.importing as importing_mod # noqa: PLC0415 + + monkeypatch.setattr(importing_mod, "MAX_UPLOAD_SIZE_BYTES", 4) + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files=file_bytes(b"way too many bytes here", "big.md", "text/markdown"), + data={"plugin_name": "simple_import", "title": "Big"}, + ) + assert resp.status_code == 413, resp.text + + +async def test_file_import_dotdot_filename_falls_back( + client: AsyncClient, library: dict +) -> None: + """A filename of '..' sanitizes to the 'unnamed' fallback and still imports.""" + lib_id = library["id"] + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=file_bytes(b"body text", "..", "text/markdown"), + data={"plugin_name": "markitdown_import", "title": "DotDot"}, + ) + # markitdown_import accepts unknown extensions only if listed; '..' has no + # extension, so the extension guard rejects it. Either way the safe_filename + # fallback line runs before that check. + assert resp.status_code in (202, 400), resp.text + + +async def test_file_import_cross_library_folder_400( + client: AsyncClient, library: dict +) -> None: + """A folder_id belonging to another library is rejected with 400.""" + # Create a second library and a folder inside it. + other = await client.post( + "/libraries", headers=AUTH_HEADERS, json=library_payload() + ) + assert other.status_code == 201, other.text + other_lib = other.json()["id"] + folder = await client.post( + f"/libraries/{other_lib}/folders", + headers=AUTH_HEADERS, + json={"name": "elsewhere"}, + ) + assert folder.status_code == 201, folder.text + foreign_folder = folder.json()["id"] + + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files=text_file("x", "a.md"), + data={ + "plugin_name": "simple_import", + "title": "Cross", + "folder_id": foreign_folder, + }, + ) + assert resp.status_code == 400, resp.text + assert "different library" in resp.json()["detail"].lower() + + +async def test_url_import_unknown_library_404(client: AsyncClient) -> None: + """URL import into a missing library returns 404.""" + resp = await client.post( + "/libraries/nope/import/url", + headers=AUTH_HEADERS, + json={"url": "https://x.com", "plugin_name": "url_import", "title": "X"}, + ) + assert resp.status_code == 404, resp.text + + +async def test_url_import_unknown_plugin_400(client: AsyncClient, library: dict) -> None: + """URL import with an unknown plugin returns 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/url", + headers=AUTH_HEADERS, + json={"url": "https://x.com", "plugin_name": "nope", "title": "X"}, + ) + assert resp.status_code == 400, resp.text + + +async def test_url_import_wrong_source_plugin_400( + client: AsyncClient, library: dict +) -> None: + """URL import using a file-only plugin returns 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/url", + headers=AUTH_HEADERS, + json={"url": "https://x.com", "plugin_name": "simple_import", "title": "X"}, + ) + assert resp.status_code == 400, resp.text + + +async def test_url_import_invalid_folder_400(client: AsyncClient, library: dict) -> None: + """URL import into a bad folder returns the service ValueError as 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/url", + headers=AUTH_HEADERS, + json={ + "url": "https://x.com", + "plugin_name": "url_import", + "title": "X", + "folder_id": "no-such-folder", + }, + ) + assert resp.status_code == 400, resp.text + + +async def test_youtube_import_unknown_library_404(client: AsyncClient) -> None: + """YouTube import into a missing library returns 404.""" + resp = await client.post( + "/libraries/nope/import/youtube", + headers=AUTH_HEADERS, + json={ + "video_url": "https://youtube.com/watch?v=x", + "plugin_name": "youtube_transcript_import", + "title": "X", + }, + ) + assert resp.status_code == 404, resp.text + + +async def test_youtube_import_unknown_plugin_400( + client: AsyncClient, library: dict +) -> None: + """YouTube import with an unknown plugin returns 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/youtube", + headers=AUTH_HEADERS, + json={ + "video_url": "https://youtube.com/watch?v=x", + "plugin_name": "nope", + "title": "X", + }, + ) + assert resp.status_code == 400, resp.text + + +async def test_youtube_import_wrong_source_plugin_400( + client: AsyncClient, library: dict +) -> None: + """YouTube import using a file-only plugin returns 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/youtube", + headers=AUTH_HEADERS, + json={ + "video_url": "https://youtube.com/watch?v=x", + "plugin_name": "simple_import", + "title": "X", + }, + ) + assert resp.status_code == 400, resp.text + + +async def test_youtube_import_invalid_folder_400( + client: AsyncClient, library: dict +) -> None: + """YouTube import into a bad folder returns the service ValueError as 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/youtube", + headers=AUTH_HEADERS, + json={ + "video_url": "https://youtube.com/watch?v=x", + "plugin_name": "youtube_transcript_import", + "title": "X", + "folder_id": "no-such-folder", + }, + ) + assert resp.status_code == 400, resp.text diff --git a/library-manager/tests/integration/test_items_move.py b/library-manager/tests/integration/test_items_move.py new file mode 100644 index 000000000..9fda4b5ec --- /dev/null +++ b/library-manager/tests/integration/test_items_move.py @@ -0,0 +1,245 @@ +"""Item move (single + bulk) and upload-with-folder_id. + +Ports ``tests/test_items_move.py``. Uploads use the real ``simple_import`` +plugin through the running worker (fast, no mocking) and ``poll_until_ready`` +to reach the terminal status before asserting placement. + +Source: ``backend/routers/folders.py`` (move route), +``backend/routers/importing.py``, ``backend/schemas/folders.py`` +(``ItemsMoveRequest`` caps the list at 1..500). +""" + +from __future__ import annotations + +from _helpers import AUTH_HEADERS, library_payload, poll_until_ready, text_file, unique_id +from httpx import AsyncClient + + +async def _upload( + client: AsyncClient, lib_id: str, *, title: str = "Doc", folder_id: str | None = None +) -> str: + data = {"plugin_name": "simple_import", "title": title} + if folder_id is not None: + data["folder_id"] = folder_id + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files=text_file("# x", "a.md"), + data=data, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + status = await poll_until_ready(client, lib_id, item_id) + assert status == "ready", f"upload did not become ready: {status}" + return item_id + + +async def _create_folder( + client: AsyncClient, lib_id: str, name: str, parent_id: str | None = None +) -> str: + resp = await client.post( + f"/libraries/{lib_id}/folders", + headers=AUTH_HEADERS, + json={"name": name, "parent_folder_id": parent_id}, + ) + assert resp.status_code == 201, resp.text + return resp.json()["id"] + + +async def _create_library(client: AsyncClient) -> str: + resp = await client.post("/libraries", headers=AUTH_HEADERS, json=library_payload()) + assert resp.status_code == 201, resp.text + return resp.json()["id"] + + +async def _folder_id_of(client: AsyncClient, lib_id: str, item_id: str) -> str | None: + resp = await client.get(f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + return resp.json()["folder_id"] + + +# --------------------------------------------------------------------------- +# Upload with folder_id +# --------------------------------------------------------------------------- + + +async def test_upload_lands_at_root_by_default(client: AsyncClient, library: dict) -> None: + """An upload with no folder_id ends up at the library root.""" + item_id = await _upload(client, library["id"]) + assert await _folder_id_of(client, library["id"], item_id) is None + + +async def test_upload_with_folder_id(client: AsyncClient, library: dict) -> None: + """An upload with a valid folder_id lands inside that folder.""" + folder_id = await _create_folder(client, library["id"], "Q1 Research") + item_id = await _upload(client, library["id"], folder_id=folder_id) + assert await _folder_id_of(client, library["id"], item_id) == folder_id + + +async def test_upload_with_invalid_folder_400(client: AsyncClient, library: dict) -> None: + """An unknown folder_id on upload is rejected with 400.""" + resp = await client.post( + f"/libraries/{library['id']}/import/file", + headers=AUTH_HEADERS, + files=text_file("# x", "a.md"), + data={"plugin_name": "simple_import", "title": "x", "folder_id": unique_id("nope")}, + ) + assert resp.status_code == 400, resp.text + + +async def test_upload_cross_library_folder_400(client: AsyncClient) -> None: + """A folder_id from another library is rejected with 400.""" + lib_a = await _create_library(client) + lib_b = await _create_library(client) + folder_in_b = await _create_folder(client, lib_b, "Foo") + resp = await client.post( + f"/libraries/{lib_a}/import/file", + headers=AUTH_HEADERS, + files=text_file("# x", "a.md"), + data={"plugin_name": "simple_import", "title": "x", "folder_id": folder_in_b}, + ) + assert resp.status_code == 400, resp.text + + +# --------------------------------------------------------------------------- +# Bulk move +# --------------------------------------------------------------------------- + + +async def test_bulk_move_items_to_folder(client: AsyncClient, library: dict) -> None: + """Bulk-moving two items reports moved=2 and updates each item's folder.""" + folder_id = await _create_folder(client, library["id"], "Target") + item1 = await _upload(client, library["id"], title="One") + item2 = await _upload(client, library["id"], title="Two") + resp = await client.post( + f"/libraries/{library['id']}/items/move", + headers=AUTH_HEADERS, + json={"item_ids": [item1, item2], "folder_id": folder_id}, + ) + assert resp.status_code == 200, resp.text + assert resp.json() == {"moved": 2, "folder_id": folder_id} + for iid in (item1, item2): + assert await _folder_id_of(client, library["id"], iid) == folder_id + + +async def test_single_move_item(client: AsyncClient, library: dict) -> None: + """A single-item move list works the same as a bulk move.""" + folder_id = await _create_folder(client, library["id"], "Solo") + item_id = await _upload(client, library["id"]) + resp = await client.post( + f"/libraries/{library['id']}/items/move", + headers=AUTH_HEADERS, + json={"item_ids": [item_id], "folder_id": folder_id}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["moved"] == 1 + assert await _folder_id_of(client, library["id"], item_id) == folder_id + + +async def test_move_items_to_root(client: AsyncClient, library: dict) -> None: + """Moving with folder_id=null sends items back to the root.""" + folder_id = await _create_folder(client, library["id"], "T") + item_id = await _upload(client, library["id"], folder_id=folder_id) + resp = await client.post( + f"/libraries/{library['id']}/items/move", + headers=AUTH_HEADERS, + json={"item_ids": [item_id], "folder_id": None}, + ) + assert resp.status_code == 200, resp.text + assert await _folder_id_of(client, library["id"], item_id) is None + + +async def test_move_missing_library_404(client: AsyncClient) -> None: + """Moving items within an unknown library returns 404.""" + resp = await client.post( + f"/libraries/{unique_id('missing')}/items/move", + headers=AUTH_HEADERS, + json={"item_ids": [unique_id("item")], "folder_id": None}, + ) + assert resp.status_code == 404, resp.text + + +async def test_move_unknown_item_404(client: AsyncClient, library: dict) -> None: + """A move referencing an item that does not exist returns 404.""" + resp = await client.post( + f"/libraries/{library['id']}/items/move", + headers=AUTH_HEADERS, + json={"item_ids": [unique_id("ghost")], "folder_id": None}, + ) + assert resp.status_code == 404, resp.text + + +async def test_move_to_missing_folder_404(client: AsyncClient, library: dict) -> None: + """A destination folder that does not exist returns 404.""" + item_id = await _upload(client, library["id"]) + resp = await client.post( + f"/libraries/{library['id']}/items/move", + headers=AUTH_HEADERS, + json={"item_ids": [item_id], "folder_id": unique_id("nope")}, + ) + assert resp.status_code == 404, resp.text + + +async def test_move_items_cross_library_400(client: AsyncClient) -> None: + """Moving an item that belongs to another library returns 400.""" + lib_a = await _create_library(client) + lib_b = await _create_library(client) + item_in_b = await _upload(client, lib_b) + resp = await client.post( + f"/libraries/{lib_a}/items/move", + headers=AUTH_HEADERS, + json={"item_ids": [item_in_b], "folder_id": None}, + ) + assert resp.status_code == 400, resp.text + + +async def test_move_to_cross_library_folder_400(client: AsyncClient) -> None: + """Moving an item into a folder from another library returns 400.""" + lib_a = await _create_library(client) + lib_b = await _create_library(client) + item_a = await _upload(client, lib_a) + folder_b = await _create_folder(client, lib_b, "Foreign") + resp = await client.post( + f"/libraries/{lib_a}/items/move", + headers=AUTH_HEADERS, + json={"item_ids": [item_a], "folder_id": folder_b}, + ) + assert resp.status_code == 400, resp.text + + +async def test_move_empty_list_422(client: AsyncClient, library: dict) -> None: + """An empty item_ids list violates min_length=1 → 422.""" + resp = await client.post( + f"/libraries/{library['id']}/items/move", + headers=AUTH_HEADERS, + json={"item_ids": [], "folder_id": None}, + ) + assert resp.status_code == 422, resp.text + + +async def test_move_payload_cap_422(client: AsyncClient, library: dict) -> None: + """More than 500 item_ids violates max_length=500 → 422.""" + big = [f"item-{i}" for i in range(501)] + resp = await client.post( + f"/libraries/{library['id']}/items/move", + headers=AUTH_HEADERS, + json={"item_ids": big, "folder_id": None}, + ) + assert resp.status_code == 422, resp.text + + +# --------------------------------------------------------------------------- +# Folder delete reparents items +# --------------------------------------------------------------------------- + + +async def test_folder_delete_reparents_items(client: AsyncClient, library: dict) -> None: + """Deleting a folder re-homes its items under the folder's parent.""" + parent = await _create_folder(client, library["id"], "Parent") + child = await _create_folder(client, library["id"], "Child", parent) + item_id = await _upload(client, library["id"], folder_id=child) + resp = await client.delete( + f"/libraries/{library['id']}/folders/{child}", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + assert await _folder_id_of(client, library["id"], item_id) == parent diff --git a/library-manager/tests/integration/test_libraries.py b/library-manager/tests/integration/test_libraries.py new file mode 100644 index 000000000..dc952c8f4 --- /dev/null +++ b/library-manager/tests/integration/test_libraries.py @@ -0,0 +1,304 @@ +"""Library CRUD + pagination boundaries. + +Ports ``tests/test_libraries.py`` and adds pagination-boundary coverage for +``GET /libraries`` (limit bounds 1..100, offset windowing). Isolation comes +from unique library ids / per-test organization ids — never from global +counts. + +Source: ``backend/routers/libraries.py``, ``backend/schemas/libraries.py``. +""" + +from __future__ import annotations + +from _helpers import AUTH_HEADERS, library_payload, unique_id +from httpx import AsyncClient + +# --------------------------------------------------------------------------- +# Create +# --------------------------------------------------------------------------- + + +async def test_create_library(client: AsyncClient) -> None: + """POST /libraries returns 201 with id, org, and a zero item_count.""" + payload = library_payload(organization_id="org-crud") + resp = await client.post("/libraries", headers=AUTH_HEADERS, json=payload) + assert resp.status_code == 201, resp.text + data = resp.json() + assert data["id"] == payload["id"] + assert data["organization_id"] == "org-crud" + assert data["item_count"] == 0 + assert data["import_config"] is None + + +async def test_create_with_import_config(client: AsyncClient) -> None: + """An import_config supplied at creation round-trips on the response.""" + payload = library_payload(import_config={"image_descriptions": "basic"}) + resp = await client.post("/libraries", headers=AUTH_HEADERS, json=payload) + assert resp.status_code == 201, resp.text + assert resp.json()["import_config"] == {"image_descriptions": "basic"} + + +async def test_duplicate_id_rejected(client: AsyncClient) -> None: + """Reusing an existing library id returns 409.""" + payload = library_payload() + resp1 = await client.post("/libraries", headers=AUTH_HEADERS, json=payload) + assert resp1.status_code == 201, resp1.text + dup = library_payload(id=payload["id"], name="Different Name") + resp2 = await client.post("/libraries", headers=AUTH_HEADERS, json=dup) + assert resp2.status_code == 409, resp2.text + + +async def test_duplicate_name_in_org_rejected(client: AsyncClient) -> None: + """Two libraries with the same name in one org return 409 on the second.""" + org = unique_id("org") + name = "Duplicate Test" + resp1 = await client.post( + "/libraries", headers=AUTH_HEADERS, json=library_payload(organization_id=org, name=name) + ) + assert resp1.status_code == 201, resp1.text + resp2 = await client.post( + "/libraries", headers=AUTH_HEADERS, json=library_payload(organization_id=org, name=name) + ) + assert resp2.status_code == 409, resp2.text + + +async def test_same_name_different_org_allowed(client: AsyncClient) -> None: + """The name uniqueness constraint is scoped per organization.""" + name = "Shared Name" + r1 = await client.post( + "/libraries", + headers=AUTH_HEADERS, + json=library_payload(organization_id=unique_id("org"), name=name), + ) + r2 = await client.post( + "/libraries", + headers=AUTH_HEADERS, + json=library_payload(organization_id=unique_id("org"), name=name), + ) + assert r1.status_code == 201, r1.text + assert r2.status_code == 201, r2.text + + +async def test_missing_required_field_422(client: AsyncClient) -> None: + """Omitting a required field (name) fails schema validation with 422.""" + resp = await client.post( + "/libraries", + headers=AUTH_HEADERS, + json={"id": unique_id("lib"), "organization_id": "org-test"}, + ) + assert resp.status_code == 422, resp.text + + +# --------------------------------------------------------------------------- +# Get +# --------------------------------------------------------------------------- + + +async def test_get_library(client: AsyncClient, library: dict) -> None: + """GET /libraries/{id} returns details including item_count.""" + resp = await client.get(f"/libraries/{library['id']}", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["name"] == library["name"] + assert data["item_count"] == 0 + + +async def test_get_nonexistent_library_404(client: AsyncClient) -> None: + """GET on an unknown library id returns 404.""" + resp = await client.get(f"/libraries/{unique_id('missing')}", headers=AUTH_HEADERS) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# List + pagination boundaries +# --------------------------------------------------------------------------- + + +async def test_list_includes_created(client: AsyncClient, library: dict) -> None: + """Listing by org includes a library just created in that org.""" + resp = await client.get( + "/libraries", + headers=AUTH_HEADERS, + params={"organization_id": library["organization_id"]}, + ) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["total"] >= 1 + assert library["id"] in [lib["id"] for lib in data["libraries"]] + + +async def test_list_requires_organization_id_422(client: AsyncClient) -> None: + """organization_id is a required query parameter.""" + resp = await client.get("/libraries", headers=AUTH_HEADERS) + assert resp.status_code == 422, resp.text + + +async def test_list_filters_by_org(client: AsyncClient) -> None: + """Listing one org never leaks libraries from another org.""" + org_a, org_b = unique_id("org"), unique_id("org") + await client.post( + "/libraries", headers=AUTH_HEADERS, json=library_payload(organization_id=org_a) + ) + await client.post( + "/libraries", headers=AUTH_HEADERS, json=library_payload(organization_id=org_b) + ) + resp = await client.get("/libraries", headers=AUTH_HEADERS, params={"organization_id": org_a}) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["total"] == 1 + assert {lib["organization_id"] for lib in data["libraries"]} == {org_a} + + +async def test_list_limit_zero_rejected_422(client: AsyncClient) -> None: + """limit below the minimum (1) is rejected by the query validator.""" + resp = await client.get( + "/libraries", headers=AUTH_HEADERS, params={"organization_id": "org-test", "limit": 0} + ) + assert resp.status_code == 422, resp.text + + +async def test_list_limit_over_max_rejected_422(client: AsyncClient) -> None: + """limit above the maximum (100) is rejected by the query validator.""" + resp = await client.get( + "/libraries", headers=AUTH_HEADERS, params={"organization_id": "org-test", "limit": 101} + ) + assert resp.status_code == 422, resp.text + + +async def test_list_negative_offset_rejected_422(client: AsyncClient) -> None: + """offset below 0 is rejected by the query validator.""" + resp = await client.get( + "/libraries", headers=AUTH_HEADERS, params={"organization_id": "org-test", "offset": -1} + ) + assert resp.status_code == 422, resp.text + + +async def test_list_limit_boundaries_accepted(client: AsyncClient) -> None: + """limit=1 and limit=100 are both accepted (inclusive bounds).""" + org = unique_id("org") + await client.post( + "/libraries", headers=AUTH_HEADERS, json=library_payload(organization_id=org) + ) + for limit in (1, 100): + resp = await client.get( + "/libraries", + headers=AUTH_HEADERS, + params={"organization_id": org, "limit": limit}, + ) + assert resp.status_code == 200, resp.text + + +async def test_offset_pagination_windows(client: AsyncClient) -> None: + """Create 5 libs in one org; assert limit/offset returns the right slices.""" + org = unique_id("org") + created = [] + for _ in range(5): + resp = await client.post( + "/libraries", headers=AUTH_HEADERS, json=library_payload(organization_id=org) + ) + assert resp.status_code == 201, resp.text + created.append(resp.json()["id"]) + + # Full list, paged 2 + 2 + 1, must reconstruct the whole set with no overlap. + seen: list[str] = [] + for offset in (0, 2, 4): + resp = await client.get( + "/libraries", + headers=AUTH_HEADERS, + params={"organization_id": org, "limit": 2, "offset": offset}, + ) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["total"] == 5 + page = [lib["id"] for lib in data["libraries"]] + assert len(page) == (2 if offset < 4 else 1) + seen.extend(page) + + assert sorted(seen) == sorted(created) + assert len(set(seen)) == 5 # no duplicates across windows + + +async def test_offset_past_end_returns_empty(client: AsyncClient) -> None: + """An offset beyond the result set returns an empty page with the true total.""" + org = unique_id("org") + await client.post( + "/libraries", headers=AUTH_HEADERS, json=library_payload(organization_id=org) + ) + resp = await client.get( + "/libraries", + headers=AUTH_HEADERS, + params={"organization_id": org, "limit": 20, "offset": 50}, + ) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["total"] == 1 + assert data["libraries"] == [] + + +# --------------------------------------------------------------------------- +# Delete +# --------------------------------------------------------------------------- + + +async def test_delete_library(client: AsyncClient) -> None: + """DELETE removes the library; a subsequent GET 404s.""" + payload = library_payload(organization_id="org-del") + await client.post("/libraries", headers=AUTH_HEADERS, json=payload) + resp = await client.delete(f"/libraries/{payload['id']}", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + assert payload["id"] in resp.json()["message"] + follow = await client.get(f"/libraries/{payload['id']}", headers=AUTH_HEADERS) + assert follow.status_code == 404, follow.text + + +async def test_delete_nonexistent_404(client: AsyncClient) -> None: + """Deleting a library that does not exist returns 404.""" + resp = await client.delete(f"/libraries/{unique_id('missing')}", headers=AUTH_HEADERS) + assert resp.status_code == 404, resp.text + + +# --------------------------------------------------------------------------- +# Import config +# --------------------------------------------------------------------------- + + +async def test_import_config_roundtrip(client: AsyncClient, library: dict) -> None: + """PUT then GET import-config returns the stored config with a warning on PUT.""" + config = {"image_descriptions": "llm", "max_discovery_depth": 5} + put = await client.put( + f"/libraries/{library['id']}/import-config", headers=AUTH_HEADERS, json=config + ) + assert put.status_code == 200, put.text + assert "warning" in put.json() + assert put.json()["import_config"] == config + + get = await client.get( + f"/libraries/{library['id']}/import-config", headers=AUTH_HEADERS + ) + assert get.status_code == 200, get.text + assert get.json()["import_config"] == config + + +async def test_import_config_default_empty(client: AsyncClient, library: dict) -> None: + """A library created without import_config reports an empty config dict.""" + resp = await client.get( + f"/libraries/{library['id']}/import-config", headers=AUTH_HEADERS + ) + assert resp.status_code == 200, resp.text + assert resp.json()["import_config"] == {} + + +async def test_get_import_config_missing_library_404(client: AsyncClient) -> None: + """GET import-config for an unknown library returns 404.""" + resp = await client.get( + f"/libraries/{unique_id('missing')}/import-config", headers=AUTH_HEADERS + ) + assert resp.status_code == 404, resp.text + + +async def test_put_import_config_missing_library_404(client: AsyncClient) -> None: + """PUT import-config for an unknown library returns 404.""" + resp = await client.put( + f"/libraries/{unique_id('missing')}/import-config", headers=AUTH_HEADERS, json={} + ) + assert resp.status_code == 404, resp.text diff --git a/library-manager/tests/integration/test_lifespan.py b/library-manager/tests/integration/test_lifespan.py new file mode 100644 index 000000000..f1ac824fd --- /dev/null +++ b/library-manager/tests/integration/test_lifespan.py @@ -0,0 +1,86 @@ +"""Cover the FastAPI app lifespan, startup helpers, and middleware. + +Drives ``main.lifespan`` directly (exercising ensure_directories → init_db → +_discover_plugins → purge_stale_uploads → recover_stale_jobs → start_worker → +stop_worker), the ``purge_stale_uploads`` helper, the request-logging +middleware, and the production docs gating. + +Source under test: ``backend/main.py`` + ``routers/importing.py`` (purge). +""" + +from __future__ import annotations + +import os +import time + +from _helpers import AUTH_HEADERS +from httpx import ASGITransport, AsyncClient +from tasks import worker + + +async def test_lifespan_starts_and_stops_worker() -> None: + """Entering the lifespan starts the worker; exiting stops it.""" + from main import app, lifespan # noqa: PLC0415 + + # Ensure a clean baseline regardless of fixture ordering. + if worker.is_worker_running(): + await worker.stop_worker() + + assert not worker.is_worker_running() + async with lifespan(app): + assert worker.is_worker_running() + assert not worker.is_worker_running() + + +async def test_purge_stale_uploads_removes_old_files() -> None: + """purge_stale_uploads deletes temp uploads older than the max age.""" + from routers.importing import ( # noqa: PLC0415 + _UPLOAD_DIR, + _UPLOAD_MAX_AGE_SECONDS, + purge_stale_uploads, + ) + + _UPLOAD_DIR.mkdir(parents=True, exist_ok=True) + stale = _UPLOAD_DIR / "stale_upload.bin" + stale.write_bytes(b"old") + # Backdate the mtime well past the cutoff. + old = time.time() - (_UPLOAD_MAX_AGE_SECONDS + 60) + os.utime(stale, (old, old)) + + fresh = _UPLOAD_DIR / "fresh_upload.bin" + fresh.write_bytes(b"new") + + purge_stale_uploads() + + assert not stale.exists(), "stale upload should have been purged" + assert fresh.exists(), "fresh upload should remain" + fresh.unlink(missing_ok=True) + + +async def test_purge_stale_uploads_noop_without_dir() -> None: + """purge_stale_uploads returns cleanly when the upload dir is absent.""" + import shutil # noqa: PLC0415 + + from routers.importing import _UPLOAD_DIR, purge_stale_uploads # noqa: PLC0415 + + if _UPLOAD_DIR.exists(): + shutil.rmtree(_UPLOAD_DIR, ignore_errors=True) + purge_stale_uploads() # must not raise + assert not _UPLOAD_DIR.exists() + + +async def test_request_logging_middleware_passthrough(client: AsyncClient) -> None: + """A normal request flows through the log_requests middleware and returns 200.""" + resp = await client.get("/health") + assert resp.status_code == 200, resp.text + + +async def test_docs_disabled_in_non_debug_env() -> None: + """With LOG_LEVEL != DEBUG, docs are gated off and /docs returns 404.""" + from main import app # noqa: PLC0415 + + assert app.docs_url is None + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + resp = await ac.get("/docs", headers=AUTH_HEADERS) + assert resp.status_code == 404, resp.text diff --git a/library-manager/tests/integration/test_migrations.py b/library-manager/tests/integration/test_migrations.py new file mode 100644 index 000000000..5fe7d2d9a --- /dev/null +++ b/library-manager/tests/integration/test_migrations.py @@ -0,0 +1,69 @@ +"""Alembic migration round-trip. + +Runs the real migration chain (``upgrade head`` then ``downgrade base``) +against an isolated temp database via the Alembic CLI, so a broken or +non-reversible revision fails the suite. Uses a subprocess (the Alembic CLI) +against its own throwaway ``DATA_DIR`` — it does not touch the session DB. +""" + +from __future__ import annotations + +import os +import sqlite3 +import subprocess +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parents[2] / "backend" + + +def _alembic(args: list[str], data_dir: str) -> subprocess.CompletedProcess: + """Invoke the Alembic CLI in the backend dir against ``data_dir``.""" + env = {**os.environ, "DATA_DIR": data_dir, "LAMB_API_TOKEN": "test"} + return subprocess.run( + [sys.executable, "-m", "alembic", "-c", "alembic.ini", *args], + cwd=_BACKEND, + env=env, + capture_output=True, + text=True, + ) + + +def _tables(db_path: Path) -> set[str]: + """Return the set of table names in the SQLite DB at ``db_path``.""" + con = sqlite3.connect(db_path) + try: + rows = con.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall() + finally: + con.close() + return {r[0] for r in rows} + + +@pytest.mark.slow +def test_upgrade_then_downgrade_roundtrip(tmp_path): + """`alembic upgrade head` creates every table; `downgrade base` removes them.""" + data_dir = str(tmp_path) + db_path = tmp_path / "library-manager.db" + + up = _alembic(["upgrade", "head"], data_dir) + assert up.returncode == 0, up.stderr + + tables = _tables(db_path) + assert { + "organizations", + "libraries", + "content_folders", + "content_items", + "content_images", + "import_jobs", + "alembic_version", + } <= tables + + down = _alembic(["downgrade", "base"], data_dir) + assert down.returncode == 0, down.stderr + + tables = _tables(db_path) + assert "content_items" not in tables + assert "organizations" not in tables diff --git a/library-manager/tests/integration/test_plugin_params_passthrough.py b/library-manager/tests/integration/test_plugin_params_passthrough.py new file mode 100644 index 000000000..8dcc6fe98 --- /dev/null +++ b/library-manager/tests/integration/test_plugin_params_passthrough.py @@ -0,0 +1,174 @@ +"""``plugin_params`` flow from the HTTP wire to ``import_content`` kwargs. + +Ported and extended from the old flat +``tests/test_plugin_params_passthrough.py``. The router must forward every +schema-declared key unchanged (the "drop in a plugin, change nothing else" +contract), merge the YouTube top-level ``language`` field into plugin_params, +and the worker must sanitize out unknown / advanced-in-SIMPLIFIED keys before +calling the plugin. + +Source under test: ``routers/importing.py`` (param merge), +``services/import_service.py`` (sanitize + dispatch), +``plugins/base.PluginRegistry.sanitize_params``. +""" + +from __future__ import annotations + +from typing import Any + +from _helpers import AUTH_HEADERS, poll_until_ready +from httpx import AsyncClient +from plugins.base import ImportResult, PluginRegistry + + +def _spy_capturing(captured: dict[str, Any]): + """Build an ``import_content`` replacement that records its kwargs.""" + + def spy(self, source_path, *, api_keys=None, **kwargs): + captured.update(kwargs) + return ImportResult( + full_text="# stub", pages=[], images=[], metadata={}, source_ref={} + ) + + return spy + + +async def test_url_plugin_params_reach_import_content( + client: AsyncClient, library: dict, monkeypatch +) -> None: + """Schema-declared url_import params arrive unchanged at the plugin.""" + lib_id = library["id"] + captured: dict[str, Any] = {} + plugin_cls = PluginRegistry.get_plugin("url_import").__class__ + monkeypatch.setattr(plugin_cls, "import_content", _spy_capturing(captured)) + + custom = {"limit": 7, "max_discovery_depth": 3} + resp = await client.post( + f"/libraries/{lib_id}/import/url", + headers=AUTH_HEADERS, + json={ + "url": "https://example.com", + "plugin_name": "url_import", + "title": "Passthrough", + "plugin_params": custom, + }, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + assert await poll_until_ready(client, lib_id, item_id) == "ready" + + for key, value in custom.items(): + assert captured.get(key) == value, (key, captured) + + +async def test_youtube_language_from_plugin_params_wins( + client: AsyncClient, library: dict, monkeypatch +) -> None: + """plugin_params['language'] overrides the deprecated top-level field.""" + lib_id = library["id"] + captured: dict[str, Any] = {} + plugin_cls = PluginRegistry.get_plugin("youtube_transcript_import").__class__ + monkeypatch.setattr(plugin_cls, "import_content", _spy_capturing(captured)) + + resp = await client.post( + f"/libraries/{lib_id}/import/youtube", + headers=AUTH_HEADERS, + json={ + "video_url": "https://www.youtube.com/watch?v=stub", + "plugin_name": "youtube_transcript_import", + "title": "YT Params", + "language": "en", + "plugin_params": {"language": "fr"}, + }, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + assert await poll_until_ready(client, lib_id, item_id) == "ready" + assert captured.get("language") == "fr", captured + + +async def test_youtube_top_level_language_fallback( + client: AsyncClient, library: dict, monkeypatch +) -> None: + """When plugin_params omits language, the top-level field fills it in.""" + lib_id = library["id"] + captured: dict[str, Any] = {} + plugin_cls = PluginRegistry.get_plugin("youtube_transcript_import").__class__ + monkeypatch.setattr(plugin_cls, "import_content", _spy_capturing(captured)) + + resp = await client.post( + f"/libraries/{lib_id}/import/youtube", + headers=AUTH_HEADERS, + json={ + "video_url": "https://www.youtube.com/watch?v=stub", + "plugin_name": "youtube_transcript_import", + "title": "YT Fallback", + "language": "de", + }, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + assert await poll_until_ready(client, lib_id, item_id) == "ready" + assert captured.get("language") == "de", captured + + +async def test_unknown_params_sanitized_out( + client: AsyncClient, library: dict, monkeypatch +) -> None: + """Unknown keys never reach the plugin (sanitize_params strips them).""" + lib_id = library["id"] + captured: dict[str, Any] = {} + plugin_cls = PluginRegistry.get_plugin("url_import").__class__ + monkeypatch.setattr(plugin_cls, "import_content", _spy_capturing(captured)) + + resp = await client.post( + f"/libraries/{lib_id}/import/url", + headers=AUTH_HEADERS, + json={ + "url": "https://example.com", + "plugin_name": "url_import", + "title": "Sanitize", + "plugin_params": {"limit": 5, "totally_made_up": "evil", "__proto__": 1}, + }, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + assert await poll_until_ready(client, lib_id, item_id) == "ready" + + assert captured.get("limit") == 5, captured + assert "totally_made_up" not in captured, captured + assert "__proto__" not in captured, captured + + +async def test_simplified_mode_drops_advanced_params( + client: AsyncClient, library: dict, monkeypatch +) -> None: + """In SIMPLIFIED mode, advanced params are stripped before the plugin call. + + ``url_import``'s ``limit`` is an *advanced* param; ``max_discovery_depth`` + is not. With governance set to SIMPLIFIED the advanced one must not reach + the plugin while the non-advanced one still does. + """ + lib_id = library["id"] + captured: dict[str, Any] = {} + plugin_cls = PluginRegistry.get_plugin("url_import").__class__ + monkeypatch.setattr(plugin_cls, "import_content", _spy_capturing(captured)) + # sanitize_params reads PLUGIN_ from the environment at call time. + monkeypatch.setenv("PLUGIN_URL_IMPORT", "SIMPLIFIED") + + resp = await client.post( + f"/libraries/{lib_id}/import/url", + headers=AUTH_HEADERS, + json={ + "url": "https://example.com", + "plugin_name": "url_import", + "title": "Simplified", + "plugin_params": {"limit": 9, "max_discovery_depth": 4}, + }, + ) + assert resp.status_code == 202, resp.text + item_id = resp.json()["item_id"] + assert await poll_until_ready(client, lib_id, item_id) == "ready" + + assert captured.get("max_discovery_depth") == 4, captured + assert "limit" not in captured, captured diff --git a/library-manager/tests/integration/test_system.py b/library-manager/tests/integration/test_system.py new file mode 100644 index 000000000..65cf2fb6f --- /dev/null +++ b/library-manager/tests/integration/test_system.py @@ -0,0 +1,63 @@ +"""System endpoints: health check and plugin listing. + +Ports ``tests/test_system.py`` into the integration tier. ``GET /health`` +reports component health (database + worker) and needs no auth; ``GET +/plugins`` lists every registered import plugin in the documented shape. + +Source: ``backend/routers/system.py``. +""" + +from __future__ import annotations + +from _helpers import AUTH_HEADERS +from httpx import AsyncClient + +# Plugins discovered at session start (root conftest calls _discover_plugins). +_EXPECTED_PLUGINS = { + "simple_import", + "markitdown_import", + "markitdown_plus_import", + "url_import", + "youtube_transcript_import", +} + + +async def test_health_ok_with_worker_running(client: AsyncClient) -> None: + """With the worker running, /health reports status ok and both checks ok.""" + resp = await client.get("/health") + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["status"] == "ok", data + assert data["service"] == "library-manager" + assert data["version"] == "1.0.0" + assert data["checks"] == {"database": "ok", "worker": "ok"} + + +async def test_health_degraded_without_worker(client_no_worker: AsyncClient) -> None: + """Without the worker the DB is still ok but status degrades.""" + resp = await client_no_worker.get("/health") + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["status"] == "degraded", data + assert data["checks"]["database"] == "ok" + assert data["checks"]["worker"] == "error" + + +async def test_list_plugins_lists_registered(client: AsyncClient) -> None: + """/plugins returns all registered import plugins.""" + resp = await client.get("/plugins", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + plugins = resp.json()["plugins"] + names = {p["name"] for p in plugins} + assert _EXPECTED_PLUGINS.issubset(names), names + + +async def test_plugin_shape(client: AsyncClient) -> None: + """Each plugin entry carries name, description, source types, and parameters.""" + resp = await client.get("/plugins", headers=AUTH_HEADERS) + assert resp.status_code == 200, resp.text + for plugin in resp.json()["plugins"]: + assert isinstance(plugin["name"], str) and plugin["name"] + assert "description" in plugin + assert isinstance(plugin["supported_source_types"], list) + assert isinstance(plugin["parameters"], list) diff --git a/library-manager/tests/integration/test_worker.py b/library-manager/tests/integration/test_worker.py new file mode 100644 index 000000000..f3b3f4a4f --- /dev/null +++ b/library-manager/tests/integration/test_worker.py @@ -0,0 +1,674 @@ +"""Strengthened background-worker tests. + +Goes beyond the old ``tests/test_load.py`` (which only asserted that every +import completed) by measuring *peak* concurrency against +``MAX_CONCURRENT_IMPORTS``, exercising every status transition (completed / +failed / timeout), and driving ``recover_stale_jobs`` directly. + +Custom test plugins are registered into ``PluginRegistry._plugins`` and +always removed in a ``finally`` so the shared session registry never leaks. + +Source under test: ``tasks/worker.py`` + ``services/import_service.py``. +""" + +from __future__ import annotations + +import contextlib +import threading +import time +import uuid +from datetime import UTC, datetime + +import config +import pytest +from database.connection import get_session_direct +from database.models import ContentItem, ImportJob +from httpx import AsyncClient +from plugins.base import ImportResult, LibraryImportPlugin, PluginRegistry +from tasks import worker + + +@contextlib.contextmanager +def _registered(plugin_cls: type[LibraryImportPlugin]): + """Register a plugin for the duration of the block, then restore. + + Snapshots the whole ``_plugins`` dict so a test never leaks a custom + plugin into the shared session registry. + """ + snapshot = dict(PluginRegistry._plugins) + PluginRegistry._plugins[plugin_cls.name] = plugin_cls + try: + yield + finally: + PluginRegistry._plugins.clear() + PluginRegistry._plugins.update(snapshot) + + +def _queue_job_row(plugin_name: str, library_id: str, org_id: str, source_url: str) -> str: + """Insert a pending ImportJob + ContentItem row directly; return job_id.""" + item_id = str(uuid.uuid4()) + job_id = str(uuid.uuid4()) + session = get_session_direct() + try: + from services.library_service import ensure_organization # noqa: PLC0415 + + ensure_organization(session, org_id) + session.add( + ContentItem( + id=item_id, + library_id=library_id, + organization_id=org_id, + title="Worker Test Item", + source_type="url", + base_path=f"/tmp/{item_id}", + permalink_base=f"/docs/{org_id}/{library_id}/{item_id}", + import_plugin=plugin_name, + status="pending", + ) + ) + session.add( + ImportJob( + id=job_id, + content_item_id=item_id, + library_id=library_id, + organization_id=org_id, + source_type="url", + plugin_name=plugin_name, + source_url=source_url, + title="Worker Test Item", + status="pending", + ) + ) + session.commit() + finally: + session.close() + return job_id + + +# --------------------------------------------------------------------------- +# Concurrency cap +# --------------------------------------------------------------------------- + + +@pytest.mark.slow +async def test_peak_concurrency_capped_at_max( + client_no_worker: AsyncClient, library: dict +) -> None: + """Peak concurrent imports equals MAX_CONCURRENT_IMPORTS, never more.""" + lib_id = library["id"] + org_id = library["organization_id"] + max_concurrent = config.MAX_CONCURRENT_IMPORTS + assert max_concurrent == 3, "test pins the documented default of 3" + + lock = threading.Lock() + state = {"current": 0, "peak": 0} + + class SlowConcurrencyPlugin(LibraryImportPlugin): + name = "slow_concurrency_test" + description = "Sleeps while tracking peak concurrency." + supported_source_types = {"url"} + + def import_content(self, source_path, *, api_keys=None, **kwargs): + with lock: + state["current"] += 1 + state["peak"] = max(state["peak"], state["current"]) + try: + time.sleep(0.6) + finally: + with lock: + state["current"] -= 1 + return ImportResult(full_text="# done", pages=[], images=[], metadata={}, source_ref={}) + + def get_parameters(self): + return [] + + num_jobs = max_concurrent + 3 + with _registered(SlowConcurrencyPlugin): + job_ids = [ + _queue_job_row("slow_concurrency_test", lib_id, org_id, f"https://x/{i}") + for i in range(num_jobs) + ] + await worker.start_worker() + try: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + session = get_session_direct() + try: + done = ( + session.query(ImportJob) + .filter( + ImportJob.id.in_(job_ids), + ImportJob.status.in_(("completed", "failed")), + ) + .count() + ) + finally: + session.close() + if done == num_jobs: + break + import asyncio # noqa: PLC0415 + + await asyncio.sleep(0.2) + finally: + await worker.stop_worker() + + assert done == num_jobs, f"only {done}/{num_jobs} finished" + assert state["peak"] == max_concurrent, ( + f"peak concurrency was {state['peak']}, expected {max_concurrent}" + ) + + +# --------------------------------------------------------------------------- +# Status transitions +# --------------------------------------------------------------------------- + + +async def test_happy_path_status_transition( + client_no_worker: AsyncClient, library: dict +) -> None: + """A successful job lands the item+job at completed/ready.""" + lib_id = library["id"] + org_id = library["organization_id"] + + class OkPlugin(LibraryImportPlugin): + name = "ok_transition_test" + description = "Always succeeds." + supported_source_types = {"url"} + + def import_content(self, source_path, *, api_keys=None, **kwargs): + return ImportResult( + full_text="# ok", pages=[], images=[], metadata={}, source_ref={"type": "url"} + ) + + def get_parameters(self): + return [] + + with _registered(OkPlugin): + job_id = _queue_job_row("ok_transition_test", lib_id, org_id, "https://x/ok") + worker._process_job_sync(job_id) + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + item = ( + session.query(ContentItem) + .filter(ContentItem.id == job.content_item_id) + .first() + ) + assert job.status == "completed" + assert job.attempts == 1 + assert job.started_at is not None + assert job.completed_at is not None + assert item.status == "ready" + finally: + session.close() + + +async def test_failure_records_truncated_error( + client_no_worker: AsyncClient, library: dict +) -> None: + """A raising plugin marks job+item failed and truncates the message to 500.""" + lib_id = library["id"] + org_id = library["organization_id"] + long_msg = "boom-" * 300 # 1500 chars + + class FailPlugin(LibraryImportPlugin): + name = "fail_transition_test" + description = "Always raises." + supported_source_types = {"url"} + + def import_content(self, source_path, *, api_keys=None, **kwargs): + raise RuntimeError(long_msg) + + def get_parameters(self): + return [] + + with _registered(FailPlugin): + job_id = _queue_job_row("fail_transition_test", lib_id, org_id, "https://x/fail") + worker._process_job_sync(job_id) + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + item = ( + session.query(ContentItem) + .filter(ContentItem.id == job.content_item_id) + .first() + ) + assert job.status == "failed" + assert item.status == "failed" + assert len(job.error_message) == 500 + assert job.error_message == long_msg[:500] + assert item.error_message == long_msg[:500] + finally: + session.close() + + +# --------------------------------------------------------------------------- +# Timeout +# --------------------------------------------------------------------------- + + +@pytest.mark.slow +async def test_job_timeout_marks_failed( + client_no_worker: AsyncClient, library: dict, monkeypatch +) -> None: + """A job exceeding IMPORT_TASK_TIMEOUT_SECONDS is marked failed with a timeout message.""" + lib_id = library["id"] + org_id = library["organization_id"] + + # The worker reads IMPORT_TASK_TIMEOUT_SECONDS from its own module-level + # import, so patch the name on tasks.worker (and config for good measure). + monkeypatch.setattr(worker, "IMPORT_TASK_TIMEOUT_SECONDS", 1) + monkeypatch.setattr(config, "IMPORT_TASK_TIMEOUT_SECONDS", 1) + + class SleepyPlugin(LibraryImportPlugin): + name = "timeout_test" + description = "Sleeps past the timeout." + supported_source_types = {"url"} + + def import_content(self, source_path, *, api_keys=None, **kwargs): + time.sleep(5) + return ImportResult(full_text="late", pages=[], images=[], metadata={}, source_ref={}) + + def get_parameters(self): + return [] + + with _registered(SleepyPlugin): + job_id = _queue_job_row("timeout_test", lib_id, org_id, "https://x/slow") + await worker._process_job_async(job_id) + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + item = ( + session.query(ContentItem) + .filter(ContentItem.id == job.content_item_id) + .first() + ) + assert job.status == "failed" + assert "timed out" in job.error_message.lower() + assert item.status == "failed" + assert "timed out" in item.error_message.lower() + finally: + session.close() + + +# --------------------------------------------------------------------------- +# recover_stale_jobs +# --------------------------------------------------------------------------- + + +async def test_recover_stale_job_resets_to_pending( + client_no_worker: AsyncClient, library: dict +) -> None: + """A processing job below the attempt cap is reset to pending on recovery.""" + lib_id = library["id"] + org_id = library["organization_id"] + job_id = _queue_job_row("simple_import", lib_id, org_id, "https://x/stale") + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + job.status = "processing" + job.attempts = worker._MAX_ATTEMPTS - 1 + session.commit() + finally: + session.close() + + worker.recover_stale_jobs() + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + assert job.status == "pending" + finally: + session.close() + + +async def test_recover_stale_job_exceeding_attempts_fails( + client_no_worker: AsyncClient, library: dict +) -> None: + """A processing job at/over the attempt cap is marked failed (item too).""" + lib_id = library["id"] + org_id = library["organization_id"] + job_id = _queue_job_row("simple_import", lib_id, org_id, "https://x/dead") + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + item_id = job.content_item_id + job.status = "processing" + job.attempts = worker._MAX_ATTEMPTS + session.commit() + finally: + session.close() + + worker.recover_stale_jobs() + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + item = session.query(ContentItem).filter(ContentItem.id == item_id).first() + assert job.status == "failed" + assert "max attempts" in job.error_message.lower() + assert item.status == "failed" + assert "max attempts" in item.error_message.lower() + finally: + session.close() + + +# --------------------------------------------------------------------------- +# store_api_keys / _job_api_keys +# --------------------------------------------------------------------------- + + +def test_store_api_keys_holds_and_pops() -> None: + """store_api_keys puts keys in memory; the worker pops them exactly once.""" + job_id = f"job-{uuid.uuid4().hex}" + worker.store_api_keys(job_id, {"openai_vision": "sk-mem"}) + assert worker._job_api_keys[job_id] == {"openai_vision": "sk-mem"} + + popped = worker._job_api_keys.pop(job_id, {}) + assert popped == {"openai_vision": "sk-mem"} + assert job_id not in worker._job_api_keys + + +def test_store_api_keys_noop_for_empty() -> None: + """store_api_keys with None/empty stores nothing (no memory footprint).""" + job_id = f"job-{uuid.uuid4().hex}" + worker.store_api_keys(job_id, None) + assert job_id not in worker._job_api_keys + worker.store_api_keys(job_id, {}) + assert job_id not in worker._job_api_keys + + +def test_process_job_sync_missing_job_is_noop() -> None: + """Processing a non-existent job id returns cleanly without raising.""" + worker._process_job_sync(f"missing-{uuid.uuid4().hex}") + + +async def test_failure_with_missing_item_still_records_job( + client_no_worker: AsyncClient, library: dict +) -> None: + """When a job's item row is gone, the error path still marks the job failed.""" + lib_id = library["id"] + org_id = library["organization_id"] + + class FailPlugin(LibraryImportPlugin): + name = "fail_no_item_test" + description = "Raises." + supported_source_types = {"url"} + + def import_content(self, source_path, *, api_keys=None, **kwargs): + raise RuntimeError("kaboom") + + def get_parameters(self): + return [] + + with _registered(FailPlugin): + job_id = _queue_job_row("fail_no_item_test", lib_id, org_id, "https://x/noitem") + # Delete the content item so the error-recording branch finds no item. + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + item = ( + session.query(ContentItem) + .filter(ContentItem.id == job.content_item_id) + .first() + ) + session.delete(item) + session.commit() + finally: + session.close() + worker._process_job_sync(job_id) + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + assert job.status == "failed" + assert job.error_message == "kaboom" + finally: + session.close() + + +async def test_plugin_not_found_fails_job( + client_no_worker: AsyncClient, library: dict +) -> None: + """A job referencing an unregistered plugin fails with a clear message.""" + lib_id = library["id"] + org_id = library["organization_id"] + job_id = _queue_job_row("ghost_plugin_xyz", lib_id, org_id, "https://x/ghost") + + worker._process_job_sync(job_id) + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + assert job.status == "failed" + assert "Plugin not found" in job.error_message + finally: + session.close() + + +async def test_job_with_no_source_fails( + client_no_worker: AsyncClient, library: dict +) -> None: + """A job with neither source_path nor source_url fails in execute_import_job.""" + lib_id = library["id"] + org_id = library["organization_id"] + job_id = _queue_job_row("simple_import", lib_id, org_id, "https://x/clearme") + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + job.source_url = None + job.source_path = None + session.commit() + finally: + session.close() + + worker._process_job_sync(job_id) + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + assert job.status == "failed" + assert "no source_path or source_url" in job.error_message + finally: + session.close() + + +async def test_execute_fails_when_item_row_missing( + client_no_worker: AsyncClient, library: dict +) -> None: + """execute_import_job raises (job fails) if the item row vanished post-queue.""" + lib_id = library["id"] + org_id = library["organization_id"] + + class OkPlugin(LibraryImportPlugin): + name = "missing_item_ok_test" + description = "Succeeds at plugin level." + supported_source_types = {"url"} + + def import_content(self, source_path, *, api_keys=None, **kwargs): + return ImportResult( + full_text="# ok", pages=[], images=[], metadata={}, source_ref={} + ) + + def get_parameters(self): + return [] + + with _registered(OkPlugin): + job_id = _queue_job_row("missing_item_ok_test", lib_id, org_id, "https://x/gone") + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + item = ( + session.query(ContentItem) + .filter(ContentItem.id == job.content_item_id) + .first() + ) + session.delete(item) + session.commit() + finally: + session.close() + worker._process_job_sync(job_id) + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + assert job.status == "failed" + assert "not found" in job.error_message.lower() + finally: + session.close() + + +async def test_write_failure_cleans_up_item_dir( + client_no_worker: AsyncClient, library: dict, monkeypatch +) -> None: + """If content writing fails, the partial item dir is removed and job fails.""" + import services.content_service as content_service # noqa: PLC0415 + from config import CONTENT_DIR # noqa: PLC0415 + + lib_id = library["id"] + org_id = library["organization_id"] + + class OkPlugin(LibraryImportPlugin): + name = "writefail_test" + description = "Succeeds at plugin level." + supported_source_types = {"url"} + + def import_content(self, source_path, *, api_keys=None, **kwargs): + return ImportResult( + full_text="# ok", pages=[], images=[], metadata={}, source_ref={} + ) + + def get_parameters(self): + return [] + + def boom(*args, **kwargs): + # Create the item dir first so the cleanup branch has something to rmtree. + item_id = kwargs["item_id"] + item_dir = CONTENT_DIR / org_id / lib_id / item_id + item_dir.mkdir(parents=True, exist_ok=True) + raise RuntimeError("disk write exploded") + + monkeypatch.setattr(content_service, "write_structured_content", boom) + + with _registered(OkPlugin): + job_id = _queue_job_row("writefail_test", lib_id, org_id, "https://x/wf") + worker._process_job_sync(job_id) + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + item_id = job.content_item_id + assert job.status == "failed" + assert "exploded" in job.error_message + finally: + session.close() + assert not (CONTENT_DIR / org_id / lib_id / item_id).exists() + + +async def test_recover_stale_job_over_cap_with_missing_item( + client_no_worker: AsyncClient, library: dict +) -> None: + """Over-cap recovery still marks the job failed when its item row is gone.""" + lib_id = library["id"] + org_id = library["organization_id"] + job_id = _queue_job_row("simple_import", lib_id, org_id, "https://x/orphan") + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + item = ( + session.query(ContentItem) + .filter(ContentItem.id == job.content_item_id) + .first() + ) + job.status = "processing" + job.attempts = worker._MAX_ATTEMPTS + session.delete(item) + session.commit() + finally: + session.close() + + worker.recover_stale_jobs() + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + assert job.status == "failed" + finally: + session.close() + + +async def test_stop_worker_without_start_is_clean() -> None: + """stop_worker is safe to call when no executor was created.""" + if worker.is_worker_running(): + await worker.stop_worker() + # Force the no-executor branch. + worker._executor = None + await worker.stop_worker() + assert not worker.is_worker_running() + + +async def test_write_failure_without_item_dir( + client_no_worker: AsyncClient, library: dict, monkeypatch +) -> None: + """Write failure before any item dir exists still fails the job cleanly.""" + import services.content_service as content_service # noqa: PLC0415 + + lib_id = library["id"] + org_id = library["organization_id"] + + class OkPlugin(LibraryImportPlugin): + name = "writefail_nodir_test" + description = "Succeeds at plugin level." + supported_source_types = {"url"} + + def import_content(self, source_path, *, api_keys=None, **kwargs): + return ImportResult( + full_text="# ok", pages=[], images=[], metadata={}, source_ref={} + ) + + def get_parameters(self): + return [] + + def boom(*args, **kwargs): + # Do NOT create the item dir — exercises the ``if item_dir.exists()`` False branch. + raise RuntimeError("write failed before any dir") + + monkeypatch.setattr(content_service, "write_structured_content", boom) + + with _registered(OkPlugin): + job_id = _queue_job_row("writefail_nodir_test", lib_id, org_id, "https://x/nodir") + worker._process_job_sync(job_id) + + session = get_session_direct() + try: + job = session.query(ImportJob).filter(ImportJob.id == job_id).first() + assert job.status == "failed" + assert "write failed" in job.error_message + finally: + session.close() + + +def test_recover_stale_jobs_noop_when_none() -> None: + """recover_stale_jobs is a no-op when no jobs are stuck processing. + + Marks any pre-existing processing rows so this assertion is about the + 'nothing to recover' branch, not other tests' leftovers. + """ + # Drain any stray processing rows first so the no-op branch is exercised. + session = get_session_direct() + try: + stuck = session.query(ImportJob).filter(ImportJob.status == "processing").all() + for j in stuck: + j.status = "completed" + j.completed_at = datetime.now(UTC) + session.commit() + finally: + session.close() + worker.recover_stale_jobs() # no exception, no commit diff --git a/library-manager/tests/test_capabilities.py b/library-manager/tests/test_capabilities.py deleted file mode 100644 index 37d48a959..000000000 --- a/library-manager/tests/test_capabilities.py +++ /dev/null @@ -1,241 +0,0 @@ -"""Tests for the capability plugin system.""" - -from __future__ import annotations - -import asyncio -import io -import json -import time -from pathlib import Path - -import pytest -from httpx import AsyncClient - -AUTH_HEADERS = {"Authorization": "Bearer test-token"} -_POLL_TIMEOUT = 15 - - -async def _wait_for_ready(client, lib_id, item_id, timeout=_POLL_TIMEOUT): - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - resp = await client.get(f"/libraries/{lib_id}/items/{item_id}/status", headers=AUTH_HEADERS) - if resp.json()["status"] in ("ready", "failed"): - return resp.json()["status"] - await asyncio.sleep(0.5) - return "timeout" - - -async def _upload_md(client, lib_id, content, title="Test Doc", filename="test.md"): - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": (filename, io.BytesIO(content.encode()), "text/markdown")}, - data={"plugin_name": "simple_import", "title": title}, - ) - item_id = resp.json()["item_id"] - await _wait_for_ready(client, lib_id, item_id) - return item_id - - -# --------------------------------------------------------------------------- -# Registry tests (no client needed) -# --------------------------------------------------------------------------- - - -def test_builtin_handlers_in_registry(): - """All three built-in handlers should be auto-registered at startup.""" - from plugins.content_handlers.capability import ( # noqa: PLC0415 - Capability, - CapabilityRegistry, - ) - - registered = CapabilityRegistry.registered_capabilities() - assert Capability.TEXT in registered - assert Capability.PAGES in registered - assert Capability.IMAGES in registered - - -def test_text_handler_roundtrip(tmp_path: Path): - """TextHandler should return the contents of content/full.md.""" - from plugins.content_handlers.capability import ( # noqa: PLC0415 - Capability, - CapabilityRegistry, - ) - - item_dir = tmp_path / "item" - (item_dir / "content").mkdir(parents=True) - (item_dir / "content" / "full.md").write_text("# Hello\n\nWorld.", encoding="utf-8") - - handler = CapabilityRegistry.get(Capability.TEXT) - assert handler is not None - - payload = handler.get(item_dir) - assert payload.mime == "text/markdown" - assert payload.body == "# Hello\n\nWorld." - - -def test_pages_handler_roundtrip(tmp_path: Path): - """PagesHandler should return per-page entries in numeric order.""" - from plugins.content_handlers.capability import ( # noqa: PLC0415 - Capability, - CapabilityRegistry, - ) - - pages_dir = tmp_path / "item" / "content" / "pages" - pages_dir.mkdir(parents=True) - (pages_dir / "page_001.md").write_text("first", encoding="utf-8") - (pages_dir / "page_002.md").write_text("second", encoding="utf-8") - (pages_dir / "page_010.md").write_text("tenth", encoding="utf-8") - - handler = CapabilityRegistry.get(Capability.PAGES) - payload = handler.get(tmp_path / "item") - assert payload.mime == "application/json" - assert [p["page"] for p in payload.body] == [1, 2, 10] - assert payload.body[0]["markdown"] == "first" - - -def test_images_handler_roundtrip(tmp_path: Path): - """ImagesHandler should list files with URLs and MIME types.""" - from plugins.content_handlers.capability import ( # noqa: PLC0415 - Capability, - CapabilityRegistry, - ) - - # path layout matches CONTENT_DIR/{org}/{library}/{item}/content/images/ - item_dir = tmp_path / "org-x" / "lib-y" / "item-z" - images_dir = item_dir / "content" / "images" - images_dir.mkdir(parents=True) - (images_dir / "img_001.png").write_bytes(b"\x89PNG\r\n\x1a\n") - (images_dir / "img_002.jpg").write_bytes(b"\xff\xd8\xff") - - handler = CapabilityRegistry.get(Capability.IMAGES) - payload = handler.get(item_dir) - assert payload.mime == "application/json" - assert len(payload.body) == 2 - assert payload.body[0]["filename"] == "img_001.png" - assert payload.body[0]["mime"] == "image/png" - assert "/libraries/lib-y/items/item-z/content/images/file/img_001.png" in payload.body[0]["url"] - - -def test_handler_unavailable_raised_when_no_file(tmp_path: Path): - """Each handler should raise HandlerUnavailable on a missing folder/file.""" - from plugins.content_handlers.capability import ( # noqa: PLC0415 - Capability, - CapabilityRegistry, - HandlerUnavailable, - ) - - for cap in (Capability.TEXT, Capability.PAGES, Capability.IMAGES): - handler = CapabilityRegistry.get(cap) - with pytest.raises(HandlerUnavailable): - handler.get(tmp_path) - - -# --------------------------------------------------------------------------- -# HTTP endpoint tests -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_capabilities_endpoint_lists_handlers(client: AsyncClient): - """GET /capabilities returns all registered handlers.""" - resp = await client.get("/capabilities", headers=AUTH_HEADERS) - assert resp.status_code == 200 - names = {row["capability"] for row in resp.json()["capabilities"]} - assert {"text", "pages", "images"}.issubset(names) - - -@pytest.mark.asyncio -async def test_item_capabilities_reflects_metadata(client: AsyncClient, library: dict): - """Item /capabilities reads from metadata.json (TEXT-only for a markdown upload).""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "# Hello\n\nBody.") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/capabilities", headers=AUTH_HEADERS - ) - assert resp.status_code == 200 - data = resp.json() - assert data["item_id"] == item_id - assert data["capabilities"] == ["text"] - - -@pytest.mark.asyncio -async def test_item_content_text_returns_markdown(client: AsyncClient, library: dict): - """GET /content/text returns the full markdown body.""" - lib_id = library["id"] - body = "# Title\n\nThe quick brown fox." - item_id = await _upload_md(client, lib_id, body) - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content/text", headers=AUTH_HEADERS - ) - assert resp.status_code == 200 - assert "text/markdown" in resp.headers["content-type"] - assert resp.text == body - - -@pytest.mark.asyncio -async def test_item_content_unknown_capability_404(client: AsyncClient, library: dict): - """Unknown capability value returns 404.""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "x") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content/audio", headers=AUTH_HEADERS - ) - # AUDIO is in the enum but has no handler in this fixture, so 404. - assert resp.status_code == 404 - - -@pytest.mark.asyncio -async def test_item_content_missing_capability_for_item_404(client: AsyncClient, library: dict): - """Asking for IMAGES on a text-only item returns 404 (HandlerUnavailable).""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "no images here") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content/images", headers=AUTH_HEADERS - ) - # NOTE: this path collides with the legacy /content/images list endpoint - # which returns 200 with an empty array. Verify either behaviour is - # acceptable (no 5xx). - assert resp.status_code in (200, 404) - - -@pytest.mark.asyncio -async def test_legacy_item_fallback_to_text(client: AsyncClient, library: dict): - """Items without a `capabilities` field in metadata.json default to [TEXT].""" - from services import content_service # noqa: PLC0415 - - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "# Legacy item.") - - # Simulate a legacy item by stripping the 'capabilities' field from - # the on-disk metadata.json. - resp = await client.get(f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS) - item = resp.json() - org_id = ( - item["organization_id"] if "organization_id" in item else library.get("organization_id") - ) - base = content_service.get_item_base_path(org_id, lib_id, item_id) - meta_path = base / "metadata.json" - raw = json.loads(meta_path.read_text(encoding="utf-8")) - raw.pop("capabilities", None) - meta_path.write_text(json.dumps(raw), encoding="utf-8") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/capabilities", headers=AUTH_HEADERS - ) - assert resp.status_code == 200 - assert resp.json()["capabilities"] == ["text"] - - -@pytest.mark.asyncio -async def test_plugins_metadata_includes_produces_capabilities(client: AsyncClient): - """/plugins response now includes produces_capabilities per plugin.""" - resp = await client.get("/plugins", headers=AUTH_HEADERS) - assert resp.status_code == 200 - plugins = {p["name"]: p for p in resp.json()["plugins"]} - assert "simple_import" in plugins - assert "text" in plugins["simple_import"]["produces_capabilities"] diff --git a/library-manager/tests/test_capabilities_branches.py b/library-manager/tests/test_capabilities_branches.py new file mode 100644 index 000000000..6c9aec169 --- /dev/null +++ b/library-manager/tests/test_capabilities_branches.py @@ -0,0 +1,135 @@ +"""Error/dispatch branch tests for ``routers.capabilities``. + +Complements ``test_capabilities.py`` (registry + handler unit tests) by +exercising the HTTP router branches: 404s for missing items / unknown +capabilities, the capability-content dispatch for a real item, and the +HandlerUnavailable path. +""" + +from __future__ import annotations + +import io + +import pytest +from httpx import AsyncClient + +AUTH_HEADERS = {"Authorization": "Bearer test-token"} + + +async def _upload_md(client, lib_id, content="# Hi\n\nbody", title="Doc"): + import asyncio + import time + + resp = await client.post( + f"/libraries/{lib_id}/import/file", + headers=AUTH_HEADERS, + files={"file": ("t.md", io.BytesIO(content.encode()), "text/markdown")}, + data={"plugin_name": "simple_import", "title": title}, + ) + item_id = resp.json()["item_id"] + deadline = time.monotonic() + 15 + while time.monotonic() < deadline: + st = await client.get( + f"/libraries/{lib_id}/items/{item_id}/status", headers=AUTH_HEADERS + ) + if st.json()["status"] in ("ready", "failed"): + break + await asyncio.sleep(0.3) + return item_id + + +# --------------------------------------------------------------------------- +# 404 branches (no real item needed) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_item_capabilities_not_found(client: AsyncClient, library): + lib_id = library["id"] + resp = await client.get( + f"/libraries/{lib_id}/items/does-not-exist/capabilities", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_capability_content_unknown_capability(client: AsyncClient, library): + lib_id = library["id"] + item_id = await _upload_md(client, lib_id) + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/bogus-capability", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 404 + assert "Unknown capability" in resp.json()["detail"] + + +@pytest.mark.asyncio +async def test_capability_content_item_not_found(client: AsyncClient, library): + lib_id = library["id"] + resp = await client.get( + f"/libraries/{lib_id}/items/missing/content/text", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_image_raw_item_not_found(client: AsyncClient, library): + lib_id = library["id"] + resp = await client.get( + f"/libraries/{lib_id}/items/missing/content/images/file/x.png", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 404 + + +# --------------------------------------------------------------------------- +# real-item dispatch +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_item_capabilities_lists_text(client: AsyncClient, library): + lib_id = library["id"] + item_id = await _upload_md(client, lib_id) + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/capabilities", headers=AUTH_HEADERS + ) + assert resp.status_code == 200 + assert "text" in resp.json()["capabilities"] + + +@pytest.mark.asyncio +async def test_capability_content_text_dispatch(client: AsyncClient, library): + lib_id = library["id"] + item_id = await _upload_md(client, lib_id, content="# Title\n\nHello world") + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/text", headers=AUTH_HEADERS + ) + assert resp.status_code == 200 + assert "Hello world" in resp.text + + +@pytest.mark.asyncio +async def test_capability_content_pages_unavailable(client: AsyncClient, library): + # A simple_import .md item has no per-page split -> pages handler reports + # the capability as unavailable -> 404. + lib_id = library["id"] + item_id = await _upload_md(client, lib_id) + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/pages", headers=AUTH_HEADERS + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_image_raw_image_not_found(client: AsyncClient, library): + lib_id = library["id"] + item_id = await _upload_md(client, lib_id) + resp = await client.get( + f"/libraries/{lib_id}/items/{item_id}/content/images/file/missing.png", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 404 diff --git a/library-manager/tests/test_capabilities_router_dispatch.py b/library-manager/tests/test_capabilities_router_dispatch.py new file mode 100644 index 000000000..20ddb375b --- /dev/null +++ b/library-manager/tests/test_capabilities_router_dispatch.py @@ -0,0 +1,136 @@ +"""Capabilities-router dispatch branches: JSON payload, raw image serve, +missing-content 404, handler-exception 500, plus the importing-router +ValueError→400 paths (invalid folder_id).""" + +from __future__ import annotations + +import uuid + +import pytest +from httpx import AsyncClient + +from .conftest import AUTH_HEADERS + + +async def _mk_library(client: AsyncClient) -> str: + lib_id = f"lib-{uuid.uuid4().hex[:8]}" + resp = await client.post( + "/libraries", headers=AUTH_HEADERS, + json={"id": lib_id, "organization_id": "org-cap", "name": f"Cap {lib_id[-6:]}"}, + ) + assert resp.status_code == 201 + return lib_id + + +def _seed_item_with_content(lib_id: str, *, with_content: bool) -> str: + """Insert a ContentItem row; optionally write a full structured layout + (text + a page + an image) to disk.""" + from database.connection import get_session_direct + from database.models import ContentItem + from services import content_service as cs + + item_id = uuid.uuid4().hex + org_id = "org-cap" + base = cs.get_item_base_path(org_id, lib_id, item_id) + if with_content: + from plugins.base import ExtractedImage, PageContent + + cs.write_structured_content( + item_id=item_id, library_id=lib_id, organization_id=org_id, + title="Doc", full_text="# Full\n\nbody", + pages=[PageContent(page_number=1, text="p1")], + images=[ExtractedImage(filename="img_001.png", data=b"PNGBYTES")], + item_metadata={}, source_ref={"type": "file"}, + ) + + session = get_session_direct() + try: + session.add(ContentItem( + id=item_id, library_id=lib_id, organization_id=org_id, title="Doc", + source_type="file", base_path=str(base), + permalink_base=f"/library/{org_id}/{lib_id}/{item_id}", + import_plugin="simple_import", status="ready", + )) + session.commit() + finally: + session.close() + return item_id + + +@pytest.mark.asyncio +async def test_capability_content_json_payload(client: AsyncClient): + lib = await _mk_library(client) + item = _seed_item_with_content(lib, with_content=True) + # images handler returns application/json -> JSONResponse branch. + resp = await client.get( + f"/libraries/{lib}/items/{item}/content/images", headers=AUTH_HEADERS + ) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("application/json") + + +@pytest.mark.asyncio +async def test_image_raw_file_served(client: AsyncClient): + lib = await _mk_library(client) + item = _seed_item_with_content(lib, with_content=True) + resp = await client.get( + f"/libraries/{lib}/items/{item}/content/images/file/img_001.png", + headers=AUTH_HEADERS, + ) + assert resp.status_code == 200 + assert resp.content == b"PNGBYTES" + + +@pytest.mark.asyncio +async def test_capability_content_missing_on_disk_404(client: AsyncClient): + lib = await _mk_library(client) + item = _seed_item_with_content(lib, with_content=False) # row but no disk + resp = await client.get( + f"/libraries/{lib}/items/{item}/content/text", headers=AUTH_HEADERS + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_capability_handler_exception_500(client: AsyncClient, monkeypatch): + lib = await _mk_library(client) + item = _seed_item_with_content(lib, with_content=True) + + import plugins.content_handlers.capability as capmod + + class _BoomHandler: + def get(self, item_path): + raise RuntimeError("handler kaboom") + + monkeypatch.setattr( + capmod.CapabilityRegistry, "get", classmethod(lambda cls, c: _BoomHandler()) + ) + resp = await client.get( + f"/libraries/{lib}/items/{item}/content/text", headers=AUTH_HEADERS + ) + assert resp.status_code == 500 + + +# --------------------------------------------------------------------------- +# importing router: ValueError -> 400 (invalid folder_id) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_import_url_bad_folder_400(client: AsyncClient): + lib = await _mk_library(client) + resp = await client.post( + f"/libraries/{lib}/import/url", headers=AUTH_HEADERS, + json={"url": "https://example.com", "title": "T", "folder_id": "missing-folder"}, + ) + assert resp.status_code == 400 + + +@pytest.mark.asyncio +async def test_import_youtube_bad_folder_400(client: AsyncClient): + lib = await _mk_library(client) + resp = await client.post( + f"/libraries/{lib}/import/youtube", headers=AUTH_HEADERS, + json={"video_url": "https://youtu.be/abc12345678", "title": "T", "folder_id": "missing-folder"}, + ) + assert resp.status_code == 400 diff --git a/library-manager/tests/test_content_service_branches.py b/library-manager/tests/test_content_service_branches.py new file mode 100644 index 000000000..4584b9e92 --- /dev/null +++ b/library-manager/tests/test_content_service_branches.py @@ -0,0 +1,67 @@ +"""Branch tests for ``services.content_service`` added by the feature. + +Covers the pages/images write loops, list_pages/list_images, capability +detection from disk, and the mkdir PermissionError / OSError handlers. +""" + +from __future__ import annotations + +from pathlib import Path +from uuid import uuid4 + +import pytest + +from plugins.base import ExtractedImage, PageContent +from services import content_service as cs + + +def _ids(): + return uuid4().hex, uuid4().hex, "org-cs" + + +def test_write_pages_images_then_list_and_detect(): + item_id, lib_id, org_id = _ids() + base = cs.write_structured_content( + item_id=item_id, + library_id=lib_id, + organization_id=org_id, + title="Doc", + full_text="# Full\n\nbody text", + pages=[PageContent(page_number=1, text="page one"), + PageContent(page_number=2, text="page two")], + images=[ExtractedImage(filename="img_001.png", data=b"PNGBYTES")], + item_metadata={"k": "v"}, + source_ref={"type": "file"}, + ) + assert cs.list_pages(org_id, lib_id, item_id) == ["page_001.md", "page_002.md"] + assert cs.list_images(org_id, lib_id, item_id) == ["img_001.png"] + caps = cs.detect_capabilities(base) + assert {"text", "pages", "images"} <= set(caps) + + +def test_write_permission_error_is_translated(monkeypatch): + def boom_mkdir(self, *a, **k): + raise PermissionError("not writable") + + monkeypatch.setattr(Path, "mkdir", boom_mkdir) + item_id, lib_id, org_id = _ids() + with pytest.raises(RuntimeError, match="not writable"): + cs.write_structured_content( + item_id=item_id, library_id=lib_id, organization_id=org_id, + title="T", full_text="x", pages=[], images=[], + item_metadata={}, source_ref={}, + ) + + +def test_write_oserror_is_translated(monkeypatch): + def boom_mkdir(self, *a, **k): + raise OSError("disk full") + + monkeypatch.setattr(Path, "mkdir", boom_mkdir) + item_id, lib_id, org_id = _ids() + with pytest.raises(RuntimeError, match="storage is"): + cs.write_structured_content( + item_id=item_id, library_id=lib_id, organization_id=org_id, + title="T", full_text="x", pages=[], images=[], + item_metadata={}, source_ref={}, + ) diff --git a/library-manager/tests/test_content_serving.py b/library-manager/tests/test_content_serving.py deleted file mode 100644 index 7acab579b..000000000 --- a/library-manager/tests/test_content_serving.py +++ /dev/null @@ -1,236 +0,0 @@ -"""Tests for content serving endpoints: pages, images, metadata, source_ref, original file.""" - -import asyncio -import io -import time - -import pytest -from httpx import AsyncClient - -AUTH_HEADERS = {"Authorization": "Bearer test-token"} -_POLL_TIMEOUT = 15 - - -async def _wait_for_ready(client, lib_id, item_id, timeout=_POLL_TIMEOUT): - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/status", headers=AUTH_HEADERS - ) - if resp.json()["status"] in ("ready", "failed"): - return resp.json()["status"] - await asyncio.sleep(0.5) - return "timeout" - - -async def _upload_md(client, lib_id, content, title="Test Doc", filename="test.md"): - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": (filename, io.BytesIO(content.encode()), "text/markdown")}, - data={"plugin_name": "simple_import", "title": title}, - ) - item_id = resp.json()["item_id"] - await _wait_for_ready(client, lib_id, item_id) - return item_id - - -@pytest.mark.asyncio -async def test_get_metadata_has_permalinks(client: AsyncClient, library: dict): - """metadata.json should contain permalink URLs for all content pieces.""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "# Test\n\nContent here.") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/metadata", headers=AUTH_HEADERS - ) - assert resp.status_code == 200 - meta = resp.json() - assert "permalinks" in meta - assert meta["permalinks"]["full_markdown"].endswith("/content/full.md") - assert meta["item_id"] == item_id - - -@pytest.mark.asyncio -async def test_get_source_ref(client: AsyncClient, library: dict): - """source_ref.json should describe the import source.""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "# Source Ref Test") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/source_ref", headers=AUTH_HEADERS - ) - assert resp.status_code == 200 - src = resp.json() - assert src["type"] == "file" - assert "original_filename" in src - - -@pytest.mark.asyncio -async def test_get_original_file(client: AsyncClient, library: dict): - """Original file should be servable via /original/{filename}.""" - lib_id = library["id"] - content = "# Original File Test\n\nContent." - item_id = await _upload_md(client, lib_id, content, filename="original-test.md") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/original/original-test.md", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - assert "Original File Test" in resp.text - - -@pytest.mark.asyncio -async def test_content_format_text(client: AsyncClient, library: dict): - """?format=text should return plain text.""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "# Text Format\n\nParagraph.") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content", - headers=AUTH_HEADERS, - params={"format": "text"}, - ) - assert resp.status_code == 200 - assert "text/plain" in resp.headers["content-type"] - - -@pytest.mark.asyncio -async def test_pages_unavailable_for_text(client: AsyncClient, library: dict): - """A plain text file exposes no ``pages`` capability — the handler 404s.""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "No pages here.") - - # The capability handler at ``/content/pages`` raises HandlerUnavailable - # → 404 when the item has no ``content/pages/`` directory on disk. The - # frontend never asks for a capability that isn't in the item's - # ``/capabilities`` list, so this is the correct contract. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content/pages", headers=AUTH_HEADERS - ) - assert resp.status_code == 404 - - # And ``capabilities`` shouldn't list it either. - caps = await client.get( - f"/libraries/{lib_id}/items/{item_id}/capabilities", headers=AUTH_HEADERS - ) - assert caps.status_code == 200 - assert "pages" not in caps.json()["capabilities"] - - -@pytest.mark.asyncio -async def test_images_unavailable_for_text(client: AsyncClient, library: dict): - """A plain text file exposes no ``images`` capability — the handler 404s.""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "No images here.") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content/images", headers=AUTH_HEADERS - ) - assert resp.status_code == 404 - - caps = await client.get( - f"/libraries/{lib_id}/items/{item_id}/capabilities", headers=AUTH_HEADERS - ) - assert caps.status_code == 200 - assert "images" not in caps.json()["capabilities"] - - -@pytest.mark.asyncio -async def test_nonexistent_item_returns_404(client: AsyncClient, library: dict): - """Requesting a non-existent item should return 404.""" - lib_id = library["id"] - resp = await client.get( - f"/libraries/{lib_id}/items/fake-item-id", headers=AUTH_HEADERS - ) - assert resp.status_code == 404 - - -@pytest.mark.asyncio -async def test_nonexistent_page_returns_404(client: AsyncClient, library: dict): - """Requesting a non-existent page should return 404.""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "No pages.") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content/pages/page_999", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 404 - - -@pytest.mark.asyncio -async def test_nonexistent_image_returns_404(client: AsyncClient, library: dict): - """Requesting a non-existent image should return 404.""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "No images.") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content/images/fake.png", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 404 - - -@pytest.mark.asyncio -async def test_nonexistent_original_returns_404(client: AsyncClient, library: dict): - """Requesting a non-existent original file should return 404.""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "Content.") - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/original/nonexistent.pdf", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 404 - - -@pytest.mark.asyncio -async def test_delete_item_removes_content(client: AsyncClient, library: dict): - """Deleting an item should make its content inaccessible.""" - lib_id = library["id"] - item_id = await _upload_md(client, lib_id, "# Will Be Deleted") - - resp = await client.delete( - f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS - ) - assert resp.status_code == 200 - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}", headers=AUTH_HEADERS - ) - assert resp.status_code == 404 - - -@pytest.mark.asyncio -async def test_items_filter_by_status(client: AsyncClient, library: dict): - """Listing items with status filter should work.""" - lib_id = library["id"] - await _upload_md(client, lib_id, "# Ready Item") - - resp = await client.get( - f"/libraries/{lib_id}/items", - headers=AUTH_HEADERS, - params={"status": "ready"}, - ) - assert resp.status_code == 200 - assert all(i["status"] == "ready" for i in resp.json()["items"]) - - -@pytest.mark.asyncio -async def test_items_filter_by_ids(client: AsyncClient, library: dict): - """Listing items with ids filter should return only matching items.""" - lib_id = library["id"] - item1 = await _upload_md(client, lib_id, "# Item 1", title="Item 1") - await _upload_md(client, lib_id, "# Item 2", title="Item 2") - - resp = await client.get( - f"/libraries/{lib_id}/items", - headers=AUTH_HEADERS, - params={"ids": item1}, - ) - assert resp.status_code == 200 - data = resp.json() - assert data["total"] == 1 - assert data["items"][0]["id"] == item1 diff --git a/library-manager/tests/test_coverage_gaps.py b/library-manager/tests/test_coverage_gaps.py deleted file mode 100644 index 055d59e18..000000000 --- a/library-manager/tests/test_coverage_gaps.py +++ /dev/null @@ -1,742 +0,0 @@ -"""Tests targeting coverage gaps in schemas/common.py, url_import.py, -youtube_transcript_import.py, and markitdown_plus_import.py.""" - -from unittest.mock import MagicMock, patch - -import pytest -from schemas.common import MessageResponse, PaginationParams - -# ----------------------------------------------------------------------- -# schemas/common.py -# ----------------------------------------------------------------------- - - -class TestPaginationParams: - def test_defaults(self): - p = PaginationParams() - assert p.limit == 20 - assert p.offset == 0 - - def test_custom_values(self): - p = PaginationParams(limit=50, offset=10) - assert p.limit == 50 - assert p.offset == 10 - - def test_limit_min_boundary(self): - p = PaginationParams(limit=1) - assert p.limit == 1 - - def test_limit_max_boundary(self): - p = PaginationParams(limit=100) - assert p.limit == 100 - - def test_limit_below_min_rejected(self): - with pytest.raises(Exception): - PaginationParams(limit=0) - - def test_limit_above_max_rejected(self): - with pytest.raises(Exception): - PaginationParams(limit=101) - - def test_offset_min_boundary(self): - p = PaginationParams(offset=0) - assert p.offset == 0 - - def test_offset_negative_rejected(self): - with pytest.raises(Exception): - PaginationParams(offset=-1) - - -class TestMessageResponse: - def test_basic(self): - r = MessageResponse(message="hello") - assert r.message == "hello" - - def test_missing_message_rejected(self): - with pytest.raises(Exception): - MessageResponse() - - -# ----------------------------------------------------------------------- -# url_import.py — validation, SSRF, _safe_int, get_parameters -# ----------------------------------------------------------------------- - - -class TestUrlImportPluginValidation: - def setup_method(self): - from plugins.url_import import UrlImportPlugin - self.plugin = UrlImportPlugin() - - def test_invalid_url_no_scheme(self): - with pytest.raises(ValueError, match="Invalid URL"): - self.plugin.import_content("example.com") - - def test_invalid_url_no_netloc(self): - with pytest.raises(ValueError, match="Invalid URL"): - self.plugin.import_content("http://") - - def test_unsupported_scheme(self): - with pytest.raises(ValueError, match="Unsupported URL scheme"): - self.plugin.import_content("ftp://example.com/file") - - def test_ssrf_localhost(self): - with pytest.raises(ValueError, match="not allowed"): - self.plugin.import_content("http://localhost/page") - - def test_ssrf_127(self): - with pytest.raises(ValueError, match="not allowed"): - self.plugin.import_content("http://127.0.0.1/page") - - def test_ssrf_0000(self): - with pytest.raises(ValueError, match="not allowed"): - self.plugin.import_content("http://0.0.0.0/page") - - def test_ssrf_ipv6_loopback(self): - with pytest.raises(ValueError, match="not allowed"): - self.plugin.import_content("http://[::1]/page") - - def test_ssrf_metadata_endpoint(self): - with pytest.raises(ValueError, match="not allowed"): - self.plugin.import_content("http://169.254.169.254/latest/meta-data") - - def test_ssrf_google_metadata(self): - with pytest.raises(ValueError, match="not allowed"): - self.plugin.import_content("http://metadata.google.internal/computeMetadata") - - def test_get_parameters(self): - params = self.plugin.get_parameters() - assert len(params) == 5 - names = [p.name for p in params] - assert "max_discovery_depth" in names - assert "limit" in names - assert "timeout" in names - assert "description" in names - assert "citation" in names - - def test_get_parameters_defaults(self): - params = self.plugin.get_parameters() - depth_param = next(p for p in params if p.name == "max_discovery_depth") - assert depth_param.default == 2 - assert depth_param.min_value == 1 - assert depth_param.max_value == 10 - - def test_markitdown_fallback_path(self): - mock_result = MagicMock() - mock_result.text_content = "# Hello World" - with patch("markitdown.MarkItDown") as MockMD: - MockMD.return_value.convert.return_value = mock_result - result = self.plugin.import_content("https://example.com/page") - assert result.full_text == "# Hello World" - assert result.metadata["fetch_method"] == "markitdown" - assert result.metadata["source_url"] == "https://example.com/page" - assert result.source_ref["type"] == "url" - - def test_markitdown_fallback_empty_content(self): - mock_result = MagicMock() - mock_result.text_content = "" - with patch("markitdown.MarkItDown") as MockMD: - MockMD.return_value.convert.return_value = mock_result - result = self.plugin.import_content("https://example.com/empty") - assert "No text content" in result.full_text - - def test_markitdown_fallback_with_description_and_citation(self): - mock_result = MagicMock() - mock_result.text_content = "Content" - with patch("markitdown.MarkItDown") as MockMD: - MockMD.return_value.convert.return_value = mock_result - result = self.plugin.import_content( - "https://example.com/page", - description="A test page", - citation="Test 2024", - ) - assert result.metadata["description"] == "A test page" - assert result.metadata["citation"] == "Test 2024" - - def test_markitdown_fallback_fetch_failure(self): - with patch("markitdown.MarkItDown") as MockMD: - MockMD.return_value.convert.side_effect = Exception("Connection refused") - with pytest.raises(RuntimeError, match="Failed to fetch"): - self.plugin.import_content("https://example.com/fail") - - -class TestSafeInt: - def test_none_returns_default(self): - from plugins.url_import import _safe_int - assert _safe_int(None, 42) == 42 - - def test_valid_int_string(self): - from plugins.url_import import _safe_int - assert _safe_int("10", 0) == 10 - - def test_valid_int(self): - from plugins.url_import import _safe_int - assert _safe_int(25, 0) == 25 - - def test_invalid_string_returns_default(self): - from plugins.url_import import _safe_int - assert _safe_int("abc", 99) == 99 - - def test_float_string_returns_default(self): - from plugins.url_import import _safe_int - assert _safe_int("3.14", 7) == 7 - - -class TestFirecrawlPath: - def setup_method(self): - from plugins.url_import import UrlImportPlugin - self.plugin = UrlImportPlugin() - - def test_firecrawl_no_data_returns_placeholder(self): - scrape_dict = {} - with patch("firecrawl.FirecrawlApp") as MockApp: - MockApp.return_value.scrape.return_value = scrape_dict - result = self.plugin.import_content( - "https://example.com", - api_keys={"firecrawl_key": "test-key"}, - ) - assert "No content" in result.full_text - assert result.metadata["pages_crawled"] == 0 - - def test_firecrawl_with_markdown_object(self): - mock_meta = MagicMock() - mock_meta.url = "https://example.com" - mock_meta.source_url = "https://example.com" - mock_meta.title = "Example" - mock_scrape = MagicMock() - mock_scrape.markdown = "# Example Content" - mock_scrape.metadata = mock_meta - mock_scrape.url = "https://example.com" - with patch("firecrawl.FirecrawlApp") as MockApp: - MockApp.return_value.scrape.return_value = mock_scrape - result = self.plugin.import_content( - "https://example.com", - api_keys={"firecrawl_key": "test-key"}, - ) - assert "# Example Content" in result.full_text - assert result.metadata["pages_crawled"] == 1 - assert result.metadata["fetch_method"] == "firecrawl" - - def test_firecrawl_with_dict_result(self): - scrape_dict = { - "markdown": "# Dict Content", - "metadata": {"sourceURL": "https://example.com", "title": "Dict Page"}, - } - with patch("firecrawl.FirecrawlApp") as MockApp: - MockApp.return_value.scrape.return_value = scrape_dict - result = self.plugin.import_content( - "https://example.com", - api_keys={"firecrawl_key": "test-key"}, - ) - assert "# Dict Content" in result.full_text - - def test_firecrawl_with_dict_no_markdown(self): - scrape_dict = {"markdown": "", "metadata": {}} - with patch("firecrawl.FirecrawlApp") as MockApp: - MockApp.return_value.scrape.return_value = scrape_dict - result = self.plugin.import_content( - "https://example.com", - api_keys={"firecrawl_key": "test-key"}, - ) - assert "No content" in result.full_text - - def test_firecrawl_scrape_failure(self): - with patch("firecrawl.FirecrawlApp") as MockApp: - MockApp.return_value.scrape.side_effect = Exception("API error") - with pytest.raises(RuntimeError, match="Firecrawl scrape failed"): - self.plugin.import_content( - "https://example.com", - api_keys={"firecrawl_key": "test-key"}, - ) - - def test_firecrawl_metadata_object_with_dict_meta(self): - mock_scrape = MagicMock() - mock_scrape.markdown = "# Content" - mock_scrape.metadata = {"sourceURL": "https://example.com", "title": "Page"} - mock_scrape.url = "https://example.com" - with patch("firecrawl.FirecrawlApp") as MockApp: - MockApp.return_value.scrape.return_value = mock_scrape - result = self.plugin.import_content( - "https://example.com", - api_keys={"firecrawl_key": "test-key"}, - ) - assert "# Content" in result.full_text - - def test_firecrawl_with_description_and_citation(self): - mock_scrape = MagicMock() - mock_scrape.markdown = "# Content" - mock_scrape.metadata = {"sourceURL": "https://example.com", "title": ""} - mock_scrape.url = "https://example.com" - with patch("firecrawl.FirecrawlApp") as MockApp: - MockApp.return_value.scrape.return_value = mock_scrape - result = self.plugin.import_content( - "https://example.com", - api_keys={"firecrawl_key": "test-key"}, - description="Test desc", - citation="Test cite", - ) - assert result.metadata["description"] == "Test desc" - assert result.metadata["citation"] == "Test cite" - - def test_firecrawl_custom_api_url(self): - mock_scrape = MagicMock() - mock_scrape.markdown = "# Content" - mock_scrape.metadata = {"sourceURL": "https://example.com", "title": ""} - mock_scrape.url = "https://example.com" - with patch("firecrawl.FirecrawlApp") as MockApp: - MockApp.return_value.scrape.return_value = mock_scrape - self.plugin.import_content( - "https://example.com", - api_keys={"firecrawl_key": "test-key", "firecrawl_url": "https://custom.api.com"}, - ) - MockApp.assert_called_once_with(api_key="test-key", api_url="https://custom.api.com") - - -# ----------------------------------------------------------------------- -# youtube_transcript_import.py — helpers and get_parameters -# ----------------------------------------------------------------------- - - -class TestYouTubeHelpers: - def test_parse_youtube_url_standard(self): - from plugins.youtube_transcript_import import _parse_youtube_url - assert _parse_youtube_url("https://www.youtube.com/watch?v=dQw4w9WgXcQ") == "dQw4w9WgXcQ" - - def test_parse_youtube_url_short(self): - from plugins.youtube_transcript_import import _parse_youtube_url - assert _parse_youtube_url("https://youtu.be/dQw4w9WgXcQ") == "dQw4w9WgXcQ" - - def test_parse_youtube_url_short_with_path(self): - from plugins.youtube_transcript_import import _parse_youtube_url - assert _parse_youtube_url("https://youtu.be/dQw4w9WgXcQ/extra") == "dQw4w9WgXcQ" - - def test_parse_youtube_url_invalid(self): - from plugins.youtube_transcript_import import _parse_youtube_url - assert _parse_youtube_url("https://example.com/video") is None - - def test_parse_youtube_url_youtu_be_empty(self): - from plugins.youtube_transcript_import import _parse_youtube_url - assert _parse_youtube_url("https://youtu.be/") is None - - def test_parse_youtube_url_no_v_param(self): - from plugins.youtube_transcript_import import _parse_youtube_url - assert _parse_youtube_url("https://www.youtube.com/channel/abc") is None - - def test_seconds_to_timestamp_no_hours(self): - from plugins.youtube_transcript_import import _seconds_to_timestamp - assert _seconds_to_timestamp(125) == "02:05" - - def test_seconds_to_timestamp_with_hours(self): - from plugins.youtube_transcript_import import _seconds_to_timestamp - assert _seconds_to_timestamp(3661) == "01:01:01" - - def test_seconds_to_timestamp_zero(self): - from plugins.youtube_transcript_import import _seconds_to_timestamp - assert _seconds_to_timestamp(0) == "00:00" - - def test_build_source_ref(self): - from plugins.youtube_transcript_import import _build_source_ref - ref = _build_source_ref("https://www.youtube.com/watch?v=abc123", "abc123", "en") - assert ref["type"] == "youtube" - assert ref["video_id"] == "abc123" - assert ref["language"] == "en" - assert ref["video_url"] == "https://www.youtube.com/watch?v=abc123" - - def test_get_parameters(self): - from plugins.youtube_transcript_import import YouTubeTranscriptImportPlugin - plugin = YouTubeTranscriptImportPlugin() - params = plugin.get_parameters() - assert len(params) == 2 - names = [p.name for p in params] - assert "language" in names - assert "proxy_url" in names - lang_param = next(p for p in params if p.name == "language") - assert lang_param.default == "en" - - def test_invalid_youtube_url_raises(self): - from plugins.youtube_transcript_import import YouTubeTranscriptImportPlugin - plugin = YouTubeTranscriptImportPlugin() - with pytest.raises(ValueError, match="Could not extract video ID"): - plugin.import_content("https://example.com/not-youtube") - - -class TestResolveLanguageKey: - def test_exact_match(self): - from plugins.youtube_transcript_import import _resolve_language_key - subs = {"en": [{"ext": "srt"}], "es": [{"ext": "srt"}]} - assert _resolve_language_key(subs, "en") == "en" - - def test_locale_suffix_match(self): - from plugins.youtube_transcript_import import _resolve_language_key - subs = {"en-US": [{"ext": "srt"}], "es": [{"ext": "srt"}]} - assert _resolve_language_key(subs, "en") == "en-US" - - def test_complex_variant_match(self): - from plugins.youtube_transcript_import import _resolve_language_key - subs = {"en-nP7-2PuUl7o": [{"ext": "srt"}]} - assert _resolve_language_key(subs, "en") == "en-nP7-2PuUl7o" - - def test_no_match(self): - from plugins.youtube_transcript_import import _resolve_language_key - subs = {"fr": [{"ext": "srt"}]} - assert _resolve_language_key(subs, "en") is None - - def test_empty_map(self): - from plugins.youtube_transcript_import import _resolve_language_key - assert _resolve_language_key({}, "en") is None - - def test_none_map(self): - from plugins.youtube_transcript_import import _resolve_language_key - assert _resolve_language_key(None, "en") is None - - -class TestParseSrtContent: - def test_basic_srt(self): - from plugins.youtube_transcript_import import _parse_srt_content - srt = ( - "1\n" - "00:00:01,000 --> 00:00:04,000\n" - "Hello world\n" - "\n" - "2\n" - "00:00:05,000 --> 00:00:08,000\n" - "Second line\n" - ) - pieces = _parse_srt_content(srt) - assert len(pieces) == 2 - assert pieces[0]["text"] == "Hello world" - assert pieces[0]["start"] == 1.0 - assert pieces[0]["duration"] == 3.0 - assert pieces[1]["text"] == "Second line" - - def test_srt_with_noise_removal(self): - from plugins.youtube_transcript_import import _parse_srt_content - srt = ( - "1\n" - "00:00:01,000 --> 00:00:04,000\n" - "[applause] Hello [music]\n" - ) - pieces = _parse_srt_content(srt) - assert len(pieces) == 1 - assert pieces[0]["text"] == "Hello" - - def test_srt_empty_blocks_skipped(self): - from plugins.youtube_transcript_import import _parse_srt_content - srt = ( - "1\n" - "00:00:01,000 --> 00:00:04,000\n" - "[noise only]\n" - "\n" - "2\n" - "00:00:05,000 --> 00:00:08,000\n" - "Real text\n" - ) - pieces = _parse_srt_content(srt) - assert len(pieces) == 1 - assert pieces[0]["text"] == "Real text" - - def test_srt_with_hours(self): - from plugins.youtube_transcript_import import _parse_srt_content - srt = ( - "1\n" - "01:30:00,000 --> 01:30:05,000\n" - "Late content\n" - ) - pieces = _parse_srt_content(srt) - assert len(pieces) == 1 - assert pieces[0]["start"] == 5400.0 - - def test_srt_empty_string(self): - from plugins.youtube_transcript_import import _parse_srt_content - assert _parse_srt_content("") == [] - - def test_srt_no_timestamp(self): - from plugins.youtube_transcript_import import _parse_srt_content - srt = "Just some text\nwithout timestamps" - assert _parse_srt_content(srt) == [] - - def test_srt_multiline_text(self): - from plugins.youtube_transcript_import import _parse_srt_content - srt = ( - "1\n" - "00:00:01,000 --> 00:00:04,000\n" - "Line one\n" - "Line two\n" - ) - pieces = _parse_srt_content(srt) - assert len(pieces) == 1 - assert pieces[0]["text"] == "Line one Line two" - - -class TestFetchTranscript: - def test_successful_manual_subtitle(self): - import yt_dlp - from plugins.youtube_transcript_import import _fetch_transcript - mock_info = {"subtitles": {"en": [{"ext": "srt"}]}, "automatic_captions": {}} - mock_ydl = MagicMock() - mock_ydl.extract_info.return_value = mock_info - mock_ydl.__enter__ = MagicMock(return_value=mock_ydl) - mock_ydl.__exit__ = MagicMock(return_value=False) - with ( - patch.object(yt_dlp, "YoutubeDL", return_value=mock_ydl), - patch("plugins.youtube_transcript_import._download_and_parse") as mock_dl, - ): - mock_dl.return_value = [{"text": "Hello", "start": 0, "duration": 1}] - pieces, source = _fetch_transcript("abc123", "en", None) - assert len(pieces) == 1 - assert source == "manual" - - -class TestDownloadAndParse: - def test_ytdlp_error_prefix_stripped(self): - import yt_dlp - from plugins.youtube_transcript_import import _download_and_parse - mock_ydl = MagicMock() - mock_ydl.download.side_effect = Exception("ERROR: Video unavailable") - mock_ydl.__enter__ = MagicMock(return_value=mock_ydl) - mock_ydl.__exit__ = MagicMock(return_value=False) - with ( - patch.object(yt_dlp, "YoutubeDL", return_value=mock_ydl), - pytest.raises(RuntimeError, match="Video unavailable"), - ): - _download_and_parse( - url="https://youtube.com/watch?v=abc", - language_key="en", - source_label="manual", - proxy_url=None, - ) - - def test_no_subtitle_files_found(self): - import yt_dlp - from plugins.youtube_transcript_import import _download_and_parse - mock_ydl = MagicMock() - mock_ydl.download.return_value = None - mock_ydl.__enter__ = MagicMock(return_value=mock_ydl) - mock_ydl.__exit__ = MagicMock(return_value=False) - with ( - patch.object(yt_dlp, "YoutubeDL", return_value=mock_ydl), - patch("os.listdir", return_value=["video.mp4"]), - ): - result = _download_and_parse( - url="https://youtube.com/watch?v=abc", - language_key="en", - source_label="manual", - proxy_url=None, - ) - assert result == [] - - def test_with_proxy(self): - import yt_dlp - from plugins.youtube_transcript_import import _download_and_parse - mock_ydl = MagicMock() - mock_ydl.download.return_value = None - mock_ydl.__enter__ = MagicMock(return_value=mock_ydl) - mock_ydl.__exit__ = MagicMock(return_value=False) - with ( - patch.object(yt_dlp, "YoutubeDL", return_value=mock_ydl) as MockYDL, - patch("os.listdir", return_value=[]), - ): - _download_and_parse( - url="https://youtube.com/watch?v=abc", - language_key="en", - source_label="auto", - proxy_url="http://proxy:8080", - ) - call_args = MockYDL.call_args[0][0] - assert call_args["proxy"] == "http://proxy:8080" - assert call_args["writeautomaticsub"] is True - assert call_args["writesubtitles"] is False - - -# ----------------------------------------------------------------------- -# markitdown_plus_import.py — get_parameters, _split_into_pages, _image_mime -# ----------------------------------------------------------------------- - - -class TestMarkitdownPlusParameters: - def test_get_parameters(self): - from plugins.markitdown_plus_import import MarkItDownPlusPlugin - plugin = MarkItDownPlusPlugin() - params = plugin.get_parameters() - assert len(params) == 3 - names = [p.name for p in params] - assert "image_descriptions" in names - assert "description" in names - assert "citation" in names - - def test_image_descriptions_choices(self): - from plugins.markitdown_plus_import import MarkItDownPlusPlugin - plugin = MarkItDownPlusPlugin() - params = plugin.get_parameters() - img_param = next(p for p in params if p.name == "image_descriptions") - assert img_param.choices == ["none", "basic", "llm"] - assert img_param.default == "basic" - - def test_produces_capabilities(self): - from plugins.content_handlers.capability import Capability - from plugins.markitdown_plus_import import MarkItDownPlusPlugin - assert Capability.TEXT in MarkItDownPlusPlugin.produces_capabilities - assert Capability.PAGES in MarkItDownPlusPlugin.produces_capabilities - assert Capability.IMAGES in MarkItDownPlusPlugin.produces_capabilities - - def test_file_extensions(self): - from plugins.markitdown_plus_import import MarkItDownPlusPlugin - assert "pdf" in MarkItDownPlusPlugin.file_extensions - assert "docx" in MarkItDownPlusPlugin.file_extensions - assert "pptx" in MarkItDownPlusPlugin.file_extensions - - -class TestSplitIntoPages: - def test_non_page_aware_type(self): - from plugins.markitdown_plus_import import _split_into_pages - assert _split_into_pages("some content", "html") == [] - - def test_no_page_breaks(self): - from plugins.markitdown_plus_import import _split_into_pages - assert _split_into_pages("just text", "pdf") == [] - - def test_horizontal_rule_page_breaks(self): - from plugins.markitdown_plus_import import _split_into_pages - content = "Page 1 content\n\n---\n\nPage 2 content\n\n---\n\nPage 3 content" - pages = _split_into_pages(content, "pdf") - assert len(pages) == 3 - assert pages[0].text == "Page 1 content" - assert pages[1].text == "Page 2 content" - assert pages[2].text == "Page 3 content" - assert pages[0].page_number == 1 - assert pages[2].page_number == 3 - - def test_form_feed_page_breaks(self): - from plugins.markitdown_plus_import import _split_into_pages - content = "Page 1\n\fPage 2\n\fPage 3" - pages = _split_into_pages(content, "docx") - assert len(pages) == 3 - - def test_html_comment_page_breaks(self): - from plugins.markitdown_plus_import import _split_into_pages - content = "Page 1\n\nPage 2" - pages = _split_into_pages(content, "pptx") - assert len(pages) == 2 - - def test_empty_pages_filtered(self): - from plugins.markitdown_plus_import import _split_into_pages - content = "Page 1\n\n---\n\n\n\n---\n\nPage 3" - pages = _split_into_pages(content, "pdf") - assert len(pages) == 2 - - def test_single_page_returns_empty(self): - from plugins.markitdown_plus_import import _split_into_pages - content = "Just one page\n\n---\n\n" - pages = _split_into_pages(content, "pdf") - assert pages == [] - - -class TestImageMime: - def test_png(self): - from plugins.markitdown_plus_import import _image_mime - assert _image_mime("png") == "image/png" - - def test_jpg(self): - from plugins.markitdown_plus_import import _image_mime - assert _image_mime("jpg") == "image/jpeg" - - def test_jpeg(self): - from plugins.markitdown_plus_import import _image_mime - assert _image_mime("jpeg") == "image/jpeg" - - def test_gif(self): - from plugins.markitdown_plus_import import _image_mime - assert _image_mime("gif") == "image/gif" - - def test_bmp(self): - from plugins.markitdown_plus_import import _image_mime - assert _image_mime("bmp") == "image/bmp" - - def test_webp(self): - from plugins.markitdown_plus_import import _image_mime - assert _image_mime("webp") == "image/webp" - - def test_tiff(self): - from plugins.markitdown_plus_import import _image_mime - assert _image_mime("tiff") == "image/tiff" - - def test_svg(self): - from plugins.markitdown_plus_import import _image_mime - assert _image_mime("svg") == "image/svg+xml" - - def test_unknown_defaults_to_png(self): - from plugins.markitdown_plus_import import _image_mime - assert _image_mime("xyz") == "image/png" - - def test_with_dot_prefix(self): - from plugins.markitdown_plus_import import _image_mime - assert _image_mime(".png") == "image/png" - - def test_uppercase(self): - from plugins.markitdown_plus_import import _image_mime - assert _image_mime("PNG") == "image/png" - - -class TestDescribeImage: - def test_basic_mode(self): - from plugins.markitdown_plus_import import _describe_image - desc = _describe_image(b"fake", "img_001.png", "png", "basic", {}, {}) - assert desc == "Image: img_001.png" - - def test_unknown_mode_returns_none(self): - from plugins.markitdown_plus_import import _describe_image - desc = _describe_image(b"fake", "img_001.png", "png", "none", {}, {}) - assert desc is None - - def test_llm_mode_no_api_key(self): - from plugins.markitdown_plus_import import _describe_image - desc = _describe_image(b"fake", "img_001.png", "png", "llm", {}, {}) - assert desc == "Image: img_001.png" - - def test_llm_mode_success(self): - from plugins.markitdown_plus_import import _describe_image - mock_response = MagicMock() - mock_response.choices = [MagicMock()] - mock_response.choices[0].message.content = "A blue sky" - mock_response.usage = MagicMock() - mock_response.usage.total_tokens = 50 - - stats: dict = {"images_with_llm_descriptions": 0, "llm_calls": []} - with patch("openai.OpenAI") as MockOpenAI: - MockOpenAI.return_value.chat.completions.create.return_value = mock_response - desc = _describe_image( - b"fake", "img_001.png", "png", "llm", {"openai_vision": "sk-test"}, stats - ) - assert desc == "A blue sky" - assert stats["images_with_llm_descriptions"] == 1 - assert stats["llm_calls"][0]["success"] is True - - def test_llm_mode_failure(self): - from plugins.markitdown_plus_import import _describe_image - stats: dict = {"images_with_llm_descriptions": 0, "llm_calls": []} - with patch("openai.OpenAI") as MockOpenAI: - MockOpenAI.return_value.chat.completions.create.side_effect = Exception("API error") - desc = _describe_image( - b"fake", "img_001.png", "png", "llm", {"openai_vision": "sk-test"}, stats - ) - assert desc == "Image: img_001.png" - assert stats["llm_calls"][0]["success"] is False - - -class TestExtractImages: - def test_pymupdf_not_installed(self): - from pathlib import Path - - from plugins.markitdown_plus_import import _extract_images - with ( - patch.dict("sys.modules", {"fitz": None}), - patch("builtins.__import__", side_effect=ImportError("No module named 'fitz'")), - ): - result = _extract_images(Path("/fake.pdf"), "basic", {}, {"images_extracted": 0}) - assert result == [] - - def test_non_pdf_returns_empty(self): - from pathlib import Path - - from plugins.markitdown_plus_import import _extract_images - result = _extract_images(Path("/fake.docx"), "basic", {}, {}) - assert result == [] diff --git a/library-manager/tests/test_edge_cases.py b/library-manager/tests/test_edge_cases.py deleted file mode 100644 index 462ed5c7b..000000000 --- a/library-manager/tests/test_edge_cases.py +++ /dev/null @@ -1,403 +0,0 @@ -"""Tests for edge cases, error paths, and coverage gaps. - -Covers: upload size limits, Unicode filenames, path traversal attempts, -empty files, malformed JSON params, stale job recovery, and source type -validation. -""" - -import asyncio -import io -import time - -import pytest -from httpx import AsyncClient - -AUTH_HEADERS = {"Authorization": "Bearer test-token"} -_POLL_TIMEOUT = 15 - - -async def _wait_for_ready(client, lib_id, item_id, timeout=_POLL_TIMEOUT): - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/status", headers=AUTH_HEADERS - ) - if resp.json()["status"] in ("ready", "failed"): - return resp.json()["status"] - await asyncio.sleep(0.5) - return "timeout" - - -# ----------------------------------------------------------------------- -# Upload size limits -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_upload_exceeds_size_limit(client: AsyncClient, library: dict): - """Uploading a file larger than MAX_UPLOAD_SIZE_BYTES should return 413.""" - import routers.importing as importing_mod # noqa: PLC0415 - - original = importing_mod.MAX_UPLOAD_SIZE_BYTES - importing_mod.MAX_UPLOAD_SIZE_BYTES = 100 # 100 bytes - - try: - content = "x" * 200 # 200 bytes > 100 limit - resp = await client.post( - f"/libraries/{library['id']}/import/file", - headers=AUTH_HEADERS, - files={"file": ("big.md", io.BytesIO(content.encode()), "text/markdown")}, - data={"plugin_name": "simple_import", "title": "Too Big"}, - ) - assert resp.status_code == 413 - finally: - importing_mod.MAX_UPLOAD_SIZE_BYTES = original - - -# ----------------------------------------------------------------------- -# Unicode filenames -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_unicode_filename(client: AsyncClient, library: dict): - """Files with Unicode names should be imported successfully.""" - lib_id = library["id"] - content = "# Documento académico\n\nContenido en español con acentos: é è ê ë." - - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("documento_académico.md", io.BytesIO(content.encode()), "text/markdown")}, - data={"plugin_name": "simple_import", "title": "Documento Académico"}, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - status = await _wait_for_ready(client, lib_id, item_id) - assert status == "ready" - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content", headers=AUTH_HEADERS - ) - assert "académico" in resp.text - - -@pytest.mark.asyncio -async def test_filename_with_spaces(client: AsyncClient, library: dict): - """Files with spaces in names should be imported successfully.""" - lib_id = library["id"] - - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("my document file.md", io.BytesIO(b"# Spaced"), "text/markdown")}, - data={"plugin_name": "simple_import", "title": "Spaced Filename"}, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - status = await _wait_for_ready(client, lib_id, item_id) - assert status == "ready" - - -# ----------------------------------------------------------------------- -# Path traversal attempts -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_path_traversal_in_page_name(client: AsyncClient, library: dict): - """Requesting a page with ../ should not escape the content directory.""" - lib_id = library["id"] - content = "# Safe Content" - - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("safe.md", io.BytesIO(content.encode()), "text/markdown")}, - data={"plugin_name": "simple_import", "title": "Safe"}, - ) - item_id = resp.json()["item_id"] - await _wait_for_ready(client, lib_id, item_id) - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content/pages/../../metadata.json", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 404 - - -@pytest.mark.asyncio -async def test_path_traversal_in_image_name(client: AsyncClient, library: dict): - """Requesting an image with ../ should not escape the content directory.""" - lib_id = library["id"] - - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("safe.md", io.BytesIO(b"# Content"), "text/markdown")}, - data={"plugin_name": "simple_import", "title": "Safe"}, - ) - item_id = resp.json()["item_id"] - await _wait_for_ready(client, lib_id, item_id) - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content/images/../../../metadata.json", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 404 - - -@pytest.mark.asyncio -async def test_path_traversal_in_original_name(client: AsyncClient, library: dict): - """Requesting an original file with ../ should not escape.""" - lib_id = library["id"] - - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("safe.md", io.BytesIO(b"# Content"), "text/markdown")}, - data={"plugin_name": "simple_import", "title": "Safe"}, - ) - item_id = resp.json()["item_id"] - await _wait_for_ready(client, lib_id, item_id) - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/original/../../metadata.json", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 404 - - -# ----------------------------------------------------------------------- -# Empty file -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_empty_file_import_rejected(client: AsyncClient, library: dict): - """An empty (0-byte) upload must be rejected with HTTP 400. - - Regression test for defect D2 (lifecycle 2026-05-03): a 0-byte file - previously succeeded and ended up as a ``ready`` item with file_size=0 - and an empty full.md, polluting the library. The router must refuse it - before any DB row or job is queued. - """ - lib_id = library["id"] - - # Snapshot item count before — must be unchanged afterwards. - items_before = await client.get( - f"/libraries/{lib_id}/items", - headers=AUTH_HEADERS, - ) - assert items_before.status_code == 200 - count_before = len(items_before.json().get("items", items_before.json())) - - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("empty.md", io.BytesIO(b""), "text/markdown")}, - data={"plugin_name": "simple_import", "title": "Empty Doc"}, - ) - assert resp.status_code == 400 - detail = resp.json()["detail"].lower() - assert "empty" in detail or "0 bytes" in detail - - # No row should have been inserted in the ContentItem table. - items_after = await client.get( - f"/libraries/{lib_id}/items", - headers=AUTH_HEADERS, - ) - assert items_after.status_code == 200 - count_after = len(items_after.json().get("items", items_after.json())) - assert count_after == count_before, ( - f"Empty-file upload created an item ({count_before} -> {count_after})" - ) - - -# ----------------------------------------------------------------------- -# Malformed JSON in form fields -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_malformed_plugin_params_json(client: AsyncClient, library: dict): - """Invalid JSON in plugin_params should return 400.""" - resp = await client.post( - f"/libraries/{library['id']}/import/file", - headers=AUTH_HEADERS, - files={"file": ("test.md", io.BytesIO(b"# Test"), "text/markdown")}, - data={ - "plugin_name": "simple_import", - "title": "Test", - "plugin_params": "{not valid json", - }, - ) - assert resp.status_code == 400 - assert "plugin_params" in resp.json()["detail"].lower() - - -@pytest.mark.asyncio -async def test_malformed_api_keys_json(client: AsyncClient, library: dict): - """Invalid JSON in api_keys should return 400.""" - resp = await client.post( - f"/libraries/{library['id']}/import/file", - headers=AUTH_HEADERS, - files={"file": ("test.md", io.BytesIO(b"# Test"), "text/markdown")}, - data={ - "plugin_name": "simple_import", - "title": "Test", - "api_keys": "not json", - }, - ) - assert resp.status_code == 400 - assert "api_keys" in resp.json()["detail"].lower() - - -# ----------------------------------------------------------------------- -# Source type validation (cross-plugin) -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_file_plugin_rejects_url_import(client: AsyncClient, library: dict): - """A file-only plugin should not accept URL imports.""" - resp = await client.post( - f"/libraries/{library['id']}/import/url", - headers=AUTH_HEADERS, - json={ - "url": "https://example.com", - "title": "Bad", - "plugin_name": "markitdown_import", - }, - ) - assert resp.status_code == 400 - - -@pytest.mark.asyncio -async def test_url_plugin_rejects_file_upload(client: AsyncClient, library: dict): - """A URL-only plugin should not accept file uploads.""" - resp = await client.post( - f"/libraries/{library['id']}/import/file", - headers=AUTH_HEADERS, - files={"file": ("test.md", io.BytesIO(b"# Test"), "text/markdown")}, - data={"plugin_name": "url_import", "title": "Bad"}, - ) - assert resp.status_code == 400 - - -# ----------------------------------------------------------------------- -# Stale job recovery -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_stale_job_recovery(client: AsyncClient, library: dict): - """Jobs stuck in 'processing' should be recovered on startup. - - Simulates a crash by directly writing a 'processing' job to the DB, - then calling recover_stale_jobs and verifying it resets to 'pending'. - """ - from database.connection import get_session_direct # noqa: PLC0415 - from database.models import ImportJob # noqa: PLC0415 - from tasks.worker import recover_stale_jobs # noqa: PLC0415 - - db = get_session_direct() - try: - stale_job = ImportJob( - id="stale-job-001", - content_item_id="fake-item", - library_id=library["id"], - organization_id="org-test", - source_type="file", - plugin_name="simple_import", - title="Stale Job", - status="processing", - attempts=1, - ) - db.add(stale_job) - db.commit() - - recover_stale_jobs() - - db.expire_all() - refreshed = db.query(ImportJob).filter(ImportJob.id == "stale-job-001").first() - assert refreshed is not None - assert refreshed.status == "pending" - finally: - db.query(ImportJob).filter(ImportJob.id == "stale-job-001").delete() - db.commit() - db.close() - - -@pytest.mark.asyncio -async def test_stale_job_exceeds_max_attempts(client: AsyncClient, library: dict): - """Jobs exceeding max attempts should be marked failed, not retried.""" - from database.connection import get_session_direct # noqa: PLC0415 - from database.models import ImportJob # noqa: PLC0415 - from tasks.worker import recover_stale_jobs # noqa: PLC0415 - - db = get_session_direct() - try: - stale_job = ImportJob( - id="stale-job-002", - content_item_id="fake-item", - library_id=library["id"], - organization_id="org-test", - source_type="file", - plugin_name="simple_import", - title="Stale Exhausted Job", - status="processing", - attempts=5, - ) - db.add(stale_job) - db.commit() - - recover_stale_jobs() - - db.expire_all() - refreshed = db.query(ImportJob).filter(ImportJob.id == "stale-job-002").first() - assert refreshed is not None - assert refreshed.status == "failed" - assert "max attempts" in refreshed.error_message.lower() - finally: - db.query(ImportJob).filter(ImportJob.id == "stale-job-002").delete() - db.commit() - db.close() - - -# ----------------------------------------------------------------------- -# Nonexistent plugin -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_nonexistent_plugin(client: AsyncClient, library: dict): - """Using a plugin name that doesn't exist should return 400.""" - resp = await client.post( - f"/libraries/{library['id']}/import/file", - headers=AUTH_HEADERS, - files={"file": ("test.md", io.BytesIO(b"# Test"), "text/markdown")}, - data={"plugin_name": "nonexistent_plugin", "title": "Bad"}, - ) - assert resp.status_code == 400 - assert "not found" in resp.json()["detail"].lower() - - -# ----------------------------------------------------------------------- -# Wrong file extension for plugin -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_wrong_extension_for_plugin(client: AsyncClient, library: dict): - """Uploading a .pdf to simple_import (which only supports txt/md/html) should fail.""" - resp = await client.post( - f"/libraries/{library['id']}/import/file", - headers=AUTH_HEADERS, - files={"file": ("fake.pdf", io.BytesIO(b"not a real pdf"), "application/pdf")}, - data={"plugin_name": "simple_import", "title": "Wrong Extension"}, - ) - assert resp.status_code == 400 - assert "does not support" in resp.json()["detail"].lower() diff --git a/library-manager/tests/test_export_import.py b/library-manager/tests/test_export_import.py deleted file mode 100644 index e37617f3e..000000000 --- a/library-manager/tests/test_export_import.py +++ /dev/null @@ -1,141 +0,0 @@ -"""Tests for library export and import via ZIP files.""" - -import asyncio -import io -import json -import time -import zipfile - -import pytest -from httpx import AsyncClient - -AUTH_HEADERS = {"Authorization": "Bearer test-token"} - -_POLL_TIMEOUT = 15 - - -async def _wait_for_ready(client, lib_id, item_id, timeout=_POLL_TIMEOUT): - """Poll until item is ready or failed.""" - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/status", - headers=AUTH_HEADERS, - ) - status = resp.json()["status"] - if status in ("ready", "failed"): - return status - await asyncio.sleep(0.5) - return "timeout" - - -@pytest.mark.asyncio -async def test_export_library(client: AsyncClient, library: dict): - """Exporting a library should produce a valid ZIP with manifest + content.""" - lib_id = library["id"] - - # Upload a file first so the export has content. - content = "# Export Test\n\nThis content should appear in the ZIP." - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("export-test.md", io.BytesIO(content.encode()), "text/markdown")}, - data={"plugin_name": "simple_import", "title": "Export Test Doc"}, - ) - item_id = resp.json()["item_id"] - await _wait_for_ready(client, lib_id, item_id) - - # Export. - resp = await client.get(f"/libraries/{lib_id}/export", headers=AUTH_HEADERS) - assert resp.status_code == 200 - assert "application/zip" in resp.headers.get("content-type", "") - - # Validate ZIP structure. - zf = zipfile.ZipFile(io.BytesIO(resp.content)) - names = zf.namelist() - assert "manifest.json" in names - - manifest = json.loads(zf.read("manifest.json")) - assert manifest["format_version"] == "1.0" - assert manifest["type"] == "library_export" - assert len(manifest["items"]) == 1 - assert manifest["items"][0]["title"] == "Export Test Doc" - - # Check content files are present. - content_files = [n for n in names if n.startswith("content/")] - assert len(content_files) >= 2 # At least metadata.json + full.md - - -@pytest.mark.asyncio -async def test_import_library_from_zip(client: AsyncClient, library: dict): - """Importing a library from ZIP should create a new library with new IDs.""" - lib_id = library["id"] - - # Upload content. - content = "# Import Roundtrip\n\nThis should survive export+import." - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("roundtrip.md", io.BytesIO(content.encode()), "text/markdown")}, - data={"plugin_name": "simple_import", "title": "Roundtrip Doc"}, - ) - item_id = resp.json()["item_id"] - await _wait_for_ready(client, lib_id, item_id) - - # Export. - resp = await client.get(f"/libraries/{lib_id}/export", headers=AUTH_HEADERS) - zip_data = resp.content - - # Import into a different org. - resp = await client.post( - "/libraries/import", - headers=AUTH_HEADERS, - params={"organization_id": "org-imported"}, - files={"file": ("export.zip", io.BytesIO(zip_data), "application/zip")}, - ) - assert resp.status_code == 201 - result = resp.json() - assert result["item_count"] == 1 - new_lib_id = result["library_id"] - assert new_lib_id != lib_id # New ID generated - - # Verify imported content is readable. - resp = await client.get( - f"/libraries/{new_lib_id}/items", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - items = resp.json()["items"] - assert len(items) == 1 - new_item_id = items[0]["id"] - assert new_item_id != item_id # New ID generated - assert items[0]["status"] == "ready" - - # Read markdown from imported item. - resp = await client.get( - f"/libraries/{new_lib_id}/items/{new_item_id}/content", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - assert "Import Roundtrip" in resp.text - - # Verify permalinks point to the NEW IDs. - resp = await client.get( - f"/libraries/{new_lib_id}/items/{new_item_id}/metadata", - headers=AUTH_HEADERS, - ) - meta = resp.json() - assert new_lib_id in meta["permalinks"]["full_markdown"] - assert new_item_id in meta["permalinks"]["full_markdown"] - - -@pytest.mark.asyncio -async def test_import_invalid_zip(client: AsyncClient): - """Importing a non-ZIP file should return 400.""" - resp = await client.post( - "/libraries/import", - headers=AUTH_HEADERS, - params={"organization_id": "org-bad"}, - files={"file": ("not-a-zip.zip", io.BytesIO(b"not a zip"), "application/zip")}, - ) - assert resp.status_code == 400 diff --git a/library-manager/tests/test_folders.py b/library-manager/tests/test_folders.py deleted file mode 100644 index a82ccb08b..000000000 --- a/library-manager/tests/test_folders.py +++ /dev/null @@ -1,327 +0,0 @@ -"""Tests for the folder + tree endpoints. - -Covers: -- Folder CRUD (create / rename / move / delete). -- Cycle prevention server-side. -- Unique sibling-name enforcement. -- Delete reparents items + subfolders to the deleted folder's parent. -- Folder name validation (length, forbidden characters). -- Cross-library FK rejection. -""" - -from __future__ import annotations - -import uuid - -import pytest -from httpx import AsyncClient - -from .conftest import AUTH_HEADERS - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -async def _create_folder( - client: AsyncClient, lib_id: str, name: str, parent_id: str | None = None -) -> dict: - resp = await client.post( - f"/libraries/{lib_id}/folders", - headers=AUTH_HEADERS, - json={"name": name, "parent_folder_id": parent_id}, - ) - assert resp.status_code == 201, resp.text - return resp.json() - - -async def _create_library(client: AsyncClient) -> dict: - lib_id = f"lib-{uuid.uuid4().hex[:8]}" - resp = await client.post( - "/libraries", - headers=AUTH_HEADERS, - json={ - "id": lib_id, - "organization_id": "org-test", - "name": f"Lib {lib_id[-8:]}", - }, - ) - assert resp.status_code == 201 - return resp.json() - - -# --------------------------------------------------------------------------- -# CRUD -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_create_folder_at_root(client: AsyncClient, library: dict) -> None: - folder = await _create_folder(client, library["id"], "Q1 Research") - assert folder["name"] == "Q1 Research" - assert folder["parent_folder_id"] is None - assert "id" in folder - - -@pytest.mark.asyncio -async def test_create_folder_nested(client: AsyncClient, library: dict) -> None: - parent = await _create_folder(client, library["id"], "Papers") - child = await _create_folder(client, library["id"], "Biology", parent["id"]) - assert child["parent_folder_id"] == parent["id"] - - -@pytest.mark.asyncio -async def test_unique_sibling_names_rejected(client: AsyncClient, library: dict) -> None: - await _create_folder(client, library["id"], "Drafts") - resp = await client.post( - f"/libraries/{library['id']}/folders", - headers=AUTH_HEADERS, - json={"name": "Drafts"}, - ) - assert resp.status_code == 409 - - -@pytest.mark.asyncio -async def test_same_name_in_different_parents_is_allowed( - client: AsyncClient, library: dict -) -> None: - a = await _create_folder(client, library["id"], "A") - b = await _create_folder(client, library["id"], "B") - # Both can have a child named "Drafts" - await _create_folder(client, library["id"], "Drafts", a["id"]) - await _create_folder(client, library["id"], "Drafts", b["id"]) - - -@pytest.mark.asyncio -async def test_rename_folder(client: AsyncClient, library: dict) -> None: - folder = await _create_folder(client, library["id"], "Old") - resp = await client.put( - f"/libraries/{library['id']}/folders/{folder['id']}", - headers=AUTH_HEADERS, - json={"name": "New"}, - ) - assert resp.status_code == 200 - assert resp.json()["name"] == "New" - - -@pytest.mark.asyncio -async def test_rename_collision_rejected(client: AsyncClient, library: dict) -> None: - await _create_folder(client, library["id"], "Existing") - other = await _create_folder(client, library["id"], "Other") - resp = await client.put( - f"/libraries/{library['id']}/folders/{other['id']}", - headers=AUTH_HEADERS, - json={"name": "Existing"}, - ) - assert resp.status_code == 409 - - -@pytest.mark.asyncio -async def test_move_folder(client: AsyncClient, library: dict) -> None: - a = await _create_folder(client, library["id"], "A") - b = await _create_folder(client, library["id"], "B") - resp = await client.put( - f"/libraries/{library['id']}/folders/{b['id']}/move", - headers=AUTH_HEADERS, - json={"parent_folder_id": a["id"]}, - ) - assert resp.status_code == 200 - assert resp.json()["parent_folder_id"] == a["id"] - - -@pytest.mark.asyncio -async def test_move_folder_to_root(client: AsyncClient, library: dict) -> None: - parent = await _create_folder(client, library["id"], "P") - child = await _create_folder(client, library["id"], "C", parent["id"]) - resp = await client.put( - f"/libraries/{library['id']}/folders/{child['id']}/move", - headers=AUTH_HEADERS, - json={"parent_folder_id": None}, - ) - assert resp.status_code == 200 - assert resp.json()["parent_folder_id"] is None - - -# --------------------------------------------------------------------------- -# Cycle prevention -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_cannot_move_folder_into_self(client: AsyncClient, library: dict) -> None: - folder = await _create_folder(client, library["id"], "Self") - resp = await client.put( - f"/libraries/{library['id']}/folders/{folder['id']}/move", - headers=AUTH_HEADERS, - json={"parent_folder_id": folder["id"]}, - ) - assert resp.status_code == 400 - - -@pytest.mark.asyncio -async def test_cannot_move_folder_into_descendant(client: AsyncClient, library: dict) -> None: - a = await _create_folder(client, library["id"], "A") - b = await _create_folder(client, library["id"], "B", a["id"]) - c = await _create_folder(client, library["id"], "C", b["id"]) - # Moving A under C would create a cycle. - resp = await client.put( - f"/libraries/{library['id']}/folders/{a['id']}/move", - headers=AUTH_HEADERS, - json={"parent_folder_id": c["id"]}, - ) - assert resp.status_code == 400 - - -# --------------------------------------------------------------------------- -# Delete (reparents items + subfolders) -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_delete_folder_reparents_subfolders( - client: AsyncClient, library: dict -) -> None: - a = await _create_folder(client, library["id"], "A") - b = await _create_folder(client, library["id"], "B", a["id"]) - c = await _create_folder(client, library["id"], "C", b["id"]) - - resp = await client.delete( - f"/libraries/{library['id']}/folders/{b['id']}", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - - # After deleting B, C should be a child of A. - tree = (await client.get(f"/libraries/{library['id']}/tree", headers=AUTH_HEADERS)).json() - folders_by_id = {f["id"]: f for f in tree["folders"]} - assert c["id"] in folders_by_id - assert folders_by_id[c["id"]]["parent_folder_id"] == a["id"] - assert b["id"] not in folders_by_id - - -@pytest.mark.asyncio -async def test_delete_folder_collision_renames_subfolder( - client: AsyncClient, library: dict -) -> None: - """If a moved-up subfolder collides with an existing sibling, suffix it.""" - a = await _create_folder(client, library["id"], "A") - drafts_at_root = await _create_folder(client, library["id"], "Drafts") # noqa: F841 - drafts_in_a = await _create_folder(client, library["id"], "Drafts", a["id"]) # noqa: F841 - - # Delete A: its child "Drafts" moves up to root where another "Drafts" lives. - resp = await client.delete( - f"/libraries/{library['id']}/folders/{a['id']}", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - - tree = (await client.get(f"/libraries/{library['id']}/tree", headers=AUTH_HEADERS)).json() - root_folder_names = sorted( - f["name"] for f in tree["folders"] if f["parent_folder_id"] is None - ) - assert "Drafts" in root_folder_names - assert any(n.startswith("Drafts (") for n in root_folder_names) - - -# --------------------------------------------------------------------------- -# Validation -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "bad_name", - [ - "", # empty - " ", # whitespace only - "x" * 200, # too long - "with/slash", - "with\\backslash", - "with\x00null", - "with\nnewline", - ], -) -async def test_invalid_folder_names_rejected( - client: AsyncClient, library: dict, bad_name: str -) -> None: - resp = await client.post( - f"/libraries/{library['id']}/folders", - headers=AUTH_HEADERS, - json={"name": bad_name}, - ) - # Pydantic 422 on min_length/max_length failures, 400 on custom validator - assert resp.status_code in (400, 422) - - -@pytest.mark.asyncio -async def test_cross_library_parent_rejected(client: AsyncClient) -> None: - lib_a = await _create_library(client) - lib_b = await _create_library(client) - folder_in_b = await _create_folder(client, lib_b["id"], "Foo") - - resp = await client.post( - f"/libraries/{lib_a['id']}/folders", - headers=AUTH_HEADERS, - json={"name": "Bar", "parent_folder_id": folder_in_b["id"]}, - ) - assert resp.status_code == 400 - - -# --------------------------------------------------------------------------- -# Tree response -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_empty_library_tree(client: AsyncClient, library: dict) -> None: - resp = await client.get(f"/libraries/{library['id']}/tree", headers=AUTH_HEADERS) - assert resp.status_code == 200 - tree = resp.json() - assert tree["library_id"] == library["id"] - assert tree["folders"] == [] - assert tree["items"] == [] - - -@pytest.mark.asyncio -async def test_tree_returns_flat_lists(client: AsyncClient, library: dict) -> None: - a = await _create_folder(client, library["id"], "A") - b = await _create_folder(client, library["id"], "B", a["id"]) - c = await _create_folder(client, library["id"], "C") - - resp = await client.get(f"/libraries/{library['id']}/tree", headers=AUTH_HEADERS) - assert resp.status_code == 200 - tree = resp.json() - assert len(tree["folders"]) == 3 - ids = {f["id"] for f in tree["folders"]} - assert {a["id"], b["id"], c["id"]} == ids - - -# --------------------------------------------------------------------------- -# Tree 404 -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_tree_on_missing_library_returns_404(client: AsyncClient) -> None: - resp = await client.get("/libraries/nonexistent/tree", headers=AUTH_HEADERS) - assert resp.status_code == 404 - - -# --------------------------------------------------------------------------- -# Auth -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_tree_requires_auth(client: AsyncClient, library: dict) -> None: - resp = await client.get(f"/libraries/{library['id']}/tree") - assert resp.status_code in (401, 403) - - -@pytest.mark.asyncio -async def test_create_folder_requires_auth(client: AsyncClient, library: dict) -> None: - resp = await client.post( - f"/libraries/{library['id']}/folders", - json={"name": "NoAuth"}, - ) - assert resp.status_code in (401, 403) diff --git a/library-manager/tests/test_folders_branches.py b/library-manager/tests/test_folders_branches.py new file mode 100644 index 000000000..d7650803b --- /dev/null +++ b/library-manager/tests/test_folders_branches.py @@ -0,0 +1,275 @@ +"""Error/edge-path tests for the folder service + router. + +Router pre-checks shadow some service not-found guards, so those are driven +by direct service calls; the rest go through the HTTP router to also cover +its error→status mapping. +""" + +from __future__ import annotations + +import io +import uuid + +import pytest +from httpx import AsyncClient + +from .conftest import AUTH_HEADERS + + +async def _mk_library(client: AsyncClient) -> str: + lib_id = f"lib-{uuid.uuid4().hex[:8]}" + resp = await client.post( + "/libraries", headers=AUTH_HEADERS, + json={"id": lib_id, "organization_id": "org-fb", "name": f"L {lib_id[-6:]}"}, + ) + assert resp.status_code == 201 + return lib_id + + +async def _mk_folder(client, lib_id, name, parent=None) -> dict: + resp = await client.post( + f"/libraries/{lib_id}/folders", headers=AUTH_HEADERS, + json={"name": name, "parent_folder_id": parent}, + ) + assert resp.status_code == 201, resp.text + return resp.json() + + +# --------------------------------------------------------------------------- +# create_folder branches +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_create_in_missing_library_404(client: AsyncClient): + resp = await client.post( + "/libraries/nope/folders", headers=AUTH_HEADERS, json={"name": "X"} + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_create_with_missing_parent_404(client: AsyncClient): + lib = await _mk_library(client) + resp = await client.post( + f"/libraries/{lib}/folders", headers=AUTH_HEADERS, + json={"name": "X", "parent_folder_id": "missing-parent"}, + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_create_duplicate_sibling_conflict(client: AsyncClient): + lib = await _mk_library(client) + await _mk_folder(client, lib, "Dup") + resp = await client.post( + f"/libraries/{lib}/folders", headers=AUTH_HEADERS, json={"name": "Dup"} + ) + assert resp.status_code == 409 + + +# --------------------------------------------------------------------------- +# rename / move / delete via router (router 404 mappings) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_rename_missing_folder_404(client: AsyncClient): + lib = await _mk_library(client) + resp = await client.put( + f"/libraries/{lib}/folders/missing", headers=AUTH_HEADERS, + json={"name": "New"}, + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_rename_same_name_noop(client: AsyncClient): + lib = await _mk_library(client) + f = await _mk_folder(client, lib, "Keep") + resp = await client.put( + f"/libraries/{lib}/folders/{f['id']}", headers=AUTH_HEADERS, + json={"name": "Keep"}, + ) + assert resp.status_code == 200 + assert resp.json()["name"] == "Keep" + + +@pytest.mark.asyncio +async def test_move_missing_folder_404(client: AsyncClient): + lib = await _mk_library(client) + resp = await client.put( + f"/libraries/{lib}/folders/missing/move", headers=AUTH_HEADERS, + json={"parent_folder_id": None}, + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_move_into_itself_cycle(client: AsyncClient): + lib = await _mk_library(client) + f = await _mk_folder(client, lib, "Self") + resp = await client.put( + f"/libraries/{lib}/folders/{f['id']}/move", headers=AUTH_HEADERS, + json={"parent_folder_id": f["id"]}, + ) + assert resp.status_code == 400 # FolderCycleError + + +@pytest.mark.asyncio +async def test_move_missing_destination_404(client: AsyncClient): + lib = await _mk_library(client) + f = await _mk_folder(client, lib, "F") + resp = await client.put( + f"/libraries/{lib}/folders/{f['id']}/move", headers=AUTH_HEADERS, + json={"parent_folder_id": "missing-dest"}, + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_move_into_descendant_cycle(client: AsyncClient): + lib = await _mk_library(client) + parent = await _mk_folder(client, lib, "Parent") + child = await _mk_folder(client, lib, "Child", parent=parent["id"]) + resp = await client.put( + f"/libraries/{lib}/folders/{parent['id']}/move", headers=AUTH_HEADERS, + json={"parent_folder_id": child["id"]}, + ) + assert resp.status_code == 400 + + +@pytest.mark.asyncio +async def test_move_across_libraries_rejected(client: AsyncClient): + lib_a = await _mk_library(client) + lib_b = await _mk_library(client) + fa = await _mk_folder(client, lib_a, "A") + fb = await _mk_folder(client, lib_b, "B") + # Move A under B's folder — cross-library. + resp = await client.put( + f"/libraries/{lib_a}/folders/{fa['id']}/move", headers=AUTH_HEADERS, + json={"parent_folder_id": fb["id"]}, + ) + assert resp.status_code in (400, 404) + + +@pytest.mark.asyncio +async def test_delete_missing_folder_404(client: AsyncClient): + lib = await _mk_library(client) + resp = await client.delete( + f"/libraries/{lib}/folders/missing", headers=AUTH_HEADERS + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_delete_reparents_with_name_collision(client: AsyncClient): + # Root already has "Drafts"; folder X also contains "Drafts". Deleting X + # reparents its "Drafts" to root, which collides -> suffixed "Drafts (2)". + lib = await _mk_library(client) + await _mk_folder(client, lib, "Drafts") + x = await _mk_folder(client, lib, "X") + await _mk_folder(client, lib, "Drafts", parent=x["id"]) + resp = await client.delete( + f"/libraries/{lib}/folders/{x['id']}", headers=AUTH_HEADERS + ) + assert resp.status_code == 200 + tree = (await client.get(f"/libraries/{lib}/tree", headers=AUTH_HEADERS)).json() + names = sorted(f["name"] for f in tree["folders"]) + assert "Drafts" in names and "Drafts (2)" in names + + +# --------------------------------------------------------------------------- +# move_items branches +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_move_items_missing_library_404(client: AsyncClient): + resp = await client.post( + "/libraries/nope/items/move", headers=AUTH_HEADERS, + json={"item_ids": ["x"], "folder_id": None}, + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_move_items_missing_destination_404(client: AsyncClient): + lib = await _mk_library(client) + resp = await client.post( + f"/libraries/{lib}/items/move", headers=AUTH_HEADERS, + json={"item_ids": ["x"], "folder_id": "missing-dest"}, + ) + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_move_items_unknown_item_404(client: AsyncClient): + lib = await _mk_library(client) + resp = await client.post( + f"/libraries/{lib}/items/move", headers=AUTH_HEADERS, + json={"item_ids": ["unknown-item"], "folder_id": None}, + ) + assert resp.status_code == 404 + + +# --------------------------------------------------------------------------- +# Direct service calls for the not-found guards the router pre-check shadows +# --------------------------------------------------------------------------- + + +def test_service_not_found_guards(): + from database.connection import get_session_direct + from services import folder_service as fs + + db = get_session_direct() + try: + with pytest.raises(fs.FolderNotFoundError): + fs.rename_folder(db, "missing", "x") + with pytest.raises(fs.FolderNotFoundError): + fs.move_folder(db, "missing", None) + with pytest.raises(fs.FolderNotFoundError): + fs.delete_folder(db, "missing") + # Empty item list short-circuits to 0 (schema blocks this via HTTP). + assert fs.move_items(db, "any-lib", [], None) == 0 + finally: + db.close() + + +@pytest.mark.asyncio +async def test_delete_reparents_with_double_name_collision(client: AsyncClient): + # Root has both "Notes" and "Notes (2)"; folder X also contains "Notes". + # Deleting X reparents its "Notes" to root, where "Notes" AND "Notes (2)" + # are taken -> the suffix loop must advance to "Notes (3)" (folder_service + # _next_available_name `except FolderConflictError: n += 1`). + lib = await _mk_library(client) + await _mk_folder(client, lib, "Notes") + await _mk_folder(client, lib, "Notes (2)") + x = await _mk_folder(client, lib, "X") + await _mk_folder(client, lib, "Notes", parent=x["id"]) + resp = await client.delete( + f"/libraries/{lib}/folders/{x['id']}", headers=AUTH_HEADERS + ) + assert resp.status_code == 200 + tree = (await client.get(f"/libraries/{lib}/tree", headers=AUTH_HEADERS)).json() + names = sorted(f["name"] for f in tree["folders"]) + assert "Notes (3)" in names + + +@pytest.mark.asyncio +async def test_delete_folder_service_error_mapped(client: AsyncClient, monkeypatch): + # Router pre-check passes (folder exists), but delete_folder raises a + # FolderError -> router's `except FolderError` translates it to a status. + lib = await _mk_library(client) + f = await _mk_folder(client, lib, "Boom") + + from services import folder_service as fs + + def _raise(db, folder_id): + raise fs.FolderConflictError("synthetic") + + monkeypatch.setattr(fs, "delete_folder", _raise) + resp = await client.delete( + f"/libraries/{lib}/folders/{f['id']}", headers=AUTH_HEADERS + ) + assert resp.status_code == 409 diff --git a/library-manager/tests/test_import_plugins.py b/library-manager/tests/test_import_plugins.py deleted file mode 100644 index 31e624ad3..000000000 --- a/library-manager/tests/test_import_plugins.py +++ /dev/null @@ -1,465 +0,0 @@ -"""End-to-end tests for each import plugin. - -Each test: -1. Creates a library -2. Uploads/imports content using the specific plugin -3. Polls until status is "ready" (or "failed") -4. Verifies markdown content is available -5. Verifies metadata.json has correct structure and permalinks -6. Verifies source_ref.json is correct -7. Verifies original file is served (for file-based plugins) - -Plugins that require external services (url_import → Firecrawl) test the -error path instead, since the external service is not available in CI. -""" - -import asyncio -import io -import time - -import pytest -from httpx import AsyncClient - -AUTH_HEADERS = {"Authorization": "Bearer test-token"} - -# Maximum time (seconds) to wait for an import job to complete. -_POLL_TIMEOUT = 20 -_POLL_INTERVAL = 0.5 - - -async def _wait_for_ready( - client: AsyncClient, - lib_id: str, - item_id: str, - timeout: float = _POLL_TIMEOUT, -) -> str: - """Poll item status until it leaves 'pending'/'processing'. - - Args: - client: HTTP client. - lib_id: Library ID. - item_id: Content item ID. - timeout: Max seconds to wait. - - Returns: - Final status string ('ready' or 'failed'). - """ - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/status", - headers=AUTH_HEADERS, - ) - status = resp.json()["status"] - if status in ("ready", "failed"): - return status - await asyncio.sleep(_POLL_INTERVAL) - return "timeout" - - -# ----------------------------------------------------------------------- -# Plugin 1: simple_import — plain text/markdown/html -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_simple_import_markdown(client: AsyncClient, library: dict): - """simple_import should import a .md file and serve it back as markdown.""" - lib_id = library["id"] - content = ( - "# Biology Notes\n\n" - "## Cell Structure\n\n" - "The cell is the basic unit of life.\n\n" - "- Nucleus\n" - "- Mitochondria\n" - "- Cell membrane\n" - ) - - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("biology.md", io.BytesIO(content.encode()), "text/markdown")}, - data={"plugin_name": "simple_import", "title": "Biology Notes"}, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - status = await _wait_for_ready(client, lib_id, item_id) - assert status == "ready", f"Expected ready, got {status}" - - # Verify markdown content. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content", - headers=AUTH_HEADERS, - params={"format": "markdown"}, - ) - assert resp.status_code == 200 - assert "# Biology Notes" in resp.text - assert "Cell membrane" in resp.text - - # Verify HTML rendering. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content", - headers=AUTH_HEADERS, - params={"format": "html"}, - ) - assert resp.status_code == 200 - assert "" - "

    HTML Test

    " - "

    Paragraph with bold text.

    " - "" - ) - - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("test.html", io.BytesIO(content.encode()), "text/html")}, - data={"plugin_name": "simple_import", "title": "HTML Test"}, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - status = await _wait_for_ready(client, lib_id, item_id) - assert status == "ready" - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content", - headers=AUTH_HEADERS, - params={"format": "markdown"}, - ) - assert resp.status_code == 200 - assert "HTML Test" in resp.text - - -@pytest.mark.asyncio -async def test_simple_import_txt(client: AsyncClient, library: dict): - """simple_import should import a .txt file.""" - lib_id = library["id"] - content = "This is a plain text file.\nLine two.\nLine three." - - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("notes.txt", io.BytesIO(content.encode()), "text/plain")}, - data={"plugin_name": "simple_import", "title": "Plain Text"}, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - status = await _wait_for_ready(client, lib_id, item_id) - assert status == "ready" - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - assert "plain text file" in resp.text - - -# ----------------------------------------------------------------------- -# Plugin 2: markitdown_import — MarkItDown conversion -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_markitdown_import(client: AsyncClient, library: dict): - """markitdown_import should convert a markdown file via MarkItDown.""" - lib_id = library["id"] - content = ( - "# MarkItDown Test\n\n" - "| Column A | Column B |\n" - "|----------|----------|\n" - "| Cell 1 | Cell 2 |\n" - "| Cell 3 | Cell 4 |\n" - ) - - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("table.md", io.BytesIO(content.encode()), "text/markdown")}, - data={"plugin_name": "markitdown_import", "title": "MarkItDown Table Test"}, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - status = await _wait_for_ready(client, lib_id, item_id) - assert status == "ready" - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - text = resp.text - assert "MarkItDown Test" in text or "Column A" in text - - # Verify item detail. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - detail = resp.json() - assert detail["import_plugin"] == "markitdown_import" - assert detail["status"] == "ready" - assert detail["metadata"]["permalinks"]["full_markdown"] is not None - - -# ----------------------------------------------------------------------- -# Plugin 3: markitdown_plus_import — enhanced with image/page awareness -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_markitdown_plus_import(client: AsyncClient, library: dict): - """markitdown_plus_import should convert with enhanced features.""" - lib_id = library["id"] - content = ( - "" - "

    Enhanced Import

    " - "

    Section 1

    " - "

    First section content.

    " - "

    Section 2

    " - "

    Second section content with a formula.

    " - "" - ) - - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("enhanced.html", io.BytesIO(content.encode()), "text/html")}, - data={ - "plugin_name": "markitdown_plus_import", - "title": "Enhanced Import Test", - "plugin_params": '{"image_descriptions": "none"}', - }, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - status = await _wait_for_ready(client, lib_id, item_id) - assert status == "ready" - - # Verify markdown. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - assert "Enhanced Import" in resp.text or "Section 1" in resp.text - - # Verify metadata. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/metadata", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - meta = resp.json() - assert meta["import_plugin"] == "markitdown_plus_import" - assert "permalinks" in meta - - # HTML imports don't produce pages or images on disk. The capability - # handler returns 404 (HandlerUnavailable) and the item's ``capabilities`` - # list excludes both — that's the correct contract for the new tabs UI. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content/pages", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 404 - - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content/images", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 404 - - caps = await client.get( - f"/libraries/{lib_id}/items/{item_id}/capabilities", headers=AUTH_HEADERS - ) - assert caps.status_code == 200 - cap_list = caps.json()["capabilities"] - assert "pages" not in cap_list - assert "images" not in cap_list - - -# ----------------------------------------------------------------------- -# Plugin 4: url_import — Firecrawl (tests error path without service) -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_url_import_no_firecrawl(client: AsyncClient, library: dict): - """url_import should fail gracefully when Firecrawl is unavailable.""" - lib_id = library["id"] - - resp = await client.post( - f"/libraries/{lib_id}/import/url", - headers=AUTH_HEADERS, - json={ - "url": "https://example.com", - "title": "Example.com", - "plugin_name": "url_import", - "api_keys": {"firecrawl_key": "dummy", "firecrawl_url": "http://localhost:3002"}, - }, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - status = await _wait_for_ready(client, lib_id, item_id, timeout=15) - assert status == "failed", "Expected failure without Firecrawl service" - - # Verify error is recorded. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/status", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - assert resp.json()["error_message"] is not None - - -@pytest.mark.asyncio -async def test_url_import_source_type_validation(client: AsyncClient, library: dict): - """Attempting to use a file plugin for URL import should fail.""" - lib_id = library["id"] - - resp = await client.post( - f"/libraries/{lib_id}/import/url", - headers=AUTH_HEADERS, - json={ - "url": "https://example.com", - "title": "Bad Plugin", - "plugin_name": "simple_import", - }, - ) - assert resp.status_code == 400 - - -# ----------------------------------------------------------------------- -# Plugin 5: youtube_transcript_import -# ----------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_youtube_import(client: AsyncClient, library: dict): - """youtube_transcript_import should download and convert a transcript. - - Uses a well-known video with reliable subtitles. - """ - lib_id = library["id"] - - resp = await client.post( - f"/libraries/{lib_id}/import/youtube", - headers=AUTH_HEADERS, - json={ - "video_url": "https://www.youtube.com/watch?v=dQw4w9WgXcQ", - "language": "en", - "title": "Never Gonna Give You Up", - "plugin_name": "youtube_transcript_import", - }, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - # YouTube downloads can take time — allow up to 30s. - status = await _wait_for_ready(client, lib_id, item_id, timeout=30) - - if status == "failed": - # YouTube may rate-limit in CI — check that error was recorded properly. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/status", - headers=AUTH_HEADERS, - ) - error = resp.json().get("error_message", "") - pytest.skip(f"YouTube import failed (likely rate-limited): {error[:100]}") - - assert status == "ready" - - # Verify markdown content exists. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/content", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - text = resp.text - - # YouTube may rate-limit subtitle requests — if no transcript was - # obtained, the plugin returns a placeholder. Both cases are valid. - if "No transcript available" in text: - pytest.skip("YouTube subtitles unavailable (likely rate-limited)") - - assert "Transcript" in text or "**[" in text # Timestamped entries - - # Verify source_ref. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/source_ref", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - src = resp.json() - assert src["type"] == "youtube" - assert src["video_id"] == "dQw4w9WgXcQ" - - # Verify metadata. - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/metadata", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - meta = resp.json() - assert meta["import_plugin"] == "youtube_transcript_import" - - -@pytest.mark.asyncio -async def test_youtube_source_type_validation(client: AsyncClient, library: dict): - """Attempting to use a file plugin for YouTube import should fail.""" - lib_id = library["id"] - - resp = await client.post( - f"/libraries/{lib_id}/import/youtube", - headers=AUTH_HEADERS, - json={ - "video_url": "https://www.youtube.com/watch?v=dQw4w9WgXcQ", - "language": "en", - "title": "Bad Plugin", - "plugin_name": "simple_import", - }, - ) - assert resp.status_code == 400 diff --git a/library-manager/tests/test_items_move.py b/library-manager/tests/test_items_move.py deleted file mode 100644 index cac4debd7..000000000 --- a/library-manager/tests/test_items_move.py +++ /dev/null @@ -1,262 +0,0 @@ -"""Tests for item move (single + bulk) and upload-with-folder_id.""" - -from __future__ import annotations - -import asyncio -import io -import time -import uuid - -import pytest -from httpx import AsyncClient - -from .conftest import AUTH_HEADERS - -_POLL_TIMEOUT = 15 - - -# --------------------------------------------------------------------------- -# Helpers (mirror the patterns in test_content_serving.py) -# --------------------------------------------------------------------------- - - -async def _wait_for_ready(client: AsyncClient, lib_id: str, item_id: str) -> str: - deadline = time.monotonic() + _POLL_TIMEOUT - while time.monotonic() < deadline: - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/status", headers=AUTH_HEADERS - ) - if resp.json()["status"] in ("ready", "failed"): - return resp.json()["status"] - await asyncio.sleep(0.5) - return "timeout" - - -async def _upload( - client: AsyncClient, - lib_id: str, - title: str = "Doc", - folder_id: str | None = None, -) -> str: - data = {"plugin_name": "simple_import", "title": title} - if folder_id is not None: - data["folder_id"] = folder_id - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": ("a.md", io.BytesIO(b"# x"), "text/markdown")}, - data=data, - ) - assert resp.status_code == 202, resp.text - item_id = resp.json()["item_id"] - status = await _wait_for_ready(client, lib_id, item_id) - assert status == "ready", f"upload did not become ready: {status}" - return item_id - - -async def _create_folder( - client: AsyncClient, lib_id: str, name: str, parent_id: str | None = None -) -> str: - resp = await client.post( - f"/libraries/{lib_id}/folders", - headers=AUTH_HEADERS, - json={"name": name, "parent_folder_id": parent_id}, - ) - assert resp.status_code == 201, resp.text - return resp.json()["id"] - - -async def _create_library(client: AsyncClient) -> str: - lib_id = f"lib-{uuid.uuid4().hex[:8]}" - resp = await client.post( - "/libraries", - headers=AUTH_HEADERS, - json={ - "id": lib_id, - "organization_id": "org-test", - "name": f"Lib {lib_id[-8:]}", - }, - ) - assert resp.status_code == 201 - return resp.json()["id"] - - -# --------------------------------------------------------------------------- -# Upload with folder_id -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_upload_lands_at_root_by_default(client: AsyncClient, library: dict) -> None: - item_id = await _upload(client, library["id"]) - resp = await client.get( - f"/libraries/{library['id']}/items/{item_id}", headers=AUTH_HEADERS - ) - assert resp.status_code == 200 - assert resp.json()["folder_id"] is None - - -@pytest.mark.asyncio -async def test_upload_with_folder_id(client: AsyncClient, library: dict) -> None: - folder_id = await _create_folder(client, library["id"], "Q1 Research") - item_id = await _upload(client, library["id"], folder_id=folder_id) - resp = await client.get( - f"/libraries/{library['id']}/items/{item_id}", headers=AUTH_HEADERS - ) - assert resp.json()["folder_id"] == folder_id - - -@pytest.mark.asyncio -async def test_upload_with_invalid_folder_rejected( - client: AsyncClient, library: dict -) -> None: - resp = await client.post( - f"/libraries/{library['id']}/import/file", - headers=AUTH_HEADERS, - files={"file": ("a.md", io.BytesIO(b"# x"), "text/markdown")}, - data={ - "plugin_name": "simple_import", - "title": "x", - "folder_id": "nonexistent-folder-id", - }, - ) - assert resp.status_code == 400 - - -@pytest.mark.asyncio -async def test_upload_cross_library_folder_rejected(client: AsyncClient) -> None: - lib_a = await _create_library(client) - lib_b = await _create_library(client) - folder_in_b = await _create_folder(client, lib_b, "Foo") - - resp = await client.post( - f"/libraries/{lib_a}/import/file", - headers=AUTH_HEADERS, - files={"file": ("a.md", io.BytesIO(b"# x"), "text/markdown")}, - data={ - "plugin_name": "simple_import", - "title": "x", - "folder_id": folder_in_b, - }, - ) - assert resp.status_code == 400 - - -# --------------------------------------------------------------------------- -# Bulk item move -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_bulk_move_items_to_folder(client: AsyncClient, library: dict) -> None: - folder_id = await _create_folder(client, library["id"], "Target") - item1 = await _upload(client, library["id"], title="One") - item2 = await _upload(client, library["id"], title="Two") - - resp = await client.post( - f"/libraries/{library['id']}/items/move", - headers=AUTH_HEADERS, - json={"item_ids": [item1, item2], "folder_id": folder_id}, - ) - assert resp.status_code == 200 - assert resp.json()["moved"] == 2 - - for iid in (item1, item2): - d = await client.get( - f"/libraries/{library['id']}/items/{iid}", headers=AUTH_HEADERS - ) - assert d.json()["folder_id"] == folder_id - - -@pytest.mark.asyncio -async def test_move_items_to_root(client: AsyncClient, library: dict) -> None: - folder_id = await _create_folder(client, library["id"], "T") - item_id = await _upload(client, library["id"], folder_id=folder_id) - - resp = await client.post( - f"/libraries/{library['id']}/items/move", - headers=AUTH_HEADERS, - json={"item_ids": [item_id], "folder_id": None}, - ) - assert resp.status_code == 200 - d = await client.get( - f"/libraries/{library['id']}/items/{item_id}", headers=AUTH_HEADERS - ) - assert d.json()["folder_id"] is None - - -@pytest.mark.asyncio -async def test_move_items_cross_library_rejected(client: AsyncClient) -> None: - lib_a = await _create_library(client) - lib_b = await _create_library(client) - item_in_b = await _upload(client, lib_b) - - # Try to move an item that belongs to lib_b through lib_a's endpoint. - resp = await client.post( - f"/libraries/{lib_a}/items/move", - headers=AUTH_HEADERS, - json={"item_ids": [item_in_b], "folder_id": None}, - ) - assert resp.status_code == 400 - - -@pytest.mark.asyncio -async def test_move_to_cross_library_folder_rejected(client: AsyncClient) -> None: - lib_a = await _create_library(client) - lib_b = await _create_library(client) - item_a = await _upload(client, lib_a) - folder_b = await _create_folder(client, lib_b, "Foreign") - - resp = await client.post( - f"/libraries/{lib_a}/items/move", - headers=AUTH_HEADERS, - json={"item_ids": [item_a], "folder_id": folder_b}, - ) - assert resp.status_code == 400 - - -@pytest.mark.asyncio -async def test_move_empty_item_list_rejected(client: AsyncClient, library: dict) -> None: - resp = await client.post( - f"/libraries/{library['id']}/items/move", - headers=AUTH_HEADERS, - json={"item_ids": [], "folder_id": None}, - ) - # Pydantic min_length=1 → 422 - assert resp.status_code == 422 - - -@pytest.mark.asyncio -async def test_bulk_move_payload_cap(client: AsyncClient, library: dict) -> None: - """Request bodies with > 500 item_ids should be rejected by the schema.""" - big = [f"item-{i}" for i in range(501)] - resp = await client.post( - f"/libraries/{library['id']}/items/move", - headers=AUTH_HEADERS, - json={"item_ids": big, "folder_id": None}, - ) - # Pydantic max_length=500 → 422 - assert resp.status_code == 422 - - -# --------------------------------------------------------------------------- -# Folder delete reparents items -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_folder_delete_reparents_items(client: AsyncClient, library: dict) -> None: - parent = await _create_folder(client, library["id"], "Parent") - child = await _create_folder(client, library["id"], "Child", parent) - item_id = await _upload(client, library["id"], folder_id=child) - - # Delete the child folder; its item should reparent to Parent. - resp = await client.delete( - f"/libraries/{library['id']}/folders/{child}", headers=AUTH_HEADERS - ) - assert resp.status_code == 200 - - d = await client.get( - f"/libraries/{library['id']}/items/{item_id}", headers=AUTH_HEADERS - ) - assert d.json()["folder_id"] == parent diff --git a/library-manager/tests/test_libraries.py b/library-manager/tests/test_libraries.py deleted file mode 100644 index 42bc051e4..000000000 --- a/library-manager/tests/test_libraries.py +++ /dev/null @@ -1,122 +0,0 @@ -"""Tests for library CRUD operations.""" - -import uuid - -import pytest -from httpx import AsyncClient - -AUTH_HEADERS = {"Authorization": "Bearer test-token"} - - -@pytest.mark.asyncio -async def test_create_library(client: AsyncClient): - """Creating a library should return 201 with correct fields.""" - lib_id = f"lib-{uuid.uuid4().hex[:8]}" - resp = await client.post( - "/libraries", - headers=AUTH_HEADERS, - json={ - "id": lib_id, - "organization_id": "org-crud", - "name": f"CRUD Test {lib_id[-8:]}", - }, - ) - assert resp.status_code == 201 - data = resp.json() - assert data["id"] == lib_id - assert data["organization_id"] == "org-crud" - assert data["item_count"] == 0 - - -@pytest.mark.asyncio -async def test_get_library(client: AsyncClient, library: dict): - """Getting a library should return its details.""" - resp = await client.get( - f"/libraries/{library['id']}", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - assert resp.json()["name"] == library["name"] - - -@pytest.mark.asyncio -async def test_get_nonexistent_library(client: AsyncClient): - """Getting a non-existent library should return 404.""" - resp = await client.get("/libraries/fake-id", headers=AUTH_HEADERS) - assert resp.status_code == 404 - - -@pytest.mark.asyncio -async def test_list_libraries(client: AsyncClient, library: dict): - """Listing libraries should include the created library.""" - resp = await client.get( - "/libraries", - headers=AUTH_HEADERS, - params={"organization_id": library["organization_id"]}, - ) - assert resp.status_code == 200 - data = resp.json() - assert data["total"] >= 1 - ids = [lib["id"] for lib in data["libraries"]] - assert library["id"] in ids - - -@pytest.mark.asyncio -async def test_duplicate_name_rejected(client: AsyncClient): - """Creating two libraries with the same name in one org should fail.""" - org_id = f"org-{uuid.uuid4().hex[:8]}" - name = "Duplicate Test" - - resp1 = await client.post( - "/libraries", - headers=AUTH_HEADERS, - json={"id": f"lib-{uuid.uuid4().hex[:8]}", "organization_id": org_id, "name": name}, - ) - assert resp1.status_code == 201 - - resp2 = await client.post( - "/libraries", - headers=AUTH_HEADERS, - json={"id": f"lib-{uuid.uuid4().hex[:8]}", "organization_id": org_id, "name": name}, - ) - assert resp2.status_code == 409 - - -@pytest.mark.asyncio -async def test_delete_library(client: AsyncClient): - """Deleting a library should remove it from the database.""" - lib_id = f"lib-{uuid.uuid4().hex[:8]}" - await client.post( - "/libraries", - headers=AUTH_HEADERS, - json={"id": lib_id, "organization_id": "org-del", "name": f"Del {lib_id[-6:]}"}, - ) - - resp = await client.delete(f"/libraries/{lib_id}", headers=AUTH_HEADERS) - assert resp.status_code == 200 - - resp = await client.get(f"/libraries/{lib_id}", headers=AUTH_HEADERS) - assert resp.status_code == 404 - - -@pytest.mark.asyncio -async def test_import_config_roundtrip(client: AsyncClient, library: dict): - """Setting and reading import config should produce consistent results.""" - config = {"image_descriptions": "llm", "max_discovery_depth": 5} - - resp = await client.put( - f"/libraries/{library['id']}/import-config", - headers=AUTH_HEADERS, - json=config, - ) - assert resp.status_code == 200 - assert "warning" in resp.json() - - resp = await client.get( - f"/libraries/{library['id']}/import-config", - headers=AUTH_HEADERS, - ) - assert resp.status_code == 200 - data = resp.json()["import_config"] - assert data["image_descriptions"] == "llm" - assert data["max_discovery_depth"] == 5 diff --git a/library-manager/tests/test_load.py b/library-manager/tests/test_load.py deleted file mode 100644 index 215fa3cb6..000000000 --- a/library-manager/tests/test_load.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Load test: verify concurrent import processing respects the semaphore. - -Uploads 20 files simultaneously and verifies: -1. All 20 are accepted (HTTP 202) -2. At most MAX_CONCURRENT_IMPORTS run at the same time -3. All 20 eventually reach 'ready' status -4. Total time is reasonable (not serialized) -""" - -import asyncio -import io -import time - -import pytest -from httpx import AsyncClient - -AUTH_HEADERS = {"Authorization": "Bearer test-token"} - -_NUM_FILES = 20 -_POLL_TIMEOUT = 60 - - -async def _wait_all_ready( - client: AsyncClient, - lib_id: str, - item_ids: list[str], - timeout: float = _POLL_TIMEOUT, -) -> dict[str, str]: - """Poll until all items leave pending/processing state. - - Returns: - Dict mapping item_id → final status. - """ - statuses = {iid: "pending" for iid in item_ids} - deadline = time.monotonic() + timeout - - while time.monotonic() < deadline: - pending = [iid for iid, s in statuses.items() if s in ("pending", "processing")] - if not pending: - break - for iid in pending: - resp = await client.get( - f"/libraries/{lib_id}/items/{iid}/status", headers=AUTH_HEADERS - ) - statuses[iid] = resp.json()["status"] - await asyncio.sleep(0.5) - - return statuses - - -@pytest.mark.asyncio -async def test_concurrent_imports(client: AsyncClient, library: dict): - """Upload 20 files concurrently and verify all complete successfully.""" - lib_id = library["id"] - - t0 = time.monotonic() - item_ids = [] - for i in range(_NUM_FILES): - content = f"# Document {i}\n\nThis is test document number {i}.\n" * 10 - resp = await client.post( - f"/libraries/{lib_id}/import/file", - headers=AUTH_HEADERS, - files={"file": (f"doc_{i:03d}.md", io.BytesIO(content.encode()), "text/markdown")}, - data={"plugin_name": "simple_import", "title": f"Load Test Doc {i}"}, - ) - assert resp.status_code == 202, f"File {i} rejected: {resp.text}" - item_ids.append(resp.json()["item_id"]) - - statuses = await _wait_all_ready(client, lib_id, item_ids) - - total_time = time.monotonic() - t0 - ready_count = sum(1 for s in statuses.values() if s == "ready") - failed_count = sum(1 for s in statuses.values() if s == "failed") - pending_count = sum(1 for s in statuses.values() if s in ("pending", "processing")) - - assert ready_count == _NUM_FILES, ( - f"Expected {_NUM_FILES} ready, got {ready_count} ready, " - f"{failed_count} failed, {pending_count} still processing. " - f"Total time: {total_time:.1f}s" - ) - - resp = await client.get( - f"/libraries/{lib_id}/items", headers=AUTH_HEADERS, params={"limit": 100} - ) - assert resp.json()["total"] == _NUM_FILES - - for iid in item_ids: - resp = await client.get( - f"/libraries/{lib_id}/items/{iid}/content", headers=AUTH_HEADERS - ) - assert resp.status_code == 200 - assert len(resp.text) > 0 diff --git a/library-manager/tests/test_markitdown_errors.py b/library-manager/tests/test_markitdown_errors.py deleted file mode 100644 index 3c12ae709..000000000 --- a/library-manager/tests/test_markitdown_errors.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Unit tests for the user-friendly markitdown error translator.""" - -from __future__ import annotations - -import pytest -from plugins._markitdown_errors import humanize_markitdown_error - - -class _Fake(Exception): - """Generic exception used to spoof markitdown's exception classes by name.""" - - -class MissingDependencyException(Exception): # noqa: N818 — mirrors markitdown's name - pass - - -class UnsupportedFormatException(Exception): # noqa: N818 - pass - - -class FileConversionException(Exception): # noqa: N818 - pass - - -def test_missing_dependency_with_ext_hint_in_message(): - exc = MissingDependencyException( - "PdfConverter recognized the input as a potential .pdf file, but the " - "dependencies needed to read .pdf files have not been installed." - ) - msg = humanize_markitdown_error(exc, "report.pdf") - assert "report.pdf" in msg - assert ".pdf" in msg - assert "not installed" in msg - # Should NOT include pip install hints. - assert "pip install" not in msg - - -def test_missing_dependency_falls_back_to_filename_extension(): - exc = MissingDependencyException("no readers available") - msg = humanize_markitdown_error(exc, "song.mp3") - assert "song.mp3" in msg - assert ".mp3" in msg - - -def test_unsupported_format(): - exc = UnsupportedFormatException("nothing matched") - msg = humanize_markitdown_error(exc, "blob.bin") - assert "blob.bin" in msg - assert "not supported" in msg.lower() - - -def test_file_conversion_corrupted(): - exc = FileConversionException("Stream error / corrupted") - msg = humanize_markitdown_error(exc, "broken.docx") - assert "broken.docx" in msg - # Should mention readability / corruption (one of: unreadable, corrupted, password) - assert any(w in msg.lower() for w in ("unreadable", "corrupted", "password")) - - -def test_generic_fallback_does_not_leak_class_name(): - exc = ValueError("Random internal error xyz") - msg = humanize_markitdown_error(exc, "doc.txt") - assert "doc.txt" in msg - # Generic message: shouldn't echo the original raw error string verbatim - assert "Random internal error xyz" not in msg - assert "ValueError" not in msg - - -def test_empty_file_detection(): - exc = Exception("File is empty") - msg = humanize_markitdown_error(exc, "x.pdf") - assert "empty" in msg.lower() - assert "x.pdf" in msg - - -@pytest.mark.parametrize("ext,expected", [ - ("report.pdf", ".pdf"), - ("photo.JPG", ".jpg"), - ("noext", ""), -]) -def test_ext_label(ext, expected): - from plugins._markitdown_errors import _ext_label - - assert _ext_label(ext) == expected diff --git a/library-manager/tests/test_markitdown_plus_plugin.py b/library-manager/tests/test_markitdown_plus_plugin.py new file mode 100644 index 000000000..b92016574 --- /dev/null +++ b/library-manager/tests/test_markitdown_plus_plugin.py @@ -0,0 +1,326 @@ +"""Unit tests for ``plugins.markitdown_plus_import``. + +MarkItDown, PyMuPDF (``fitz``) and OpenAI are all mocked so the tests run +offline. Covers the conversion happy/error paths, page splitting, MIME maps, +image extraction (bitmap + rasterize fallback + no-fitz), and the three +``_describe_image`` modes. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path + +import pytest + +import plugins.markitdown_plus_import as mp +from plugins.markitdown_plus_import import ( + MarkItDownPlusPlugin, + _describe_image, + _guess_mime, + _image_mime, + _split_into_pages, +) + + +# --------------------------------------------------------------------------- +# pure helpers +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "ext,expected", + [("png", "image/png"), (".jpg", "image/jpeg"), ("JPEG", "image/jpeg"), + ("webp", "image/webp"), ("unknown", "image/png")], +) +def test_image_mime(ext, expected): + assert _image_mime(ext) == expected + + +@pytest.mark.parametrize( + "ext,expected", + [(".pdf", "application/pdf"), (".csv", "text/csv"), + (".mp3", "audio/mpeg"), (".zzz", "application/octet-stream")], +) +def test_guess_mime(ext, expected): + assert _guess_mime(ext) == expected + + +def test_split_into_pages_non_page_aware_returns_empty(): + assert _split_into_pages("a\n---\nb", "html") == [] + + +def test_split_into_pages_with_breaks(): + content = "Page one\n\n---\n\nPage two\n\n---\n\nPage three" + pages = _split_into_pages(content, "pdf") + assert len(pages) == 3 + assert pages[0].text == "Page one" + + +def test_split_into_pages_no_breaks_returns_empty(): + assert _split_into_pages("single page, no breaks", "pdf") == [] + + +def test_get_parameters(): + names = {p.name for p in MarkItDownPlusPlugin().get_parameters()} + assert "image_descriptions" in names + + +# --------------------------------------------------------------------------- +# _describe_image +# --------------------------------------------------------------------------- + + +def test_describe_image_basic(): + assert _describe_image(b"x", "i.png", "png", "basic", {}, {}) == "Image: i.png" + + +def test_describe_image_none_mode(): + assert _describe_image(b"x", "i.png", "png", "none", {}, {}) is None + + +def test_describe_image_llm_no_key(): + stats = {"llm_calls": []} + assert ( + _describe_image(b"x", "i.png", "png", "llm", {}, stats) == "Image: i.png" + ) + + +def test_describe_image_llm_success(monkeypatch): + import openai + + class _Msg: + content = " A blue diagram. " + + class _Choice: + message = _Msg() + + class _Usage: + total_tokens = 42 + + class _Resp: + choices = [_Choice()] + usage = _Usage() + + class FakeOpenAI: + def __init__(self, **kwargs): + self.chat = types.SimpleNamespace( + completions=types.SimpleNamespace(create=lambda **kw: _Resp()) + ) + + monkeypatch.setattr(openai, "OpenAI", FakeOpenAI) + stats = {"llm_calls": [], "images_with_llm_descriptions": 0} + out = _describe_image(b"x", "i.png", "png", "llm", {"openai_vision": "sk"}, stats) + assert out == "A blue diagram." + assert stats["images_with_llm_descriptions"] == 1 + assert stats["llm_calls"][0]["success"] is True + + +def test_describe_image_llm_failure(monkeypatch): + import openai + + class FakeOpenAI: + def __init__(self, **kwargs): + raise RuntimeError("api down") + + monkeypatch.setattr(openai, "OpenAI", FakeOpenAI) + stats = {"llm_calls": []} + out = _describe_image(b"x", "i.png", "png", "llm", {"openai_vision": "sk"}, stats) + assert out == "Image: i.png" + assert stats["llm_calls"][0]["success"] is False + + +# --------------------------------------------------------------------------- +# _extract_images (fitz mocked) +# --------------------------------------------------------------------------- + + +def _install_fake_fitz(monkeypatch, *, pages): + fake = types.ModuleType("fitz") + + class FakeDoc: + def __init__(self, pages): + self._pages = pages + + def __len__(self): + return len(self._pages) + + def __getitem__(self, i): + return self._pages[i] + + def extract_image(self, xref): + return self._extract_map.get(xref) + + def close(self): + pass + + doc = FakeDoc(pages) + doc._extract_map = {} + for p in pages: + doc._extract_map.update(getattr(p, "_extract_map", {})) + fake.open = lambda path: doc + monkeypatch.setitem(sys.modules, "fitz", fake) + return doc + + +class _FakePage: + def __init__(self, images=None, extract_map=None): + self._images = images or [] + self._extract_map = extract_map or {} + + def get_images(self, full=True): + return self._images + + def get_pixmap(self, dpi=144): + return types.SimpleNamespace(tobytes=lambda fmt: b"PNGBYTES") + + +def test_extract_images_no_fitz(monkeypatch): + monkeypatch.setitem(sys.modules, "fitz", None) + assert mp._extract_images(Path("x.pdf"), "basic", {}, {}) == [] + + +def test_extract_images_non_pdf_suffix(monkeypatch): + _install_fake_fitz(monkeypatch, pages=[]) + assert mp._extract_images(Path("x.docx"), "basic", {}, {}) == [] + + +def test_extract_images_bitmap(monkeypatch): + page = _FakePage( + images=[(7,)], + extract_map={7: {"image": b"IMGDATA", "ext": "png"}}, + ) + _install_fake_fitz(monkeypatch, pages=[page]) + out = mp._extract_images(Path("x.pdf"), "basic", {}, {"llm_calls": []}) + assert len(out) == 1 + assert out[0].filename == "img_001.png" + assert out[0].page_number == 1 + + +def test_extract_images_open_failure(monkeypatch): + fake = types.ModuleType("fitz") + + def _open(path): + raise RuntimeError("cannot open") + + fake.open = _open + monkeypatch.setitem(sys.modules, "fitz", fake) + assert mp._extract_images(Path("x.pdf"), "basic", {}, {}) == [] + + +def test_extract_images_skips_bad_xref_and_empty(monkeypatch): + # xref 1 raises in extract_image (skipped); xref 2 returns empty (skipped). + class _BadPage(_FakePage): + def __init__(self): + super().__init__(images=[(1,), (2,)]) + + page = _BadPage() + + class _Doc: + def __len__(self): + return 1 + + def __getitem__(self, i): + return page + + def extract_image(self, xref): + if xref == 1: + raise RuntimeError("corrupt") + return {"image": b""} # empty -> skipped + + def close(self): + pass + + fake = types.ModuleType("fitz") + fake.open = lambda path: _Doc() + monkeypatch.setitem(sys.modules, "fitz", fake) + # No bitmap images survive -> falls back to rasterizing the single page. + out = mp._extract_images(Path("x.pdf"), "basic", {}, {"llm_calls": []}) + assert out[0].filename == "page_001.png" + + +def test_extract_images_rasterize_pixmap_failure(monkeypatch): + class _BoomPage(_FakePage): + def get_pixmap(self, dpi=144): + raise RuntimeError("render failed") + + _install_fake_fitz(monkeypatch, pages=[_BoomPage(images=[])]) + # Pixmap render fails -> page skipped -> no images. + assert mp._extract_images(Path("x.pdf"), "basic", {}, {"llm_calls": []}) == [] + + +def test_extract_images_rasterize_fallback(monkeypatch): + # No bitmap images -> render each page as PNG. + page = _FakePage(images=[]) + stats = {"llm_calls": []} + _install_fake_fitz(monkeypatch, pages=[page, _FakePage(images=[])]) + out = mp._extract_images(Path("x.pdf"), "basic", {}, stats) + assert len(out) == 2 + assert out[0].filename == "page_001.png" + assert stats["rendered_page_fallback"] == 2 + + +# --------------------------------------------------------------------------- +# import_content +# --------------------------------------------------------------------------- + + +def _fake_markitdown(monkeypatch, *, text="converted", raises=None): + import markitdown + + class FakeMD: + def convert(self, path): + if raises is not None: + raise raises + return types.SimpleNamespace(text_content=text) + + monkeypatch.setattr(markitdown, "MarkItDown", FakeMD) + + +def test_import_content_file_not_found(): + with pytest.raises(FileNotFoundError): + MarkItDownPlusPlugin().import_content("/nonexistent/file.pdf") + + +def test_import_content_success_html(tmp_path, monkeypatch): + # .html is not page-aware -> no image extraction, no pages. + f = tmp_path / "doc.html" + f.write_text("") + _fake_markitdown(monkeypatch, text="# Title\n\nbody") + result = MarkItDownPlusPlugin().import_content( + str(f), description="d", citation="c" + ) + assert "Title" in result.full_text + assert result.metadata["import_plugin"] == "markitdown_plus_import" + assert result.metadata["description"] == "d" + assert result.pages == [] + assert result.images == [] + + +def test_import_content_conversion_error_humanized(tmp_path, monkeypatch): + f = tmp_path / "doc.html" + f.write_text("x") + _fake_markitdown(monkeypatch, raises=RuntimeError("MissingDependencyException")) + with pytest.raises(RuntimeError): + MarkItDownPlusPlugin().import_content(str(f)) + + +def test_import_content_pdf_with_images(tmp_path, monkeypatch): + f = tmp_path / "doc.pdf" + f.write_bytes(b"%PDF-1.4 fake") + _fake_markitdown(monkeypatch, text="page one\n\n---\n\npage two") + page = _FakePage(images=[(1,)], extract_map={1: {"image": b"I", "ext": "png"}}) + _install_fake_fitz(monkeypatch, pages=[page]) + result = MarkItDownPlusPlugin().import_content(str(f), image_descriptions="basic") + assert result.metadata["image_count"] == 1 + assert result.metadata["page_count"] == 2 + + +def test_import_content_image_mode_none_skips_extraction(tmp_path, monkeypatch): + f = tmp_path / "doc.pdf" + f.write_bytes(b"%PDF fake") + _fake_markitdown(monkeypatch, text="no breaks here") + result = MarkItDownPlusPlugin().import_content(str(f), image_descriptions="none") + assert result.metadata["image_count"] == 0 + assert result.metadata["image_descriptions_mode"] == "none" diff --git a/library-manager/tests/test_plugin_discovery.py b/library-manager/tests/test_plugin_discovery.py deleted file mode 100644 index 804da314f..000000000 --- a/library-manager/tests/test_plugin_discovery.py +++ /dev/null @@ -1,189 +0,0 @@ -"""Auto-discovery regression test for library-manager plugins. - -The product thesis is **zero-touch plugin drop-in**: a fresh ``.py`` file -in ``backend/plugins/`` must be picked up at startup with no code edits -and no config edits. This test exercises ``_discover_plugins`` against a -synthetic plugin module placed inside a tmpdir and confirms the module -self-registered via ``@PluginRegistry.register``. - -Pattern: - 1. Create a tmpdir laid out as a Python package mirror of the real - ``plugins/`` package (``plugins/__init__.py`` + a fresh module). - 2. ``monkeypatch.syspath_prepend`` so ``import plugins`` resolves there. - 3. Pop ``plugins`` (and any subpackage) out of ``sys.modules`` so the - new path is used on next import. - 4. Call the production ``_discover_plugins`` helper. - 5. Assert the synthetic plugin name is in ``PluginRegistry``. - 6. Cleanup: restore the original ``plugins`` package on teardown so - subsequent tests don't see the synthetic plugin. -""" - -from __future__ import annotations - -import sys -import textwrap -from pathlib import Path - -import pytest -from main import _discover_plugins -from plugins.base import PluginRegistry - -_DROPIN_NAME = "test_dropin_import_h" - - -_PLUGIN_SOURCE = textwrap.dedent( - f""" - from plugins.base import ( - ImportResult, - LibraryImportPlugin, - PluginParameter, - PluginRegistry, - ) - - - @PluginRegistry.register - class TestDropinPlugin(LibraryImportPlugin): - name = "{_DROPIN_NAME}" - description = "Synthetic plugin used by test_plugin_discovery." - supported_source_types = ["file"] - file_extensions = ["xyz"] - human_label = "Drop-in plugin for discovery test" - - def get_parameters(self): - return [] - - def import_content(self, source_path, *, api_keys=None, **kwargs): - return ImportResult(full_text="") - """ -).strip() - - -@pytest.fixture -def isolate_plugins_package(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): - """Mirror the real plugins package into a tmpdir and shadow it on sys.path. - - Copying every module is overkill for this test — instead we re-export - the real package by re-using its filesystem path AND adding the tmpdir - onto its ``__path__``. That way the real registry survives, AND - ``pkgutil.iter_modules`` picks up our synthetic module on the next scan. - """ - import plugins as plugins_pkg # noqa: PLC0415 - - # Snapshot for teardown. - original_paths = list(plugins_pkg.__path__) - original_registry = dict(PluginRegistry._plugins) - - plugins_pkg.__path__.insert(0, str(tmp_path)) - - yield tmp_path - - # Restore __path__ exactly as we found it. - plugins_pkg.__path__[:] = original_paths - # Restore registry exactly as we found it (drop synthetic plugin). - PluginRegistry._plugins.clear() - PluginRegistry._plugins.update(original_registry) - # Forget our synthetic module so it is GC'd and a re-run is clean. - sys.modules.pop(f"plugins.{_DROPIN_NAME}", None) - - -def test_dropin_plugin_is_auto_discovered(isolate_plugins_package: Path) -> None: - """A fresh .py dropped into the plugins folder appears in the registry.""" - plugin_file = isolate_plugins_package / f"{_DROPIN_NAME}.py" - plugin_file.write_text(_PLUGIN_SOURCE, encoding="utf-8") - - # Sanity: not already registered before discovery. - assert _DROPIN_NAME not in PluginRegistry._plugins - - _discover_plugins() - - assert _DROPIN_NAME in PluginRegistry._plugins, ( - "Drop-in plugin was not picked up by _discover_plugins. " - f"Registered: {sorted(PluginRegistry._plugins)}" - ) - - -def test_broken_plugin_does_not_block_discovery( - isolate_plugins_package: Path, caplog: pytest.LogCaptureFixture -) -> None: - """A plugin with a syntax/import error must not break other plugins. - - The discovery helper wraps every ``importlib.import_module`` in a - try/except so a single broken file degrades to a warning, never a - startup failure. We prove that by dropping in TWO files — one broken, - one valid — and asserting the valid one still registers. - """ - broken = isolate_plugins_package / "broken_test_dropin.py" - broken.write_text("raise RuntimeError('intentional')\n", encoding="utf-8") - - valid = isolate_plugins_package / f"{_DROPIN_NAME}.py" - valid.write_text(_PLUGIN_SOURCE, encoding="utf-8") - - with caplog.at_level("WARNING"): - _discover_plugins() - - assert _DROPIN_NAME in PluginRegistry._plugins - # Some logger message should reference the broken module name. - assert any("broken_test_dropin" in (rec.getMessage() or "") for rec in caplog.records), ( - "Expected a warning log mentioning the broken plugin file." - ) - - -def test_subpackage_modules_are_recursed( - isolate_plugins_package: Path, -) -> None: - """Plugins inside a sub-package (e.g. content_handlers/) are also imported. - - The capability plugins live under ``plugins/content_handlers/`` and - must keep working. This drops a fresh module into a sub-folder and - asserts ``_discover_plugins`` recurses into it. - """ - from plugins.content_handlers.capability import CapabilityRegistry # noqa: PLC0415 - - snapshot = dict(CapabilityRegistry._handlers) - - subpkg = isolate_plugins_package / "content_handlers" - subpkg.mkdir() - # We mirror the real sub-package by giving it an __init__ that imports - # the real one — but the cleanest approach is just to extend the real - # sub-package's __path__ in the same way as the parent. - import plugins.content_handlers as ch_pkg # noqa: PLC0415 - - original_ch_paths = list(ch_pkg.__path__) - ch_pkg.__path__.insert(0, str(subpkg)) - - synthetic = subpkg / "synthetic_handler_h.py" - synthetic.write_text( - textwrap.dedent( - """ - from pathlib import Path - from plugins.content_handlers.capability import ( - Capability, - CapabilityPayload, - CapabilityRegistry, - ContentHandler, - ) - - - @CapabilityRegistry.register - class SyntheticAudioHandler(ContentHandler): - capability = Capability.AUDIO - description = "Synthetic handler for discovery test." - - def get(self, item_path: Path) -> CapabilityPayload: - return CapabilityPayload(mime="audio/mpeg", body=b"") - """ - ).strip(), - encoding="utf-8", - ) - - try: - _discover_plugins() - assert any( - getattr(h, "__name__", "") == "SyntheticAudioHandler" - for h in CapabilityRegistry._handlers.values() - ), "Sub-package plugin was not discovered." - finally: - ch_pkg.__path__[:] = original_ch_paths - CapabilityRegistry._handlers.clear() - CapabilityRegistry._handlers.update(snapshot) - sys.modules.pop("plugins.content_handlers.synthetic_handler_h", None) diff --git a/library-manager/tests/test_plugin_params_passthrough.py b/library-manager/tests/test_plugin_params_passthrough.py deleted file mode 100644 index 9516811cb..000000000 --- a/library-manager/tests/test_plugin_params_passthrough.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Verify ``plugin_params`` flow from the HTTP layer to ``import_content`` kwargs. - -Each import-route accepts a ``plugin_params`` dict that mirrors the -plugin's declared parameter schema. The router must forward every key -unchanged so a drop-in plugin's parameters arrive at the plugin without -any router-side knowledge of them. - -This contract is what makes the "add a plugin file, no other code -changes" promise hold across the renderer → service → router → worker -chain. -""" - -import asyncio -import time -from typing import Any - -import pytest -from httpx import AsyncClient -from plugins.base import PluginRegistry - -AUTH_HEADERS = {"Authorization": "Bearer test-token"} - - -async def _wait_for_done( - client: AsyncClient, lib_id: str, item_id: str, timeout: float = 10.0 -) -> str: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - resp = await client.get( - f"/libraries/{lib_id}/items/{item_id}/status", - headers=AUTH_HEADERS, - ) - status = resp.json()["status"] - if status in ("ready", "failed"): - return status - await asyncio.sleep(0.3) - return "timeout" - - -@pytest.mark.asyncio -async def test_url_import_forwards_plugin_params(client: AsyncClient, library: dict, monkeypatch): - """POST /import/url with plugin_params must reach the plugin's import_content kwargs. - - We monkey-patch ``UrlImportPlugin.import_content`` to capture the kwargs - it receives, then submit a request with a contrived param dict. - """ - lib_id = library["id"] - captured: dict[str, Any] = {} - - plugin = PluginRegistry.get_plugin("url_import") - assert plugin is not None - plugin_cls = plugin.__class__ - original = plugin_cls.import_content - - def spy(self, source_path, *, api_keys=None, **kwargs): - captured.update(kwargs) - # Return a minimal valid ImportResult so the worker can finish. - from plugins.base import ImportResult - return ImportResult( - full_text="# stub", pages=[], images=[], metadata={}, source_ref={} - ) - - monkeypatch.setattr(plugin_cls, "import_content", spy) - - # All keys are schema-declared params from url_import's get_parameters(). - # The router/worker sanitizes against the plugin's declared schema and - # drops unknown keys — so the test pins schema-declared keys to confirm - # the wire format actually delivers them end-to-end. - custom_params = { - "limit": 7, - "max_discovery_depth": 3, - } - resp = await client.post( - f"/libraries/{lib_id}/import/url", - headers=AUTH_HEADERS, - json={ - "url": "https://example.com", - "plugin_name": "url_import", - "title": "Passthrough Test", - "plugin_params": custom_params, - }, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - final = await _wait_for_done(client, lib_id, item_id) - assert final == "ready", f"worker did not finish: {final}" - - # Every schema-declared key the caller supplied must appear in the - # plugin's kwargs unchanged. - for key, value in custom_params.items(): - assert captured.get(key) == value, ( - f"plugin_params['{key}'] did not reach the plugin (got {captured})" - ) - - # Restore — pytest's monkeypatch cleans up, but be explicit anyway. - monkeypatch.setattr(plugin_cls, "import_content", original) - - -@pytest.mark.asyncio -async def test_youtube_import_language_via_plugin_params( - client: AsyncClient, library: dict, monkeypatch -): - """``plugin_params['language']`` overrides the top-level ``language`` field. - - This is the schema-driven path the UI now uses; the top-level - ``language`` field is kept only for backward compat. When both are - sent, ``plugin_params['language']`` wins. - """ - lib_id = library["id"] - captured: dict[str, Any] = {} - - plugin = PluginRegistry.get_plugin("youtube_transcript_import") - assert plugin is not None - plugin_cls = plugin.__class__ - - def spy(self, source_path, *, api_keys=None, **kwargs): - captured.update(kwargs) - from plugins.base import ImportResult - return ImportResult( - full_text="# stub", pages=[], images=[], metadata={}, source_ref={} - ) - - monkeypatch.setattr(plugin_cls, "import_content", spy) - - resp = await client.post( - f"/libraries/{lib_id}/import/youtube", - headers=AUTH_HEADERS, - json={ - "video_url": "https://www.youtube.com/watch?v=stub", - "plugin_name": "youtube_transcript_import", - "title": "YT Plugin Params Test", - # Both transports set — plugin_params wins. - "language": "en", - "plugin_params": {"language": "fr"}, - }, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - final = await _wait_for_done(client, lib_id, item_id) - assert final == "ready" - assert captured.get("language") == "fr" - - -@pytest.mark.asyncio -async def test_youtube_import_top_level_language_fallback( - client: AsyncClient, library: dict, monkeypatch -): - """When ``plugin_params`` omits ``language``, the top-level field fills the gap. - - Preserves backward compat for API clients that haven't migrated to - sending all knobs via ``plugin_params``. - """ - lib_id = library["id"] - captured: dict[str, Any] = {} - - plugin = PluginRegistry.get_plugin("youtube_transcript_import") - plugin_cls = plugin.__class__ - - def spy(self, source_path, *, api_keys=None, **kwargs): - captured.update(kwargs) - from plugins.base import ImportResult - return ImportResult( - full_text="# stub", pages=[], images=[], metadata={}, source_ref={} - ) - - monkeypatch.setattr(plugin_cls, "import_content", spy) - - resp = await client.post( - f"/libraries/{lib_id}/import/youtube", - headers=AUTH_HEADERS, - json={ - "video_url": "https://www.youtube.com/watch?v=stub", - "plugin_name": "youtube_transcript_import", - "title": "YT Fallback Test", - "language": "de", - # No plugin_params at all. - }, - ) - assert resp.status_code == 202 - item_id = resp.json()["item_id"] - - final = await _wait_for_done(client, lib_id, item_id) - assert final == "ready" - assert captured.get("language") == "de" diff --git a/library-manager/tests/test_remaining_branches.py b/library-manager/tests/test_remaining_branches.py new file mode 100644 index 000000000..8fda017b4 --- /dev/null +++ b/library-manager/tests/test_remaining_branches.py @@ -0,0 +1,251 @@ +"""Coverage for the remaining small branch-added lines across library-manager: +the folder_id migration, plugin-discovery failure, markitdown error humanizer, +capability registry edge cases, handler 'unavailable' guards, and the simple/ +markitdown import error paths. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path + +import pytest + + +# --------------------------------------------------------------------------- +# database._apply_lightweight_migrations (adds content_items.folder_id) +# --------------------------------------------------------------------------- + + +def test_apply_migration_adds_folder_id(tmp_path): + import sqlite3 + + from sqlalchemy import create_engine + + from database.connection import _apply_lightweight_migrations + + db = tmp_path / "legacy.db" + conn = sqlite3.connect(str(db)) + conn.executescript( + "CREATE TABLE content_items (id TEXT PRIMARY KEY, library_id TEXT);" + ) + conn.commit() + conn.close() + + engine = create_engine(f"sqlite:///{db}") + _apply_lightweight_migrations(engine) # should ADD COLUMN folder_id + engine.dispose() + + conn = sqlite3.connect(str(db)) + cols = {row[1] for row in conn.execute("PRAGMA table_info(content_items)")} + conn.close() + assert "folder_id" in cols + + +# --------------------------------------------------------------------------- +# main._discover_plugins import failure +# --------------------------------------------------------------------------- + + +def test_discover_plugins_bad_package(): + import main + + main._discover_plugins("nonexistent_pkg_zzz") # logged + return, no raise + + +# --------------------------------------------------------------------------- +# _markitdown_errors fallback +# --------------------------------------------------------------------------- + + +def test_humanize_missing_dependency_no_format(): + from plugins._markitdown_errors import humanize_markitdown_error + + exc = type("MissingDependencyException", (Exception,), {})("opaque failure") + # filename without a recognisable extension -> generic reader fallback. + msg = humanize_markitdown_error(exc, "documentwithoutext") + assert "required reader" in msg + + +def test_humanize_generic_fallback(): + from plugins._markitdown_errors import humanize_markitdown_error + + msg = humanize_markitdown_error(RuntimeError("weird"), "x.bin") + assert "x.bin" in msg + + +# --------------------------------------------------------------------------- +# CapabilityRegistry edge cases +# --------------------------------------------------------------------------- + + +def test_capability_registry_register_validation_and_dup(caplog): + from plugins.content_handlers.capability import ( + Capability, + CapabilityRegistry, + ContentHandler, + ) + + saved = dict(CapabilityRegistry._handlers) + try: + # Handler without a Capability attribute -> ValueError. + class _BadHandler(ContentHandler): + def get(self, item_path): + raise NotImplementedError + + with pytest.raises(ValueError, match="must set"): + CapabilityRegistry.register(_BadHandler) + + # Duplicate capability -> replace-with-warning branch. + class _H1(ContentHandler): + capability = Capability.TEXT + + def get(self, item_path): + raise NotImplementedError + + class _H2(ContentHandler): + capability = Capability.TEXT + + def get(self, item_path): + raise NotImplementedError + + CapabilityRegistry.register(_H1) + CapabilityRegistry.register(_H2) # logs "already registered ... replacing" + finally: + CapabilityRegistry._handlers.clear() + CapabilityRegistry._handlers.update(saved) + + +def test_capability_registry_get_unknown_string(): + from plugins.content_handlers.capability import CapabilityRegistry + + # A string that isn't a valid Capability -> None. + assert CapabilityRegistry.get("not-a-capability") is None + + +def test_capability_registry_reset(): + from plugins.content_handlers.capability import CapabilityRegistry + + saved = dict(CapabilityRegistry._handlers) + try: + CapabilityRegistry._reset() + assert CapabilityRegistry._handlers == {} + finally: + CapabilityRegistry._handlers.update(saved) + + +# --------------------------------------------------------------------------- +# Handler "unavailable" guards (dir exists but empty) +# --------------------------------------------------------------------------- + + +def test_images_handler_unavailable_when_empty(tmp_path): + from plugins.content_handlers.capability import HandlerUnavailable + from plugins.content_handlers.images_handler import ImagesHandler + + (tmp_path / "content" / "images").mkdir(parents=True) # exists but empty + with pytest.raises(HandlerUnavailable): + ImagesHandler().get(tmp_path) + + +def test_pages_handler_unavailable_when_empty(tmp_path): + from plugins.content_handlers.capability import HandlerUnavailable + from plugins.content_handlers.pages_handler import PagesHandler + + (tmp_path / "content" / "pages").mkdir(parents=True) # exists but empty + with pytest.raises(HandlerUnavailable): + PagesHandler().get(tmp_path) + + +# --------------------------------------------------------------------------- +# import plugin error paths +# --------------------------------------------------------------------------- + + +def test_simple_import_non_utf8(tmp_path): + from plugins.simple_import import SimpleImportPlugin + + f = tmp_path / "bad.md" + f.write_bytes(b"\xff\xfe invalid utf-8 \x80") + with pytest.raises(ValueError, match="not valid UTF-8"): + SimpleImportPlugin().import_content(str(f)) + + +def test_markitdown_import_conversion_error(tmp_path, monkeypatch): + import markitdown + + class FakeMD: + def convert(self, path): + raise RuntimeError("boom") + + monkeypatch.setattr(markitdown, "MarkItDown", FakeMD) + f = tmp_path / "doc.docx" + f.write_bytes(b"x") + from plugins.markitdown_import import MarkItDownImportPlugin + + with pytest.raises(RuntimeError): + MarkItDownImportPlugin().import_content(str(f)) + + +# --------------------------------------------------------------------------- +# schemas.folders name validation +# --------------------------------------------------------------------------- + + +def test_folder_name_too_long_rejected(): + from pydantic import ValidationError + + from schemas.folders import FolderCreateRequest + + with pytest.raises(ValidationError): + FolderCreateRequest(name="x" * 1000) + + +# --------------------------------------------------------------------------- +# url_import firecrawl object-doc with dict metadata +# --------------------------------------------------------------------------- + + +def test_validate_folder_name_too_long_direct(): + # The schema's Field(max_length=128) rejects over-long names before the + # validator runs, so the validator's own length guard is only reachable + # by calling it directly. + from schemas.folders import _validate_folder_name + + with pytest.raises(ValueError, match="cannot exceed"): + _validate_folder_name("x" * 200) + + +def test_detect_capabilities_no_content_dir(tmp_path): + # No content/ directory -> early return with an empty capability list. + from services.content_service import detect_capabilities + + assert detect_capabilities(tmp_path) == [] + + +def test_url_import_firecrawl_object_doc_dict_metadata(monkeypatch): + import firecrawl + + class _Doc: + markdown = "# page\n\nbody" + metadata = {"sourceURL": "https://x/p", "title": "P"} # dict metadata + + class _Result: + data = [_Doc()] + + class FakeApp: + def __init__(self, **kw): + pass + + def crawl(self, url, **kw): + return _Result() + + monkeypatch.setattr(firecrawl, "FirecrawlApp", FakeApp) + from plugins.url_import import UrlImportPlugin + + result = UrlImportPlugin().import_content( + "https://x", api_keys={"firecrawl_key": "fc"} + ) + assert "## P" in result.full_text + assert "https://x/p" in result.full_text diff --git a/library-manager/tests/test_system.py b/library-manager/tests/test_system.py deleted file mode 100644 index 807a6b183..000000000 --- a/library-manager/tests/test_system.py +++ /dev/null @@ -1,59 +0,0 @@ -"""Tests for system endpoints: health check, auth, and plugin listing.""" - -import pytest -from httpx import AsyncClient - -AUTH_HEADERS = {"Authorization": "Bearer test-token"} - - -@pytest.mark.asyncio -async def test_health_no_auth(client: AsyncClient): - """Health endpoint should respond without authentication and show component health.""" - resp = await client.get("/health") - assert resp.status_code == 200 - data = resp.json() - assert data["status"] in ("ok", "degraded") - assert data["service"] == "library-manager" - assert "checks" in data - assert data["checks"]["database"] == "ok" - - -@pytest.mark.asyncio -async def test_unauthenticated_rejected(client: AsyncClient): - """Protected endpoints should reject requests without a valid token.""" - resp = await client.get("/plugins") - assert resp.status_code in (401, 403) - - -@pytest.mark.asyncio -async def test_wrong_token_rejected(client: AsyncClient): - """Protected endpoints should reject requests with a wrong token.""" - resp = await client.get( - "/plugins", - headers={"Authorization": "Bearer wrong-token"}, - ) - assert resp.status_code == 401 - - -@pytest.mark.asyncio -async def test_list_plugins(client: AsyncClient): - """Plugin listing should return all registered import plugins.""" - resp = await client.get("/plugins", headers=AUTH_HEADERS) - assert resp.status_code == 200 - - data = resp.json() - plugins = data["plugins"] - names = {p["name"] for p in plugins} - - assert "simple_import" in names - assert "markitdown_import" in names - assert "markitdown_plus_import" in names - assert "url_import" in names - assert "youtube_transcript_import" in names - - # Each plugin should have required fields. - for plugin in plugins: - assert "name" in plugin - assert "description" in plugin - assert "supported_source_types" in plugin - assert "parameters" in plugin diff --git a/library-manager/tests/test_url_import_plugin.py b/library-manager/tests/test_url_import_plugin.py new file mode 100644 index 000000000..be166c942 --- /dev/null +++ b/library-manager/tests/test_url_import_plugin.py @@ -0,0 +1,215 @@ +"""Unit tests for ``plugins.url_import`` (Firecrawl + MarkItDown fallback). + +Both backends are import-mocked so the tests never hit the network. We cover +the URL-validation guard, both crawl paths, the various Firecrawl document +shapes (object metadata, dict metadata, empty pages), error wrapping, and the +``_safe_int`` / ``_safe_bool`` helpers. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from plugins.url_import import UrlImportPlugin, _safe_bool, _safe_int + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "value,default,expected", + [(None, 7, 7), ("3", 0, 3), ("x", 9, 9), (5, 0, 5), (1.5, 0, 1)], +) +def test_safe_int(value, default, expected): + assert _safe_int(value, default) == expected + + +@pytest.mark.parametrize( + "value,default,expected", + [ + (None, True, True), + (True, False, True), + ("false", True, False), + ("0", True, False), + ("yes", False, True), + (1, False, True), + ], +) +def test_safe_bool(value, default, expected): + assert _safe_bool(value, default) is expected + + +def test_get_parameters(): + names = {p.name for p in UrlImportPlugin().get_parameters()} + assert {"max_discovery_depth", "limit", "crawl_entire_domain", "timeout"} <= names + + +def test_invalid_url_raises(): + with pytest.raises(ValueError): + UrlImportPlugin().import_content("not-a-url") + + +# --------------------------------------------------------------------------- +# Firecrawl path +# --------------------------------------------------------------------------- + + +def _install_fake_firecrawl(monkeypatch, *, crawl_result=None, raises=None): + import firecrawl + + class FakeApp: + def __init__(self, **kwargs): + self.init_kwargs = kwargs + + def crawl(self, url, **kwargs): + if raises is not None: + raise raises + return crawl_result + + monkeypatch.setattr(firecrawl, "FirecrawlApp", FakeApp) + + +class _Doc: + def __init__(self, markdown, metadata=None, url=""): + self.markdown = markdown + if metadata is not None: + self.metadata = metadata + self.url = url + + +class _Meta: + def __init__(self, url="", title=""): + self.url = url + self.title = title + + +class _CrawlResult: + def __init__(self, data): + self.data = data + + +def test_firecrawl_crawl_failure_wrapped(monkeypatch): + _install_fake_firecrawl(monkeypatch, raises=RuntimeError("boom")) + with pytest.raises(RuntimeError, match="Firecrawl crawl failed"): + UrlImportPlugin().import_content( + "https://example.com", api_keys={"firecrawl_key": "fc-key"} + ) + + +def test_firecrawl_no_pages_returns_placeholder(monkeypatch): + _install_fake_firecrawl(monkeypatch, crawl_result=_CrawlResult(data=[])) + result = UrlImportPlugin().import_content( + "https://example.com", api_keys={"firecrawl_key": "fc-key"} + ) + assert "No content could be crawled" in result.full_text + assert result.metadata["pages_crawled"] == 0 + + +def test_firecrawl_object_docs_with_object_metadata(monkeypatch): + docs = [ + _Doc("# Page one\n\ntext", metadata=_Meta(url="https://example.com/a", + title="Page A")), + _Doc(" ", metadata=_Meta()), # blank -> skipped + ] + _install_fake_firecrawl(monkeypatch, crawl_result=_CrawlResult(data=docs)) + result = UrlImportPlugin().import_content( + "https://example.com", + api_keys={"firecrawl_key": "fc-key"}, + description="desc", citation="cite", + ) + assert "## Page A" in result.full_text + assert "Source: https://example.com/a" in result.full_text + assert result.metadata["fetch_method"] == "firecrawl" + assert result.metadata["description"] == "desc" + assert result.metadata["citation"] == "cite" + assert result.metadata["pages_crawled"] == 2 + + +def test_firecrawl_dict_docs(monkeypatch): + docs = [ + {"markdown": "dict page", "metadata": {"sourceURL": "https://x/b", + "title": "B"}}, + ] + _install_fake_firecrawl(monkeypatch, crawl_result=_CrawlResult(data=docs)) + result = UrlImportPlugin().import_content( + "https://x", api_keys={"firecrawl_key": "fc-key"} + ) + assert "## B" in result.full_text + assert "dict page" in result.full_text + + +def test_firecrawl_dict_crawl_result(monkeypatch): + # crawl returns a plain dict (no .data attribute). + _install_fake_firecrawl( + monkeypatch, crawl_result={"data": [{"markdown": "hello"}]} + ) + result = UrlImportPlugin().import_content( + "https://x", api_keys={"firecrawl_key": "fc-key"} + ) + assert "hello" in result.full_text + + +# --------------------------------------------------------------------------- +# MarkItDown fallback path +# --------------------------------------------------------------------------- + + +def _install_fake_markitdown(monkeypatch, *, text="", raises=None): + import markitdown + + class FakeResult: + def __init__(self, t): + self.text_content = t + + class FakeMD: + def convert(self, url): + if raises is not None: + raise raises + return FakeResult(text) + + monkeypatch.setattr(markitdown, "MarkItDown", FakeMD) + + +def test_markitdown_success(monkeypatch): + _install_fake_markitdown(monkeypatch, text="# Converted\n\nbody") + result = UrlImportPlugin().import_content( + "https://example.com", description="d", citation="c" + ) + assert "Converted" in result.full_text + assert result.metadata["fetch_method"] == "markitdown" + assert result.metadata["description"] == "d" + + +def test_markitdown_empty_text_placeholder(monkeypatch): + _install_fake_markitdown(monkeypatch, text=" ") + result = UrlImportPlugin().import_content("https://example.com") + assert "No text content could be extracted" in result.full_text + + +def test_markitdown_convert_failure_wrapped(monkeypatch): + _install_fake_markitdown(monkeypatch, raises=RuntimeError("bad")) + with pytest.raises(RuntimeError, match="Failed to fetch and convert"): + UrlImportPlugin().import_content("https://example.com") + + +def test_markitdown_not_installed(monkeypatch): + # Simulate the markitdown import failing. + monkeypatch.setitem(sys.modules, "markitdown", None) + with pytest.raises(RuntimeError, match="markitdown is not installed"): + UrlImportPlugin().import_content("https://example.com") + + +def test_progress_callback_is_invoked(monkeypatch): + _install_fake_markitdown(monkeypatch, text="body") + events = [] + UrlImportPlugin().import_content( + "https://example.com", + progress_callback=lambda c, t, m: events.append((c, t, m)), + ) + assert events # callback fired for each step + assert events[-1][0] == events[-1][1] # last step is "complete" diff --git a/library-manager/tests/test_youtube_plugin.py b/library-manager/tests/test_youtube_plugin.py new file mode 100644 index 000000000..7d2d5d866 --- /dev/null +++ b/library-manager/tests/test_youtube_plugin.py @@ -0,0 +1,270 @@ +"""Unit tests for ``plugins.youtube_transcript_import``. + +``yt_dlp`` is import-mocked so nothing touches YouTube. Covers URL parsing, +SRT/VTT parsing, timestamp formatting, language-key resolution, the +fetch/download orchestration (manual vs auto, rate-limit, all-fail), and the +``import_content`` entry point. +""" + +from __future__ import annotations + +import os +import sys +import types + +import pytest + +import plugins.youtube_transcript_import as yt +from plugins.youtube_transcript_import import ( + YouTubeTranscriptImportPlugin, + _build_source_ref, + _parse_srt_content, + _parse_youtube_url, + _resolve_language_key, + _seconds_to_timestamp, +) + + +# --------------------------------------------------------------------------- +# pure helpers +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "url,expected", + [ + ("https://www.youtube.com/watch?v=abc12345678", "abc12345678"), + ("https://youtu.be/xyz98765432", "xyz98765432"), + ("https://example.com/watch?v=x", None), + ("https://www.youtube.com/watch", None), + ], +) +def test_parse_youtube_url(url, expected): + assert _parse_youtube_url(url) == expected + + +@pytest.mark.parametrize( + "seconds,expected", + [(0, "00:00"), (65, "01:05"), (3661, "01:01:01")], +) +def test_seconds_to_timestamp(seconds, expected): + assert _seconds_to_timestamp(seconds) == expected + + +def test_build_source_ref(): + ref = _build_source_ref("https://youtu.be/v", "v", "en") + assert ref["type"] == "youtube" + assert ref["video_id"] == "v" + assert ref["video_url"].endswith("watch?v=v") + + +@pytest.mark.parametrize( + "subs,lang,expected", + [ + ({"en": []}, "en", "en"), # exact + ({"en-US": []}, "en", "en-US"), # prefix + ({"de": []}, "en", None), # no match + ({}, "en", None), # empty + ], +) +def test_resolve_language_key(subs, lang, expected): + assert _resolve_language_key(subs, lang) == expected + + +def test_parse_srt_content(): + srt = ( + "1\n" + "00:00:01,000 --> 00:00:04,000\n" + "Hello world\n\n" + "2\n" + "00:00:05,000 --> 00:00:08,000\n" + "[music] Second line\n\n" + "3\n" # block with no timestamp -> skipped + "just one line\n" + ) + pieces = _parse_srt_content(srt) + assert len(pieces) == 2 + assert pieces[0]["text"] == "Hello world" + assert pieces[0]["start"] == 1.0 + assert pieces[1]["text"] == "Second line" # [music] noise stripped + + +def test_parse_srt_content_empty(): + assert _parse_srt_content("") == [] + + +def test_get_parameters(): + names = {p.name for p in YouTubeTranscriptImportPlugin().get_parameters()} + assert names == {"language", "proxy_url"} + + +# --------------------------------------------------------------------------- +# yt_dlp mock + fetch orchestration +# --------------------------------------------------------------------------- + +_SRT = ( + "1\n00:00:01,000 --> 00:00:03,000\nhello\n\n" + "2\n00:00:04,000 --> 00:00:06,000\nworld\n" +) + + +def _real_fetch(): + """The conftest installs a file-cache wrapper around ``_fetch_transcript``; + these tests target the real implementation underneath it.""" + return getattr(yt._fetch_transcript, "__wrapped__", yt._fetch_transcript) + + +def _install_fake_ytdlp( + monkeypatch, *, info=None, extract_raises=None, + download_raises=None, write_file=None, +): + fake = types.ModuleType("yt_dlp") + + class FakeYDL: + def __init__(self, opts): + self.opts = opts + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def extract_info(self, url, download=False): + if extract_raises is not None: + raise extract_raises + return info + + def download(self, urls): + if download_raises is not None: + raise download_raises + if write_file is not None: + fname, content = write_file + outdir = os.path.dirname(self.opts["outtmpl"]) + with open(os.path.join(outdir, fname), "w", encoding="utf-8") as fh: + fh.write(content) + + fake.YoutubeDL = FakeYDL + monkeypatch.setitem(sys.modules, "yt_dlp", fake) + + +def test_fetch_transcript_extract_info_failure(monkeypatch): + _install_fake_ytdlp(monkeypatch, extract_raises=RuntimeError("network")) + with pytest.raises(RuntimeError, match="Failed to fetch video info"): + _real_fetch()("vid12345678", "en", None) + + +def test_fetch_transcript_no_subtitles(monkeypatch): + _install_fake_ytdlp(monkeypatch, info={"subtitles": {}, "automatic_captions": {}}) + with pytest.raises(RuntimeError, match="No subtitles available"): + _real_fetch()("vid12345678", "en", None) + + +def test_fetch_transcript_manual_success(monkeypatch): + info = {"subtitles": {"en": [{}]}, "automatic_captions": {}} + _install_fake_ytdlp(monkeypatch, info=info) + # Isolate from _download_and_parse by stubbing it. + monkeypatch.setattr( + yt, "_download_and_parse", + lambda **kw: [{"text": "hi", "start": 0.0, "duration": 1.0}], + ) + pieces, label = _real_fetch()("vid12345678", "en", "http://proxy") + assert label == "manual" + assert pieces[0]["text"] == "hi" + + +def test_fetch_transcript_falls_back_to_auto(monkeypatch): + info = {"subtitles": {}, "automatic_captions": {"en": [{}]}} + _install_fake_ytdlp(monkeypatch, info=info) + monkeypatch.setattr( + yt, "_download_and_parse", + lambda **kw: [{"text": "auto", "start": 0.0, "duration": 1.0}], + ) + pieces, label = _real_fetch()("vid12345678", "en", None) + assert label == "auto" + + +def test_fetch_transcript_rate_limited(monkeypatch): + info = {"subtitles": {"en": [{}]}, "automatic_captions": {}} + _install_fake_ytdlp(monkeypatch, info=info) + + def boom(**kw): + raise RuntimeError("HTTP Error 429: Too Many Requests") + + monkeypatch.setattr(yt, "_download_and_parse", boom) + with pytest.raises(RuntimeError, match="rate-limited"): + _real_fetch()("vid12345678", "en", None) + + +def test_fetch_transcript_all_attempts_fail(monkeypatch): + info = {"subtitles": {"en": [{}]}, "automatic_captions": {}} + _install_fake_ytdlp(monkeypatch, info=info) + monkeypatch.setattr( + yt, "_download_and_parse", + lambda **kw: (_ for _ in ()).throw(RuntimeError("generic failure")), + ) + with pytest.raises(RuntimeError, match="Could not download subtitles"): + _real_fetch()("vid12345678", "en", None) + + +# --------------------------------------------------------------------------- +# _download_and_parse +# --------------------------------------------------------------------------- + + +def test_download_and_parse_success(monkeypatch): + _install_fake_ytdlp(monkeypatch, write_file=("video.en.srt", _SRT)) + pieces = yt._download_and_parse( + url="https://youtu.be/v", language_key="en", + source_label="manual", proxy_url=None, + ) + assert [p["text"] for p in pieces] == ["hello", "world"] + + +def test_download_and_parse_no_subtitle_files(monkeypatch): + # download writes nothing -> no .srt/.vtt files found. + _install_fake_ytdlp(monkeypatch, write_file=None) + out = yt._download_and_parse( + url="https://youtu.be/v", language_key="en", + source_label="auto", proxy_url=None, + ) + assert out == [] + + +def test_download_and_parse_download_error_strips_prefix(monkeypatch): + _install_fake_ytdlp( + monkeypatch, download_raises=RuntimeError("ERROR: video unavailable") + ) + with pytest.raises(RuntimeError) as exc: + yt._download_and_parse( + url="https://youtu.be/v", language_key="en", + source_label="manual", proxy_url="http://proxy", + ) + assert "video unavailable" in str(exc.value) + assert "ERROR:" not in str(exc.value) + + +# --------------------------------------------------------------------------- +# import_content +# --------------------------------------------------------------------------- + + +def test_import_content_invalid_url(): + with pytest.raises(ValueError, match="Could not extract video ID"): + YouTubeTranscriptImportPlugin().import_content("https://example.com/no-video") + + +def test_import_content_success(monkeypatch): + monkeypatch.setattr( + yt, "_fetch_transcript", + lambda vid, lang, proxy: ( + [{"text": "hello", "start": 1.0, "duration": 2.0}], "manual" + ), + ) + result = YouTubeTranscriptImportPlugin().import_content( + "https://www.youtube.com/watch?v=abc12345678", language="en" + ) + assert "Transcript:" in result.full_text + assert "**[00:01]** hello" in result.full_text + assert result.metadata["subtitle_source"] == "manual" + assert result.metadata["video_id"] == "abc12345678" diff --git a/library-manager/tests/unit/__init__.py b/library-manager/tests/unit/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/library-manager/tests/unit/conftest.py b/library-manager/tests/unit/conftest.py new file mode 100644 index 000000000..2328b5ae6 --- /dev/null +++ b/library-manager/tests/unit/conftest.py @@ -0,0 +1,66 @@ +"""Unit-tier fixtures: direct module calls, no FastAPI app, no HTTP. + +Unit tests exercise one module's logic in isolation. They get: + +* ``tmp_storage`` — a throwaway directory for filesystem-touching services. +* ``db_session`` — a real SQLAlchemy session (the session DB is initialised + once by the root conftest); the caller never has to manage HTTP. +* ``fresh_plugin_registry`` / ``fresh_capability_registry`` — snapshot the + class-level registry dicts and restore them after the test, so tests that + register/disable plugins cannot leak state into their neighbours. +""" + +from __future__ import annotations + +import shutil +import tempfile +from collections.abc import Iterator +from pathlib import Path + +import pytest + + +@pytest.fixture +def tmp_storage() -> Iterator[Path]: + """A per-test temporary directory, removed on teardown.""" + path = Path(tempfile.mkdtemp(prefix="lm-unit-")) + try: + yield path + finally: + shutil.rmtree(path, ignore_errors=True) + + +@pytest.fixture +def db_session() -> Iterator: + """A direct SQLAlchemy session bound to the session test database.""" + from database.connection import get_session_direct # noqa: PLC0415 + + session = get_session_direct() + try: + yield session + finally: + session.close() + + +@pytest.fixture +def fresh_plugin_registry() -> Iterator: + """Snapshot and restore ``PluginRegistry._plugins`` around a test.""" + from plugins.base import PluginRegistry # noqa: PLC0415 + + saved = dict(PluginRegistry._plugins) + try: + yield PluginRegistry + finally: + PluginRegistry._plugins = saved + + +@pytest.fixture +def fresh_capability_registry() -> Iterator: + """Snapshot and restore ``CapabilityRegistry._handlers`` around a test.""" + from plugins.content_handlers.capability import CapabilityRegistry # noqa: PLC0415 + + saved = dict(CapabilityRegistry._handlers) + try: + yield CapabilityRegistry + finally: + CapabilityRegistry._handlers = saved diff --git a/library-manager/tests/unit/test_config.py b/library-manager/tests/unit/test_config.py new file mode 100644 index 000000000..a535d5100 --- /dev/null +++ b/library-manager/tests/unit/test_config.py @@ -0,0 +1,193 @@ +"""Unit tests for ``backend/config.py`` env parsing, defaults, and helpers. + +``config`` reads ``os.environ`` once at import time into module-level +constants. To test parsing of *different* env values we reload the module +under ``monkeypatch.setenv`` and always reload it back to the real +(session-conftest) environment afterwards, so neighbouring tests that import +``config`` see the canonical values. +""" + +from __future__ import annotations + +import importlib +from pathlib import Path + +import config +import pytest + +# Constants these tests reload/patch and that other modules bind by value +# (e.g. ``from config import CONTENT_DIR``). They MUST be returned to their +# canonical session values after every test. +_CONFIG_CONSTANTS = ( + "HOST", + "PORT", + "LOG_LEVEL", + "LAMB_API_TOKEN", + "DATA_DIR", + "CONTENT_DIR", + "DB_PATH", + "MAX_CONCURRENT_IMPORTS", + "IMPORT_TASK_TIMEOUT_SECONDS", + "MAX_UPLOAD_SIZE_BYTES", + "MAX_ZIP_IMPORT_SIZE_BYTES", + "PERMALINK_PREFIX", +) + + +@pytest.fixture(autouse=True) +def _isolate_config(): + """Snapshot config's constants and restore them by value after each test. + + ``importlib.reload(config)`` mutates the module in place, and the in-test + ``finally: _reload_config()`` runs *before* monkeypatch restores the + environment — so without this fixture a reload under a patched/deleted env + would leak the wrong ``config.CONTENT_DIR`` (etc.) into every later test in + the session. Restoring by value here is independent of env-teardown + ordering and keeps the session canonical. + """ + saved = {k: getattr(config, k) for k in _CONFIG_CONSTANTS} + yield + for k, v in saved.items(): + setattr(config, k, v) + + +def _reload_config(): + """Reload the config module and return the fresh module object.""" + return importlib.reload(config) + + +class TestDefaults: + """Defaults apply when the corresponding env vars are unset.""" + + def test_server_defaults(self, monkeypatch): + """HOST/PORT/LOG_LEVEL fall back to documented defaults.""" + for var in ("HOST", "PORT", "LOG_LEVEL"): + monkeypatch.delenv(var, raising=False) + try: + cfg = _reload_config() + assert cfg.HOST == "0.0.0.0", "default host" + assert cfg.PORT == 9091, "default port" + assert cfg.LOG_LEVEL == "INFO", "default log level" + finally: + _reload_config() + + def test_task_and_upload_defaults(self, monkeypatch): + """Task concurrency/timeout and upload-size defaults match source.""" + for var in ( + "MAX_CONCURRENT_IMPORTS", + "IMPORT_TASK_TIMEOUT_SECONDS", + "MAX_UPLOAD_SIZE_BYTES", + "MAX_ZIP_IMPORT_SIZE_BYTES", + "PERMALINK_PREFIX", + ): + monkeypatch.delenv(var, raising=False) + try: + cfg = _reload_config() + assert cfg.MAX_CONCURRENT_IMPORTS == 3 + assert cfg.IMPORT_TASK_TIMEOUT_SECONDS == 600 + assert cfg.MAX_UPLOAD_SIZE_BYTES == 500 * 1024 * 1024 + assert cfg.MAX_ZIP_IMPORT_SIZE_BYTES == 200 * 1024 * 1024 + assert cfg.PERMALINK_PREFIX == "/docs" + finally: + _reload_config() + + +class TestEnvParsing: + """Environment values override the defaults and are typed correctly.""" + + def test_int_vars_parsed_from_strings(self, monkeypatch): + """Integer env vars are coerced from their string representation.""" + monkeypatch.setenv("PORT", "1234") + monkeypatch.setenv("MAX_CONCURRENT_IMPORTS", "7") + monkeypatch.setenv("IMPORT_TASK_TIMEOUT_SECONDS", "42") + monkeypatch.setenv("MAX_UPLOAD_SIZE_BYTES", "999") + monkeypatch.setenv("MAX_ZIP_IMPORT_SIZE_BYTES", "888") + try: + cfg = _reload_config() + assert cfg.PORT == 1234 and isinstance(cfg.PORT, int) + assert cfg.MAX_CONCURRENT_IMPORTS == 7 + assert cfg.IMPORT_TASK_TIMEOUT_SECONDS == 42 + assert cfg.MAX_UPLOAD_SIZE_BYTES == 999 + assert cfg.MAX_ZIP_IMPORT_SIZE_BYTES == 888 + finally: + _reload_config() + + def test_log_level_uppercased(self, monkeypatch): + """LOG_LEVEL is upper-cased regardless of the env value's case.""" + monkeypatch.setenv("LOG_LEVEL", "debug") + try: + cfg = _reload_config() + assert cfg.LOG_LEVEL == "DEBUG", "log level should be upper-cased" + finally: + _reload_config() + + def test_token_and_permalink_from_env(self, monkeypatch): + """LAMB_API_TOKEN and PERMALINK_PREFIX read raw string env values.""" + monkeypatch.setenv("LAMB_API_TOKEN", "secret-xyz") + monkeypatch.setenv("PERMALINK_PREFIX", "/files") + try: + cfg = _reload_config() + assert cfg.LAMB_API_TOKEN == "secret-xyz" + assert cfg.PERMALINK_PREFIX == "/files" + finally: + _reload_config() + + def test_token_defaults_empty(self, monkeypatch): + """LAMB_API_TOKEN defaults to an empty string when unset.""" + monkeypatch.delenv("LAMB_API_TOKEN", raising=False) + try: + cfg = _reload_config() + assert cfg.LAMB_API_TOKEN == "" + finally: + _reload_config() + + +class TestDerivedPaths: + """DATA_DIR drives the derived CONTENT_DIR and DB_PATH paths.""" + + def test_data_dir_from_env_and_derived_paths(self, monkeypatch): + """CONTENT_DIR and DB_PATH are derived under the configured DATA_DIR.""" + monkeypatch.setenv("DATA_DIR", "/tmp/lm-cfg-test") + try: + cfg = _reload_config() + assert Path("/tmp/lm-cfg-test") == cfg.DATA_DIR + assert Path("/tmp/lm-cfg-test/content") == cfg.CONTENT_DIR + assert Path("/tmp/lm-cfg-test/library-manager.db") == cfg.DB_PATH + finally: + _reload_config() + + def test_default_data_dir_relative(self, monkeypatch): + """DATA_DIR defaults to the relative ``data`` directory.""" + monkeypatch.delenv("DATA_DIR", raising=False) + try: + cfg = _reload_config() + assert Path("data") == cfg.DATA_DIR + assert Path("data/content") == cfg.CONTENT_DIR + assert Path("data/library-manager.db") == cfg.DB_PATH + finally: + _reload_config() + + +class TestEnsureDirectories: + """``ensure_directories`` creates DATA_DIR and CONTENT_DIR idempotently.""" + + def test_creates_data_and_content(self, monkeypatch, tmp_storage): + """Both DATA_DIR and CONTENT_DIR are created when missing.""" + data = tmp_storage / "newdata" + content = data / "content" + monkeypatch.setattr(config, "DATA_DIR", data) + monkeypatch.setattr(config, "CONTENT_DIR", content) + assert not data.exists() + config.ensure_directories() + assert data.is_dir(), "DATA_DIR should be created" + assert content.is_dir(), "CONTENT_DIR should be created" + + def test_idempotent(self, monkeypatch, tmp_storage): + """Calling twice with existing dirs does not raise (exist_ok=True).""" + data = tmp_storage / "d" + content = data / "content" + monkeypatch.setattr(config, "DATA_DIR", data) + monkeypatch.setattr(config, "CONTENT_DIR", content) + config.ensure_directories() + config.ensure_directories() + assert content.is_dir() diff --git a/library-manager/tests/unit/test_content_handlers.py b/library-manager/tests/unit/test_content_handlers.py new file mode 100644 index 000000000..cb871024d --- /dev/null +++ b/library-manager/tests/unit/test_content_handlers.py @@ -0,0 +1,357 @@ +"""Unit tests for the content-handler capability layer. + +Covers ``plugins.content_handlers.capability`` (enum, payload, registry) +plus the three concrete handlers (text, pages, images). Registry-mutating +tests use ``fresh_capability_registry`` so handler registrations never +leak into other tests. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from plugins.content_handlers.capability import ( + Capability, + CapabilityPayload, + ContentHandler, + HandlerUnavailable, +) +from plugins.content_handlers.images_handler import ImagesHandler +from plugins.content_handlers.pages_handler import PagesHandler +from plugins.content_handlers.text_handler import TextHandler + + +def _item_dir(tmp_storage: Path) -> Path: + """Create a CONTENT_DIR/{org}/{library}/{item} style item directory.""" + item = tmp_storage / "org1" / "lib1" / "item1" + (item / "content").mkdir(parents=True) + return item + + +# --------------------------------------------------------------------------- +# Capability enum + CapabilityPayload +# --------------------------------------------------------------------------- + + +def test_capability_enum_values(): + """Capability is a str enum with the documented controlled vocabulary.""" + assert Capability.TEXT == "text" + assert Capability.PAGES == "pages" + assert Capability.IMAGES == "images" + assert Capability.AUDIO == "audio" + assert Capability.TRANSCRIPT == "transcript" + assert {c.value for c in Capability} == { + "text", "pages", "images", "audio", "transcript" + } + + +def test_capability_payload_fields(): + """CapabilityPayload stores mime and arbitrary body.""" + payload = CapabilityPayload(mime="text/markdown", body="# hi") + assert payload.mime == "text/markdown" + assert payload.body == "# hi" + + +# --------------------------------------------------------------------------- +# CapabilityRegistry.register +# --------------------------------------------------------------------------- + + +def test_register_rejects_non_enum_capability(fresh_capability_registry): + """register raises ValueError when capability is not a Capability enum.""" + + class BadHandler(ContentHandler): + capability = "text" # plain string, not the enum + + def get(self, item_path): + return CapabilityPayload(mime="text/plain", body="") + + with pytest.raises(ValueError, match="Capability enum member"): + fresh_capability_registry.register(BadHandler) + + +def test_register_warns_and_replaces_on_duplicate(fresh_capability_registry, caplog): + """register warns and replaces when a capability is registered twice.""" + + class FirstAudio(ContentHandler): + capability = Capability.AUDIO + + def get(self, item_path): + return CapabilityPayload(mime="audio/mpeg", body=b"") + + class SecondAudio(ContentHandler): + capability = Capability.AUDIO + + def get(self, item_path): + return CapabilityPayload(mime="audio/mpeg", body=b"") + + fresh_capability_registry.register(FirstAudio) + with caplog.at_level("WARNING"): + fresh_capability_registry.register(SecondAudio) + + assert fresh_capability_registry._handlers[Capability.AUDIO] is SecondAudio + assert any("already registered" in (r.getMessage() or "") for r in caplog.records) + + +def test_register_same_class_twice_no_warning(fresh_capability_registry, caplog): + """Re-registering the identical class is idempotent and does not warn.""" + + class SameAudio(ContentHandler): + capability = Capability.AUDIO + + def get(self, item_path): + return CapabilityPayload(mime="audio/mpeg", body=b"") + + fresh_capability_registry.register(SameAudio) + with caplog.at_level("WARNING"): + fresh_capability_registry.register(SameAudio) + assert not any("already registered" in (r.getMessage() or "") for r in caplog.records) + + +def test_register_returns_class(fresh_capability_registry): + """register returns the handler class unchanged (decorator-friendly).""" + + class RetAudio(ContentHandler): + capability = Capability.AUDIO + + def get(self, item_path): + return CapabilityPayload(mime="audio/mpeg", body=b"") + + assert fresh_capability_registry.register(RetAudio) is RetAudio + + +# --------------------------------------------------------------------------- +# CapabilityRegistry.get +# --------------------------------------------------------------------------- + + +def test_get_by_enum_returns_fresh_instance(fresh_capability_registry): + """get returns a fresh handler instance keyed by enum value.""" + + class TranscriptH(ContentHandler): + capability = Capability.TRANSCRIPT + + def get(self, item_path): + return CapabilityPayload(mime="text/plain", body="") + + fresh_capability_registry.register(TranscriptH) + a = fresh_capability_registry.get(Capability.TRANSCRIPT) + b = fresh_capability_registry.get(Capability.TRANSCRIPT) + assert isinstance(a, TranscriptH) + assert a is not b + + +def test_get_by_string(fresh_capability_registry): + """get accepts the string form of a capability.""" + + class TranscriptH(ContentHandler): + capability = Capability.TRANSCRIPT + + def get(self, item_path): + return CapabilityPayload(mime="text/plain", body="") + + fresh_capability_registry.register(TranscriptH) + assert isinstance(fresh_capability_registry.get("transcript"), TranscriptH) + + +def test_get_unknown_string_returns_none(fresh_capability_registry): + """get returns None for a string that is not a Capability value.""" + assert fresh_capability_registry.get("nonsense") is None + + +def test_get_unregistered_capability_returns_none(fresh_capability_registry): + """get returns None when nothing is registered for a valid capability.""" + fresh_capability_registry._handlers.clear() + assert fresh_capability_registry.get(Capability.AUDIO) is None + + +# --------------------------------------------------------------------------- +# list_handlers / registered_capabilities / _reset +# --------------------------------------------------------------------------- + + +def test_list_handlers_sorted(fresh_capability_registry): + """list_handlers returns rows sorted by capability value.""" + fresh_capability_registry._handlers.clear() + + class ImagesH(ContentHandler): + capability = Capability.IMAGES + description = "imgs" + + def get(self, item_path): + return CapabilityPayload(mime="application/json", body=[]) + + class AudioH(ContentHandler): + capability = Capability.AUDIO + description = "aud" + + def get(self, item_path): + return CapabilityPayload(mime="audio/mpeg", body=b"") + + fresh_capability_registry.register(ImagesH) + fresh_capability_registry.register(AudioH) + rows = fresh_capability_registry.list_handlers() + caps = [r["capability"] for r in rows] + assert caps == sorted(caps) + assert {"capability": "audio", "description": "aud"} in rows + + +def test_registered_capabilities(fresh_capability_registry): + """registered_capabilities lists the currently registered enum keys.""" + fresh_capability_registry._handlers.clear() + + class AudioH(ContentHandler): + capability = Capability.AUDIO + + def get(self, item_path): + return CapabilityPayload(mime="audio/mpeg", body=b"") + + fresh_capability_registry.register(AudioH) + assert fresh_capability_registry.registered_capabilities() == [Capability.AUDIO] + + +def test_reset_clears_handlers(fresh_capability_registry): + """_reset empties the handler registry.""" + + class AudioH(ContentHandler): + capability = Capability.AUDIO + + def get(self, item_path): + return CapabilityPayload(mime="audio/mpeg", body=b"") + + fresh_capability_registry.register(AudioH) + fresh_capability_registry._reset() + assert fresh_capability_registry._handlers == {} + + +# --------------------------------------------------------------------------- +# TextHandler +# --------------------------------------------------------------------------- + + +def test_text_handler_returns_full_md(tmp_storage): + """TextHandler reads content/full.md and returns it as text/markdown.""" + item = _item_dir(tmp_storage) + (item / "content" / "full.md").write_text("# Title\nBody", encoding="utf-8") + payload = TextHandler().get(item) + assert payload.mime == "text/markdown" + assert payload.body == "# Title\nBody" + + +def test_text_handler_unavailable_without_file(tmp_storage): + """TextHandler raises HandlerUnavailable when full.md is missing.""" + item = _item_dir(tmp_storage) + with pytest.raises(HandlerUnavailable): + TextHandler().get(item) + + +# --------------------------------------------------------------------------- +# PagesHandler +# --------------------------------------------------------------------------- + + +def test_pages_handler_returns_sorted_pages(tmp_storage): + """PagesHandler lists content/pages/*.md ordered by page number.""" + item = _item_dir(tmp_storage) + pages = item / "content" / "pages" + pages.mkdir() + (pages / "page_002.md").write_text("two", encoding="utf-8") + (pages / "page_010.md").write_text("ten", encoding="utf-8") + (pages / "page_001.md").write_text("one", encoding="utf-8") + (pages / "notes.txt").write_text("ignored", encoding="utf-8") # non-md ignored + payload = PagesHandler().get(item) + assert payload.mime == "application/json" + assert payload.body == [ + {"page": 1, "markdown": "one"}, + {"page": 2, "markdown": "two"}, + {"page": 10, "markdown": "ten"}, + ] + + +def test_pages_handler_filename_without_number_falls_back_to_zero(tmp_storage): + """A page file with no embedded number is treated as page 0 (sorts first).""" + item = _item_dir(tmp_storage) + pages = item / "content" / "pages" + pages.mkdir() + (pages / "intro.md").write_text("intro", encoding="utf-8") + (pages / "page_001.md").write_text("one", encoding="utf-8") + payload = PagesHandler().get(item) + assert payload.body == [ + {"page": 0, "markdown": "intro"}, + {"page": 1, "markdown": "one"}, + ] + + +def test_pages_handler_unavailable_without_dir(tmp_storage): + """PagesHandler raises HandlerUnavailable when the pages dir is absent.""" + item = _item_dir(tmp_storage) + with pytest.raises(HandlerUnavailable): + PagesHandler().get(item) + + +def test_pages_handler_unavailable_when_empty(tmp_storage): + """PagesHandler raises HandlerUnavailable when no .md page files exist.""" + item = _item_dir(tmp_storage) + pages = item / "content" / "pages" + pages.mkdir() + (pages / "readme.txt").write_text("x", encoding="utf-8") + with pytest.raises(HandlerUnavailable): + PagesHandler().get(item) + + +# --------------------------------------------------------------------------- +# ImagesHandler +# --------------------------------------------------------------------------- + + +def test_images_handler_returns_gallery(tmp_storage): + """ImagesHandler lists images with filename, public URL and MIME type.""" + item = _item_dir(tmp_storage) + images = item / "content" / "images" + images.mkdir() + (images / "img_001.png").write_bytes(b"\x89PNG") + (images / "photo.JPG").write_bytes(b"\xff\xd8") + (images / "notes.txt").write_text("ignored", encoding="utf-8") # non-image ignored + payload = ImagesHandler().get(item) + assert payload.mime == "application/json" + # Sorted by filename: img_001.png before photo.JPG. + assert payload.body == [ + { + "filename": "img_001.png", + "url": "/libraries/lib1/items/item1/content/images/file/img_001.png", + "mime": "image/png", + }, + { + "filename": "photo.JPG", + "url": "/libraries/lib1/items/item1/content/images/file/photo.JPG", + "mime": "image/jpeg", + }, + ] + + +def test_images_handler_recognised_ext_uses_mapped_mime(tmp_storage): + """A recognised image extension resolves to its mapped MIME type.""" + item = _item_dir(tmp_storage) + images = item / "content" / "images" + images.mkdir() + (images / "scan.pnm").write_bytes(b"P6") + payload = ImagesHandler().get(item) + assert payload.body[0]["mime"] == "image/pnm" + + +def test_images_handler_unavailable_without_dir(tmp_storage): + """ImagesHandler raises HandlerUnavailable when the images dir is absent.""" + item = _item_dir(tmp_storage) + with pytest.raises(HandlerUnavailable): + ImagesHandler().get(item) + + +def test_images_handler_unavailable_when_no_images(tmp_storage): + """ImagesHandler raises HandlerUnavailable when only non-image files exist.""" + item = _item_dir(tmp_storage) + images = item / "content" / "images" + images.mkdir() + (images / "readme.txt").write_text("x", encoding="utf-8") + with pytest.raises(HandlerUnavailable): + ImagesHandler().get(item) diff --git a/library-manager/tests/unit/test_database.py b/library-manager/tests/unit/test_database.py new file mode 100644 index 000000000..f63ecdb2a --- /dev/null +++ b/library-manager/tests/unit/test_database.py @@ -0,0 +1,96 @@ +"""Unit tests for ``backend/database/connection.py``. + +The session database is already initialised by the root conftest (which also +holds the single-instance file lock). These tests therefore never call +``init_db`` again; they exercise the session factories, the WAL/foreign-key +pragmas, and the not-initialised RuntimeError branches by temporarily +nulling the module-level session factory under try/finally. +""" + +from __future__ import annotations + +import sqlite3 + +import pytest +from database import connection +from sqlalchemy import text +from sqlalchemy.orm import Session + + +class TestSessionFactories: + """``get_session`` / ``get_session_direct`` yield usable sessions.""" + + def test_get_session_direct_returns_session(self): + """``get_session_direct`` returns a live Session the caller closes.""" + session = connection.get_session_direct() + try: + assert isinstance(session, Session) + assert session.execute(text("SELECT 1")).scalar() == 1 + finally: + session.close() + + def test_get_session_generator_yields_and_closes(self): + """``get_session`` yields a working Session and closes it on exit.""" + gen = connection.get_session() + session = next(gen) + assert isinstance(session, Session) + assert session.execute(text("SELECT 1")).scalar() == 1 + # Exhausting the generator runs the finally: close(). + with pytest.raises(StopIteration): + next(gen) + + +class TestNotInitialisedBranches: + """Both accessors raise RuntimeError when the factory is unset.""" + + def test_get_session_direct_raises_when_uninitialised(self): + """``get_session_direct`` raises RuntimeError if init_db not called.""" + saved = connection._SessionLocal + connection._SessionLocal = None + try: + with pytest.raises(RuntimeError, match="Database not initialized"): + connection.get_session_direct() + finally: + connection._SessionLocal = saved + + def test_get_session_raises_when_uninitialised(self): + """``get_session`` raises RuntimeError if init_db not called.""" + saved = connection._SessionLocal + connection._SessionLocal = None + try: + gen = connection.get_session() + with pytest.raises(RuntimeError, match="Database not initialized"): + next(gen) + finally: + connection._SessionLocal = saved + + +class TestPragmas: + """WAL journal mode and foreign-key enforcement are active.""" + + def test_pragmas_on_live_engine_connection(self): + """A connection from the real engine reports WAL + foreign_keys=ON.""" + with connection._engine.connect() as conn: + journal = conn.execute(text("PRAGMA journal_mode")).scalar() + fk = conn.execute(text("PRAGMA foreign_keys")).scalar() + assert journal.lower() == "wal", "journal_mode should be WAL" + assert fk == 1, "foreign_keys enforcement should be ON" + + def test_enable_sqlite_wal_applies_pragmas(self): + """``_enable_sqlite_wal`` sets WAL + foreign_keys on a raw connection.""" + # Use a real on-disk DB: WAL is not available on :memory: connections. + import tempfile # noqa: PLC0415 + from pathlib import Path # noqa: PLC0415 + + tmp = Path(tempfile.mkdtemp(prefix="lm-wal-")) / "wal.db" + raw = sqlite3.connect(str(tmp)) + try: + connection._enable_sqlite_wal(raw, None) + cur = raw.cursor() + journal = cur.execute("PRAGMA journal_mode").fetchone()[0] + fk = cur.execute("PRAGMA foreign_keys").fetchone()[0] + cur.close() + assert journal.lower() == "wal" + assert fk == 1 + finally: + raw.close() diff --git a/library-manager/tests/unit/test_dependencies.py b/library-manager/tests/unit/test_dependencies.py new file mode 100644 index 000000000..c5a33455a --- /dev/null +++ b/library-manager/tests/unit/test_dependencies.py @@ -0,0 +1,47 @@ +"""Unit tests for ``backend/dependencies.py`` bearer-token verification.""" + +from __future__ import annotations + +import dependencies +import pytest +from fastapi import HTTPException, status +from fastapi.security import HTTPAuthorizationCredentials + + +def _creds(token: str) -> HTTPAuthorizationCredentials: + """Build a Bearer credentials object carrying ``token``.""" + return HTTPAuthorizationCredentials(scheme="Bearer", credentials=token) + + +class TestVerifyToken: + """``verify_token`` accepts the configured token and rejects others.""" + + async def test_correct_token_passes(self): + """A matching token is returned unchanged.""" + result = await dependencies.verify_token(_creds(dependencies.LAMB_API_TOKEN)) + assert result == dependencies.LAMB_API_TOKEN + + async def test_known_token_value(self): + """The session token (``test-token``) is the one accepted here.""" + # The root conftest sets LAMB_API_TOKEN=test-token before import. + result = await dependencies.verify_token(_creds("test-token")) + assert result == "test-token" + + async def test_wrong_token_raises_401(self): + """A non-matching token raises HTTPException 401.""" + with pytest.raises(HTTPException) as exc_info: + await dependencies.verify_token(_creds("wrong-token")) + assert exc_info.value.status_code == status.HTTP_401_UNAUTHORIZED + assert exc_info.value.detail == "Invalid service token." + + async def test_empty_token_raises_401(self): + """An empty credentials string does not match the real token.""" + with pytest.raises(HTTPException) as exc_info: + await dependencies.verify_token(_creds("")) + assert exc_info.value.status_code == status.HTTP_401_UNAUTHORIZED + + async def test_prefix_of_token_rejected(self): + """A prefix of the valid token is rejected (compare_digest is exact).""" + prefix = dependencies.LAMB_API_TOKEN[:-1] + with pytest.raises(HTTPException): + await dependencies.verify_token(_creds(prefix)) diff --git a/library-manager/tests/unit/test_markitdown_errors.py b/library-manager/tests/unit/test_markitdown_errors.py new file mode 100644 index 000000000..cc2c8c729 --- /dev/null +++ b/library-manager/tests/unit/test_markitdown_errors.py @@ -0,0 +1,161 @@ +"""Unit tests for the user-friendly markitdown error translator. + +Covers ``humanize_markitdown_error`` across each exception/message branch, +plus the ``_ext_label`` and ``_guess_format_from_message`` helpers. +""" + +from __future__ import annotations + +import pytest +from plugins._markitdown_errors import ( + _ext_label, + _guess_format_from_message, + humanize_markitdown_error, +) + + +class MissingDependencyException(Exception): # noqa: N818 — mirrors markitdown's name + pass + + +class UnsupportedFormatException(Exception): # noqa: N818 + pass + + +class FileConversionException(Exception): # noqa: N818 + pass + + +# --------------------------------------------------------------------------- +# humanize_markitdown_error +# --------------------------------------------------------------------------- + + +def test_missing_dependency_with_ext_hint_in_message(): + """MissingDependency with a .pdf hint names the format and hides pip hints.""" + exc = MissingDependencyException( + "PdfConverter recognized the input as a potential .pdf file, but the " + "dependencies needed to read .pdf files have not been installed." + ) + msg = humanize_markitdown_error(exc, "report.pdf") + assert "report.pdf" in msg + assert ".pdf" in msg + assert "not installed" in msg + assert "pip install" not in msg + + +def test_missing_dependency_falls_back_to_filename_extension(): + """Without a message hint, the filename extension supplies the format.""" + exc = MissingDependencyException("no readers available") + msg = humanize_markitdown_error(exc, "song.mp3") + assert "song.mp3" in msg + assert ".mp3" in msg + + +def test_missing_dependency_no_ext_anywhere_uses_generic_reader_message(): + """No message hint and no filename extension yields the generic reader text.""" + exc = MissingDependencyException("no readers available") + msg = humanize_markitdown_error(exc, "noextfile") + assert "noextfile" in msg + assert "required reader for this file" in msg + + +def test_missing_dependency_detected_by_message_substring(): + """Detection also works when only the message mentions the exception name.""" + exc = Exception("wrapped: MissingDependencyException for .docx") + msg = humanize_markitdown_error(exc, "doc.docx") + assert "doc.docx" in msg + assert "not installed" in msg + + +def test_unsupported_format(): + """UnsupportedFormat yields the 'not supported / try a different plugin' text.""" + exc = UnsupportedFormatException("nothing matched") + msg = humanize_markitdown_error(exc, "blob.bin") + assert "blob.bin" in msg + assert "not supported" in msg.lower() + + +def test_unsupported_format_detected_by_message(): + """UnsupportedFormat is detected via the 'unsupportedformat' message token.""" + exc = Exception("internal unsupportedformat raised") + msg = humanize_markitdown_error(exc, "x.bin") + assert "not supported" in msg.lower() + + +def test_file_conversion_corrupted(): + """FileConversion yields the unreadable/corrupted/password text.""" + exc = FileConversionException("Stream error / corrupted") + msg = humanize_markitdown_error(exc, "broken.docx") + assert "broken.docx" in msg + assert any(w in msg.lower() for w in ("unreadable", "corrupted", "password")) + + +def test_file_conversion_detected_by_message(): + """FileConversion is detected via the 'fileconversion' message token.""" + exc = Exception("a fileconversion problem occurred") + msg = humanize_markitdown_error(exc, "y.docx") + assert any(w in msg.lower() for w in ("unreadable", "corrupted", "password")) + + +def test_empty_file_detection(): + """An 'empty' message produces the dedicated empty-file message.""" + exc = Exception("File is empty") + msg = humanize_markitdown_error(exc, "x.pdf") + assert "empty" in msg.lower() + assert "x.pdf" in msg + + +def test_zero_bytes_detection(): + """A '0 bytes' message is also classified as an empty file.""" + exc = Exception("file has 0 bytes") + msg = humanize_markitdown_error(exc, "z.pdf") + assert "empty" in msg.lower() + + +def test_generic_fallback_does_not_leak_class_name(): + """An unrecognised error falls back without leaking the raw string or class.""" + exc = ValueError("Random internal error xyz") + msg = humanize_markitdown_error(exc, "doc.txt") + assert "doc.txt" in msg + assert "Random internal error xyz" not in msg + assert "ValueError" not in msg + assert "could not be converted to" in msg + + +# --------------------------------------------------------------------------- +# _ext_label +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "filename,expected", + [ + ("report.pdf", ".pdf"), + ("photo.JPG", ".jpg"), + ("noext", ""), + ("archive.tar.gz", ".gz"), + ("trailingdot.", ""), + ], +) +def test_ext_label(filename, expected): + """_ext_label returns the lowercase dotted extension or '' when absent.""" + assert _ext_label(filename) == expected + + +# --------------------------------------------------------------------------- +# _guess_format_from_message +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "message,expected", + [ + ("recognized the input as a potential .pdf file", ".pdf"), + ("as a potential .DOCX file, but", ".docx"), + ("no format hint here", ""), + ], +) +def test_guess_format_from_message(message, expected): + """_guess_format_from_message extracts the '.ext' hint or returns ''.""" + assert _guess_format_from_message(message) == expected diff --git a/library-manager/tests/unit/test_mime.py b/library-manager/tests/unit/test_mime.py new file mode 100644 index 000000000..31634a39f --- /dev/null +++ b/library-manager/tests/unit/test_mime.py @@ -0,0 +1,53 @@ +"""Unit tests for ``backend/plugins/_mime.py`` extension-to-MIME mapping.""" + +from __future__ import annotations + +import pytest +from plugins._mime import guess_mime + + +class TestGuessMime: + """``guess_mime`` maps known extensions and falls back otherwise.""" + + @pytest.mark.parametrize( + ("ext", "expected"), + [ + (".pdf", "application/pdf"), + ( + ".docx", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ), + ( + ".pptx", + "application/vnd.openxmlformats-officedocument.presentationml.presentation", + ), + ( + ".xlsx", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + ), + (".xls", "application/vnd.ms-excel"), + (".html", "text/html"), + (".csv", "text/csv"), + (".json", "application/json"), + (".xml", "application/xml"), + (".zip", "application/zip"), + (".epub", "application/epub+zip"), + (".txt", "text/plain"), + (".md", "text/markdown"), + (".mp3", "audio/mpeg"), + (".wav", "audio/wav"), + ], + ) + def test_known_extensions(self, ext, expected): + """Each known extension maps to its documented MIME type.""" + assert guess_mime(ext) == expected + + def test_case_insensitive(self): + """Extension matching is case-insensitive (lower-cased internally).""" + assert guess_mime(".PDF") == "application/pdf" + assert guess_mime(".Md") == "text/markdown" + + @pytest.mark.parametrize("ext", [".unknown", "", ".xyz", "pdf", ".tar.gz"]) + def test_unknown_falls_back(self, ext): + """Unknown or malformed extensions fall back to octet-stream.""" + assert guess_mime(ext) == "application/octet-stream" diff --git a/library-manager/tests/unit/test_models.py b/library-manager/tests/unit/test_models.py new file mode 100644 index 000000000..2f5606e4a --- /dev/null +++ b/library-manager/tests/unit/test_models.py @@ -0,0 +1,233 @@ +"""Unit tests for ``backend/database/models.py`` ORM behavior. + +The session database is shared across the whole run, so every test creates +rows with unique ids/names and tears down what it created. Foreign keys are +enforced (``PRAGMA foreign_keys=ON``), so cascade / SET NULL behavior is real. +""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest +from _helpers import unique_id +from database.models import ( + ContentFolder, + ContentImage, + ContentItem, + Library, + Organization, + _utcnow, +) +from sqlalchemy import select +from sqlalchemy.exc import IntegrityError + + +def _make_org(db) -> Organization: + """Persist and return a fresh organization.""" + org = Organization(id=unique_id("org"), name="Org " + unique_id("n")) + db.add(org) + db.commit() + return org + + +def _make_library(db, org: Organization, name: str | None = None) -> Library: + """Persist and return a fresh library under ``org``.""" + lib = Library( + id=unique_id("lib"), + organization_id=org.id, + name=name or ("Lib " + unique_id("ln")), + ) + db.add(lib) + db.commit() + return lib + + +def _make_item(db, lib: Library, org: Organization, folder_id=None) -> ContentItem: + """Persist and return a minimal content item under ``lib``.""" + item = ContentItem( + id=unique_id("item"), + library_id=lib.id, + organization_id=org.id, + folder_id=folder_id, + title="T " + unique_id("t"), + source_type="file", + base_path="base/path", + permalink_base="/docs/x", + import_plugin="simple", + ) + db.add(item) + db.commit() + return item + + +class TestUtcNow: + """``_utcnow`` returns a timezone-aware UTC datetime.""" + + def test_returns_aware_utc(self): + """Result is tz-aware and in UTC.""" + now = _utcnow() + assert isinstance(now, datetime) + assert now.tzinfo is not None, "must be timezone-aware" + assert now.utcoffset() == UTC.utcoffset(now), "must be UTC offset" + + +class TestUniqueConstraints: + """Composite unique constraints reject duplicate rows.""" + + def test_library_org_name_unique(self, db_session): + """Two libraries with the same org+name violate uq_library_org_name.""" + org = _make_org(db_session) + name = "Dup " + unique_id("d") + _make_library(db_session, org, name=name) + dup = Library(id=unique_id("lib"), organization_id=org.id, name=name) + db_session.add(dup) + with pytest.raises(IntegrityError): + db_session.commit() + db_session.rollback() + + def test_folder_sibling_name_unique(self, db_session): + """Two sibling folders with the same name violate uq_folder_sibling_name. + + SQLite treats NULL as distinct in unique constraints, so the + duplicates must share a non-NULL parent for the constraint to bite. + """ + org = _make_org(db_session) + lib = _make_library(db_session, org) + parent = ContentFolder(id=unique_id("fld"), library_id=lib.id, name="P") + db_session.add(parent) + db_session.commit() + name = "Folder " + unique_id("f") + f1 = ContentFolder( + id=unique_id("fld"), + library_id=lib.id, + parent_folder_id=parent.id, + name=name, + ) + db_session.add(f1) + db_session.commit() + f2 = ContentFolder( + id=unique_id("fld"), + library_id=lib.id, + parent_folder_id=parent.id, + name=name, + ) + db_session.add(f2) + with pytest.raises(IntegrityError): + db_session.commit() + db_session.rollback() + + def test_same_name_different_parent_allowed(self, db_session): + """Same folder name is allowed under different parents.""" + org = _make_org(db_session) + lib = _make_library(db_session, org) + parent = ContentFolder(id=unique_id("fld"), library_id=lib.id, name="P") + db_session.add(parent) + db_session.commit() + name = "Child " + unique_id("c") + top = ContentFolder(id=unique_id("fld"), library_id=lib.id, name=name) + nested = ContentFolder( + id=unique_id("fld"), + library_id=lib.id, + parent_folder_id=parent.id, + name=name, + ) + db_session.add_all([top, nested]) + db_session.commit() # must not raise + assert top.id != nested.id + + +class TestCascades: + """Foreign-key cascade and SET NULL rules are enforced.""" + + def test_delete_library_cascades_items_and_folders_and_images(self, db_session): + """Deleting a library removes its folders, items, and images.""" + org = _make_org(db_session) + lib = _make_library(db_session, org) + folder = ContentFolder(id=unique_id("fld"), library_id=lib.id, name="F") + db_session.add(folder) + db_session.commit() + item = _make_item(db_session, lib, org, folder_id=folder.id) + img = ContentImage( + id=unique_id("img"), content_item_id=item.id, image_path="p.png" + ) + db_session.add(img) + db_session.commit() + item_id, folder_id, img_id = item.id, folder.id, img.id + + db_session.delete(lib) + db_session.commit() + # The library row deletion fires DB-level ON DELETE CASCADE for folders + # (no ORM relationship covers them); expire the identity map so reads + # hit the database rather than returning stale cached objects. + db_session.expire_all() + + assert db_session.get(ContentItem, item_id) is None, "item cascaded" + assert db_session.get(ContentFolder, folder_id) is None, "folder cascaded" + assert db_session.get(ContentImage, img_id) is None, "image cascaded" + + def test_delete_org_cascades_libraries(self, db_session): + """Deleting an org removes its libraries (ORM cascade).""" + org = _make_org(db_session) + lib = _make_library(db_session, org) + lib_id = lib.id + db_session.delete(org) + db_session.commit() + assert db_session.get(Library, lib_id) is None, "library cascaded with org" + + def test_delete_folder_sets_item_folder_id_null(self, db_session): + """Deleting a folder SETs NULL on referencing items' folder_id.""" + org = _make_org(db_session) + lib = _make_library(db_session, org) + folder = ContentFolder(id=unique_id("fld"), library_id=lib.id, name="F") + db_session.add(folder) + db_session.commit() + item = _make_item(db_session, lib, org, folder_id=folder.id) + item_id = item.id + + # Delete folder via raw SQL so the DB-level SET NULL fires (ORM-level + # delete-orphan on the library relationship is not involved here). + from sqlalchemy import text # noqa: PLC0415 + + db_session.execute( + text("DELETE FROM content_folders WHERE id = :i"), {"i": folder.id} + ) + db_session.commit() + db_session.expire_all() + + refreshed = db_session.get(ContentItem, item_id) + assert refreshed is not None, "item must survive folder deletion" + assert refreshed.folder_id is None, "folder_id set NULL on folder delete" + + # cleanup + db_session.delete(lib) + db_session.commit() + + +class TestMetadataColumn: + """The ``metadata_`` attribute maps to the ``metadata`` DB column.""" + + def test_metadata_attr_maps_to_metadata_column(self, db_session): + """Writing ``metadata_`` round-trips via the ``metadata`` column.""" + org = _make_org(db_session) + lib = _make_library(db_session, org) + item = _make_item(db_session, lib, org) + item.metadata_ = '{"k": "v"}' + db_session.commit() + + from sqlalchemy import text # noqa: PLC0415 + + raw = db_session.execute( + text("SELECT metadata FROM content_items WHERE id = :i"), {"i": item.id} + ).scalar() + assert raw == '{"k": "v"}', "metadata_ stored in the metadata column" + + # And reads back through the ORM attribute. + db_session.expire_all() + again = db_session.execute( + select(ContentItem).where(ContentItem.id == item.id) + ).scalar_one() + assert again.metadata_ == '{"k": "v"}' + + db_session.delete(lib) + db_session.commit() diff --git a/library-manager/tests/unit/test_plugin_discovery.py b/library-manager/tests/unit/test_plugin_discovery.py new file mode 100644 index 000000000..bb79a17f1 --- /dev/null +++ b/library-manager/tests/unit/test_plugin_discovery.py @@ -0,0 +1,140 @@ +"""Unit tests for ``main._discover_plugins`` (zero-touch plugin drop-in). + +A fresh ``.py`` file dropped into ``backend/plugins/`` must be picked up +at startup with no code/config edits. We synthesise plugin modules inside +a tmpdir appended to the real package ``__path__`` and assert discovery +self-registers them via the decorators. ``fresh_plugin_registry`` / +``fresh_capability_registry`` guarantee no registration leaks out. +""" + +from __future__ import annotations + +import sys +import textwrap + +import pytest +from main import _discover_plugins + +_DROPIN_NAME = "test_unit_dropin_import" + +_PLUGIN_SOURCE = textwrap.dedent( + f""" + from plugins.base import ( + ImportResult, + LibraryImportPlugin, + PluginParameter, + PluginRegistry, + ) + + + @PluginRegistry.register + class UnitDropinPlugin(LibraryImportPlugin): + name = "{_DROPIN_NAME}" + description = "Synthetic plugin used by unit discovery test." + supported_source_types = ["file"] + file_extensions = ["xyz"] + human_label = "Unit drop-in plugin" + + def get_parameters(self): + return [] + + def import_content(self, source_path, *, api_keys=None, **kwargs): + return ImportResult(full_text="") + """ +).strip() + + +@pytest.fixture +def plugins_dropin_path(tmp_path, monkeypatch, fresh_plugin_registry): + """Append a tmpdir onto the real plugins package __path__ for discovery. + + ``fresh_plugin_registry`` already snapshots/restores the registry; this + fixture only manages the package ``__path__`` and the synthetic modules + in ``sys.modules`` so the next run is clean. + """ + import plugins as plugins_pkg # noqa: PLC0415 + + original_paths = list(plugins_pkg.__path__) + plugins_pkg.__path__.insert(0, str(tmp_path)) + try: + yield tmp_path + finally: + plugins_pkg.__path__[:] = original_paths + sys.modules.pop(f"plugins.{_DROPIN_NAME}", None) + sys.modules.pop("plugins.broken_unit_dropin", None) + + +def test_dropin_plugin_is_auto_discovered(plugins_dropin_path, fresh_plugin_registry): + """A fresh .py dropped into the plugins folder appears in the registry.""" + (plugins_dropin_path / f"{_DROPIN_NAME}.py").write_text(_PLUGIN_SOURCE, encoding="utf-8") + assert _DROPIN_NAME not in fresh_plugin_registry._plugins + + _discover_plugins() + + assert _DROPIN_NAME in fresh_plugin_registry._plugins, ( + f"Drop-in not discovered. Registered: {sorted(fresh_plugin_registry._plugins)}" + ) + + +def test_broken_plugin_does_not_block_discovery( + plugins_dropin_path, fresh_plugin_registry, caplog +): + """A plugin with an import error degrades to a warning, never blocks others.""" + (plugins_dropin_path / "broken_unit_dropin.py").write_text( + "raise RuntimeError('intentional')\n", encoding="utf-8" + ) + (plugins_dropin_path / f"{_DROPIN_NAME}.py").write_text(_PLUGIN_SOURCE, encoding="utf-8") + + with caplog.at_level("WARNING"): + _discover_plugins() + + assert _DROPIN_NAME in fresh_plugin_registry._plugins + assert any( + "broken_unit_dropin" in (rec.getMessage() or "") for rec in caplog.records + ), "Expected a warning log mentioning the broken plugin file." + + +def test_subpackage_modules_are_recursed( + plugins_dropin_path, fresh_plugin_registry, fresh_capability_registry +): + """Plugins inside a sub-package (content_handlers/) are also imported.""" + import plugins.content_handlers as ch_pkg # noqa: PLC0415 + from plugins.content_handlers.capability import CapabilityRegistry # noqa: PLC0415 + + subpkg = plugins_dropin_path / "content_handlers" + subpkg.mkdir() + (subpkg / "synthetic_unit_handler.py").write_text( + textwrap.dedent( + """ + from pathlib import Path + from plugins.content_handlers.capability import ( + Capability, + CapabilityPayload, + CapabilityRegistry, + ContentHandler, + ) + + + @CapabilityRegistry.register + class SyntheticUnitAudioHandler(ContentHandler): + capability = Capability.AUDIO + description = "Synthetic handler for unit discovery test." + + def get(self, item_path: Path) -> CapabilityPayload: + return CapabilityPayload(mime="audio/mpeg", body=b"") + """ + ).strip(), + encoding="utf-8", + ) + + original_ch_paths = list(ch_pkg.__path__) + ch_pkg.__path__.insert(0, str(subpkg)) + try: + _discover_plugins() + assert any( + getattr(h, "__name__", "") == "SyntheticUnitAudioHandler" + for h in CapabilityRegistry._handlers.values() + ), "Sub-package plugin was not discovered." + finally: + ch_pkg.__path__[:] = original_ch_paths + sys.modules.pop("plugins.content_handlers.synthetic_unit_handler", None) diff --git a/library-manager/tests/unit/test_plugin_markitdown.py b/library-manager/tests/unit/test_plugin_markitdown.py new file mode 100644 index 000000000..ae1738abb --- /dev/null +++ b/library-manager/tests/unit/test_plugin_markitdown.py @@ -0,0 +1,139 @@ +"""Unit tests for ``MarkItDownImportPlugin`` with MarkItDown mocked at the boundary. + +The ``markitdown`` SDK is replaced via ``patch_markitdown`` so the plugin's own +logic (page splitting, metadata building, empty-content placeholder, error +humanisation) runs for real. Files are created under ``tmp_storage`` and only +need to exist with the right extension — their bytes are never read because +MarkItDown is faked. +""" + +from __future__ import annotations + +import pytest +from _fakes import patch_markitdown +from plugins.base import ImportResult +from plugins.markitdown_import import MarkItDownImportPlugin +from plugins.markitdown_plus_import import _split_into_pages + + +def _touch(path): + """Create an empty file (content irrelevant — MarkItDown is faked).""" + path.write_bytes(b"") + return path + + +class TestImportContent: + """``MarkItDownImportPlugin.import_content`` with a faked converter.""" + + def setup_method(self): + """Fresh plugin per test.""" + self.plugin = MarkItDownImportPlugin() + + def test_returns_converted_text(self, tmp_storage): + """The converter's text_content becomes ``full_text``.""" + path = _touch(tmp_storage / "doc.docx") + with patch_markitdown(text="# Converted\n\nbody"): + result = self.plugin.import_content(str(path)) + + assert isinstance(result, ImportResult) + assert result.full_text == "# Converted\n\nbody" + assert result.metadata["import_plugin"] == "markitdown_import" + + def test_pdf_with_page_breaks_splits_pages(self, tmp_storage): + """A page-aware ext (.pdf) with ``---`` markers yields multiple pages.""" + path = _touch(tmp_storage / "doc.pdf") + text = "Page one\n\n---\n\nPage two\n\n---\n\nPage three" + with patch_markitdown(text=text): + result = self.plugin.import_content(str(path)) + + assert len(result.pages) == 3 + assert result.metadata["page_count"] == 3 + assert result.pages[0].text == "Page one" + + def test_non_page_aware_ext_has_no_pages(self, tmp_storage): + """A non-page-aware ext (.html) never splits even with ``---``.""" + path = _touch(tmp_storage / "doc.html") + with patch_markitdown(text="A\n\n---\n\nB"): + result = self.plugin.import_content(str(path)) + + assert result.pages == [] + assert result.metadata["page_count"] == 0 + + def test_empty_content_placeholder(self, tmp_storage): + """Blank conversion output is replaced with a no-text placeholder.""" + path = _touch(tmp_storage / "blank.pdf") + with patch_markitdown(text=" \n "): + result = self.plugin.import_content(str(path)) + + assert "No extractable text found" in result.full_text + assert "blank.pdf" in result.full_text + + def test_metadata_and_source_ref_shape(self, tmp_storage): + """Metadata carries mime/size/char count; source_ref is a file ref.""" + path = _touch(tmp_storage / "sheet.xlsx") + with patch_markitdown(text="data"): + result = self.plugin.import_content(str(path)) + + meta = result.metadata + assert meta["original_filename"] == "sheet.xlsx" + assert meta["content_type"].endswith("spreadsheetml.sheet") + assert meta["character_count"] == len("data") + assert result.source_ref == { + "type": "file", + "original_filename": "sheet.xlsx", + "content_type": meta["content_type"], + } + + def test_description_and_citation_passthrough(self, tmp_storage): + """Optional description/citation kwargs land in metadata.""" + path = _touch(tmp_storage / "doc.docx") + with patch_markitdown(text="body"): + result = self.plugin.import_content( + str(path), description="A doc", citation="Ref 2024" + ) + + assert result.metadata["description"] == "A doc" + assert result.metadata["citation"] == "Ref 2024" + + def test_missing_file_raises(self, tmp_storage): + """A nonexistent source path raises FileNotFoundError.""" + with pytest.raises(FileNotFoundError, match="Source file not found"): + self.plugin.import_content(str(tmp_storage / "absent.pdf")) + + def test_conversion_error_humanized(self, tmp_storage): + """A converter exception is translated into a RuntimeError.""" + path = _touch(tmp_storage / "doc.pdf") + with ( + patch_markitdown(raises=ValueError("boom internal")), + pytest.raises(RuntimeError) as exc, + ): + self.plugin.import_content(str(path)) + # The raw internal message is not surfaced verbatim. + assert "boom internal" not in str(exc.value) + + +class TestSplitIntoPages: + """``_split_into_pages`` direct calls (shared with markitdown_plus).""" + + def test_non_page_aware_returns_empty(self): + """A non-page-aware extension returns no pages.""" + assert _split_into_pages("a\n\n---\n\nb", "html") == [] + + def test_no_markers_returns_empty(self): + """Page-aware ext with no break markers returns no pages.""" + assert _split_into_pages("just text", "pdf") == [] + + def test_multiple_pages(self): + """``---`` markers split a PDF into numbered pages.""" + pages = _split_into_pages("one\n\n---\n\ntwo", "pdf") + assert [p.page_number for p in pages] == [1, 2] + assert [p.text for p in pages] == ["one", "two"] + + +class TestGetParameters: + """``get_parameters`` contract.""" + + def test_description_and_citation(self): + """Exposes optional description and citation parameters.""" + names = [p.name for p in MarkItDownImportPlugin().get_parameters()] + assert names == ["description", "citation"] diff --git a/library-manager/tests/unit/test_plugin_markitdown_plus.py b/library-manager/tests/unit/test_plugin_markitdown_plus.py new file mode 100644 index 000000000..dd2e47de9 --- /dev/null +++ b/library-manager/tests/unit/test_plugin_markitdown_plus.py @@ -0,0 +1,381 @@ +"""Unit tests for ``MarkItDownPlusPlugin`` — the richest import plugin. + +All third-party SDKs are mocked at their boundaries: + +* ``markitdown.MarkItDown`` via ``patch_markitdown`` +* ``fitz.open`` via ``patch_fitz`` (PyMuPDF) +* ``openai.OpenAI`` via ``patch_openai_vision`` + +so the plugin's own logic — page splitting, image extraction, page +rasterisation, MIME mapping, description modes, and stats accounting — runs +for real. Files live under ``tmp_storage`` and only need to exist with the +correct extension. +""" + +from __future__ import annotations + +from unittest import mock + +import pytest +from _fakes import FakeFitzDoc, patch_fitz, patch_markitdown, patch_openai_vision +from plugins.base import ImportResult +from plugins.markitdown_plus_import import ( + _EXT_NORMALIZE, + MarkItDownPlusPlugin, + _describe_image, + _extract_images, + _image_mime, + _split_into_pages, +) + + +def _touch(path): + """Create an empty file (content irrelevant — SDKs are faked).""" + path.write_bytes(b"") + return path + + +# --------------------------------------------------------------------------- +# _split_into_pages — all three break patterns +# --------------------------------------------------------------------------- + + +class TestSplitIntoPages: + """Every page-break pattern plus the non-splitting edge cases.""" + + def test_horizontal_rule(self): + """``---`` horizontal rules split page-aware content.""" + pages = _split_into_pages("one\n\n---\n\ntwo\n\n---\n\nthree", "pdf") + assert [p.text for p in pages] == ["one", "two", "three"] + assert [p.page_number for p in pages] == [1, 2, 3] + + def test_form_feed(self): + """Form-feed (\\f) markers split content.""" + pages = _split_into_pages("one\n\ftwo\n\fthree", "docx") + assert len(pages) == 3 + + def test_html_comment_marker(self): + """```` comments split content (case-insensitive).""" + pages = _split_into_pages("one\n\ntwo", "pptx") + assert len(pages) == 2 + + def test_non_page_aware_returns_empty(self): + """Non-page-aware extensions never split.""" + assert _split_into_pages("a\n\n---\n\nb", "html") == [] + + def test_single_page_returns_empty(self): + """Markers that yield fewer than two non-empty pages return empty.""" + assert _split_into_pages("only one page\n\n---\n\n", "pdf") == [] + + def test_no_markers_returns_empty(self): + """No markers at all returns empty.""" + assert _split_into_pages("plain text", "pdf") == [] + + def test_empty_pages_filtered(self): + """Blank sections between markers are dropped.""" + pages = _split_into_pages("Page 1\n\n---\n\n\n\n---\n\nPage 3", "pdf") + assert len(pages) == 2 + + +# --------------------------------------------------------------------------- +# _image_mime + _EXT_NORMALIZE +# --------------------------------------------------------------------------- + + +class TestImageMime: + """``_image_mime`` mapping including dot/case handling and the fallback.""" + + @pytest.mark.parametrize( + ("ext", "expected"), + [ + ("png", "image/png"), + ("jpg", "image/jpeg"), + ("jpeg", "image/jpeg"), + ("gif", "image/gif"), + ("bmp", "image/bmp"), + ("webp", "image/webp"), + ("tiff", "image/tiff"), + ("svg", "image/svg+xml"), + ], + ) + def test_known_extensions(self, ext, expected): + """Each known extension maps to its MIME type.""" + assert _image_mime(ext) == expected + + def test_unknown_defaults_to_png(self): + """Unknown extensions fall back to image/png.""" + assert _image_mime("xyz") == "image/png" + + def test_dot_prefix_stripped(self): + """A leading dot is stripped before lookup.""" + assert _image_mime(".jpeg") == "image/jpeg" + + def test_uppercase_normalized(self): + """Lookup is case-insensitive.""" + assert _image_mime("PNG") == "image/png" + + +class TestExtNormalize: + """``_EXT_NORMALIZE`` rewrites exotic PyMuPDF extensions to viewable ones.""" + + def test_known_normalizations(self): + """Each exotic format normalizes to a browser-friendly equivalent.""" + assert _EXT_NORMALIZE["tif"] == "tiff" + assert _EXT_NORMALIZE["jpx"] == "jpg" + assert _EXT_NORMALIZE["jb2"] == "png" + + def test_unmapped_passthrough(self): + """An extension not in the map is left to the caller (no key).""" + assert "png" not in _EXT_NORMALIZE + + +# --------------------------------------------------------------------------- +# _describe_image — all four modes +# --------------------------------------------------------------------------- + + +class TestDescribeImage: + """``_describe_image`` across basic / none / llm modes.""" + + def test_basic_mode(self): + """basic mode returns a filename-style description.""" + assert _describe_image(b"x", "img_001.png", "png", "basic", {}, {}) == ( + "Image: img_001.png" + ) + + def test_none_mode(self): + """none (and any non-basic/non-llm) mode returns None.""" + assert _describe_image(b"x", "img_001.png", "png", "none", {}, {}) is None + + def test_llm_mode_without_key_falls_back(self): + """llm mode with no openai_vision key falls back to filename text.""" + assert _describe_image(b"x", "img_001.png", "png", "llm", {}, {}) == ( + "Image: img_001.png" + ) + + def test_llm_mode_success_records_stats(self): + """llm mode with a key returns the model description and records stats.""" + stats = {"images_with_llm_descriptions": 0, "llm_calls": []} + with patch_openai_vision(description="A red square."): + desc = _describe_image( + b"x", "img_001.png", "png", "llm", {"openai_vision": "sk-1"}, stats + ) + assert desc == "A red square." + assert stats["images_with_llm_descriptions"] == 1 + assert stats["llm_calls"][0]["success"] is True + assert stats["llm_calls"][0]["tokens_used"] == 42 + + def test_llm_mode_failure_falls_back_and_records(self): + """A vision-call exception falls back to filename text and logs failure.""" + stats = {"images_with_llm_descriptions": 0, "llm_calls": []} + with patch_openai_vision(raises=RuntimeError("api down")): + desc = _describe_image( + b"x", "img_001.png", "png", "llm", {"openai_vision": "sk-1"}, stats + ) + assert desc == "Image: img_001.png" + assert stats["llm_calls"][0]["success"] is False + assert "api down" in stats["llm_calls"][0]["error"] + + +# --------------------------------------------------------------------------- +# _extract_images — PyMuPDF mocked +# --------------------------------------------------------------------------- + + +class TestExtractImages: + """``_extract_images`` against a fake PyMuPDF document.""" + + def test_non_pdf_returns_empty(self, tmp_storage): + """A non-PDF suffix short-circuits to no images.""" + path = _touch(tmp_storage / "doc.docx") + assert _extract_images(path, "basic", {}, {}) == [] + + def test_open_failure_returns_empty(self, tmp_storage): + """If fitz.open raises, extraction returns no images.""" + path = _touch(tmp_storage / "doc.pdf") + with patch_fitz(raises=RuntimeError("corrupt pdf")): + assert _extract_images(path, "basic", {}, {}) == [] + + def test_extracts_bitmaps_and_rasterized_pages(self, tmp_storage): + """Each embedded image plus one rasterized PNG per page is returned.""" + path = _touch(tmp_storage / "doc.pdf") + stats = {} + doc = FakeFitzDoc(pages=2, images_per_page=1) + with patch_fitz(doc=doc): + images = _extract_images(path, "basic", {}, stats) + + # 2 embedded bitmaps + 2 rasterized pages. + assert len(images) == 4 + bitmaps = [i for i in images if i.filename.startswith("img_")] + renders = [i for i in images if i.filename.startswith("page_")] + assert [i.filename for i in bitmaps] == ["img_001.png", "img_002.png"] + assert [i.filename for i in renders] == ["page_001.png", "page_002.png"] + assert [i.page_number for i in renders] == [1, 2] + # basic mode attaches filename-style descriptions. + assert images[0].description == "Image: img_001.png" + assert stats["rendered_page_fallback"] == 2 + assert doc.closed is True + + def test_fitz_not_installed_returns_empty(self, tmp_storage): + """If PyMuPDF (fitz) cannot be imported, extraction returns empty.""" + import builtins + + path = _touch(tmp_storage / "doc.pdf") + real_import = builtins.__import__ + + def _no_fitz(name, *args, **kwargs): + if name == "fitz": + raise ImportError("No module named 'fitz'") + return real_import(name, *args, **kwargs) + + with mock.patch("builtins.__import__", side_effect=_no_fitz): + assert _extract_images(path, "basic", {}, {}) == [] + + def test_extract_image_failure_skips_bitmap(self, tmp_storage): + """A bitmap whose extract_image raises is skipped; pages still render.""" + path = _touch(tmp_storage / "doc.pdf") + + class _Doc(FakeFitzDoc): + def extract_image(self, xref): + raise RuntimeError("bad xref") + + with patch_fitz(doc=_Doc(pages=1, images_per_page=1)): + images = _extract_images(path, "basic", {}, {}) + + # No bitmaps survived; the single rasterized page remains. + assert [i.filename for i in images] == ["page_001.png"] + + def test_empty_image_payload_skipped(self, tmp_storage): + """A bitmap with no image bytes is skipped.""" + path = _touch(tmp_storage / "doc.pdf") + + class _Doc(FakeFitzDoc): + def extract_image(self, xref): + return {"image": b"", "ext": "png"} + + with patch_fitz(doc=_Doc(pages=1, images_per_page=1)): + images = _extract_images(path, "basic", {}, {}) + + assert [i.filename for i in images] == ["page_001.png"] + + def test_rasterize_failure_skips_page(self, tmp_storage): + """A page whose get_pixmap raises is skipped during rasterisation.""" + from _fakes import _FakeFitzPage + + path = _touch(tmp_storage / "doc.pdf") + + class _BadPage(_FakeFitzPage): + def get_pixmap(self, *a, **k): + raise RuntimeError("render failed") + + class _Doc(FakeFitzDoc): + def __init__(self): + super().__init__(pages=1, images_per_page=0) + self._pages = [_BadPage([], b"")] + + doc = _Doc() + with patch_fitz(doc=doc): + images = _extract_images(path, "basic", {}, {}) + + # No bitmaps, and the only page failed to rasterize. + assert images == [] + + def test_exotic_extension_normalized_in_filename(self, tmp_storage): + """A PyMuPDF 'jpx' image is normalized to a 'jpg' filename.""" + path = _touch(tmp_storage / "doc.pdf") + doc = FakeFitzDoc(pages=1, images_per_page=1, extension="jpx") + with patch_fitz(doc=doc): + images = _extract_images(path, "basic", {}, {}) + + assert images[0].filename == "img_001.jpg" + + +# --------------------------------------------------------------------------- +# import_content — full pipeline +# --------------------------------------------------------------------------- + + +class TestImportContent: + """``MarkItDownPlusPlugin.import_content`` end to end with faked SDKs.""" + + def setup_method(self): + """Fresh plugin per test.""" + self.plugin = MarkItDownPlusPlugin() + + def test_pdf_with_images_and_pages(self, tmp_storage): + """A faked PDF yields full text, pages, images, and processing stats.""" + path = _touch(tmp_storage / "report.pdf") + text = "Page one\n\n---\n\nPage two" + doc = FakeFitzDoc(pages=2, images_per_page=1) + with patch_markitdown(text=text), patch_fitz(doc=doc): + result = self.plugin.import_content(str(path)) + + assert isinstance(result, ImportResult) + assert result.full_text == text + assert len(result.pages) == 2 + # 2 bitmaps + 2 rasterized pages. + assert len(result.images) == 4 + assert result.metadata["image_count"] == 4 + assert result.metadata["page_count"] == 2 + assert result.metadata["image_descriptions_mode"] == "basic" + assert result.metadata["import_plugin"] == "markitdown_plus_import" + + def test_image_mode_none_skips_extraction(self, tmp_storage): + """image_descriptions='none' skips image extraction entirely.""" + path = _touch(tmp_storage / "report.pdf") + with patch_markitdown(text="body"): + result = self.plugin.import_content(str(path), image_descriptions="none") + + assert result.images == [] + assert result.metadata["image_count"] == 0 + + def test_description_and_citation_passthrough(self, tmp_storage): + """Optional description/citation kwargs land in metadata.""" + path = _touch(tmp_storage / "report.pdf") + with patch_markitdown(text="body"): + result = self.plugin.import_content( + str(path), + image_descriptions="none", + description="A report", + citation="Cite 2024", + ) + + assert result.metadata["description"] == "A report" + assert result.metadata["citation"] == "Cite 2024" + + def test_non_pdf_extracts_no_bitmaps(self, tmp_storage): + """A .docx is page-aware but yields no images (PyMuPDF only reads PDFs).""" + path = _touch(tmp_storage / "doc.docx") + with patch_markitdown(text="A\n\f\nB"): + result = self.plugin.import_content(str(path)) + + assert result.images == [] + assert len(result.pages) == 2 + + def test_conversion_error_humanized(self, tmp_storage): + """A converter exception becomes a RuntimeError without leaking internals.""" + path = _touch(tmp_storage / "report.pdf") + with ( + patch_markitdown(raises=ValueError("internal boom")), + pytest.raises(RuntimeError) as exc, + ): + self.plugin.import_content(str(path)) + assert "internal boom" not in str(exc.value) + + def test_missing_file_raises(self, tmp_storage): + """A nonexistent source path raises FileNotFoundError.""" + with pytest.raises(FileNotFoundError, match="Source file not found"): + self.plugin.import_content(str(tmp_storage / "absent.pdf")) + + +class TestGetParameters: + """``get_parameters`` contract.""" + + def test_parameter_names_and_choices(self): + """Exposes image_descriptions (enum) plus description/citation.""" + params = MarkItDownPlusPlugin().get_parameters() + names = [p.name for p in params] + assert names == ["image_descriptions", "description", "citation"] + img = next(p for p in params if p.name == "image_descriptions") + assert img.choices == ["none", "basic", "llm"] + assert img.default == "basic" diff --git a/library-manager/tests/unit/test_plugin_registry.py b/library-manager/tests/unit/test_plugin_registry.py new file mode 100644 index 000000000..6b77bc6f2 --- /dev/null +++ b/library-manager/tests/unit/test_plugin_registry.py @@ -0,0 +1,324 @@ +"""Unit tests for ``plugins.base``: dataclasses and ``PluginRegistry``. + +These exercise the registry classmethods directly (no FastAPI, no HTTP) +and use the ``fresh_plugin_registry`` fixture to snapshot/restore the +class-level ``_plugins`` dict so no registration leaks into other tests. +""" + +from __future__ import annotations + +import pytest +from plugins.base import ( + ExtractedImage, + ImportResult, + LibraryImportPlugin, + PageContent, + PluginParameter, +) + + +def _make_plugin_cls( + plugin_name: str, + *, + parameters: list[PluginParameter] | None = None, + source_types: set[str] | None = None, + extensions: list[str] | None = None, + capabilities: list | None = None, + label: str = "", + desc: str = "A throwaway plugin", + keys: list[str] | None = None, +): + """Build a concrete throwaway ``LibraryImportPlugin`` subclass for tests.""" + params = parameters or [] + + class _Throwaway(LibraryImportPlugin): + name = plugin_name + description = desc + supported_source_types = source_types if source_types is not None else {"file"} + required_keys = keys if keys is not None else [] + produces_capabilities = capabilities or [] + file_extensions = extensions or [] + human_label = label + + def import_content(self, source_path, *, api_keys=None, **kwargs): + return ImportResult(full_text="") + + def get_parameters(self): + return params + + return _Throwaway + + +# --------------------------------------------------------------------------- +# Dataclasses +# --------------------------------------------------------------------------- + + +def test_page_content_fields(): + """PageContent stores page number and markdown text.""" + page = PageContent(page_number=3, text="# Hi") + assert page.page_number == 3 + assert page.text == "# Hi" + + +def test_extracted_image_defaults(): + """ExtractedImage defaults page_number and description to None.""" + img = ExtractedImage(filename="img_001.png", data=b"\x89PNG") + assert img.filename == "img_001.png" + assert img.data == b"\x89PNG" + assert img.page_number is None + assert img.description is None + + +def test_import_result_default_factories_are_independent(): + """ImportResult mutable defaults come from independent factories.""" + a = ImportResult(full_text="a") + b = ImportResult(full_text="b") + assert a.pages == [] and a.images == [] and a.metadata == {} and a.source_ref == {} + a.pages.append(PageContent(1, "x")) + a.metadata["k"] = "v" + # b must not see a's mutations — proves default_factory, not shared default. + assert b.pages == [] + assert b.metadata == {} + + +def test_plugin_parameter_defaults(): + """PluginParameter applies the documented defaults for optional fields.""" + param = PluginParameter(name="chunk", type="int") + assert param.description == "" + assert param.default is None + assert param.required is False + assert param.choices is None + assert param.min_value is None + assert param.max_value is None + assert param.advanced is False + + +# --------------------------------------------------------------------------- +# LibraryImportPlugin.report_progress +# --------------------------------------------------------------------------- + + +def test_report_progress_invokes_callable_callback(): + """report_progress forwards (current, total, message) to a callable callback.""" + plugin = _make_plugin_cls("rp_a")() + calls = [] + kwargs = {"progress_callback": lambda c, t, m: calls.append((c, t, m))} + plugin.report_progress(kwargs, 2, 5, "halfway") + assert calls == [(2, 5, "halfway")] + + +def test_report_progress_noop_when_no_callback(): + """report_progress is a silent no-op when no callback is supplied.""" + plugin = _make_plugin_cls("rp_b")() + plugin.report_progress({}, 1, 1, "done") # must not raise + + +def test_report_progress_noop_when_callback_not_callable(): + """report_progress ignores a non-callable progress_callback value.""" + plugin = _make_plugin_cls("rp_c")() + plugin.report_progress({"progress_callback": "nope"}, 1, 1, "x") # must not raise + + +# --------------------------------------------------------------------------- +# PluginRegistry.register +# --------------------------------------------------------------------------- + + +def test_register_adds_plugin_and_returns_class(fresh_plugin_registry): + """register stores the class and returns it unchanged (decorator-friendly).""" + cls = _make_plugin_cls("reg_default") + returned = fresh_plugin_registry.register(cls) + assert returned is cls + assert fresh_plugin_registry._plugins["reg_default"] is cls + + +def test_register_respects_disable_env(fresh_plugin_registry, monkeypatch): + """register skips storage when PLUGIN_=DISABLE but still returns the class.""" + monkeypatch.setenv("PLUGIN_REG_DISABLED", "DISABLE") + cls = _make_plugin_cls("reg_disabled") + returned = fresh_plugin_registry.register(cls) + assert returned is cls + assert "reg_disabled" not in fresh_plugin_registry._plugins + + +# --------------------------------------------------------------------------- +# PluginRegistry.get_plugin +# --------------------------------------------------------------------------- + + +def test_get_plugin_returns_fresh_instance(fresh_plugin_registry): + """get_plugin instantiates a new object each call.""" + cls = _make_plugin_cls("gp_fresh") + fresh_plugin_registry.register(cls) + a = fresh_plugin_registry.get_plugin("gp_fresh") + b = fresh_plugin_registry.get_plugin("gp_fresh") + assert isinstance(a, cls) + assert a is not b + + +def test_get_plugin_none_when_missing(fresh_plugin_registry): + """get_plugin returns None for an unregistered name.""" + assert fresh_plugin_registry.get_plugin("does_not_exist") is None + + +def test_get_plugin_none_when_disabled(fresh_plugin_registry, monkeypatch): + """A DISABLE'd plugin never enters the registry, so get_plugin is None.""" + monkeypatch.setenv("PLUGIN_GP_DISABLED", "DISABLE") + fresh_plugin_registry.register(_make_plugin_cls("gp_disabled")) + assert fresh_plugin_registry.get_plugin("gp_disabled") is None + + +# --------------------------------------------------------------------------- +# PluginRegistry.get_plugin_mode +# --------------------------------------------------------------------------- + + +def test_get_plugin_mode_defaults_advanced(fresh_plugin_registry, monkeypatch): + """Mode defaults to ADVANCED when no env var is set.""" + monkeypatch.delenv("PLUGIN_MODE_X", raising=False) + assert fresh_plugin_registry.get_plugin_mode("mode_x") == "ADVANCED" + + +@pytest.mark.parametrize("value", ["DISABLE", "SIMPLIFIED", "ADVANCED"]) +def test_get_plugin_mode_reads_env(fresh_plugin_registry, monkeypatch, value): + """Mode reads PLUGIN_ and recognises the three valid values.""" + monkeypatch.setenv("PLUGIN_MODE_Y", value.lower()) + assert fresh_plugin_registry.get_plugin_mode("mode_y") == value + + +def test_get_plugin_mode_unrecognized_falls_back_advanced(fresh_plugin_registry, monkeypatch): + """An unrecognised env value falls back to ADVANCED.""" + monkeypatch.setenv("PLUGIN_MODE_Z", "bogus") + assert fresh_plugin_registry.get_plugin_mode("mode_z") == "ADVANCED" + + +# --------------------------------------------------------------------------- +# PluginRegistry.list_plugins +# --------------------------------------------------------------------------- + + +def test_list_plugins_metadata_shape(fresh_plugin_registry, monkeypatch): + """list_plugins returns the documented keys with derived/normalised values.""" + from plugins.content_handlers.capability import Capability # noqa: PLC0415 + + monkeypatch.delenv("PLUGIN_LP_SHAPE", raising=False) + cls = _make_plugin_cls( + "lp_shape", + parameters=[PluginParameter(name="depth", type="int")], + source_types={"url", "file"}, + extensions=[".PDF", "Docx"], + capabilities=[Capability.TEXT, Capability.PAGES], + label="Nice Label", + keys=["openai_vision"], + ) + # mutate the dict directly so the fixture's attribute rebind restores cleanly + fresh_plugin_registry._plugins["lp_shape"] = cls + + rows = [r for r in fresh_plugin_registry.list_plugins() if r["name"] == "lp_shape"] + assert len(rows) == 1 + row = rows[0] + assert set(row) == { + "name", "description", "source_type", "supported_source_types", + "required_keys", "produces_capabilities", "file_extensions", + "human_label", "mode", "parameters", + } + # source_type is the first sorted supported source type. + assert row["supported_source_types"] == ["file", "url"] + assert row["source_type"] == "file" + # extensions normalised: lowercase, no leading dot. + assert row["file_extensions"] == ["pdf", "docx"] + # capabilities serialised to their string values. + assert row["produces_capabilities"] == ["text", "pages"] + assert row["human_label"] == "Nice Label" + assert row["required_keys"] == ["openai_vision"] + assert row["mode"] == "ADVANCED" + assert row["parameters"][0]["name"] == "depth" + + +def test_list_plugins_human_label_falls_back_to_name(fresh_plugin_registry): + """An empty human_label falls back to the plugin name.""" + fresh_plugin_registry._plugins["lp_nolabel"] = _make_plugin_cls("lp_nolabel", label="") + row = next(r for r in fresh_plugin_registry.list_plugins() if r["name"] == "lp_nolabel") + assert row["human_label"] == "lp_nolabel" + + +def test_list_plugins_no_source_types_defaults_file(fresh_plugin_registry): + """A plugin with no supported source types reports source_type 'file'.""" + fresh_plugin_registry._plugins["lp_nosrc"] = _make_plugin_cls( + "lp_nosrc", source_types=set() + ) + row = next(r for r in fresh_plugin_registry.list_plugins() if r["name"] == "lp_nosrc") + assert row["supported_source_types"] == [] + assert row["source_type"] == "file" + + +def test_list_plugins_simplified_strips_advanced_params(fresh_plugin_registry, monkeypatch): + """In SIMPLIFIED mode advanced parameters are dropped from the listing.""" + monkeypatch.setenv("PLUGIN_LP_SIMPLE", "SIMPLIFIED") + fresh_plugin_registry._plugins["lp_simple"] = _make_plugin_cls( + "lp_simple", + parameters=[ + PluginParameter(name="basic", type="string"), + PluginParameter(name="expert", type="int", advanced=True), + ], + ) + row = next(r for r in fresh_plugin_registry.list_plugins() if r["name"] == "lp_simple") + assert row["mode"] == "SIMPLIFIED" + names = [p["name"] for p in row["parameters"]] + assert names == ["basic"] + + +# --------------------------------------------------------------------------- +# PluginRegistry.sanitize_params +# --------------------------------------------------------------------------- + + +def test_sanitize_params_drops_unknown_always(fresh_plugin_registry, monkeypatch): + """sanitize_params strips parameters the plugin does not declare.""" + monkeypatch.delenv("PLUGIN_SP_ADV", raising=False) + fresh_plugin_registry._plugins["sp_adv"] = _make_plugin_cls( + "sp_adv", + parameters=[PluginParameter(name="known", type="string")], + ) + out = fresh_plugin_registry.sanitize_params( + "sp_adv", {"known": 1, "bogus": 2} + ) + assert out == {"known": 1} + + +def test_sanitize_params_drops_advanced_in_simplified(fresh_plugin_registry, monkeypatch): + """In SIMPLIFIED mode advanced params are stripped as well as unknown ones.""" + monkeypatch.setenv("PLUGIN_SP_SIMPLE", "SIMPLIFIED") + fresh_plugin_registry._plugins["sp_simple"] = _make_plugin_cls( + "sp_simple", + parameters=[ + PluginParameter(name="basic", type="string"), + PluginParameter(name="expert", type="int", advanced=True), + ], + ) + out = fresh_plugin_registry.sanitize_params( + "sp_simple", {"basic": 1, "expert": 2, "bogus": 3} + ) + assert out == {"basic": 1} + + +def test_sanitize_params_keeps_advanced_in_advanced(fresh_plugin_registry, monkeypatch): + """In ADVANCED mode advanced params survive sanitisation.""" + monkeypatch.delenv("PLUGIN_SP_KEEP", raising=False) + fresh_plugin_registry._plugins["sp_keep"] = _make_plugin_cls( + "sp_keep", + parameters=[ + PluginParameter(name="basic", type="string"), + PluginParameter(name="expert", type="int", advanced=True), + ], + ) + out = fresh_plugin_registry.sanitize_params( + "sp_keep", {"basic": 1, "expert": 2} + ) + assert out == {"basic": 1, "expert": 2} + + +def test_sanitize_params_unknown_plugin_returns_empty(fresh_plugin_registry): + """sanitize_params returns {} for a plugin that is not registered.""" + assert fresh_plugin_registry.sanitize_params("nope_plugin", {"a": 1}) == {} diff --git a/library-manager/tests/unit/test_plugin_simple.py b/library-manager/tests/unit/test_plugin_simple.py new file mode 100644 index 000000000..4196af7e7 --- /dev/null +++ b/library-manager/tests/unit/test_plugin_simple.py @@ -0,0 +1,145 @@ +"""Unit tests for ``SimpleImportPlugin`` — direct calls, no HTTP, no worker. + +Exercises the plain-text import path end to end against the real filesystem +(files created under ``tmp_storage``) plus the ``_mime_for_extension`` map and +the plugin's class-level capability/extension declarations. +""" + +from __future__ import annotations + +import pytest +from plugins.base import ImportResult +from plugins.content_handlers.capability import Capability +from plugins.simple_import import SimpleImportPlugin, _mime_for_extension + + +class TestImportContent: + """``SimpleImportPlugin.import_content`` reading real files.""" + + def setup_method(self): + """Build a fresh plugin instance per test.""" + self.plugin = SimpleImportPlugin() + + def test_reads_markdown_file(self, tmp_storage): + """A .md file is read verbatim into ``full_text``.""" + path = tmp_storage / "notes.md" + body = "# Title\n\nSome **markdown** body.\n" + path.write_text(body, encoding="utf-8") + + result = self.plugin.import_content(str(path)) + + assert isinstance(result, ImportResult) + assert result.full_text == body + assert result.pages == [] + assert result.images == [] + + def test_reads_txt_file(self, tmp_storage): + """A .txt file is read as UTF-8 text.""" + path = tmp_storage / "plain.txt" + path.write_text("line one\nline two", encoding="utf-8") + + result = self.plugin.import_content(str(path)) + + assert result.full_text == "line one\nline two" + + def test_reads_html_file(self, tmp_storage): + """A .html file is read as-is (no conversion).""" + path = tmp_storage / "page.html" + body = "

    hi

    " + path.write_text(body, encoding="utf-8") + + result = self.plugin.import_content(str(path)) + + assert result.full_text == body + + def test_metadata_shape(self, tmp_storage): + """Metadata records filename, mime, size, and character count.""" + path = tmp_storage / "doc.md" + body = "héllo" # multibyte to make byte-size differ from char count + path.write_text(body, encoding="utf-8") + + result = self.plugin.import_content(str(path)) + + meta = result.metadata + assert meta["original_filename"] == "doc.md" + assert meta["content_type"] == "text/markdown" + assert meta["file_size"] == path.stat().st_size + assert meta["character_count"] == len(body) + + def test_source_ref_shape(self, tmp_storage): + """source_ref is a file reference echoing filename and content type.""" + path = tmp_storage / "doc.txt" + path.write_text("x", encoding="utf-8") + + result = self.plugin.import_content(str(path)) + + assert result.source_ref == { + "type": "file", + "original_filename": "doc.txt", + "content_type": "text/plain", + } + + def test_missing_file_raises(self, tmp_storage): + """A nonexistent source path raises FileNotFoundError.""" + with pytest.raises(FileNotFoundError, match="Source file not found"): + self.plugin.import_content(str(tmp_storage / "absent.txt")) + + def test_invalid_utf8_raises_value_error(self, tmp_storage): + """Bytes that are not valid UTF-8 raise a humanized ValueError.""" + path = tmp_storage / "bad.txt" + path.write_bytes(b"\xff\xfe\x00invalid") + + with pytest.raises(ValueError, match="not valid UTF-8"): + self.plugin.import_content(str(path)) + + +class TestGetParameters: + """``get_parameters`` contract.""" + + def test_returns_empty_list(self): + """The simple plugin exposes no configurable parameters.""" + assert SimpleImportPlugin().get_parameters() == [] + + +class TestMimeForExtension: + """``_mime_for_extension`` mapping for every supported extension.""" + + def test_txt(self): + """.txt maps to text/plain.""" + assert _mime_for_extension(".txt") == "text/plain" + + def test_md(self): + """.md maps to text/markdown.""" + assert _mime_for_extension(".md") == "text/markdown" + + def test_html(self): + """.html maps to text/html.""" + assert _mime_for_extension(".html") == "text/html" + + def test_uppercase_normalized(self): + """Extension matching is case-insensitive.""" + assert _mime_for_extension(".MD") == "text/markdown" + + def test_unknown_defaults_to_plain(self): + """Unknown extensions fall back to text/plain.""" + assert _mime_for_extension(".xyz") == "text/plain" + + +class TestClassAttributes: + """Class-level declarations used by routing/UI.""" + + def test_supported_source_types(self): + """Only file sources are handled.""" + assert SimpleImportPlugin.supported_source_types == {"file"} + + def test_file_extensions(self): + """Declares the plain-text extensions it accepts.""" + assert SimpleImportPlugin.file_extensions == ["txt", "md", "html", "htm"] + + def test_capabilities(self): + """Produces only the TEXT capability.""" + assert SimpleImportPlugin.produces_capabilities == [Capability.TEXT] + + def test_name(self): + """Registered under the stable plugin name.""" + assert SimpleImportPlugin.name == "simple_import" diff --git a/library-manager/tests/unit/test_plugin_url.py b/library-manager/tests/unit/test_plugin_url.py new file mode 100644 index 000000000..314c9f55a --- /dev/null +++ b/library-manager/tests/unit/test_plugin_url.py @@ -0,0 +1,264 @@ +"""Unit tests for ``UrlImportPlugin`` with HTTP backends mocked at the boundary. + +``firecrawl.FirecrawlApp`` and ``markitdown.MarkItDown`` are replaced via the +``patch_firecrawl`` / ``patch_markitdown`` context managers so the plugin's URL +validation, SSRF blocklist, backend-selection branch, metadata building, and +error humanisation all run for real with no network access. +""" + +from __future__ import annotations + +from unittest import mock + +import pytest +from _fakes import patch_firecrawl, patch_markitdown +from plugins.base import ImportResult +from plugins.url_import import UrlImportPlugin, _safe_int + +FIRECRAWL_KEYS = {"firecrawl_key": "fc-test"} + + +class TestUrlValidation: + """URL scheme/host validation and the SSRF blocklist.""" + + def setup_method(self): + """Fresh plugin per test.""" + self.plugin = UrlImportPlugin() + + def test_no_scheme_raises(self): + """A URL without a scheme is rejected.""" + with pytest.raises(ValueError, match="Invalid URL"): + self.plugin.import_content("example.com") + + def test_no_netloc_raises(self): + """A scheme with no host is rejected.""" + with pytest.raises(ValueError, match="Invalid URL"): + self.plugin.import_content("http://") + + def test_non_http_scheme_raises(self): + """Non-http(s) schemes are rejected.""" + with pytest.raises(ValueError, match="Unsupported URL scheme"): + self.plugin.import_content("ftp://example.com/x") + + @pytest.mark.parametrize( + "url", + [ + "http://localhost/p", + "http://127.0.0.1/p", + "http://0.0.0.0/p", + "http://[::1]/p", + "http://169.254.169.254/latest/meta-data", + "http://metadata.google.internal/computeMetadata", + ], + ) + def test_blocked_hosts(self, url): + """Each SSRF-sensitive host is rejected.""" + with pytest.raises(ValueError, match="not allowed"): + self.plugin.import_content(url) + + +class TestMarkitdownFallback: + """No firecrawl_key → direct MarkItDown fetch path.""" + + def setup_method(self): + """Fresh plugin per test.""" + self.plugin = UrlImportPlugin() + + def test_returns_markitdown_result(self): + """Without a key, content is fetched via MarkItDown.""" + with patch_markitdown(text="# Hello"): + result = self.plugin.import_content("https://example.com/page") + + assert isinstance(result, ImportResult) + assert result.full_text == "# Hello" + assert result.metadata["fetch_method"] == "markitdown" + assert result.metadata["source_url"] == "https://example.com/page" + assert result.source_ref["type"] == "url" + + def test_empty_text_placeholder(self): + """Blank conversion output yields a no-content placeholder.""" + with patch_markitdown(text=""): + result = self.plugin.import_content("https://example.com/empty") + + assert "No text content could be extracted" in result.full_text + + def test_fetch_failure_raises_runtime(self): + """A converter exception becomes a RuntimeError.""" + with ( + patch_markitdown(raises=Exception("connection refused")), + pytest.raises(RuntimeError, match="Failed to fetch"), + ): + self.plugin.import_content("https://example.com/fail") + + def test_description_and_citation_passthrough(self): + """Optional description/citation kwargs reach metadata.""" + with patch_markitdown(text="body"): + result = self.plugin.import_content( + "https://example.com/p", description="A page", citation="Ref" + ) + + assert result.metadata["description"] == "A page" + assert result.metadata["citation"] == "Ref" + + +class TestFirecrawlPath: + """firecrawl_key present → Firecrawl scrape path.""" + + def setup_method(self): + """Fresh plugin per test.""" + self.plugin = UrlImportPlugin() + + def test_returns_firecrawl_result(self): + """With a key, content is scraped via Firecrawl.""" + meta = {"title": "Example", "sourceURL": "https://example.com"} + with patch_firecrawl(markdown="# Example body", metadata=meta): + result = self.plugin.import_content( + "https://example.com", api_keys=FIRECRAWL_KEYS + ) + + assert "# Example body" in result.full_text + assert result.metadata["fetch_method"] == "firecrawl" + assert result.metadata["pages_crawled"] == 1 + assert result.source_ref["type"] == "url" + + def test_empty_markdown_yields_empty_text(self): + """A doc with blank markdown contributes no text parts. + + The fake doc still exposes a ``.markdown`` attribute, so it counts as + one crawled page, but its empty body produces an empty ``full_text``. + """ + with patch_firecrawl(markdown="", metadata={}): + result = self.plugin.import_content( + "https://example.com", api_keys=FIRECRAWL_KEYS + ) + + assert result.full_text == "" + assert result.metadata["pages_crawled"] == 1 + + def test_scrape_failure_raises_runtime(self): + """A scrape exception becomes a RuntimeError.""" + with ( + patch_firecrawl(raises=Exception("API error")), + pytest.raises(RuntimeError, match="Firecrawl scrape failed"), + ): + self.plugin.import_content("https://example.com", api_keys=FIRECRAWL_KEYS) + + def test_object_metadata_branch(self): + """Metadata supplied as an object (not a dict) is read via attributes.""" + + class _Meta: + url = "https://example.com/from-attr" + source_url = "" + title = "Attr Title" + + with patch_firecrawl(markdown="# body", metadata=_Meta()): + result = self.plugin.import_content( + "https://example.com", api_keys=FIRECRAWL_KEYS + ) + + # Title header and source note derived from the object attributes. + assert "## Attr Title" in result.full_text + assert "https://example.com/from-attr" in result.full_text + + def test_description_and_citation_passthrough(self): + """Optional description/citation kwargs reach metadata.""" + meta = {"title": "", "sourceURL": "https://example.com"} + with patch_firecrawl(markdown="# body", metadata=meta): + result = self.plugin.import_content( + "https://example.com", + api_keys=FIRECRAWL_KEYS, + description="Desc", + citation="Cite", + ) + + assert result.metadata["description"] == "Desc" + assert result.metadata["citation"] == "Cite" + + +class TestFirecrawlDictResults: + """Firecrawl returning plain dicts instead of doc objects. + + The shared ``patch_firecrawl`` fake always yields an object with a + ``.markdown`` attribute, so the dict-shaped branches are driven by patching + ``firecrawl.FirecrawlApp`` directly at the same boundary. + """ + + def setup_method(self): + """Fresh plugin per test.""" + self.plugin = UrlImportPlugin() + + def test_dict_with_markdown(self): + """A dict result carrying markdown is parsed via the dict branch.""" + doc = {"markdown": "# Dict body", "metadata": {"sourceURL": "u", "title": "T"}} + with mock.patch("firecrawl.FirecrawlApp") as App: + App.return_value.scrape.return_value = doc + result = self.plugin.import_content( + "https://example.com", api_keys=FIRECRAWL_KEYS + ) + + assert "# Dict body" in result.full_text + assert result.metadata["pages_crawled"] == 1 + + def test_dict_without_markdown_is_no_content(self): + """An empty dict (no markdown) produces the no-content placeholder.""" + with mock.patch("firecrawl.FirecrawlApp") as App: + App.return_value.scrape.return_value = {} + result = self.plugin.import_content( + "https://example.com", api_keys=FIRECRAWL_KEYS + ) + + assert "No content could be scraped" in result.full_text + assert result.metadata["pages_crawled"] == 0 + + +class TestMarkitdownNotInstalled: + """The markitdown-fallback path when the SDK is unavailable.""" + + def test_import_error_raises_runtime(self): + """A missing markitdown package surfaces an install hint.""" + plugin = UrlImportPlugin() + real_import = __import__ + + def _no_md(name, *args, **kwargs): + if name == "markitdown": + raise ImportError("No module named 'markitdown'") + return real_import(name, *args, **kwargs) + + with ( + mock.patch("builtins.__import__", side_effect=_no_md), + pytest.raises(RuntimeError, match="markitdown is not installed"), + ): + plugin.import_content("https://example.com/p") + + +class TestSafeInt: + """``_safe_int`` conversion helper.""" + + def test_none_returns_default(self): + """None yields the default.""" + assert _safe_int(None, 42) == 42 + + def test_valid_int_string(self): + """A numeric string is parsed.""" + assert _safe_int("10", 0) == 10 + + def test_valid_int(self): + """An int passes through.""" + assert _safe_int(25, 0) == 25 + + def test_invalid_string_returns_default(self): + """A non-numeric string yields the default.""" + assert _safe_int("abc", 99) == 99 + + def test_float_string_returns_default(self): + """A float-shaped string is not an int and yields the default.""" + assert _safe_int("3.14", 7) == 7 + + +class TestGetParameters: + """``get_parameters`` contract.""" + + def test_parameter_names(self): + """Exposes crawl knobs plus description/citation.""" + names = [p.name for p in UrlImportPlugin().get_parameters()] + assert names == ["max_discovery_depth", "limit", "timeout", "description", "citation"] diff --git a/library-manager/tests/unit/test_plugin_youtube.py b/library-manager/tests/unit/test_plugin_youtube.py new file mode 100644 index 000000000..776a09900 --- /dev/null +++ b/library-manager/tests/unit/test_plugin_youtube.py @@ -0,0 +1,578 @@ +"""Unit tests for ``YouTubeTranscriptImportPlugin`` and its pure helpers. + +The transcript fetch is patched at session scope by ``tests/youtube_cache.py`` +(installed in the root conftest), so ``import_content`` for the cached video +ids (``dQw4w9WgXcQ`` and the synthetic ``abc123``) runs fully offline — no +yt-dlp, no network. The URL/timestamp/SRT helpers are exercised directly. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +from plugins.base import ImportResult +from plugins.youtube_transcript_import import ( + YouTubeTranscriptImportPlugin, + _build_source_ref, + _parse_srt_content, + _parse_youtube_url, + _resolve_language_key, + _seconds_to_timestamp, +) + +# --------------------------------------------------------------------------- +# _parse_youtube_url +# --------------------------------------------------------------------------- + + +class TestParseYoutubeUrl: + """Video-id extraction across the URL shapes the plugin actually supports.""" + + def test_watch_query(self): + """A standard watch?v= URL yields the id.""" + assert _parse_youtube_url( + "https://www.youtube.com/watch?v=dQw4w9WgXcQ" + ) == "dQw4w9WgXcQ" + + def test_youtu_be_short(self): + """A youtu.be short link yields the id from the path.""" + assert _parse_youtube_url("https://youtu.be/dQw4w9WgXcQ") == "dQw4w9WgXcQ" + + def test_youtu_be_with_trailing_path(self): + """youtu.be takes the first path segment as the id.""" + assert _parse_youtube_url("https://youtu.be/dQw4w9WgXcQ/x") == "dQw4w9WgXcQ" + + def test_youtu_be_empty_path(self): + """An empty youtu.be path yields None.""" + assert _parse_youtube_url("https://youtu.be/") is None + + def test_watch_without_v_param(self): + """A youtube.com URL with no v= query yields None.""" + assert _parse_youtube_url("https://www.youtube.com/channel/abc") is None + + def test_embed_url_yields_none(self): + """An /embed/ URL has no v= query, so the parser returns None. + + (Documented behavior: the parser only reads the v= query param and the + youtu.be path — it does NOT special-case /embed/ or /shorts/.) + """ + assert _parse_youtube_url("https://www.youtube.com/embed/dQw4w9WgXcQ") is None + + def test_shorts_url_yields_none(self): + """A /shorts/ URL likewise yields None (no v= query).""" + assert _parse_youtube_url("https://www.youtube.com/shorts/dQw4w9WgXcQ") is None + + def test_non_youtube_url(self): + """A non-YouTube host yields None.""" + assert _parse_youtube_url("https://example.com/video") is None + + +# --------------------------------------------------------------------------- +# _seconds_to_timestamp +# --------------------------------------------------------------------------- + + +class TestSecondsToTimestamp: + """Timestamp formatting with and without an hours component.""" + + @pytest.mark.parametrize( + ("seconds", "expected"), + [ + (0, "00:00"), + (59, "00:59"), + (60, "01:00"), + (125, "02:05"), + (3600, "01:00:00"), + (3661, "01:01:01"), + ], + ) + def test_formats(self, seconds, expected): + """Seconds render as MM:SS, switching to HH:MM:SS past an hour.""" + assert _seconds_to_timestamp(seconds) == expected + + def test_fractional_seconds_floored(self): + """Fractional seconds are floored to whole seconds.""" + assert _seconds_to_timestamp(5.9) == "00:05" + + +# --------------------------------------------------------------------------- +# _resolve_language_key +# --------------------------------------------------------------------------- + + +class TestResolveLanguageKey: + """Resolving a requested language to the variant the video exposes.""" + + def test_exact_match(self): + """An exact key wins.""" + assert _resolve_language_key({"en": [], "es": []}, "en") == "en" + + def test_locale_suffix_match(self): + """A locale-suffixed variant (en-US) matches the base code.""" + assert _resolve_language_key({"en-US": [], "es": []}, "en") == "en-US" + + def test_translation_chain_variant(self): + """A translation-chain variant matches on the prefix before the dash.""" + assert _resolve_language_key({"en-nP7-2PuUl7o": []}, "en") == "en-nP7-2PuUl7o" + + def test_no_match(self): + """No matching key yields None.""" + assert _resolve_language_key({"fr": []}, "en") is None + + def test_empty_map(self): + """An empty map yields None.""" + assert _resolve_language_key({}, "en") is None + + def test_none_map(self): + """A None map yields None.""" + assert _resolve_language_key(None, "en") is None + + +# --------------------------------------------------------------------------- +# _parse_srt_content (handles SRT and VTT alike) +# --------------------------------------------------------------------------- + + +class TestParseSrtContent: + """SRT/VTT parsing into timed text pieces.""" + + def test_basic_srt(self): + """Two cues parse into two pieces with start/duration.""" + srt = ( + "1\n00:00:01,000 --> 00:00:04,000\nHello world\n\n" + "2\n00:00:05,000 --> 00:00:08,000\nSecond line\n" + ) + pieces = _parse_srt_content(srt) + assert len(pieces) == 2 + assert pieces[0]["text"] == "Hello world" + assert pieces[0]["start"] == 1.0 + assert pieces[0]["duration"] == 3.0 + + def test_vtt_dot_separator(self): + """VTT-style ``.`` millisecond separators parse the same as SRT commas.""" + vtt = "00:00:01.000 --> 00:00:02.500\nVTT cue\n" + pieces = _parse_srt_content(vtt) + assert len(pieces) == 1 + assert pieces[0]["start"] == 1.0 + assert pieces[0]["duration"] == 1.5 + + def test_noise_removed(self): + """Bracketed/parenthesised noise is stripped from cue text.""" + srt = "1\n00:00:01,000 --> 00:00:04,000\n[applause] Hello [music]\n" + pieces = _parse_srt_content(srt) + assert pieces[0]["text"] == "Hello" + + def test_noise_only_block_skipped(self): + """A cue whose text is entirely noise is dropped.""" + srt = ( + "1\n00:00:01,000 --> 00:00:04,000\n[noise]\n\n" + "2\n00:00:05,000 --> 00:00:08,000\nReal\n" + ) + pieces = _parse_srt_content(srt) + assert len(pieces) == 1 + assert pieces[0]["text"] == "Real" + + def test_multiline_cue_joined(self): + """Multiple text lines in a cue are joined with spaces.""" + srt = "1\n00:00:01,000 --> 00:00:04,000\nLine one\nLine two\n" + assert _parse_srt_content(srt)[0]["text"] == "Line one Line two" + + def test_hours_component(self): + """Hour-level timestamps are summed into seconds.""" + srt = "1\n01:30:00,000 --> 01:30:05,000\nLate\n" + assert _parse_srt_content(srt)[0]["start"] == 5400.0 + + def test_empty_string(self): + """Empty input yields no pieces.""" + assert _parse_srt_content("") == [] + + def test_no_timestamp(self): + """A block with no timestamp line is skipped.""" + assert _parse_srt_content("just text\nno timestamps") == [] + + +# --------------------------------------------------------------------------- +# _build_source_ref +# --------------------------------------------------------------------------- + + +class TestBuildSourceRef: + """source_ref construction for YouTube imports.""" + + def test_shape(self): + """source_ref carries type, id, canonical url, and language.""" + ref = _build_source_ref( + "https://www.youtube.com/watch?v=abc123", "abc123", "en" + ) + assert ref == { + "type": "youtube", + "source_url": "https://www.youtube.com/watch?v=abc123", + "video_id": "abc123", + "video_url": "https://www.youtube.com/watch?v=abc123", + "language": "en", + } + + +# --------------------------------------------------------------------------- +# import_content (offline via the session YouTube cache) +# --------------------------------------------------------------------------- + + +class TestImportContent: + """Full import driven by the offline transcript cache.""" + + def setup_method(self): + """Fresh plugin per test.""" + self.plugin = YouTubeTranscriptImportPlugin() + + def test_invalid_url_raises(self): + """A non-YouTube URL raises ValueError before any fetch.""" + with pytest.raises(ValueError, match="Could not extract video ID"): + self.plugin.import_content("https://example.com/not-youtube") + + def test_cached_real_transcript(self): + """The cached Rick Astley video imports into timestamped Markdown.""" + url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ" + result = self.plugin.import_content(url) + + assert isinstance(result, ImportResult) + assert result.full_text.startswith(f"# Transcript: {url}") + assert "**[" in result.full_text # timestamped entries + assert result.metadata["video_id"] == "dQw4w9WgXcQ" + assert result.metadata["subtitle_source"] == "manual" + assert result.metadata["import_plugin"] == "youtube_transcript_import" + assert result.source_ref["type"] == "youtube" + assert result.source_ref["video_id"] == "dQw4w9WgXcQ" + + def test_cached_synthetic_transcript(self): + """The synthetic ``abc123`` cache entry imports a single timed piece.""" + url = "https://www.youtube.com/watch?v=abc123" + result = self.plugin.import_content(url) + + assert "**[00:00]** Hello" in result.full_text + assert result.metadata["transcript_pieces"] == 1 + assert result.metadata["video_id"] == "abc123" + + +# --------------------------------------------------------------------------- +# _fetch_transcript / _download_and_parse — yt-dlp mocked at the boundary +# --------------------------------------------------------------------------- + + +def _mock_ydl(**attrs): + """Build a context-manager-capable MagicMock standing in for YoutubeDL().""" + inst = MagicMock() + for name, value in attrs.items(): + setattr(inst, name, value) + inst.__enter__ = MagicMock(return_value=inst) + inst.__exit__ = MagicMock(return_value=False) + return inst + + +class TestFetchTranscript: + """``_fetch_transcript`` source-selection logic with yt-dlp mocked. + + The session cache patches the module-level ``_fetch_transcript`` to read + from disk; the real implementation is preserved on its ``__wrapped__`` + attribute, which we call here to exercise the live probe/select code. + """ + + def setup_method(self): + """Resolve the unwrapped (real) implementation under test.""" + import plugins.youtube_transcript_import as yt_mod + + self.fetch = getattr( + yt_mod._fetch_transcript, "__wrapped__", yt_mod._fetch_transcript + ) + + def test_prefers_manual_subtitle(self): + """When manual subtitles exist, they are chosen over auto-captions.""" + import yt_dlp + + info = {"subtitles": {"en": [{"ext": "srt"}]}, "automatic_captions": {}} + ydl = _mock_ydl(extract_info=MagicMock(return_value=info)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + patch( + "plugins.youtube_transcript_import._download_and_parse", + return_value=[{"text": "Hi", "start": 0, "duration": 1}], + ), + ): + pieces, source = self.fetch("abc123", "en", None) + + assert source == "manual" + assert len(pieces) == 1 + + def test_falls_back_to_auto(self): + """With only auto-captions, the auto source label is returned.""" + import yt_dlp + + info = {"subtitles": {}, "automatic_captions": {"en": [{"ext": "srt"}]}} + ydl = _mock_ydl(extract_info=MagicMock(return_value=info)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + patch( + "plugins.youtube_transcript_import._download_and_parse", + return_value=[{"text": "Hi", "start": 0, "duration": 1}], + ), + ): + _pieces, source = self.fetch("abc123", "en", None) + + assert source == "auto" + + def test_no_subtitles_raises(self): + """No tracks in the requested language raises a RuntimeError.""" + import yt_dlp + + info = {"subtitles": {"fr": [{}]}, "automatic_captions": {}} + ydl = _mock_ydl(extract_info=MagicMock(return_value=info)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + pytest.raises(RuntimeError, match="No subtitles available"), + ): + self.fetch("abc123", "en", None) + + def test_proxy_passed_to_probe(self): + """A proxy_url is threaded into the yt-dlp probe options.""" + import yt_dlp + + info = {"subtitles": {"en": [{"ext": "srt"}]}, "automatic_captions": {}} + ydl = _mock_ydl(extract_info=MagicMock(return_value=info)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl) as mk, + patch( + "plugins.youtube_transcript_import._download_and_parse", + return_value=[{"text": "Hi", "start": 0, "duration": 1}], + ), + ): + self.fetch("abc123", "en", "http://proxy:8080") + + assert mk.call_args_list[0][0][0]["proxy"] == "http://proxy:8080" + + def test_extract_info_failure_raises(self): + """A yt-dlp probe failure raises a RuntimeError naming the video.""" + import yt_dlp + + ydl = _mock_ydl(extract_info=MagicMock(side_effect=Exception("blocked"))) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + pytest.raises(RuntimeError, match="Failed to fetch video info"), + ): + self.fetch("abc123", "en", None) + + def test_rate_limit_message(self): + """A 429 during download surfaces the rate-limit message.""" + import yt_dlp + + info = {"subtitles": {"en": [{"ext": "srt"}]}, "automatic_captions": {}} + ydl = _mock_ydl(extract_info=MagicMock(return_value=info)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + patch( + "plugins.youtube_transcript_import._download_and_parse", + side_effect=Exception("HTTP Error 429: Too Many Requests"), + ), + pytest.raises(RuntimeError, match="rate-limited"), + ): + self.fetch("abc123", "en", None) + + def test_non_429_failure_raises_generic(self): + """A non-rate-limit download failure yields the generic error.""" + import yt_dlp + + info = {"subtitles": {"en": [{"ext": "srt"}]}, "automatic_captions": {}} + ydl = _mock_ydl(extract_info=MagicMock(return_value=info)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + patch( + "plugins.youtube_transcript_import._download_and_parse", + side_effect=Exception("network blip"), + ), + pytest.raises(RuntimeError, match="Could not download subtitles"), + ): + self.fetch("abc123", "en", None) + + def test_all_attempts_empty_raises(self): + """When every download returns no pieces, a generic error is raised.""" + import yt_dlp + + info = {"subtitles": {"en": [{"ext": "srt"}]}, "automatic_captions": {}} + ydl = _mock_ydl(extract_info=MagicMock(return_value=info)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + patch( + "plugins.youtube_transcript_import._download_and_parse", + return_value=[], + ), + pytest.raises(RuntimeError, match="Could not download subtitles"), + ): + self.fetch("abc123", "en", None) + + +class TestDownloadAndParse: + """``_download_and_parse`` temp-file handling with yt-dlp mocked.""" + + def test_error_prefix_stripped(self): + """A yt-dlp ``ERROR:`` prefix is trimmed from the bubbled-up message.""" + import yt_dlp + from plugins.youtube_transcript_import import _download_and_parse + + ydl = _mock_ydl(download=MagicMock(side_effect=Exception("ERROR: Video unavailable"))) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + pytest.raises(RuntimeError, match="^Video unavailable"), + ): + _download_and_parse( + url="https://youtube.com/watch?v=abc", + language_key="en", + source_label="manual", + proxy_url=None, + ) + + def test_error_without_prefix_preserved(self): + """A download error lacking the ``ERROR:`` prefix is preserved as-is.""" + import yt_dlp + from plugins.youtube_transcript_import import _download_and_parse + + ydl = _mock_ydl(download=MagicMock(side_effect=Exception("plain failure"))) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + pytest.raises(RuntimeError, match="^plain failure$"), + ): + _download_and_parse( + url="https://youtube.com/watch?v=abc", + language_key="en", + source_label="manual", + proxy_url=None, + ) + + def test_no_subtitle_files_returns_empty(self): + """A temp dir with no .srt/.vtt files yields no pieces.""" + import yt_dlp + from plugins.youtube_transcript_import import _download_and_parse + + ydl = _mock_ydl(download=MagicMock(return_value=None)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + patch("os.listdir", return_value=["video.mp4"]), + ): + assert ( + _download_and_parse( + url="https://youtube.com/watch?v=abc", + language_key="en", + source_label="manual", + proxy_url=None, + ) + == [] + ) + + def test_listdir_oserror_returns_empty(self): + """If listing the temp dir fails, no pieces are returned.""" + import yt_dlp + from plugins.youtube_transcript_import import _download_and_parse + + ydl = _mock_ydl(download=MagicMock(return_value=None)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + patch("os.listdir", side_effect=OSError("gone")), + ): + assert ( + _download_and_parse( + url="https://youtube.com/watch?v=abc", + language_key="en", + source_label="manual", + proxy_url=None, + ) + == [] + ) + + def test_open_oserror_returns_empty(self): + """If the subtitle file cannot be opened, no pieces are returned.""" + import yt_dlp + from plugins.youtube_transcript_import import _download_and_parse + + ydl = _mock_ydl(download=MagicMock(return_value=None)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + patch("os.listdir", return_value=["abc.en.srt"]), + patch("builtins.open", side_effect=OSError("denied")), + ): + assert ( + _download_and_parse( + url="https://youtube.com/watch?v=abc", + language_key="en", + source_label="manual", + proxy_url=None, + ) + == [] + ) + + def test_proxy_and_auto_options(self): + """Proxy and auto-only download options are passed to yt-dlp.""" + import yt_dlp + from plugins.youtube_transcript_import import _download_and_parse + + ydl = _mock_ydl(download=MagicMock(return_value=None)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl) as mk, + patch("os.listdir", return_value=[]), + ): + _download_and_parse( + url="https://youtube.com/watch?v=abc", + language_key="en", + source_label="auto", + proxy_url="http://proxy:8080", + ) + + opts = mk.call_args[0][0] + assert opts["proxy"] == "http://proxy:8080" + assert opts["writeautomaticsub"] is True + assert opts["writesubtitles"] is False + + def test_parses_downloaded_srt(self, tmp_storage): + """A real .srt written to the temp dir is read and parsed. + + ``_download_and_parse`` does ``import tempfile`` inside the function, + so we patch ``tempfile.TemporaryDirectory`` to hand back a directory we + pre-populated with a subtitle file matching the language key. + """ + import yt_dlp + from plugins.youtube_transcript_import import _download_and_parse + + sub = tmp_storage / "abc.en.srt" + sub.write_text("1\n00:00:01,000 --> 00:00:03,000\nReal cue\n", encoding="utf-8") + + class _TD: + def __enter__(self): + return str(tmp_storage) + + def __exit__(self, *a): + return False + + ydl = _mock_ydl(download=MagicMock(return_value=None)) + with ( + patch.object(yt_dlp, "YoutubeDL", return_value=ydl), + patch("tempfile.TemporaryDirectory", lambda: _TD()), + ): + pieces = _download_and_parse( + url="https://youtube.com/watch?v=abc", + language_key="en", + source_label="manual", + proxy_url=None, + ) + + assert len(pieces) == 1 + assert pieces[0]["text"] == "Real cue" + + +class TestGetParameters: + """``get_parameters`` contract.""" + + def test_includes_language_and_proxy(self): + """Exposes a language param (default 'en') and an advanced proxy_url.""" + params = YouTubeTranscriptImportPlugin().get_parameters() + names = [p.name for p in params] + assert names == ["language", "proxy_url"] + lang = next(p for p in params if p.name == "language") + assert lang.default == "en" diff --git a/library-manager/tests/unit/test_schemas.py b/library-manager/tests/unit/test_schemas.py new file mode 100644 index 000000000..ea78f56f4 --- /dev/null +++ b/library-manager/tests/unit/test_schemas.py @@ -0,0 +1,252 @@ +"""Unit tests for the Pydantic schemas under ``backend/schemas/``. + +Covers ``common``, ``libraries``, ``content``, and ``folders``: valid input +constructs, and boundary / bad-type / missing-field cases raise +``pydantic.ValidationError``. +""" + +from __future__ import annotations + +from datetime import datetime + +import pytest +from pydantic import ValidationError +from schemas.common import MessageResponse, PaginationParams +from schemas.content import ( + ContentItemDetail, + ContentItemSummary, + FileImportParams, + ImportAcceptedResponse, + UrlImportRequest, + YoutubeImportRequest, +) +from schemas.folders import ( + _FOLDER_NAME_MAX_LEN, + FolderCreateRequest, + FolderMoveRequest, + FolderRenameRequest, + ItemsMoveRequest, + _validate_folder_name, +) +from schemas.libraries import LibraryCreate, LibraryResponse, LibraryUpdate + + +class TestCommon: + """``PaginationParams`` bounds and ``MessageResponse``.""" + + def test_pagination_defaults(self): + """Defaults are limit=20, offset=0.""" + p = PaginationParams() + assert p.limit == 20 and p.offset == 0 + + def test_pagination_bounds_ok(self): + """Edge values within bounds are accepted (limit 1..100, offset >=0).""" + assert PaginationParams(limit=1, offset=0).limit == 1 + assert PaginationParams(limit=100, offset=5).limit == 100 + + @pytest.mark.parametrize( + "kwargs", + [{"limit": 0}, {"limit": 101}, {"offset": -1}, {"limit": "x"}], + ) + def test_pagination_invalid(self, kwargs): + """Out-of-range or wrong-typed pagination values are rejected.""" + with pytest.raises(ValidationError): + PaginationParams(**kwargs) + + def test_message_response(self): + """MessageResponse requires a string message.""" + assert MessageResponse(message="hi").message == "hi" + with pytest.raises(ValidationError): + MessageResponse() + + +class TestLibrarySchemas: + """LibraryCreate/Update/Response validation.""" + + def test_library_create_valid(self): + """A full valid body constructs successfully.""" + m = LibraryCreate(id="lib-1", organization_id="org-1", name="Docs") + assert m.name == "Docs" and m.import_config is None + + @pytest.mark.parametrize( + "kwargs", + [ + {"organization_id": "o", "name": "n"}, # missing id + {"id": "i", "name": "n"}, # missing org + {"id": "i", "organization_id": "o"}, # missing name + {"id": "i", "organization_id": "o", "name": ""}, # empty name + {"id": "i", "organization_id": "o", "name": "x" * 256}, # too long + ], + ) + def test_library_create_invalid(self, kwargs): + """Missing required fields or name length violations are rejected.""" + with pytest.raises(ValidationError): + LibraryCreate(**kwargs) + + def test_library_update_all_optional(self): + """LibraryUpdate accepts an empty body (partial update).""" + assert LibraryUpdate().name is None + + def test_library_update_name_bounds(self): + """LibraryUpdate enforces name length when provided.""" + with pytest.raises(ValidationError): + LibraryUpdate(name="") + with pytest.raises(ValidationError): + LibraryUpdate(name="x" * 256) + + def test_library_response_from_attrs(self): + """LibraryResponse builds with defaults and a datetime.""" + m = LibraryResponse( + id="l", organization_id="o", name="n", created_at=datetime.now() + ) + assert m.item_count == 0 + + +class TestContentImportSchemas: + """File/URL/YouTube import request validation.""" + + def test_file_import_valid(self): + """A valid file-import params body constructs.""" + m = FileImportParams(plugin_name="simple", title="Doc") + assert m.api_keys is None + + @pytest.mark.parametrize( + "kwargs", + [ + {"plugin_name": "", "title": "t"}, # empty plugin + {"plugin_name": "p", "title": ""}, # empty title + {"plugin_name": "p", "title": "x" * 501}, # title too long + {"plugin_name": "p"}, # missing title + ], + ) + def test_file_import_invalid(self, kwargs): + """Empty/oversize/missing fields are rejected.""" + with pytest.raises(ValidationError): + FileImportParams(**kwargs) + + def test_url_import_valid_and_default_plugin(self): + """URL import accepts http(s) and defaults plugin_name.""" + m = UrlImportRequest(url="https://example.com", title="Page") + assert m.plugin_name == "url_import" + + @pytest.mark.parametrize("url", ["ftp://x", "example.com", "", "javascript:1"]) + def test_url_import_rejects_non_http(self, url): + """URLs not starting with http:// or https:// are rejected.""" + with pytest.raises(ValidationError): + UrlImportRequest(url=url, title="t") + + def test_youtube_import_valid(self): + """A valid YouTube URL constructs with default language/plugin.""" + m = YoutubeImportRequest( + video_url="https://youtu.be/abc", title="Vid" + ) + assert m.language == "en" + assert m.plugin_name == "youtube_transcript_import" + m2 = YoutubeImportRequest( + video_url="https://www.youtube.com/watch?v=z", title="Vid" + ) + assert m2.video_url.endswith("v=z") + + @pytest.mark.parametrize( + "url", ["https://vimeo.com/1", "https://example.com", "not-a-url"] + ) + def test_youtube_import_rejects_non_youtube(self, url): + """Non-YouTube links are rejected by the validator.""" + with pytest.raises(ValidationError): + YoutubeImportRequest(video_url=url, title="t") + + def test_import_accepted_defaults(self): + """ImportAcceptedResponse defaults status to 'processing'.""" + m = ImportAcceptedResponse(item_id="i", job_id="j") + assert m.status == "processing" + + +class TestContentItemSchemas: + """ContentItemSummary / ContentItemDetail validation.""" + + def _summary_kwargs(self): + return { + "id": "i", + "title": "T", + "source_type": "file", + "import_plugin": "simple", + "status": "ready", + "created_at": datetime.now(), + "updated_at": datetime.now(), + } + + def test_summary_valid_with_defaults(self): + """Summary constructs with required fields and defaulted counts.""" + m = ContentItemSummary(**self._summary_kwargs()) + assert m.page_count == 0 and m.image_count == 0 + assert m.source_url is None + + def test_summary_missing_required(self): + """Omitting a required field raises ValidationError.""" + kwargs = self._summary_kwargs() + del kwargs["status"] + with pytest.raises(ValidationError): + ContentItemSummary(**kwargs) + + def test_detail_extends_summary(self): + """Detail adds optional metadata/permalink fields over the summary.""" + m = ContentItemDetail( + **self._summary_kwargs(), + metadata={"a": 1}, + permalink_base="/docs/x", + ) + assert m.metadata == {"a": 1} + assert m.permalink_base == "/docs/x" + + +class TestFolderSchemas: + """Folder create/rename/move and items-move validation.""" + + def test_validate_folder_name_over_length_branch(self): + """``_validate_folder_name`` rejects names longer than the cap. + + Exercised directly because the Field's ``max_length`` would otherwise + reject the value before the validator's length check is reached. + """ + with pytest.raises(ValueError, match="cannot exceed"): + _validate_folder_name("x" * (_FOLDER_NAME_MAX_LEN + 1)) + + def test_folder_create_valid_and_trims(self): + """A valid name is accepted and surrounding whitespace trimmed.""" + m = FolderCreateRequest(name=" Reports ") + assert m.name == "Reports", "name should be trimmed" + assert m.parent_folder_id is None + + @pytest.mark.parametrize( + "name", + ["", " ", "a/b", "a\\b", "a\x00b", "a\x01b", "x" * 129], + ) + def test_folder_create_invalid_names(self, name): + """Empty, separator, control-char, and oversize names are rejected.""" + with pytest.raises(ValidationError): + FolderCreateRequest(name=name) + + def test_folder_rename_valid_and_invalid(self): + """Rename applies the same name validation as create.""" + assert FolderRenameRequest(name=" Ok ").name == "Ok" + with pytest.raises(ValidationError): + FolderRenameRequest(name="bad/name") + + def test_folder_move_optional_parent(self): + """FolderMoveRequest allows a null parent (move to root).""" + assert FolderMoveRequest().parent_folder_id is None + assert FolderMoveRequest(parent_folder_id="f1").parent_folder_id == "f1" + + def test_items_move_valid(self): + """ItemsMoveRequest requires 1..500 item ids.""" + m = ItemsMoveRequest(item_ids=["a", "b"], folder_id="f") + assert m.item_ids == ["a", "b"] + + @pytest.mark.parametrize( + "item_ids", + [[], ["x"] * 501], + ) + def test_items_move_invalid_length(self, item_ids): + """Empty or oversize item_ids lists are rejected.""" + with pytest.raises(ValidationError): + ItemsMoveRequest(item_ids=item_ids) diff --git a/library-manager/tests/unit/test_service_content.py b/library-manager/tests/unit/test_service_content.py new file mode 100644 index 000000000..f17029103 --- /dev/null +++ b/library-manager/tests/unit/test_service_content.py @@ -0,0 +1,495 @@ +"""Unit tests for ``services.content_service`` — direct function calls.""" + +from __future__ import annotations + +import json + +import pytest +from _helpers import unique_id +from database.models import ContentItem +from plugins.base import ExtractedImage, PageContent +from services import content_service +from services.library_service import create_library + + +@pytest.fixture +def storage(tmp_storage, monkeypatch): + """Point content_service at a throwaway CONTENT_DIR.""" + monkeypatch.setattr(content_service, "CONTENT_DIR", tmp_storage) + return tmp_storage + + +def _write_sample(org, lib, item, *, pages=True, images=True, original=None): + """Write a structured content item and return its base path.""" + page_list = ( + [PageContent(page_number=1, text="page one"), PageContent(page_number=2, text="page two")] + if pages + else [] + ) + image_list = ( + [ + ExtractedImage( + filename="img_001.png", data=b"\x89PNG-bytes", page_number=1, description="d" + ) + ] + if images + else [] + ) + return content_service.write_structured_content( + item_id=item, + library_id=lib, + organization_id=org, + title="Doc Title", + full_text="# Full\n\nbody", + pages=page_list, + images=image_list, + item_metadata={"language": "en", "character_count": 12}, + source_ref={"type": "file", "original_filename": "src.txt"}, + original_file_path=original, + original_filename="src.txt" if original else None, + ) + + +# --------------------------------------------------------------------------- +# write_structured_content + detect_capabilities +# --------------------------------------------------------------------------- + + +def test_write_structured_content_creates_full_tree(storage): + """write_structured_content writes metadata, source_ref, pages, images, full.md.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + base = _write_sample(org, lib, item) + assert (base / "metadata.json").is_file() + assert (base / "source_ref.json").is_file() + assert (base / "content" / "full.md").read_text() == "# Full\n\nbody" + assert (base / "content" / "pages" / "page_001.md").read_text() == "page one" + assert (base / "content" / "pages" / "page_002.md").read_text() == "page two" + assert (base / "content" / "images" / "img_001.png").read_bytes() == b"\x89PNG-bytes" + + +def test_write_structured_content_copies_original(storage, tmp_storage): + """write_structured_content copies the original file into original/.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + src = tmp_storage / "upload.txt" + src.write_text("raw original") + base = _write_sample(org, lib, item, original=src) + assert (base / "original" / "src.txt").read_text() == "raw original" + meta = json.loads((base / "metadata.json").read_text()) + assert meta["permalinks"]["original"].endswith("/original/src.txt") + + +def test_write_structured_content_metadata_has_capabilities(storage): + """metadata.json reflects the on-disk capabilities (text, pages, images).""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + base = _write_sample(org, lib, item) + meta = json.loads((base / "metadata.json").read_text()) + assert sorted(meta["capabilities"]) == ["images", "pages", "text"] + assert meta["title"] == "Doc Title" + assert meta["language"] == "en" + + +def test_write_structured_content_permission_error_humanized(storage, monkeypatch): + """A PermissionError on mkdir becomes a user-actionable RuntimeError.""" + from pathlib import Path + + def _deny(self, *args, **kwargs): + raise PermissionError("denied") + + monkeypatch.setattr(Path, "mkdir", _deny) + with pytest.raises(RuntimeError, match="not writable"): + _write_sample(unique_id("org"), unique_id("lib"), unique_id("item")) + + +def test_write_structured_content_oserror_humanized(storage, monkeypatch): + """A generic OSError on mkdir becomes a storage-unavailable RuntimeError.""" + from pathlib import Path + + def _fail(self, *args, **kwargs): + raise OSError("disk gone") + + monkeypatch.setattr(Path, "mkdir", _fail) + with pytest.raises(RuntimeError, match="storage is"): + _write_sample(unique_id("org"), unique_id("lib"), unique_id("item")) + + +def test_detect_capabilities_text_only(storage): + """detect_capabilities returns only ``text`` when there are no pages/images.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + base = _write_sample(org, lib, item, pages=False, images=False) + assert content_service.detect_capabilities(base) == ["text"] + + +def test_detect_capabilities_empty_when_no_content_dir(tmp_storage): + """detect_capabilities returns [] when no content/ directory exists.""" + assert content_service.detect_capabilities(tmp_storage / "ghost") == [] + + +def test_detect_capabilities_ignores_non_image_files(storage, tmp_storage): + """detect_capabilities does not advertise images for non-image files.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + base = _write_sample(org, lib, item, pages=False, images=False) + images_dir = base / "content" / "images" + images_dir.mkdir() + (images_dir / "notes.txt").write_text("not an image") + assert "images" not in content_service.detect_capabilities(base) + + +# --------------------------------------------------------------------------- +# Readers +# --------------------------------------------------------------------------- + + +def test_read_full_markdown(storage): + """read_full_markdown returns the full.md contents.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + assert content_service.read_full_markdown(org, lib, item) == "# Full\n\nbody" + + +def test_read_full_markdown_missing_returns_none(storage): + """read_full_markdown returns None when the file is absent.""" + assert content_service.read_full_markdown("o", "l", "missing") is None + + +def test_read_page_markdown_with_and_without_extension(storage): + """read_page_markdown accepts the page name with or without ``.md``.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + assert content_service.read_page_markdown(org, lib, item, "page_001.md") == "page one" + assert content_service.read_page_markdown(org, lib, item, "page_002") == "page two" + + +def test_read_page_markdown_missing_returns_none(storage): + """read_page_markdown returns None for a non-existent page.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + assert content_service.read_page_markdown(org, lib, item, "page_999") is None + + +def test_read_metadata_json(storage): + """read_metadata_json parses and returns metadata.json.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + meta = content_service.read_metadata_json(org, lib, item) + assert meta["item_id"] == item + + +def test_read_metadata_json_missing_returns_none(storage): + """read_metadata_json returns None when metadata.json is absent.""" + assert content_service.read_metadata_json("o", "l", "missing") is None + + +def test_read_source_ref(storage): + """read_source_ref parses and returns source_ref.json.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + ref = content_service.read_source_ref(org, lib, item) + assert ref["original_filename"] == "src.txt" + + +def test_read_source_ref_missing_returns_none(storage): + """read_source_ref returns None when source_ref.json is absent.""" + assert content_service.read_source_ref("o", "l", "missing") is None + + +def test_list_pages(storage): + """list_pages returns sorted page filenames.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + assert content_service.list_pages(org, lib, item) == ["page_001.md", "page_002.md"] + + +def test_list_pages_empty_when_no_dir(storage): + """list_pages returns [] when there is no pages directory.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item, pages=False) + assert content_service.list_pages(org, lib, item) == [] + + +def test_list_images(storage): + """list_images returns sorted image filenames.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + assert content_service.list_images(org, lib, item) == ["img_001.png"] + + +def test_list_images_empty_when_no_dir(storage): + """list_images returns [] when there is no images directory.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item, images=False) + assert content_service.list_images(org, lib, item) == [] + + +def test_get_image_path_resolves(storage): + """get_image_path returns the path to an existing image file.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + path = content_service.get_image_path(org, lib, item, "img_001.png") + assert path is not None and path.is_file() + + +def test_get_image_path_missing_returns_none(storage): + """get_image_path returns None for a non-existent image.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + assert content_service.get_image_path(org, lib, item, "nope.png") is None + + +def test_get_original_path_resolves(storage, tmp_storage): + """get_original_path returns the path to the stored original file.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + src = tmp_storage / "upload.txt" + src.write_text("raw") + _write_sample(org, lib, item, original=src) + path = content_service.get_original_path(org, lib, item, "src.txt") + assert path is not None and path.is_file() + + +def test_get_original_path_missing_returns_none(storage): + """get_original_path returns None when no original file is stored.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + assert content_service.get_original_path(org, lib, item, "src.txt") is None + + +# --------------------------------------------------------------------------- +# Path-traversal guards +# --------------------------------------------------------------------------- + + +def test_safe_resolve_blocks_traversal(tmp_storage): + """_safe_resolve returns None when the path escapes the expected parent.""" + base = tmp_storage / "item" + base.mkdir() + escaped = base / ".." / ".." / "etc" / "passwd" + assert content_service._safe_resolve(escaped, base) is None + + +def test_safe_resolve_allows_inside(tmp_storage): + """_safe_resolve returns the resolved path when it stays inside the parent.""" + base = tmp_storage / "item" + base.mkdir() + inside = base / "content" / "full.md" + resolved = content_service._safe_resolve(inside, base) + assert resolved is not None + + +def test_read_page_markdown_traversal_sanitized(storage): + """read_page_markdown sanitizes ``../`` page names to a single component.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + # Sanitization strips directory components, so it cannot escape; returns None. + assert content_service.read_page_markdown(org, lib, item, "../../../etc/passwd") is None + + +def test_get_image_path_traversal_sanitized(storage): + """get_image_path sanitizes traversal attempts in the image name.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + _write_sample(org, lib, item) + assert content_service.get_image_path(org, lib, item, "../../secret.png") is None + + +# --------------------------------------------------------------------------- +# _sanitize_filename +# --------------------------------------------------------------------------- + + +def test_sanitize_filename_strips_directories(): + """_sanitize_filename keeps only the final path component.""" + assert content_service._sanitize_filename("../../etc/passwd") == "passwd" + assert content_service._sanitize_filename("/abs/path/file.png") == "file.png" + + +def test_sanitize_filename_strips_null_bytes(): + """_sanitize_filename removes embedded null bytes.""" + assert content_service._sanitize_filename("a\x00b.txt") == "ab.txt" + + +def test_sanitize_filename_fallback_for_empty(): + """_sanitize_filename falls back to ``unnamed`` for empty/dot names.""" + assert content_service._sanitize_filename("") == "unnamed" + assert content_service._sanitize_filename(".") == "unnamed" + assert content_service._sanitize_filename("..") == "unnamed" + + +# --------------------------------------------------------------------------- +# delete_content_item +# --------------------------------------------------------------------------- + + +def test_delete_content_item_removes_row_and_disk(storage, db_session): + """delete_content_item removes the DB row and the on-disk directory.""" + org, lib, item_id = unique_id("org"), unique_id("lib"), unique_id("item") + create_library(db_session, lib, org, "Lib") + base = _write_sample(org, lib, item_id) + db_session.add( + ContentItem( + id=item_id, + library_id=lib, + organization_id=org, + title="t", + source_type="file", + base_path=str(base), + permalink_base="/docs/x", + import_plugin="simple_import", + status="ready", + ) + ) + db_session.commit() + assert content_service.delete_content_item(db_session, org, lib, item_id) is True + assert content_service.get_content_item(db_session, item_id) is None + assert not base.exists() + + +def test_delete_content_item_missing_returns_false(storage, db_session): + """delete_content_item returns False for an unknown item.""" + assert content_service.delete_content_item(db_session, "o", "l", unique_id("x")) is False + + +def test_delete_content_item_no_disk_dir(storage, db_session): + """delete_content_item still succeeds when there is no on-disk directory.""" + org, lib, item_id = unique_id("org"), unique_id("lib"), unique_id("item") + create_library(db_session, lib, org, "Lib") + db_session.add( + ContentItem( + id=item_id, + library_id=lib, + organization_id=org, + title="t", + source_type="file", + base_path="/tmp/none", + permalink_base="/docs/x", + import_plugin="simple_import", + status="ready", + ) + ) + db_session.commit() + # No directory was written for this item. + assert content_service.delete_content_item(db_session, org, lib, item_id) is True + + +def test_detect_capabilities_pages_without_full_md(storage): + """detect_capabilities reports ``pages`` even when full.md is absent.""" + org, lib, item = unique_id("org"), unique_id("lib"), unique_id("item") + base = _write_sample(org, lib, item, pages=True, images=False) + (base / "content" / "full.md").unlink() + caps = content_service.detect_capabilities(base) + assert "pages" in caps + assert "text" not in caps + + +# --------------------------------------------------------------------------- +# get_content_item / item_to_summary / list_content_items +# --------------------------------------------------------------------------- + + +def _insert_item(db, lib, org, **overrides): + """Insert a content item directly into the DB.""" + overrides.setdefault("status", "ready") + item = ContentItem( + id=unique_id("item"), + library_id=lib, + organization_id=org, + title="T", + source_type="file", + base_path="/tmp/x", + permalink_base="/docs/x", + import_plugin="simple_import", + **overrides, + ) + db.add(item) + db.commit() + return item + + +def test_get_content_item_returns_row(db_session): + """get_content_item returns the matching row.""" + org, lib = unique_id("org"), unique_id("lib") + create_library(db_session, lib, org, "Lib") + item = _insert_item(db_session, lib, org) + assert content_service.get_content_item(db_session, item.id).id == item.id + + +def test_get_content_item_missing_returns_none(db_session): + """get_content_item returns None for an unknown id.""" + assert content_service.get_content_item(db_session, unique_id("x")) is None + + +def test_item_to_summary_shape(): + """item_to_summary exposes the canonical API summary keys.""" + item = ContentItem( + id="i1", + library_id="l1", + organization_id="o1", + title="T", + source_type="url", + source_url="http://x", + original_filename="f.txt", + content_type="text/plain", + file_size=10, + import_plugin="url_import", + status="ready", + error_message=None, + page_count=3, + image_count=2, + folder_id="fold1", + base_path="/x", + permalink_base="/docs/x", + ) + summary = content_service.item_to_summary(item) + assert summary["id"] == "i1" + assert summary["page_count"] == 3 + assert summary["folder_id"] == "fold1" + assert summary["source_url"] == "http://x" + assert set(summary) == { + "id", "title", "source_type", "source_url", "original_filename", + "content_type", "file_size", "import_plugin", "status", "error_message", + "page_count", "image_count", "folder_id", "created_at", "updated_at", + } + + +def test_list_content_items_filters_by_library(db_session): + """list_content_items returns only the requested library's items.""" + org, lib = unique_id("org"), unique_id("lib") + other = unique_id("lib") + create_library(db_session, lib, org, "Lib") + create_library(db_session, other, org, "Other") + _insert_item(db_session, lib, org) + _insert_item(db_session, lib, org) + _insert_item(db_session, other, org) + items, total = content_service.list_content_items(db_session, lib) + assert total == 2 + assert all(i.library_id == lib for i in items) + + +def test_list_content_items_status_filter(db_session): + """list_content_items honors the status filter.""" + org, lib = unique_id("org"), unique_id("lib") + create_library(db_session, lib, org, "Lib") + _insert_item(db_session, lib, org, status="ready") + _insert_item(db_session, lib, org, status="failed") + items, total = content_service.list_content_items(db_session, lib, status_filter="failed") + assert total == 1 + assert items[0].status == "failed" + + +def test_list_content_items_ids_filter_and_pagination(db_session): + """list_content_items supports an id filter plus limit/offset.""" + org, lib = unique_id("org"), unique_id("lib") + create_library(db_session, lib, org, "Lib") + a = _insert_item(db_session, lib, org) + b = _insert_item(db_session, lib, org) + _insert_item(db_session, lib, org) + items, total = content_service.list_content_items( + db_session, lib, ids_filter=[a.id, b.id] + ) + assert total == 2 + page, total2 = content_service.list_content_items(db_session, lib, limit=1, offset=0) + assert total2 == 3 + assert len(page) == 1 + + +def test_get_item_base_path_layout(storage, tmp_storage): + """get_item_base_path composes {CONTENT_DIR}/{org}/{lib}/{item}.""" + path = content_service.get_item_base_path("o", "l", "i") + assert path == tmp_storage / "o" / "l" / "i" diff --git a/library-manager/tests/unit/test_service_export.py b/library-manager/tests/unit/test_service_export.py new file mode 100644 index 000000000..39a7de36e --- /dev/null +++ b/library-manager/tests/unit/test_service_export.py @@ -0,0 +1,291 @@ +"""Unit tests for ``services.export_service`` — direct function calls.""" + +from __future__ import annotations + +import json +import zipfile +from io import BytesIO + +import pytest +from _helpers import unique_id +from database.models import ContentItem +from services import export_service +from services.library_service import create_library, get_library + + +@pytest.fixture +def storage(tmp_storage, monkeypatch): + """Point export_service at a throwaway CONTENT_DIR.""" + monkeypatch.setattr(export_service, "CONTENT_DIR", tmp_storage) + return tmp_storage + + +def _write_item_dir(base, org, lib, item_id): + """Materialize a minimal structured-content directory for an item.""" + item_dir = base / org / lib / item_id + (item_dir / "content" / "pages").mkdir(parents=True) + (item_dir / "content" / "images").mkdir(parents=True) + (item_dir / "original").mkdir(parents=True) + (item_dir / "content" / "full.md").write_text("# body") + (item_dir / "content" / "pages" / "page_001.md").write_text("p1") + (item_dir / "content" / "images" / "img_001.png").write_bytes(b"PNG") + (item_dir / "original" / "src.txt").write_text("orig") + (item_dir / "metadata.json").write_text(json.dumps({"item_id": item_id})) + return item_dir + + +def _insert_item(db, lib, org, item_id, *, status="ready", **overrides): + """Insert a content item row for export.""" + item = ContentItem( + id=item_id, + library_id=lib, + organization_id=org, + title="Doc", + source_type="file", + original_filename="src.txt", + content_type="text/plain", + import_plugin="simple_import", + import_params=json.dumps({"x": 1}), + metadata_=json.dumps({"page_count": 1, "language": "en"}), + base_path="/tmp/x", + permalink_base="/docs/x", + status=status, + **overrides, + ) + db.add(item) + db.commit() + return item + + +# --------------------------------------------------------------------------- +# export_library_zip +# --------------------------------------------------------------------------- + + +def test_export_library_zip_manifest_and_content(storage, db_session): + """export_library_zip writes a v1.0 manifest plus each ready item's files.""" + org, lib = unique_id("org"), unique_id("lib") + create_library(db_session, lib, org, "Lib") + item_id = unique_id("item") + _write_item_dir(storage, org, lib, item_id) + _insert_item(db_session, lib, org, item_id) + + buf = export_service.export_library_zip( + db_session, lib, org, "Lib", {"plugin": "simple_import"}, exported_by="me@x" + ) + with zipfile.ZipFile(buf) as zf: + names = zf.namelist() + manifest = json.loads(zf.read("manifest.json")) + assert manifest["format_version"] == "1.0" + assert manifest["type"] == "library_export" + assert manifest["exported_by"] == "me@x" + assert manifest["items"][0]["id"] == item_id + assert f"content/{item_id}/content/full.md" in names + assert f"content/{item_id}/metadata.json" in names + + +def test_export_library_zip_excludes_non_ready(storage, db_session): + """export_library_zip omits items whose status is not ready.""" + org, lib = unique_id("org"), unique_id("lib") + create_library(db_session, lib, org, "Lib") + ready_id = unique_id("item") + pending_id = unique_id("item") + _write_item_dir(storage, org, lib, ready_id) + _write_item_dir(storage, org, lib, pending_id) + _insert_item(db_session, lib, org, ready_id, status="ready") + _insert_item(db_session, lib, org, pending_id, status="pending") + + buf = export_service.export_library_zip(db_session, lib, org, "Lib", None) + with zipfile.ZipFile(buf) as zf: + manifest = json.loads(zf.read("manifest.json")) + names = zf.namelist() + ids = {i["id"] for i in manifest["items"]} + assert ids == {ready_id} + assert not any(pending_id in n for n in names) + + +def test_export_library_zip_missing_dir_still_lists_item(storage, db_session): + """An item with no on-disk dir is still listed in the manifest.""" + org, lib = unique_id("org"), unique_id("lib") + create_library(db_session, lib, org, "Lib") + item_id = unique_id("item") + _insert_item(db_session, lib, org, item_id) # no dir written + buf = export_service.export_library_zip(db_session, lib, org, "Lib", None) + with zipfile.ZipFile(buf) as zf: + manifest = json.loads(zf.read("manifest.json")) + assert manifest["items"][0]["id"] == item_id + + +# --------------------------------------------------------------------------- +# import_library_zip — round trip +# --------------------------------------------------------------------------- + + +def _build_zip(manifest, files): + """Build an in-memory ZIP with manifest.json and the given arcname->bytes.""" + buf = BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr("manifest.json", json.dumps(manifest)) + for arcname, data in files.items(): + zf.writestr(arcname, data) + return buf.getvalue() + + +def test_import_library_zip_round_trip(storage, db_session): + """import_library_zip recreates a library with new IDs and ready items.""" + org, lib = unique_id("org"), unique_id("lib") + create_library(db_session, lib, org, "Export Source") + item_id = unique_id("item") + _write_item_dir(storage, org, lib, item_id) + _insert_item(db_session, lib, org, item_id) + + zip_bytes = export_service.export_library_zip( + db_session, lib, org, "Export Source", {"plugin": "simple_import"} + ).getvalue() + + target_org = unique_id("org") + result = export_service.import_library_zip(db_session, zip_bytes, target_org) + assert result["item_count"] == 1 + assert result["library_id"] != lib # New library ID generated. + assert result["library_name"] == "Export Source" + + new_lib = get_library(db_session, result["library_id"]) + assert new_lib.organization_id == target_org + new_items = ( + db_session.query(ContentItem) + .filter(ContentItem.library_id == result["library_id"]) + .all() + ) + assert len(new_items) == 1 + new_item = new_items[0] + assert new_item.id != item_id # New item ID generated. + assert new_item.status == "ready" + assert new_item.page_count == 1 + assert new_item.image_count == 1 + # Permalink regenerated under the new org/lib/item. + assert new_item.permalink_base.endswith(f"/{target_org}/{result['library_id']}/{new_item.id}") + # Files landed on disk under the new path. + new_dir = storage / target_org / result["library_id"] / new_item.id + assert (new_dir / "content" / "full.md").read_text() == "# body" + meta = json.loads((new_dir / "metadata.json").read_text()) + assert meta["item_id"] == new_item.id + assert meta["permalinks"]["original"].endswith("/original/src.txt") + assert meta["page_count"] == 1 + assert meta["language"] == "en" + + +def test_import_library_zip_bad_zip_raises(db_session): + """import_library_zip raises ValueError on a corrupt ZIP.""" + with pytest.raises(ValueError, match="Invalid ZIP"): + export_service.import_library_zip(db_session, b"not a zip", unique_id("org")) + + +def test_import_library_zip_missing_manifest_raises(db_session): + """import_library_zip raises ValueError when manifest.json is absent.""" + buf = BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr("content/foo.txt", b"x") + with pytest.raises(ValueError, match="missing manifest"): + export_service.import_library_zip(db_session, buf.getvalue(), unique_id("org")) + + +def test_import_library_zip_bad_format_version_raises(db_session): + """import_library_zip rejects an unsupported manifest format_version.""" + manifest = {"format_version": "2.0", "type": "library_export", "library": {}, "items": []} + with pytest.raises(ValueError, match="Unsupported manifest version"): + export_service.import_library_zip( + db_session, _build_zip(manifest, {}), unique_id("org") + ) + + +def test_import_library_zip_bad_type_raises(db_session): + """import_library_zip rejects an unexpected manifest type.""" + manifest = {"format_version": "1.0", "type": "not_a_library", "library": {}, "items": []} + with pytest.raises(ValueError, match="Unexpected manifest type"): + export_service.import_library_zip( + db_session, _build_zip(manifest, {}), unique_id("org") + ) + + +def test_import_library_zip_default_name_when_missing(storage, db_session): + """import_library_zip synthesizes a name when the manifest omits one.""" + manifest = { + "format_version": "1.0", + "type": "library_export", + "library": {}, + "items": [], + } + result = export_service.import_library_zip( + db_session, _build_zip(manifest, {}), unique_id("org") + ) + assert result["library_name"].startswith("Imported Library ") + assert result["item_count"] == 0 + + +def test_import_library_zip_blocks_zip_slip(storage, db_session): + """import_library_zip skips entries that try to escape via ``../``.""" + item_id = "evil-item" + manifest = { + "format_version": "1.0", + "type": "library_export", + "library": {"name": "Slip"}, + "items": [{"id": item_id, "title": "Evil"}], + } + files = { + f"content/{item_id}/content/full.md": b"safe", + f"content/{item_id}/../../../escape.txt": b"pwned", + } + target_org = unique_id("org") + result = export_service.import_library_zip( + db_session, _build_zip(manifest, files), target_org + ) + new_dir = storage / target_org / result["library_id"] + # The legitimate file landed; the escape file did not appear anywhere. + assert any(p.name == "full.md" for p in new_dir.rglob("*")) + assert not (storage / "escape.txt").exists() + assert not (storage.parent / "escape.txt").exists() + + +# --------------------------------------------------------------------------- +# _regenerate_metadata +# --------------------------------------------------------------------------- + + +def test_regenerate_metadata_builds_permalinks(tmp_storage): + """_regenerate_metadata rebuilds permalinks and merges manifest extras.""" + item_dir = tmp_storage / "item" + (item_dir / "content" / "pages").mkdir(parents=True) + (item_dir / "content" / "images").mkdir(parents=True) + (item_dir / "original").mkdir(parents=True) + (item_dir / "content" / "pages" / "page_001.md").write_text("p") + (item_dir / "content" / "images" / "i.png").write_bytes(b"x") + (item_dir / "original" / "o.txt").write_text("o") + + manifest = { + "title": "Title", + "source_type": "file", + "original_filename": "o.txt", + "content_type": "text/plain", + "import_plugin": "simple_import", + "metadata": {"page_count": 1, "language": "es", "description": "d"}, + } + export_service._regenerate_metadata(item_dir, "new-id", "/docs/o/l/new-id", manifest) + meta = json.loads((item_dir / "metadata.json").read_text()) + assert meta["item_id"] == "new-id" + assert meta["permalinks"]["full_markdown"] == "/docs/o/l/new-id/content/full.md" + assert meta["permalinks"]["pages"] == ["/docs/o/l/new-id/content/pages/page_001.md"] + assert meta["permalinks"]["images"] == ["/docs/o/l/new-id/content/images/i.png"] + assert meta["permalinks"]["original"] == "/docs/o/l/new-id/original/o.txt" + assert meta["language"] == "es" + assert meta["description"] == "d" + + +def test_regenerate_metadata_no_subdirs(tmp_storage): + """_regenerate_metadata handles items with no pages/images/original.""" + item_dir = tmp_storage / "bare" + item_dir.mkdir() + export_service._regenerate_metadata(item_dir, "id2", "/docs/o/l/id2", {"title": "T"}) + meta = json.loads((item_dir / "metadata.json").read_text()) + assert meta["permalinks"]["pages"] == [] + assert meta["permalinks"]["images"] == [] + assert meta["permalinks"]["original"] is None diff --git a/library-manager/tests/unit/test_service_folder.py b/library-manager/tests/unit/test_service_folder.py new file mode 100644 index 000000000..e1c070010 --- /dev/null +++ b/library-manager/tests/unit/test_service_folder.py @@ -0,0 +1,510 @@ +"""Unit tests for ``services.folder_service`` — direct function calls.""" + +from __future__ import annotations + +import pytest +from _helpers import unique_id +from database.models import ContentItem, Library +from services import folder_service +from services.folder_service import ( + FolderConflictError, + FolderCycleError, + FolderNotFoundError, + FolderValidationError, +) +from services.library_service import create_library + + +@pytest.fixture +def lib(db_session): + """Create a fresh library and return its id.""" + org_id = unique_id("org") + lib_id = unique_id("lib") + create_library(db_session, lib_id, org_id, "Folder Lib") + return lib_id + + +def _make_item(db, library_id, folder_id=None): + """Insert a content item, optionally inside a folder.""" + lib = db.query(Library).filter(Library.id == library_id).first() + org_id = lib.organization_id + item = ContentItem( + id=unique_id("item"), + library_id=library_id, + organization_id=org_id, + folder_id=folder_id, + title="t", + source_type="file", + base_path="/tmp/none", + permalink_base="/docs/x", + import_plugin="simple_import", + status="ready", + ) + db.add(item) + db.commit() + return item + + +# --------------------------------------------------------------------------- +# Error subclasses +# --------------------------------------------------------------------------- + + +def test_error_status_codes(): + """Each FolderError subclass carries the documented HTTP status code.""" + assert FolderNotFoundError().status_code == 404 + assert FolderConflictError().status_code == 409 + assert FolderCycleError().status_code == 400 + assert FolderValidationError().status_code == 400 + assert folder_service.FolderError().status_code == 400 + + +# --------------------------------------------------------------------------- +# create_folder +# --------------------------------------------------------------------------- + + +def test_create_folder_at_root(db_session, lib): + """create_folder creates a root-level folder with no parent.""" + f = folder_service.create_folder( + db_session, library_id=lib, name="Root", parent_folder_id=None + ) + assert f.parent_folder_id is None + assert f.library_id == lib + assert f.name == "Root" + + +def test_create_folder_nested(db_session, lib): + """create_folder nests a folder under an existing parent.""" + parent = folder_service.create_folder( + db_session, library_id=lib, name="Parent", parent_folder_id=None + ) + child = folder_service.create_folder( + db_session, library_id=lib, name="Child", parent_folder_id=parent.id + ) + assert child.parent_folder_id == parent.id + + +def test_create_folder_unknown_library_raises_notfound(db_session): + """create_folder raises FolderNotFoundError when the library is missing.""" + with pytest.raises(FolderNotFoundError): + folder_service.create_folder( + db_session, library_id=unique_id("x"), name="A", parent_folder_id=None + ) + + +def test_create_folder_unknown_parent_raises_notfound(db_session, lib): + """create_folder raises FolderNotFoundError when the parent is missing.""" + with pytest.raises(FolderNotFoundError): + folder_service.create_folder( + db_session, library_id=lib, name="A", parent_folder_id=unique_id("x") + ) + + +def test_create_folder_parent_in_other_library_raises_validation(db_session, lib): + """create_folder rejects a parent that lives in a different library.""" + other_lib = unique_id("lib") + create_library(db_session, other_lib, unique_id("org"), "Other") + foreign = folder_service.create_folder( + db_session, library_id=other_lib, name="P", parent_folder_id=None + ) + with pytest.raises(FolderValidationError): + folder_service.create_folder( + db_session, library_id=lib, name="A", parent_folder_id=foreign.id + ) + + +def test_create_folder_duplicate_sibling_raises_conflict(db_session, lib): + """create_folder raises FolderConflictError (409) on a duplicate sibling.""" + folder_service.create_folder( + db_session, library_id=lib, name="Same", parent_folder_id=None + ) + db_session.rollback() # discard nothing; just ensure clean state + with pytest.raises(FolderConflictError) as exc: + folder_service.create_folder( + db_session, library_id=lib, name="Same", parent_folder_id=None + ) + assert exc.value.status_code == 409 + + +def test_create_folder_same_name_different_parent_ok(db_session, lib): + """Identical names are allowed under different parents.""" + p1 = folder_service.create_folder( + db_session, library_id=lib, name="P1", parent_folder_id=None + ) + p2 = folder_service.create_folder( + db_session, library_id=lib, name="P2", parent_folder_id=None + ) + folder_service.create_folder( + db_session, library_id=lib, name="Dup", parent_folder_id=p1.id + ) + # No conflict — different parent. + folder_service.create_folder( + db_session, library_id=lib, name="Dup", parent_folder_id=p2.id + ) + + +# --------------------------------------------------------------------------- +# rename_folder +# --------------------------------------------------------------------------- + + +def test_rename_folder_changes_name(db_session, lib): + """rename_folder updates the folder's name.""" + f = folder_service.create_folder( + db_session, library_id=lib, name="Old", parent_folder_id=None + ) + renamed = folder_service.rename_folder(db_session, f.id, "New") + assert renamed.name == "New" + + +def test_rename_folder_noop_same_name(db_session, lib): + """rename_folder is a no-op when the name is unchanged.""" + f = folder_service.create_folder( + db_session, library_id=lib, name="Keep", parent_folder_id=None + ) + result = folder_service.rename_folder(db_session, f.id, "Keep") + assert result.name == "Keep" + + +def test_rename_folder_missing_raises_notfound(db_session): + """rename_folder raises FolderNotFoundError for an unknown folder.""" + with pytest.raises(FolderNotFoundError): + folder_service.rename_folder(db_session, unique_id("x"), "X") + + +def test_rename_folder_conflict_raises(db_session, lib): + """rename_folder raises FolderConflictError when colliding with a sibling.""" + folder_service.create_folder( + db_session, library_id=lib, name="A", parent_folder_id=None + ) + b = folder_service.create_folder( + db_session, library_id=lib, name="B", parent_folder_id=None + ) + with pytest.raises(FolderConflictError): + folder_service.rename_folder(db_session, b.id, "A") + + +# --------------------------------------------------------------------------- +# move_folder +# --------------------------------------------------------------------------- + + +def test_move_folder_reparents(db_session, lib): + """move_folder re-parents a folder under a new parent.""" + p = folder_service.create_folder( + db_session, library_id=lib, name="P", parent_folder_id=None + ) + f = folder_service.create_folder( + db_session, library_id=lib, name="F", parent_folder_id=None + ) + moved = folder_service.move_folder(db_session, f.id, p.id) + assert moved.parent_folder_id == p.id + + +def test_move_folder_to_root(db_session, lib): + """move_folder with None parent moves the folder to the library root.""" + p = folder_service.create_folder( + db_session, library_id=lib, name="P", parent_folder_id=None + ) + f = folder_service.create_folder( + db_session, library_id=lib, name="F", parent_folder_id=p.id + ) + moved = folder_service.move_folder(db_session, f.id, None) + assert moved.parent_folder_id is None + + +def test_move_folder_into_self_raises_cycle(db_session, lib): + """move_folder raises FolderCycleError when moving a folder into itself.""" + f = folder_service.create_folder( + db_session, library_id=lib, name="Self", parent_folder_id=None + ) + with pytest.raises(FolderCycleError): + folder_service.move_folder(db_session, f.id, f.id) + + +def test_move_folder_into_descendant_raises_cycle(db_session, lib): + """move_folder raises FolderCycleError when moving into a descendant.""" + a = folder_service.create_folder( + db_session, library_id=lib, name="A", parent_folder_id=None + ) + b = folder_service.create_folder( + db_session, library_id=lib, name="B", parent_folder_id=a.id + ) + c = folder_service.create_folder( + db_session, library_id=lib, name="C", parent_folder_id=b.id + ) + with pytest.raises(FolderCycleError): + folder_service.move_folder(db_session, a.id, c.id) + + +def test_move_folder_missing_raises_notfound(db_session): + """move_folder raises FolderNotFoundError for an unknown folder.""" + with pytest.raises(FolderNotFoundError): + folder_service.move_folder(db_session, unique_id("x"), None) + + +def test_move_folder_unknown_destination_raises_notfound(db_session, lib): + """move_folder raises FolderNotFoundError for an unknown destination.""" + f = folder_service.create_folder( + db_session, library_id=lib, name="F", parent_folder_id=None + ) + with pytest.raises(FolderNotFoundError): + folder_service.move_folder(db_session, f.id, unique_id("x")) + + +def test_move_folder_cross_library_raises_validation(db_session, lib): + """move_folder rejects a destination in a different library.""" + other_lib = unique_id("lib") + create_library(db_session, other_lib, unique_id("org"), "Other") + dest = folder_service.create_folder( + db_session, library_id=other_lib, name="Dest", parent_folder_id=None + ) + f = folder_service.create_folder( + db_session, library_id=lib, name="F", parent_folder_id=None + ) + with pytest.raises(FolderValidationError): + folder_service.move_folder(db_session, f.id, dest.id) + + +def test_move_folder_name_collision_raises_conflict(db_session, lib): + """move_folder raises FolderConflictError when the destination has the name.""" + dest = folder_service.create_folder( + db_session, library_id=lib, name="Dest", parent_folder_id=None + ) + folder_service.create_folder( + db_session, library_id=lib, name="Dup", parent_folder_id=dest.id + ) + f = folder_service.create_folder( + db_session, library_id=lib, name="Dup", parent_folder_id=None + ) + with pytest.raises(FolderConflictError): + folder_service.move_folder(db_session, f.id, dest.id) + + +# --------------------------------------------------------------------------- +# delete_folder +# --------------------------------------------------------------------------- + + +def test_delete_folder_reparents_items_up(db_session, lib): + """delete_folder re-homes items up to the deleted folder's parent.""" + parent = folder_service.create_folder( + db_session, library_id=lib, name="Parent", parent_folder_id=None + ) + child = folder_service.create_folder( + db_session, library_id=lib, name="Child", parent_folder_id=parent.id + ) + item = _make_item(db_session, lib, folder_id=child.id) + new_parent = folder_service.delete_folder(db_session, child.id) + assert new_parent == parent.id + db_session.refresh(item) + assert item.folder_id == parent.id + + +def test_delete_folder_reparents_subfolders_up(db_session, lib): + """delete_folder re-homes subfolders up to the deleted folder's parent.""" + parent = folder_service.create_folder( + db_session, library_id=lib, name="Parent", parent_folder_id=None + ) + mid = folder_service.create_folder( + db_session, library_id=lib, name="Mid", parent_folder_id=parent.id + ) + sub = folder_service.create_folder( + db_session, library_id=lib, name="Sub", parent_folder_id=mid.id + ) + folder_service.delete_folder(db_session, mid.id) + db_session.refresh(sub) + assert sub.parent_folder_id == parent.id + + +def test_delete_folder_root_reparents_to_none(db_session, lib): + """Deleting a root folder re-homes its children to the library root.""" + root = folder_service.create_folder( + db_session, library_id=lib, name="Root", parent_folder_id=None + ) + child = folder_service.create_folder( + db_session, library_id=lib, name="Child", parent_folder_id=root.id + ) + new_parent = folder_service.delete_folder(db_session, root.id) + assert new_parent is None + db_session.refresh(child) + assert child.parent_folder_id is None + + +def test_delete_folder_collision_renames_with_suffix(db_session, lib): + """delete_folder appends ``(2)`` when a reparented subfolder name collides. + + NOTE: the deleted folder is at the library root so its children reparent + to a NULL parent. The non-root case is covered by + test_delete_folder_collision_under_nonroot_parent. + """ + # Root already has a "Drafts". + folder_service.create_folder( + db_session, library_id=lib, name="Drafts", parent_folder_id=None + ) + mid = folder_service.create_folder( + db_session, library_id=lib, name="Mid", parent_folder_id=None + ) + # The deleted folder ALSO contains a "Drafts". + collide = folder_service.create_folder( + db_session, library_id=lib, name="Drafts", parent_folder_id=mid.id + ) + folder_service.delete_folder(db_session, mid.id) + db_session.refresh(collide) + assert collide.parent_folder_id is None + assert collide.name == "Drafts (2)" + + +def test_delete_folder_collision_increments_suffix(db_session, lib): + """_next_available_name keeps incrementing when ``(2)`` is also taken.""" + folder_service.create_folder( + db_session, library_id=lib, name="Drafts", parent_folder_id=None + ) + folder_service.create_folder( + db_session, library_id=lib, name="Drafts (2)", parent_folder_id=None + ) + mid = folder_service.create_folder( + db_session, library_id=lib, name="Mid", parent_folder_id=None + ) + collide = folder_service.create_folder( + db_session, library_id=lib, name="Drafts", parent_folder_id=mid.id + ) + folder_service.delete_folder(db_session, mid.id) + db_session.refresh(collide) + assert collide.name == "Drafts (3)" + + +def test_delete_folder_collision_under_nonroot_parent(db_session, lib): + """A reparented subfolder colliding under a NON-root parent is renamed. + + The collided 'Drafts' is renamed to 'Drafts (2)' under the parent. The + deduped name is computed before parent_folder_id is reassigned, so the + autoflush during _next_available_name never writes a colliding row. + """ + parent = folder_service.create_folder( + db_session, library_id=lib, name="Parent", parent_folder_id=None + ) + # Parent already has a "Drafts". + folder_service.create_folder( + db_session, library_id=lib, name="Drafts", parent_folder_id=parent.id + ) + mid = folder_service.create_folder( + db_session, library_id=lib, name="Mid", parent_folder_id=parent.id + ) + # Mid ALSO contains a "Drafts"; deleting Mid reparents it under Parent. + collide = folder_service.create_folder( + db_session, library_id=lib, name="Drafts", parent_folder_id=mid.id + ) + folder_service.delete_folder(db_session, mid.id) + db_session.refresh(collide) + assert collide.parent_folder_id == parent.id + assert collide.name == "Drafts (2)" + + +def test_delete_folder_missing_raises_notfound(db_session): + """delete_folder raises FolderNotFoundError for an unknown folder.""" + with pytest.raises(FolderNotFoundError): + folder_service.delete_folder(db_session, unique_id("x")) + + +# --------------------------------------------------------------------------- +# move_items +# --------------------------------------------------------------------------- + + +def test_move_items_single_into_folder(db_session, lib): + """move_items moves one item into a folder.""" + folder = folder_service.create_folder( + db_session, library_id=lib, name="Dest", parent_folder_id=None + ) + item = _make_item(db_session, lib) + n = folder_service.move_items(db_session, lib, [item.id], folder.id) + assert n == 1 + db_session.refresh(item) + assert item.folder_id == folder.id + + +def test_move_items_bulk_to_root(db_session, lib): + """move_items moves multiple items to the library root (folder_id None).""" + folder = folder_service.create_folder( + db_session, library_id=lib, name="Dest", parent_folder_id=None + ) + i1 = _make_item(db_session, lib, folder_id=folder.id) + i2 = _make_item(db_session, lib, folder_id=folder.id) + n = folder_service.move_items(db_session, lib, [i1.id, i2.id], None) + assert n == 2 + db_session.refresh(i1) + db_session.refresh(i2) + assert i1.folder_id is None + assert i2.folder_id is None + + +def test_move_items_empty_list_returns_zero(db_session, lib): + """move_items returns 0 for an empty item list without touching the DB.""" + assert folder_service.move_items(db_session, lib, [], None) == 0 + + +def test_move_items_unknown_destination_raises_notfound(db_session, lib): + """move_items raises FolderNotFoundError for an unknown destination.""" + item = _make_item(db_session, lib) + with pytest.raises(FolderNotFoundError): + folder_service.move_items(db_session, lib, [item.id], unique_id("x")) + + +def test_move_items_destination_other_library_raises_validation(db_session, lib): + """move_items rejects a destination folder in a different library.""" + other_lib = unique_id("lib") + create_library(db_session, other_lib, unique_id("org"), "Other") + dest = folder_service.create_folder( + db_session, library_id=other_lib, name="Dest", parent_folder_id=None + ) + item = _make_item(db_session, lib) + with pytest.raises(FolderValidationError): + folder_service.move_items(db_session, lib, [item.id], dest.id) + + +def test_move_items_missing_item_raises_notfound(db_session, lib): + """move_items raises FolderNotFoundError when an item id does not exist.""" + item = _make_item(db_session, lib) + with pytest.raises(FolderNotFoundError): + folder_service.move_items(db_session, lib, [item.id, unique_id("ghost")], None) + + +def test_move_items_cross_library_item_raises_validation(db_session, lib): + """move_items rejects items that belong to a different library.""" + other_lib = unique_id("lib") + create_library(db_session, other_lib, unique_id("org"), "Other") + foreign_item = _make_item(db_session, other_lib) + with pytest.raises(FolderValidationError): + folder_service.move_items(db_session, lib, [foreign_item.id], None) + + +# --------------------------------------------------------------------------- +# Queries +# --------------------------------------------------------------------------- + + +def test_get_folder_returns_none_when_missing(db_session): + """get_folder returns None for an unknown id.""" + assert folder_service.get_folder(db_session, unique_id("x")) is None + + +def test_list_folders_sorted_by_name(db_session, lib): + """list_folders returns the library's folders ordered alphabetically.""" + folder_service.create_folder( + db_session, library_id=lib, name="Zebra", parent_folder_id=None + ) + folder_service.create_folder( + db_session, library_id=lib, name="Alpha", parent_folder_id=None + ) + names = [f.name for f in folder_service.list_folders(db_session, lib)] + assert names == ["Alpha", "Zebra"] + + +def test_list_items_for_tree_returns_library_items(db_session, lib): + """list_items_for_tree returns the items belonging to the library.""" + i1 = _make_item(db_session, lib) + i2 = _make_item(db_session, lib) + ids = {i.id for i in folder_service.list_items_for_tree(db_session, lib)} + assert {i1.id, i2.id} <= ids diff --git a/library-manager/tests/unit/test_service_import.py b/library-manager/tests/unit/test_service_import.py new file mode 100644 index 000000000..3515a09c4 --- /dev/null +++ b/library-manager/tests/unit/test_service_import.py @@ -0,0 +1,489 @@ +"""Unit tests for ``services.import_service`` — direct function calls.""" + +from __future__ import annotations + +import pytest +from _fakes import patch_markitdown +from _helpers import unique_id +from database.models import ContentImage, ContentItem, ImportJob +from services import content_service, import_service +from services.folder_service import create_folder +from services.library_service import create_library +from tasks import worker + + +@pytest.fixture +def lib(db_session): + """Create a library and return (organization_id, library_id).""" + org_id = unique_id("org") + lib_id = unique_id("lib") + create_library(db_session, lib_id, org_id, "Import Lib") + return org_id, lib_id + + +@pytest.fixture +def storage(tmp_storage, monkeypatch): + """Point both import_service and content_service at a throwaway CONTENT_DIR.""" + monkeypatch.setattr(import_service, "CONTENT_DIR", tmp_storage) + monkeypatch.setattr(content_service, "CONTENT_DIR", tmp_storage) + return tmp_storage + + +def _get_item(db, item_id): + return db.query(ContentItem).filter(ContentItem.id == item_id).first() + + +def _get_job(db, job_id): + return db.query(ImportJob).filter(ImportJob.id == job_id).first() + + +# --------------------------------------------------------------------------- +# _validate_folder_id +# --------------------------------------------------------------------------- + + +def test_validate_folder_id_none_is_noop(db_session, lib): + """_validate_folder_id accepts None (library root).""" + _, lib_id = lib + import_service._validate_folder_id(db_session, lib_id, None) # no raise + + +def test_validate_folder_id_valid_passes(db_session, lib): + """_validate_folder_id accepts a folder in the same library.""" + _, lib_id = lib + folder = create_folder(db_session, library_id=lib_id, name="F", parent_folder_id=None) + import_service._validate_folder_id(db_session, lib_id, folder.id) # no raise + + +def test_validate_folder_id_missing_raises(db_session, lib): + """_validate_folder_id raises ValueError for an unknown folder.""" + _, lib_id = lib + with pytest.raises(ValueError, match="not found"): + import_service._validate_folder_id(db_session, lib_id, unique_id("ghost")) + + +def test_validate_folder_id_foreign_raises(db_session, lib): + """_validate_folder_id raises ValueError for a folder in another library.""" + _, lib_id = lib + other_lib = unique_id("lib") + create_library(db_session, other_lib, unique_id("org"), "Other") + foreign = create_folder(db_session, library_id=other_lib, name="F", parent_folder_id=None) + with pytest.raises(ValueError, match="different library"): + import_service._validate_folder_id(db_session, lib_id, foreign.id) + + +# --------------------------------------------------------------------------- +# queue_* — create pending item + pending job, no api_keys persisted +# --------------------------------------------------------------------------- + + +def test_queue_file_import_creates_pending_records(db_session, lib): + """queue_file_import inserts a pending item + job and returns their ids.""" + org_id, lib_id = lib + item_id, job_id = import_service.queue_file_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="My Doc", + plugin_name="simple_import", + file_path="/tmp/upload.txt", + original_filename="upload.txt", + content_type="text/plain", + file_size=42, + api_keys={"openai": "sk-secret"}, + ) + item = _get_item(db_session, item_id) + job = _get_job(db_session, job_id) + assert item.status == "pending" + assert item.source_type == "file" + assert item.original_filename == "upload.txt" + assert job.status == "pending" + assert job.source_path == "/tmp/upload.txt" + assert job.plugin_name == "simple_import" + + +def test_queue_does_not_persist_api_keys(db_session, lib): + """API keys never land on the ImportJob row (only in worker memory).""" + org_id, lib_id = lib + _, job_id = import_service.queue_file_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="Doc", + plugin_name="simple_import", + file_path="/tmp/u.txt", + original_filename="u.txt", + api_keys={"openai": "sk-secret"}, + ) + job = _get_job(db_session, job_id) + # No column holds the keys; the whole serialized row must not leak them. + serialized = " ".join( + str(getattr(job, c.name)) for c in ImportJob.__table__.columns + ) + assert "sk-secret" not in serialized + # They are held in the worker's in-memory dict instead. + assert worker._job_api_keys.get(job_id) == {"openai": "sk-secret"} + worker._job_api_keys.pop(job_id, None) + + +def test_queue_url_import_creates_pending(db_session, lib): + """queue_url_import records a url-source pending item + job.""" + org_id, lib_id = lib + item_id, job_id = import_service.queue_url_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="Page", + plugin_name="url_import", + url="https://example.com", + ) + item = _get_item(db_session, item_id) + job = _get_job(db_session, job_id) + assert item.source_type == "url" + assert item.source_url == "https://example.com" + assert job.source_url == "https://example.com" + assert job.source_path is None + + +def test_queue_youtube_import_creates_pending(db_session, lib): + """queue_youtube_import records a youtube-source pending item + job.""" + org_id, lib_id = lib + item_id, job_id = import_service.queue_youtube_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="Video", + plugin_name="youtube_import", + video_url="https://youtu.be/abc", + ) + item = _get_item(db_session, item_id) + assert item.source_type == "youtube" + assert item.source_url == "https://youtu.be/abc" + + +def test_queue_file_import_with_folder(db_session, lib): + """queue_file_import stores the destination folder on the item.""" + org_id, lib_id = lib + folder = create_folder(db_session, library_id=lib_id, name="Dest", parent_folder_id=None) + item_id, _ = import_service.queue_file_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="Doc", + plugin_name="simple_import", + file_path="/tmp/u.txt", + original_filename="u.txt", + folder_id=folder.id, + ) + assert _get_item(db_session, item_id).folder_id == folder.id + + +def test_queue_file_import_bad_folder_raises(db_session, lib): + """queue_file_import propagates folder validation errors.""" + org_id, lib_id = lib + with pytest.raises(ValueError): + import_service.queue_file_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="Doc", + plugin_name="simple_import", + file_path="/tmp/u.txt", + original_filename="u.txt", + folder_id=unique_id("ghost"), + ) + + +# --------------------------------------------------------------------------- +# execute_import_job — success +# --------------------------------------------------------------------------- + + +def test_execute_import_job_success(storage, db_session, lib): + """execute_import_job runs the plugin, writes content, marks item ready.""" + org_id, lib_id = lib + src = storage / "input.md" + src.write_text("# Title\n\nbody text") + item_id, job_id = import_service.queue_file_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="Title", + plugin_name="simple_import", + file_path=str(src), + original_filename="input.md", + ) + worker._job_api_keys.pop(job_id, None) + job = _get_job(db_session, job_id) + + import_service.execute_import_job(db_session, job, api_keys={}) + + item = _get_item(db_session, item_id) + assert item.status == "ready" + assert item.page_count == 0 + assert item.image_count == 0 + base = storage / org_id / lib_id / item_id + assert (base / "content" / "full.md").read_text() == "# Title\n\nbody text" + assert item.full_markdown_path == str(base / "content" / "full.md") + assert item.content_type == "text/markdown" # from plugin metadata + # Temp source file is removed after a successful file import. + assert not src.exists() + + +def test_execute_import_job_inserts_image_rows(storage, db_session, lib, monkeypatch): + """execute_import_job writes ContentImage rows for extracted images.""" + from plugins.base import ExtractedImage, ImportResult, PluginRegistry + + org_id, lib_id = lib + src = storage / "doc.pdf" + src.write_bytes(b"%PDF-fake") + item_id, job_id = import_service.queue_file_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="With Images", + plugin_name="simple_import", + file_path=str(src), + original_filename="doc.pdf", + ) + worker._job_api_keys.pop(job_id, None) + job = _get_job(db_session, job_id) + + fake_result = ImportResult( + full_text="text", + pages=[], + images=[ + ExtractedImage(filename="img_001.png", data=b"PNG", page_number=1, description="d"), + ], + metadata={"content_type": "application/pdf"}, + source_ref={"type": "file"}, + ) + + class _FakePlugin: + def import_content(self, source, *, api_keys=None, **kwargs): + return fake_result + + monkeypatch.setattr(PluginRegistry, "get_plugin", staticmethod(lambda name: _FakePlugin())) + monkeypatch.setattr(PluginRegistry, "sanitize_params", staticmethod(lambda n, p: {})) + + import_service.execute_import_job(db_session, job, api_keys={}) + + item = _get_item(db_session, item_id) + assert item.status == "ready" + assert item.image_count == 1 + imgs = db_session.query(ContentImage).filter(ContentImage.content_item_id == item_id).all() + assert len(imgs) == 1 + assert imgs[0].image_path == "content/images/img_001.png" + + +def test_execute_import_job_plugin_not_found_raises(db_session, lib): + """execute_import_job raises RuntimeError when the plugin is unknown.""" + org_id, lib_id = lib + item_id, job_id = import_service.queue_url_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="X", + plugin_name="does_not_exist", + url="https://x", + ) + worker._job_api_keys.pop(job_id, None) + job = _get_job(db_session, job_id) + with pytest.raises(RuntimeError, match="Plugin not found"): + import_service.execute_import_job(db_session, job, api_keys={}) + + +def test_execute_import_job_url_success_no_temp_cleanup(storage, db_session, lib, monkeypatch): + """A URL job writes content and skips the temp-file cleanup (no source_path).""" + from plugins.base import ImportResult, PluginRegistry + + org_id, lib_id = lib + item_id, job_id = import_service.queue_url_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="Page", + plugin_name="url_import", + url="https://example.com", + ) + worker._job_api_keys.pop(job_id, None) + job = _get_job(db_session, job_id) + + class _FakePlugin: + def import_content(self, source, *, api_keys=None, **kwargs): + return ImportResult( + full_text="scraped", + metadata={"file_size": 7}, + source_ref={"type": "url"}, + ) + + monkeypatch.setattr(PluginRegistry, "get_plugin", staticmethod(lambda name: _FakePlugin())) + monkeypatch.setattr(PluginRegistry, "sanitize_params", staticmethod(lambda n, p: {})) + + import_service.execute_import_job(db_session, job, api_keys={}) + item = _get_item(db_session, item_id) + assert item.status == "ready" + assert item.file_size == 7 + base = storage / org_id / lib_id / item_id + assert (base / "content" / "full.md").read_text() == "scraped" + + +def test_execute_import_job_missing_item_raises(storage, db_session, lib, monkeypatch): + """execute_import_job raises RuntimeError when the ContentItem vanished.""" + from plugins.base import ImportResult, PluginRegistry + + org_id, lib_id = lib + src = storage / "x.txt" + src.write_text("body") + item_id, job_id = import_service.queue_file_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="Doc", + plugin_name="simple_import", + file_path=str(src), + original_filename="x.txt", + ) + worker._job_api_keys.pop(job_id, None) + job = _get_job(db_session, job_id) + + class _FakePlugin: + def import_content(self, source, *, api_keys=None, **kwargs): + return ImportResult(full_text="t", metadata={}, source_ref={}) + + monkeypatch.setattr(PluginRegistry, "get_plugin", staticmethod(lambda name: _FakePlugin())) + monkeypatch.setattr(PluginRegistry, "sanitize_params", staticmethod(lambda n, p: {})) + + # Delete the item after queueing but before execution. + db_session.query(ContentItem).filter(ContentItem.id == item_id).delete() + db_session.commit() + + with pytest.raises(RuntimeError, match="not found"): + import_service.execute_import_job(db_session, job, api_keys={}) + + +def test_execute_import_job_temp_unlink_failure_is_tolerated( + storage, db_session, lib, monkeypatch +): + """A failure to delete the temp source file does not fail the import.""" + from pathlib import Path as _Path + + org_id, lib_id = lib + src = storage / "keep.md" + src.write_text("# body") + item_id, job_id = import_service.queue_file_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="Doc", + plugin_name="simple_import", + file_path=str(src), + original_filename="keep.md", + ) + worker._job_api_keys.pop(job_id, None) + job = _get_job(db_session, job_id) + + real_unlink = _Path.unlink + + def _boom_unlink(self, *args, **kwargs): + if self == src: + raise OSError("cannot delete") + return real_unlink(self, *args, **kwargs) + + monkeypatch.setattr(_Path, "unlink", _boom_unlink) + import_service.execute_import_job(db_session, job, api_keys={}) + assert _get_item(db_session, item_id).status == "ready" + + +def test_execute_import_job_no_source_raises(storage, db_session, lib): + """execute_import_job raises RuntimeError when the job has no source.""" + org_id, lib_id = lib + item_id, job_id = import_service.queue_url_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="X", + plugin_name="simple_import", + url="placeholder", + ) + worker._job_api_keys.pop(job_id, None) + job = _get_job(db_session, job_id) + job.source_url = None + job.source_path = None + db_session.commit() + with pytest.raises(RuntimeError, match="no source"): + import_service.execute_import_job(db_session, job, api_keys={}) + + +# --------------------------------------------------------------------------- +# execute_import_job — failure cleans up the partial dir +# --------------------------------------------------------------------------- + + +def test_execute_import_job_failure_cleans_dir(storage, db_session, lib, monkeypatch): + """A write failure rmtree's the partial item directory and re-raises.""" + from plugins.base import ImportResult, PluginRegistry + + org_id, lib_id = lib + src = storage / "doc.txt" + src.write_text("body") + item_id, job_id = import_service.queue_file_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="Doc", + plugin_name="simple_import", + file_path=str(src), + original_filename="doc.txt", + ) + worker._job_api_keys.pop(job_id, None) + job = _get_job(db_session, job_id) + + class _FakePlugin: + def import_content(self, source, *, api_keys=None, **kwargs): + return ImportResult(full_text="t", metadata={}, source_ref={}) + + monkeypatch.setattr(PluginRegistry, "get_plugin", staticmethod(lambda name: _FakePlugin())) + monkeypatch.setattr(PluginRegistry, "sanitize_params", staticmethod(lambda n, p: {})) + + # Make the disk write blow up after the directory is created. + item_dir = storage / org_id / lib_id / item_id + + def _boom(**kwargs): + item_dir.mkdir(parents=True, exist_ok=True) + (item_dir / "marker").write_text("partial") + raise RuntimeError("disk full") + + monkeypatch.setattr(content_service, "write_structured_content", _boom) + + with pytest.raises(RuntimeError, match="disk full"): + import_service.execute_import_job(db_session, job, api_keys={}) + + assert not item_dir.exists() # Partial dir removed. + + +def test_execute_import_job_markitdown_failure_propagates(storage, db_session, lib): + """A plugin that raises inside import_content propagates (item not ready).""" + org_id, lib_id = lib + src = storage / "doc.docx" + src.write_bytes(b"fake-docx") + item_id, job_id = import_service.queue_file_import( + db_session, + library_id=lib_id, + organization_id=org_id, + title="Doc", + plugin_name="markitdown_import", + file_path=str(src), + original_filename="doc.docx", + ) + worker._job_api_keys.pop(job_id, None) + job = _get_job(db_session, job_id) + + with patch_markitdown(raises=RuntimeError("conversion boom")), pytest.raises(Exception): + import_service.execute_import_job(db_session, job, api_keys={}) + + db_session.rollback() + item = _get_item(db_session, item_id) + # The item never reached ``ready`` — the service layer leaves status + # handling to the worker, so it stays at its pre-execution value. + assert item.status != "ready" diff --git a/library-manager/tests/unit/test_service_library.py b/library-manager/tests/unit/test_service_library.py new file mode 100644 index 000000000..09ecf4451 --- /dev/null +++ b/library-manager/tests/unit/test_service_library.py @@ -0,0 +1,227 @@ +"""Unit tests for ``services.library_service`` — direct function calls.""" + +from __future__ import annotations + +import json +import uuid + +import pytest +from _helpers import unique_id +from database.models import ContentItem, Library +from services import library_service +from sqlalchemy.exc import IntegrityError + + +def _make_item(db, library_id, organization_id, **overrides): + """Insert a minimal ready ContentItem for count assertions.""" + item = ContentItem( + id=unique_id("item"), + library_id=library_id, + organization_id=organization_id, + title="t", + source_type="file", + base_path="/tmp/none", + permalink_base="/docs/x", + import_plugin="simple_import", + status="ready", + **overrides, + ) + db.add(item) + db.commit() + return item + + +def test_ensure_organization_creates_when_missing(db_session): + """ensure_organization creates a new org row when it does not exist.""" + org_id = unique_id("org") + org = library_service.ensure_organization(db_session, org_id, name="My Org") + db_session.commit() + assert org.id == org_id + assert org.name == "My Org" + + +def test_ensure_organization_defaults_name_to_id(db_session): + """ensure_organization falls back to the id when no name is given.""" + org_id = unique_id("org") + org = library_service.ensure_organization(db_session, org_id) + db_session.commit() + assert org.name == org_id + + +def test_ensure_organization_idempotent_returns_existing(db_session): + """ensure_organization returns the existing row rather than duplicating.""" + org_id = unique_id("org") + first = library_service.ensure_organization(db_session, org_id, name="First") + db_session.commit() + second = library_service.ensure_organization(db_session, org_id, name="Second") + db_session.commit() + assert second.id == first.id + # Name is not overwritten on the second call. + assert second.name == "First" + + +def test_create_library_persists_row_and_config(db_session): + """create_library inserts a row and serializes import_config to JSON.""" + org_id = unique_id("org") + lib_id = unique_id("lib") + cfg = {"plugin": "simple_import", "params": {"a": 1}} + lib = library_service.create_library(db_session, lib_id, org_id, "Lib A", cfg) + assert lib.id == lib_id + assert lib.organization_id == org_id + assert json.loads(lib.import_config) == cfg + + +def test_create_library_null_config_when_omitted(db_session): + """create_library stores NULL import_config when none is supplied.""" + org_id = unique_id("org") + lib = library_service.create_library(db_session, unique_id("lib"), org_id, "Lib") + assert lib.import_config is None + + +def test_create_library_duplicate_org_name_raises_integrity_error(db_session): + """Duplicate (org, name) violates the unique constraint and raises.""" + org_id = unique_id("org") + name = f"Dup {uuid.uuid4().hex[:6]}" + library_service.create_library(db_session, unique_id("lib"), org_id, name) + with pytest.raises(IntegrityError): + library_service.create_library(db_session, unique_id("lib"), org_id, name) + db_session.rollback() + + +def test_get_library_returns_none_when_missing(db_session): + """get_library returns None for an unknown id.""" + assert library_service.get_library(db_session, unique_id("nope")) is None + + +def test_get_library_returns_row(db_session): + """get_library returns the matching Library row.""" + org_id = unique_id("org") + lib_id = unique_id("lib") + library_service.create_library(db_session, lib_id, org_id, "Lib") + found = library_service.get_library(db_session, lib_id) + assert found is not None and found.id == lib_id + + +def test_get_library_with_item_count_none_when_missing(db_session): + """get_library_with_item_count returns None for an unknown id.""" + assert library_service.get_library_with_item_count(db_session, unique_id("x")) is None + + +def test_get_library_with_item_count_reflects_items(db_session): + """item_count matches the number of items in the library.""" + org_id = unique_id("org") + lib_id = unique_id("lib") + cfg = {"k": "v"} + library_service.create_library(db_session, lib_id, org_id, "Lib", cfg) + _make_item(db_session, lib_id, org_id) + _make_item(db_session, lib_id, org_id) + info = library_service.get_library_with_item_count(db_session, lib_id) + assert info["item_count"] == 2 + assert info["import_config"] == cfg + assert info["id"] == lib_id + + +def test_list_libraries_filters_by_organization(db_session): + """list_libraries only returns libraries for the requested org.""" + org_a = unique_id("orgA") + org_b = unique_id("orgB") + library_service.create_library(db_session, unique_id("lib"), org_a, "A1") + library_service.create_library(db_session, unique_id("lib"), org_a, "A2") + library_service.create_library(db_session, unique_id("lib"), org_b, "B1") + libs, total = library_service.list_libraries(db_session, org_a) + assert total == 2 + assert {lib["organization_id"] for lib in libs} == {org_a} + + +def test_list_libraries_pagination_boundaries(db_session): + """list_libraries honors limit/offset while reporting the full total.""" + org_id = unique_id("org") + for i in range(5): + library_service.create_library(db_session, unique_id("lib"), org_id, f"L{i}") + page1, total = library_service.list_libraries(db_session, org_id, limit=2, offset=0) + page2, _ = library_service.list_libraries(db_session, org_id, limit=2, offset=2) + page3, _ = library_service.list_libraries(db_session, org_id, limit=2, offset=4) + assert total == 5 + assert len(page1) == 2 + assert len(page2) == 2 + assert len(page3) == 1 + ids = {lib["id"] for lib in page1 + page2 + page3} + assert len(ids) == 5 # No overlap across pages. + + +def test_list_libraries_includes_item_count(db_session): + """list_libraries embeds a per-library item_count.""" + org_id = unique_id("org") + lib_id = unique_id("lib") + library_service.create_library(db_session, lib_id, org_id, "Lib") + _make_item(db_session, lib_id, org_id) + libs, _ = library_service.list_libraries(db_session, org_id) + assert libs[0]["item_count"] == 1 + + +def test_delete_library_removes_row_and_cascades(db_session): + """delete_library removes the library and cascades to its items.""" + org_id = unique_id("org") + lib_id = unique_id("lib") + library_service.create_library(db_session, lib_id, org_id, "Lib") + item = _make_item(db_session, lib_id, org_id) + item_id = item.id + assert library_service.delete_library(db_session, lib_id) is True + assert library_service.get_library(db_session, lib_id) is None + assert db_session.query(ContentItem).filter(ContentItem.id == item_id).first() is None + + +def test_delete_library_returns_false_when_missing(db_session): + """delete_library returns False for an unknown id.""" + assert library_service.delete_library(db_session, unique_id("nope")) is False + + +def test_update_import_config_persists_json(db_session): + """update_import_config persists the new config as JSON.""" + org_id = unique_id("org") + lib_id = unique_id("lib") + library_service.create_library(db_session, lib_id, org_id, "Lib") + new_cfg = {"plugin": "markitdown_import", "x": [1, 2, 3]} + updated = library_service.update_import_config(db_session, lib_id, new_cfg) + assert json.loads(updated.import_config) == new_cfg + # Re-read to confirm it was committed. + reread = library_service.get_library(db_session, lib_id) + assert json.loads(reread.import_config) == new_cfg + + +def test_update_import_config_only_affects_target(db_session): + """update_import_config does not touch sibling libraries.""" + org_id = unique_id("org") + lib_a = unique_id("lib") + lib_b = unique_id("lib") + library_service.create_library(db_session, lib_a, org_id, "A", {"orig": "a"}) + library_service.create_library(db_session, lib_b, org_id, "B", {"orig": "b"}) + library_service.update_import_config(db_session, lib_a, {"changed": True}) + other = library_service.get_library(db_session, lib_b) + assert json.loads(other.import_config) == {"orig": "b"} + + +def test_update_import_config_returns_none_when_missing(db_session): + """update_import_config returns None for an unknown library.""" + assert library_service.update_import_config(db_session, unique_id("x"), {}) is None + + +def test_delete_library_removes_content_dir(db_session, tmp_storage, monkeypatch): + """delete_library rmtree's the on-disk content directory when present.""" + org_id = unique_id("org") + lib_id = unique_id("lib") + monkeypatch.setattr(library_service, "CONTENT_DIR", tmp_storage) + content_dir = tmp_storage / org_id / lib_id + content_dir.mkdir(parents=True) + (content_dir / "marker.txt").write_text("x") + library_service.create_library(db_session, lib_id, org_id, "Lib") + assert library_service.delete_library(db_session, lib_id) is True + assert not content_dir.exists() + + +def test_create_library_returns_library_instance(db_session): + """create_library returns a refreshed Library ORM instance.""" + org_id = unique_id("org") + lib = library_service.create_library(db_session, unique_id("lib"), org_id, "L") + assert isinstance(lib, Library) + assert lib.created_at is not None diff --git a/testing/playwright/global-setup.js b/testing/playwright/global-setup.js index 207e896ac..30126c0e8 100644 --- a/testing/playwright/global-setup.js +++ b/testing/playwright/global-setup.js @@ -3,13 +3,10 @@ const fs = require('fs'); const path = require('path'); module.exports = async (config) => { - const baseURL = config.projects[0]?.use?.baseURL || 'http://localhost:5173/'; + const baseURL = config.projects[0]?.use?.baseURL || 'http://localhost:9099/'; const email = process.env.LOGIN_EMAIL || 'admin@owi.com'; const password = process.env.LOGIN_PASSWORD || 'admin'; - // Derive the API base from the baseURL (frontend proxies /creator to the backend). - const apiBase = baseURL.replace(/\/$/, ''); - const authDir = path.join(__dirname, '.auth'); const statePath = path.join(authDir, 'state.json'); @@ -20,46 +17,62 @@ module.exports = async (config) => { return; } - // Call the login API directly — avoids fighting the form's missing preventDefault(). - const formData = new URLSearchParams({ email, password }); - const resp = await fetch(`${apiBase}/creator/login`, { - method: 'POST', - headers: { 'Content-Type': 'application/x-www-form-urlencoded' }, - body: formData.toString(), - }); + const browser = await chromium.launch({ headless: !!process.env.CI }); + const context = await browser.newContext(); + const page = await context.newPage(); - if (!resp.ok) { - throw new Error(`Login API returned HTTP ${resp.status}`); - } + await page.goto(baseURL); + await page.waitForLoadState('domcontentloaded'); - const json = await resp.json(); - if (!json.success || !json.data?.token) { - throw new Error(`Login failed: ${json.error || JSON.stringify(json)}`); - } + // If already authenticated in some environments, just persist storage. + const existingToken = await page.evaluate(() => localStorage.getItem('userToken')); + if (!existingToken) { + // Wait for the email input to exist in the DOM. + await page.waitForSelector('#email', { timeout: 30_000 }); + + // Wait for full load + network idle so Vite has finished streaming all JS + // chunks and SvelteKit hydration has wired up the form's onsubmit handler. + // Without this, clicking on slower machines fires a native form submission + // (URL becomes /?email=…&password=…) instead of the XHR handler. + await page.waitForLoadState('load'); + await page.waitForLoadState('networkidle'); - const token = json.data.token; + // Double-check that SvelteKit has hydrated by polling for its runtime + // marker. This is a zero-cost guard on fast machines (already true by the + // time networkidle resolves) and a reliable gate on slow ones. + await page.waitForFunction( + () => typeof window.__sveltekit_dev !== 'undefined' || typeof window.__sveltekit !== 'undefined', + { timeout: 30_000 } + ); - // Build the storageState manually with the token and all related keys - // so the Svelte app considers itself logged in from the first page load. - const storage = { - cookies: [], - origins: [ - { - origin: new URL(baseURL).origin, - localStorage: [ - { name: 'userToken', value: token }, - { name: 'userName', value: json.data.name || '' }, - { name: 'userEmail', value: json.data.email || email }, - { name: 'userData', value: JSON.stringify(json.data) }, - ], - }, - ], - }; + await page.fill('#email', email); + await page.fill('#password', password); + + // Click the submit button (selector is intentionally narrow) and + // concurrently wait for the POST /creator/login XHR response. This is + // deterministic: we know auth succeeded the moment the server replies, + // with no blind sleep required. + await Promise.all([ + page.waitForResponse( + (r) => r.url().includes('/creator/login') && r.request().method() === 'POST', + { timeout: 30_000 } + ), + page.click('form > button') + ]); + + // Poll until the SPA has stored the token in localStorage. waitForFunction + // resolves as soon as the predicate returns truthy — no fixed wait. + await page.waitForFunction(() => !!localStorage.getItem('userToken'), { timeout: 15_000 }); + } - if (json.data.launch_url) { - storage.origins[0].localStorage.push({ name: 'OWI_url', value: json.data.launch_url }); + const tokenAfter = await page.evaluate(() => localStorage.getItem('userToken')); + if (!tokenAfter) { + await browser.close(); + throw new Error( + 'Login did not produce localStorage.userToken. Check UI selectors or credentials (LOGIN_EMAIL/LOGIN_PASSWORD).' + ); } - fs.writeFileSync(statePath, JSON.stringify(storage, null, 2)); - console.log(`[global-setup] Logged in as ${email}, state saved to ${statePath}`); + await context.storageState({ path: statePath }); + await browser.close(); }; diff --git a/testing/playwright/tests/admin_and_sharing_flow.spec.js b/testing/playwright/tests/admin_and_sharing_flow.spec.js index 645caec1e..71c1b176e 100644 --- a/testing/playwright/tests/admin_and_sharing_flow.spec.js +++ b/testing/playwright/tests/admin_and_sharing_flow.spec.js @@ -822,6 +822,12 @@ test.describe.serial("Admin & Assistant Sharing Flow", () => { } } + // Re-navigate to clear any lingering success banner/overlay from the first + // disable before acting on the second user (otherwise the overlay can + // intercept the next click). + await page.goto("admin?view=users"); + await page.waitForLoadState("networkidle"); + // Disable second test user if (await searchBox.count()) { await searchBox.fill(sharingUser2Email); diff --git a/testing/playwright/tests/assistant_with_knowledge_store.spec.js b/testing/playwright/tests/assistant_with_knowledge_store.spec.js index 818052bc7..98aca1758 100644 --- a/testing/playwright/tests/assistant_with_knowledge_store.spec.js +++ b/testing/playwright/tests/assistant_with_knowledge_store.spec.js @@ -40,6 +40,7 @@ test.describe.serial("Assistant with Knowledge Store RAG (UI) @phase5-pending", let libraryId; let itemId; let knowledgeStoreId; + let knowledgeStoreId2; let knowledgeStoreName; let assistantId; let assistantName; @@ -95,6 +96,43 @@ test.describe.serial("Assistant with Knowledge Store RAG (UI) @phase5-pending", return fallback; } + /** + * Drive the AssistantForm into a state where `knowledge_store_rag` is a + * selectable RAG processor and the KS picker is mounted. + * + * In the merged UI the RAG dropdown is filtered by the selected Prompt + * Processor: `knowledge_store_rag` is only compatible with the + * `kvcache_augment` PPS. The PPS dropdown itself is only rendered in + * "Advanced Mode" during create. So the sequence is: + * 1. enable Advanced Mode + * 2. switch PPS -> kvcache_augment + * 3. select knowledge_store_rag in the RAG dropdown + */ + async function selectKnowledgeStoreRag(page) { + // Enable Advanced Mode so the Prompt Processor select is rendered. The + // toggle is an sr-only checkbox inside a