From 136664d9ffd46daa027772aaf5a9b49d922448f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A0=20Guilera?= Date: Sun, 17 May 2026 14:19:46 +0200 Subject: [PATCH 001/156] feat(kg-rag): port semantic-graph KG-RAG into new lamb-kb-server arch Ports the KG-RAG / semantic-graph stack from the legacy lamb-kb-server-stable into the new pluggable lamb-kb-server (port 9092) and wires it through Knowledge Stores instead of legacy Knowledge Bases. KB server: - New /graph and /benchmarks routers, mounted only when KG_RAG_ENABLED. - Neo4j graph_store + LLM concept extractor adapted to string collection IDs and request-scoped OpenAI credentials. - New graph_indexing helper feeds Neo4j from the ChromaDB backend after vector ingestion succeeds (fails open). - KG-RAG query augmentation hangs off query_service.query_with_plugin ('simple_query' | 'kg_rag_query'), preserving the trace contract. - Optional 'kg-rag' extra adds neo4j + openai. LAMB backend: - New knowledge_store_graph_router proxies graph + benchmark calls through KnowledgeStoreClient using per-org token/URL resolution. - KnowledgeStoreCreate threads graph_enabled to the KB server. Frontend: - New graphService.js + benchmarkService.js axios clients. - New KnowledgeStoreGraphView + KnowledgeStoreBenchmarkView Svelte components, mounted as tabs on KnowledgeStoreDetail when KG-RAG is on. Curation actions (approve / reject concepts and relationships) proxy through the LAMB router. - CreateKnowledgeWizard exposes a Graph RAG opt-in checkbox. Infra: - docker-compose.next.yaml gets an optional neo4j service under --profile kg-rag. - .env.next.example documents the KG_RAG_* block. Tests: - tests/unit/test_kg_rag.py covers config parsing, concept-extraction helpers and the plugin's three graceful-degradation paths. Quality gates: lamb-kb-server-stable untouched; KB server starts and all existing unit + integration tests pass without Neo4j; graph features are off by default per-server AND per-collection. --- .env.next.example | 31 + .../knowledge_store_client.py | 175 ++ .../knowledge_store_graph_router.py | 332 +++ .../knowledge_store_router.py | 6 + backend/creator_interface/main.py | 8 + docker-compose.next.yaml | 28 + .../knowledge/CreateKnowledgeWizard.svelte | 3 +- .../wizard/Step8_ReviewCreate.svelte | 3 +- .../knowledge/wizard/StepKSSetup.svelte | 27 +- .../KnowledgeStoreBenchmarkView.svelte | 207 ++ .../KnowledgeStoreDetail.svelte | 55 + .../KnowledgeStoreGraphView.svelte | 345 +++ .../src/lib/services/benchmarkService.js | 43 + .../src/lib/services/graphService.js | 141 + .../svelte-app/src/lib/utils/graphCuration.js | 268 ++ lamb-kb-server/backend/config.py | 79 + lamb-kb-server/backend/database/models.py | 8 + lamb-kb-server/backend/main.py | 11 + .../backend/plugins/kg_rag_query.py | 282 ++ lamb-kb-server/backend/routers/benchmarks.py | 93 + lamb-kb-server/backend/routers/graph.py | 546 ++++ lamb-kb-server/backend/schemas/benchmark.py | 109 + lamb-kb-server/backend/schemas/collection.py | 10 + lamb-kb-server/backend/schemas/graph.py | 156 ++ lamb-kb-server/backend/services/benchmark.py | 699 +++++ .../backend/services/collection_service.py | 1 + .../backend/services/concept_extraction.py | 328 +++ .../backend/services/graph_indexing.py | 180 ++ .../backend/services/graph_store.py | 2472 +++++++++++++++++ .../backend/services/ingestion_service.py | 99 + .../backend/services/query_service.py | 115 + lamb-kb-server/pyproject.toml | 8 +- lamb-kb-server/tests/unit/test_kg_rag.py | 187 ++ 33 files changed, 7051 insertions(+), 4 deletions(-) create mode 100644 backend/creator_interface/knowledge_store_graph_router.py create mode 100644 frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreBenchmarkView.svelte create mode 100644 frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte create mode 100644 frontend/svelte-app/src/lib/services/benchmarkService.js create mode 100644 frontend/svelte-app/src/lib/services/graphService.js create mode 100644 frontend/svelte-app/src/lib/utils/graphCuration.js create mode 100644 lamb-kb-server/backend/plugins/kg_rag_query.py create mode 100644 lamb-kb-server/backend/routers/benchmarks.py create mode 100644 lamb-kb-server/backend/routers/graph.py create mode 100644 lamb-kb-server/backend/schemas/benchmark.py create mode 100644 lamb-kb-server/backend/schemas/graph.py create mode 100644 lamb-kb-server/backend/services/benchmark.py create mode 100644 lamb-kb-server/backend/services/concept_extraction.py create mode 100644 lamb-kb-server/backend/services/graph_indexing.py create mode 100644 lamb-kb-server/backend/services/graph_store.py create mode 100644 lamb-kb-server/tests/unit/test_kg_rag.py diff --git a/.env.next.example b/.env.next.example index e27c95eba..ece16f648 100644 --- a/.env.next.example +++ b/.env.next.example @@ -102,3 +102,34 @@ 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 +# 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 diff --git a/backend/creator_interface/knowledge_store_client.py b/backend/creator_interface/knowledge_store_client.py index 539e9e8e5..a0c2955fd 100644 --- a/backend/creator_interface/knowledge_store_client.py +++ b/backend/creator_interface/knowledge_store_client.py @@ -253,6 +253,7 @@ async def create_collection( description: str = "", chunking_params: Dict[str, Any] = None, embedding_endpoint: str = "", + graph_enabled: bool = False, creator_user: Dict[str, Any] = None, ) -> Dict: """Create a collection on the KB Server. @@ -274,6 +275,7 @@ async def create_collection( "api_endpoint": embedding_endpoint or "", }, "vector_db_backend": vector_db_backend, + "graph_enabled": bool(graph_enabled), } return await self._request("POST", "/collections", config, json=payload) @@ -407,6 +409,179 @@ async def get_job_status(self, job_id: str, config = self._get_ks_config(creator_user) return await self._request("GET", f"/jobs/{job_id}", 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 run_benchmark( + self, + knowledge_store_id: str, + body: Dict[str, Any], + embedding_api_key: str = "", + embedding_api_endpoint: str = "", + creator_user: Dict[str, Any] = None, + ) -> Dict: + """Run a single benchmark dataset against a Knowledge Store.""" + config = self._get_ks_config(creator_user) + payload = { + **body, + "embedding_credentials": { + "api_key": embedding_api_key or "", + "api_endpoint": embedding_api_endpoint or "", + }, + } + return await self._request( + "POST", + f"/benchmarks/collections/{knowledge_store_id}/run", + config, + json=payload, + ) + + async def list_benchmark_datasets( + self, creator_user: Dict[str, Any] = None + ) -> Dict: + config = self._get_ks_config(creator_user) + return await self._request("GET", "/benchmarks/datasets", config) + + 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..969d59b72 --- /dev/null +++ b/backend/creator_interface/knowledge_store_graph_router.py @@ -0,0 +1,332 @@ +"""Creator-Interface proxy for KG-RAG / semantic-graph endpoints. + +Routes mount at ``/creator/knowledge-stores/{ks_id}/graph/...`` and +``/creator/knowledge-stores/{ks_id}/benchmarks/...``. Each call resolves +the per-org KB Server URL/token through ``KnowledgeStoreClient`` and +proxies to the new KB Server's ``/graph`` and ``/benchmarks`` routers. + +The KB Server only exposes those routers when ``KG_RAG_ENABLED=true``; if +the flag is off the proxy will return 503 from the KB Server. 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.completions.org_config_resolver import OrganizationConfigResolver +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: + 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=200), + 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, + ) + + +# ---------------------------------------------------------------------- +# Benchmarks +# ---------------------------------------------------------------------- + + +class BenchmarkRunBody(BaseModel): + dataset_id: Optional[str] = "educational" + top_k: Optional[int] = None + graph_depth: Optional[int] = None + threshold: float = 0.0 + + +@router.get("/benchmarks/datasets") +async def list_benchmark_datasets( + auth: AuthContext = Depends(get_auth_context), +): + return await _client.list_benchmark_datasets(creator_user=auth.user) + + +@router.post("/{ks_id}/benchmarks/run") +async def run_benchmark( + ks_id: str, + body: BenchmarkRunBody, + auth: AuthContext = Depends(get_auth_context), +): + _assert_ks_access(ks_id, auth) + resolver = OrganizationConfigResolver(auth.user.get("email")) + # Resolve embedding key the same way ingestion does, so the benchmark + # baseline vector pass works without the caller threading creds. + ks = _db.get_knowledge_store(ks_id) + embedding_api_key = "" + embedding_api_endpoint = "" + try: + embedding_api_key = ( + resolver.get_provider_api_key(ks.get("embedding_vendor")) or "" + ) + embedding_api_endpoint = ( + resolver.get_provider_endpoint(ks.get("embedding_vendor")) or "" + ) + except Exception: # noqa: BLE001 + embedding_api_key = "" + return await _client.run_benchmark( + knowledge_store_id=ks_id, + body=body.model_dump(exclude_none=True), + embedding_api_key=embedding_api_key, + embedding_api_endpoint=embedding_api_endpoint, + creator_user=auth.user, + ) diff --git a/backend/creator_interface/knowledge_store_router.py b/backend/creator_interface/knowledge_store_router.py index ad50ebb4c..e13992d0c 100644 --- a/backend/creator_interface/knowledge_store_router.py +++ b/backend/creator_interface/knowledge_store_router.py @@ -49,6 +49,11 @@ class KnowledgeStoreCreate(BaseModel): embedding_model: str embedding_endpoint: Optional[str] = None vector_db_backend: str + # 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 class KnowledgeStoreUpdate(BaseModel): @@ -213,6 +218,7 @@ async def create_knowledge_store( description=body.description, chunking_params=body.chunking_params, embedding_endpoint=resolved_endpoint or "", + graph_enabled=bool(body.graph_enabled), creator_user=auth.user, ) except Exception as e: diff --git a/backend/creator_interface/main.py b/backend/creator_interface/main.py index 56b25de0c..9cbed5795 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_graph_router import ( + 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 @@ -145,6 +148,11 @@ 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") +# Graph + benchmark proxy. Mounted on the same prefix so the frontend +# can keep ``/creator/knowledge-stores/...`` as the only KB-related base. +router.include_router( + knowledge_store_graph_router, prefix="/knowledge-stores" +) # REMOVED: assistant_sharing_router - functionality moved to services, accessed via /creator/lamb/* proxy diff --git a/docker-compose.next.yaml b/docker-compose.next.yaml index 4bbbbb8e7..09a27d71e 100644 --- a/docker-compose.next.yaml +++ b/docker-compose.next.yaml @@ -88,9 +88,37 @@ 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: kb-static: openwebui-data: ollama-data: + neo4j-data: + neo4j-logs: diff --git a/frontend/svelte-app/src/lib/components/knowledge/CreateKnowledgeWizard.svelte b/frontend/svelte-app/src/lib/components/knowledge/CreateKnowledgeWizard.svelte index 8d28ef45e..ab08a1d3e 100644 --- a/frontend/svelte-app/src/lib/components/knowledge/CreateKnowledgeWizard.svelte +++ b/frontend/svelte-app/src/lib/components/knowledge/CreateKnowledgeWizard.svelte @@ -89,7 +89,8 @@ embedding_vendor: '', embedding_model: '', embedding_endpoint: '', - vector_db_backend: '' + vector_db_backend: '', + graph_enabled: false }, selectedItemIds: [], selectionInitialized: false, diff --git a/frontend/svelte-app/src/lib/components/knowledge/wizard/Step8_ReviewCreate.svelte b/frontend/svelte-app/src/lib/components/knowledge/wizard/Step8_ReviewCreate.svelte index e6bc6f8fe..748f33995 100644 --- a/frontend/svelte-app/src/lib/components/knowledge/wizard/Step8_ReviewCreate.svelte +++ b/frontend/svelte-app/src/lib/components/knowledge/wizard/Step8_ReviewCreate.svelte @@ -342,7 +342,8 @@ embedding_vendor: wizardState.ksConfig.embedding_vendor, embedding_model: wizardState.ksConfig.embedding_model, embedding_endpoint: wizardState.ksConfig.embedding_endpoint || undefined, - vector_db_backend: wizardState.ksConfig.vector_db_backend + vector_db_backend: wizardState.ksConfig.vector_db_backend, + graph_enabled: !!wizardState.ksConfig.graph_enabled }); ksId = ks.id; ksName = ks.name; diff --git a/frontend/svelte-app/src/lib/components/knowledge/wizard/StepKSSetup.svelte b/frontend/svelte-app/src/lib/components/knowledge/wizard/StepKSSetup.svelte index d0a31ebc4..66ce239b4 100644 --- a/frontend/svelte-app/src/lib/components/knowledge/wizard/StepKSSetup.svelte +++ b/frontend/svelte-app/src/lib/components/knowledge/wizard/StepKSSetup.svelte @@ -101,6 +101,7 @@ let embeddingModel = $state(wizardState.ksConfig?.embedding_model || ''); let embeddingEndpoint = $state(wizardState.ksConfig?.embedding_endpoint || ''); let vectorDb = $state(wizardState.ksConfig?.vector_db_backend || ''); + let graphEnabled = $state(!!wizardState.ksConfig?.graph_enabled); // Param descriptors for the currently-selected strategy. Each entry has // { name, type, description, default, min_value, max_value }. @@ -326,7 +327,8 @@ embedding_vendor: embeddingVendor, embedding_model: embeddingModel, embedding_endpoint: embeddingEndpoint, - vector_db_backend: vectorDb + vector_db_backend: vectorDb, + graph_enabled: graphEnabled } }); }); @@ -702,6 +704,29 @@ class="mt-1 block w-full rounded-md border border-gray-300 px-3 py-2 text-sm" /> + +
+ +
{/if} diff --git a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreBenchmarkView.svelte b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreBenchmarkView.svelte new file mode 100644 index 000000000..61b17b474 --- /dev/null +++ b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreBenchmarkView.svelte @@ -0,0 +1,207 @@ + + + +
+ {#if !graphEnabled} +
+ Graph RAG is not enabled on this Knowledge Store, so the + kg_rag_query arm of the benchmark will + degrade to the vector baseline. Enable the graph first for a meaningful + comparison. +
+ {/if} + +
+ + + + + +
+ + {#if error} +
{error}
+ {/if} + + {#if result} +
+

Summary

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
P@kR@kMRRAvg vectorAvg graphAvg total
Baseline{fmt(result?.baseline?.precision_at_k)}{fmt(result?.baseline?.recall_at_k)}{fmt(result?.baseline?.mrr)}{fmtMs(result?.baseline?.avg_vector_ms)}{fmtMs(result?.baseline?.avg_graph_ms)}{fmtMs(result?.baseline?.avg_total_ms)}
KG-RAG{fmt(result?.kg_rag?.precision_at_k)}{fmt(result?.kg_rag?.recall_at_k)}{fmt(result?.kg_rag?.mrr)}{fmtMs(result?.kg_rag?.avg_vector_ms)}{fmtMs(result?.kg_rag?.avg_graph_ms)}{fmtMs(result?.kg_rag?.avg_total_ms)}
+ {#if result?.comparison} +

+ ΔP@k {fmt(result.comparison.delta_precision_at_k)} · ΔR@k {fmt( + result.comparison.delta_recall_at_k, + )} · ΔMRR {fmt(result.comparison.delta_mrr)} · graph overhead {fmtMs( + result.comparison.graph_overhead_ms, + )} +

+

{result.comparison.expected_behavior}

+ {/if} +
+ +
+

Per question

+ + + + + + + + + + + + {#each result?.results || [] as row (row.question_id)} + + + + + + + + {/each} + +
QuestionKindBaseline P@kKG-RAG P@kΔMRR
{row.question}{row.kind}{fmt(row?.baseline?.precision_at_k)}{fmt(row?.kg_rag?.precision_at_k)} + {fmt( + (row?.kg_rag?.mrr ?? 0) - (row?.baseline?.mrr ?? 0), + )} +
+
+ {/if} +
diff --git a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte index 9048c75c7..14d47fa37 100644 --- a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte +++ b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte @@ -21,6 +21,9 @@ import { _ } from '$lib/i18n'; import ConfirmationModal from '$lib/components/modals/ConfirmationModal.svelte'; import AddContentToKSModal from '$lib/components/knowledgeStores/AddContentToKSModal.svelte'; + import KnowledgeStoreGraphView from '$lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte'; + import KnowledgeStoreBenchmarkView from '$lib/components/knowledgeStores/KnowledgeStoreBenchmarkView.svelte'; + import { getGraphStatus } from '$lib/services/graphService'; /** @type {{ ksId: string }} */ let { ksId } = $props(); @@ -31,6 +34,10 @@ let error = $state(''); let successMessage = $state(''); + // KG-RAG / semantic graph state. Loaded lazily after the KS itself. + let graphStatus = $state(/** @type {any} */ (null)); + let activeTab = $state('content'); + // Edit state let editingMeta = $state(false); let editName = $state(''); @@ -115,6 +122,13 @@ editName = ks?.name ?? ''; editDescription = ks?.description ?? ''; schedulePollIfNeeded(); + // Best-effort graph status — KG-RAG is optional, so a 5xx here + // shouldn't break the detail view. + try { + graphStatus = await getGraphStatus(); + } catch (_) { + graphStatus = { enabled: false }; + } } catch (/** @type {unknown} */ err) { console.error('Error loading Knowledge Store:', err); error = err instanceof Error ? err.message : 'Failed to load Knowledge Store'; @@ -445,6 +459,46 @@ + +
+ +
+ + {#if activeTab === 'graph'} + + {:else if activeTab === 'benchmark'} + + {:else}
@@ -708,6 +762,7 @@ {/if}
+ {/if}
diff --git a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte new file mode 100644 index 000000000..d25e40cb2 --- /dev/null +++ b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte @@ -0,0 +1,345 @@ + + + +
+ {#if !graphEnabled} +
+

Graph RAG is not enabled on this Knowledge Store.

+

+ Run the migration below to extract concepts and relationships from + existing chunks. This calls the LLM extractor and writes results to + Neo4j; vector retrieval keeps working either way. +

+
+ + +
+
+ {/if} + + {#if error} +
{error}
+ {/if} + {#if success} +
{success}
+ {/if} + +
+ + + + +
+ + {#if loading} +

Loading graph…

+ {:else if snapshot} +
+
+
Concepts
+
{snapshot?.counts?.concepts ?? 0}
+
+
+
Documents
+
{snapshot?.counts?.documents ?? 0}
+
+
+
Chunks
+
{snapshot?.counts?.chunks ?? 0}
+
+
+
Edges
+
{snapshot?.counts?.edges ?? 0}
+
+
+ +
+

Concepts

+ {#if !snapshot?.nodes?.length} +

No concept nodes yet.

+ {:else} +
    + {#each snapshot.nodes.filter((/** @type {any} */ n) => n.type === 'concept') as node (node.id)} +
  • +
    +
    {node.label}
    +
    {node.data?.verification_state || 'unverified'}
    +
    +
    + + +
    +
    • + {/each} +
+ {/if} +
+ +
+

Relationships

+ {#if !snapshot?.edges?.length} +

No relationships yet.

+ {:else} +
    + {#each snapshot.edges.filter((/** @type {any} */ e) => e.type === 'RELATES_TO') as edge (edge.id)} +
  • +
    + {edge.data?.source_label || edge.source} + → + {edge.data?.target_label || edge.target} + + {edge.label || edge.data?.relation} + +
    +
    + + +
    +
    • + {/each} +
+ {/if} +
+ {/if} + +
+

Recent changes

+ {#if !changes.length} +

No change events yet.

+ {:else} +
    + {#each changes as change (change.event_id)} +
  • +
    {changeOperationLabel(change)}
    +
    {changeMetaLine(change)}
    +
    {changeDetail(change)}
    +
    • + {/each} +
+ {/if} +
+
diff --git a/frontend/svelte-app/src/lib/services/benchmarkService.js b/frontend/svelte-app/src/lib/services/benchmarkService.js new file mode 100644 index 000000000..c48a27cac --- /dev/null +++ b/frontend/svelte-app/src/lib/services/benchmarkService.js @@ -0,0 +1,43 @@ +/** + * @module benchmarkService + * KG-RAG benchmark API client for Knowledge Stores. + * + * Targets the LAMB backend proxy at + * ``/creator/knowledge-stores/{ksId}/benchmarks/...``, which forwards to the + * new KB Server's ``/benchmarks`` router. Like the graph endpoints, these + * are only available when ``KG_RAG_ENABLED=true`` on the KB Server. + */ + +import axios from 'axios'; +import { browser } from '$app/environment'; +import { getApiUrl } from '$lib/config'; + +function authHeaders() { + const token = localStorage.getItem('userToken'); + if (!token) throw new Error('User not authenticated.'); + return { Authorization: `Bearer ${token}` }; +} + +/** + * List built-in benchmark datasets. + */ +export async function listDatasets() { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl('/knowledge-stores/benchmarks/datasets'); + const response = await axios.get(url, { headers: authHeaders() }); + return response.data; +} + +/** + * Run a benchmark dataset against a Knowledge Store, comparing + * ``simple_query`` (vector baseline) against ``kg_rag_query`` (vector + + * graph expansion). + * @param {string} ksId + * @param {{ dataset_id?: string, top_k?: number, graph_depth?: number, threshold?: number }} body + */ +export async function runBenchmark(ksId, body) { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl(`/knowledge-stores/${ksId}/benchmarks/run`); + const response = await axios.post(url, body, { headers: authHeaders() }); + return response.data; +} diff --git a/frontend/svelte-app/src/lib/services/graphService.js b/frontend/svelte-app/src/lib/services/graphService.js new file mode 100644 index 000000000..91de147a2 --- /dev/null +++ b/frontend/svelte-app/src/lib/services/graphService.js @@ -0,0 +1,141 @@ +/** + * @module graphService + * KG-RAG / semantic-graph API client for Knowledge Stores. + * + * Targets the LAMB backend proxy at + * ``/creator/knowledge-stores/{ksId}/graph/...``, which forwards to the new + * KB Server's ``/graph`` router. The KB Server only exposes graph + * endpoints when ``KG_RAG_ENABLED=true`` — callers should gate UI on the + * result of {@link getGraphStatus}. + */ + +import axios from 'axios'; +import { browser } from '$app/environment'; +import { getApiUrl } from '$lib/config'; + +function authHeaders() { + const token = localStorage.getItem('userToken'); + if (!token) { + throw new Error('User not authenticated.'); + } + return { Authorization: `Bearer ${token}` }; +} + +/** + * Get the KG-RAG feature-flag + Neo4j availability for the current user's + * org. Cached by the caller (it shouldn't flap inside a session). + * @returns {Promise<{enabled: boolean, index_on_ingest: boolean, neo4j_configured: boolean, neo4j_available: boolean}>} + */ +export async function getGraphStatus() { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl('/knowledge-stores/graph/status'); + const response = await axios.get(url, { headers: authHeaders() }); + return response.data; +} + +/** + * Trigger KG-RAG migration for an existing Knowledge Store. + * @param {string} ksId + * @param {{ openai_api_key?: string }} [body] + */ +export async function migrateToGraph(ksId, body = {}) { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl(`/knowledge-stores/${ksId}/graph/migrate`); + const response = await axios.post(url, body, { headers: authHeaders() }); + return response.data; +} + +/** + * Fetch the graph snapshot (nodes + edges + counts) for visualization. + * @param {string} ksId + * @param {Object} [params] + */ +export async function getGraphSnapshot(ksId, params = {}) { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl(`/knowledge-stores/${ksId}/graph/snapshot`); + const response = await axios.get(url, { headers: authHeaders(), params }); + return response.data; +} + +/** + * List recent graph change events. + * @param {string} ksId + * @param {Object} [params] + */ +export async function listGraphChanges(ksId, params = {}) { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl(`/knowledge-stores/${ksId}/graph/changes`); + const response = await axios.get(url, { headers: authHeaders(), params }); + return response.data; +} + +/** + * Rename a concept within a Knowledge Store's graph. + * @param {string} ksId + * @param {string} concept + * @param {{ new_name: string, actor?: string, reason?: string }} body + */ +export async function renameConcept(ksId, concept, body) { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl( + `/knowledge-stores/${ksId}/graph/concepts/${encodeURIComponent(concept)}/rename`, + ); + const response = await axios.patch(url, body, { headers: authHeaders() }); + return response.data; +} + +/** + * Merge concepts into a target concept. + * @param {string} ksId + * @param {{ source_names: string[], target_name: string, actor?: string, reason?: string }} body + */ +export async function mergeConcepts(ksId, body) { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl(`/knowledge-stores/${ksId}/graph/concepts/merge`); + const response = await axios.post(url, body, { headers: authHeaders() }); + return response.data; +} + +/** + * Update concept curation metadata (notes / tags / verification_state). + * Approving sets ``verification_state='verified'``, rejecting sets + * ``'rejected'`` which expunges the concept from KG-RAG retrieval but + * keeps the audit event for traceability. + * @param {string} ksId + * @param {string} concept + * @param {Object} body + */ +export async function curateConcept(ksId, concept, body) { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl( + `/knowledge-stores/${ksId}/graph/concepts/${encodeURIComponent(concept)}/curation`, + ); + const response = await axios.patch(url, body, { headers: authHeaders() }); + return response.data; +} + +/** + * Edit a relationship (rename / re-weight / annotate / verify). + * @param {string} ksId + * @param {Object} body + */ +export async function editRelationship(ksId, body) { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl(`/knowledge-stores/${ksId}/graph/relationships`); + const response = await axios.patch(url, body, { headers: authHeaders() }); + return response.data; +} + +/** + * Approve / reject / annotate a relationship without changing its type. + * @param {string} ksId + * @param {Object} body + */ +export async function curateRelationship(ksId, body) { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl( + `/knowledge-stores/${ksId}/graph/relationships/curation`, + ); + const response = await axios.patch(url, body, { headers: authHeaders() }); + return response.data; +} diff --git a/frontend/svelte-app/src/lib/utils/graphCuration.js b/frontend/svelte-app/src/lib/utils/graphCuration.js new file mode 100644 index 000000000..62113b501 --- /dev/null +++ b/frontend/svelte-app/src/lib/utils/graphCuration.js @@ -0,0 +1,268 @@ +const GENERIC_CURATION_REASON = 'Educator graph curation'; +const INTERNAL_ACTORS = new Set(['graph-curation-api', 'graph-traceability-api']); + +/** @typedef {Record} GraphData */ +/** @typedef {{ id: string, type: string, label: string, data: GraphData }} GraphNode */ +/** @typedef {{ id: string, type: string, source: string, target: string, label?: string, weight?: number, data: GraphData }} GraphEdge */ +/** @typedef {{ event_id?: string, operation?: string, actor?: string, timestamp?: string, filename?: string, document_id?: string, concepts?: string[], payload_json?: string | null }} GraphChange */ + +/** @param {GraphChange | null | undefined} change @returns {GraphData} */ +export function changePayload(change) { + try { + const payload = JSON.parse(change?.payload_json || '{}'); + return payload && typeof payload === 'object' ? payload : {}; + } catch { + return {}; + } +} + +/** @param {GraphChange} change */ +export function changeOperationLabel(change) { + const payload = changePayload(change); + const operation = change?.operation || ''; + if (operation === 'automatic_ingestion') return 'Created by ingestion'; + if (operation === 'revert_change') { + if (payload.reverted_operation === 'manual_expunge_relationship') + return 'Relationship restored'; + if (payload.reverted_operation === 'manual_expunge_concept') return 'Concept restored'; + return 'Change reverted'; + } + if (operation === 'manual_edit_relationship') return 'Relationship updated'; + if (operation === 'manual_curate_relationship') return 'Relationship status updated'; + if (operation === 'manual_expunge_relationship') return 'Relationship expunged'; + if (operation === 'manual_curate_concept') return 'Concept status updated'; + if (operation === 'manual_expunge_concept') return 'Concept expunged'; + if (operation === 'manual_rename_concept') return 'Concept renamed'; + if (operation === 'manual_merge_concepts') return 'Concepts merged'; + return operation || 'Graph change'; +} + +/** @param {GraphChange} change */ +export function changeDetail(change) { + const payload = changePayload(change); + const operation = change?.operation || ''; + const parts = []; + if (operation === 'revert_change') { + if (payload.reverted_operation === 'manual_expunge_relationship') + parts.push('Restored relationship as approved'); + else if (payload.reverted_operation === 'manual_expunge_concept') + parts.push('Restored concept as approved'); + else parts.push('Reverted previous graph change'); + if (payload.reason && !isGenericReason(payload.reason)) parts.push(String(payload.reason)); + return parts.join(' | '); + } + if (payload.reason && !isGenericReason(payload.reason)) parts.push(String(payload.reason)); + const transition = verificationTransition(payload); + if (transition) parts.push(transition); + if (operation === 'manual_edit_relationship' && payload.source && payload.target) { + parts.push( + `${payload.source} / ${payload.relation || payload.new_relation || 'related_to'} / ${payload.target}` + ); + if (payload.new_relation && payload.relation && payload.new_relation !== payload.relation) { + parts.push(`Relation ${payload.relation} -> ${payload.new_relation}`); + } + } + if (Array.isArray(payload.removed_chunk_mentions)) { + parts.push(`${payload.removed_chunk_mentions.length} mention links removed`); + } + if (Array.isArray(payload.removed_relationships)) { + parts.push(`${payload.removed_relationships.length} relationships removed`); + } + return parts.join(' | '); +} + +/** @param {unknown} value */ +export function isGenericReason(value) { + return String(value || '').trim() === GENERIC_CURATION_REASON; +} + +/** @param {GraphData} payload */ +export function verificationTransition(payload) { + if ( + !Object.prototype.hasOwnProperty.call(payload || {}, 'verification_state') || + !payload?.verification_state + ) + return ''; + const oldState = String(payload.old_verification_state || 'unverified'); + const newState = String(payload.verification_state || 'unverified'); + if (oldState === newState) return ''; + return `${stateLabel(oldState)} -> ${stateLabel(newState)}`; +} + +/** @param {GraphChange} change */ +export function changeMetaLine(change) { + const parts = []; + const source = change?.filename || change?.document_id || ''; + const actor = String(change?.actor || '').trim(); + if (source) parts.push(source); + if (actor && !INTERNAL_ACTORS.has(actor)) parts.push(actorLabel(actor)); + return parts.join(' / '); +} + +/** @param {string} actor */ +export function actorLabel(actor) { + if (actor === 'lamb-ingestion-pipeline') return 'Ingestion pipeline'; + return actor; +} + +/** @param {unknown} value */ +export function stateLabel(value) { + const state = String(value || 'unverified'); + if (state === 'unverified') return 'Unreviewed'; + if (state === 'needs_review') return 'Needs review'; + if (state === 'verified') return 'Approved'; + if (state === 'rejected') return 'Expunged'; + return state; +} + +/** @param {unknown} value */ +export function stateRank(value) { + const state = String(value || 'unverified'); + if (state === 'needs_review') return 0; + if (state === 'unverified') return 1; + if (state === 'rejected') return 2; + if (state === 'verified') return 3; + return 4; +} + +/** @param {unknown} value */ +export function stateClass(value) { + const state = String(value || 'unverified'); + if (state === 'verified') return 'bg-emerald-50 text-emerald-700 ring-emerald-200'; + if (state === 'rejected') return 'bg-red-50 text-red-700 ring-red-200'; + if (state === 'needs_review') return 'bg-amber-50 text-amber-800 ring-amber-200'; + return 'bg-slate-50 text-slate-700 ring-slate-200'; +} + +/** + * @param {GraphChange[]} sourceChanges + * @param {GraphEdge[]} activeEdges + * @returns {GraphEdge[]} + */ +export function buildExpungedRelationships(sourceChanges, activeEdges) { + const activeKeys = new Set( + (activeEdges || []).map((edge) => + relationshipIdentityKey(edge.data?.source, edge.data?.target, edge.data?.relation) + ) + ); + const seen = new Set(); + const rows = []; + for (const change of sourceChanges || []) { + const payload = changePayload(change); + const candidates = []; + if (change.operation === 'manual_expunge_relationship') { + candidates.push({ + source: payload.source, + target: payload.target, + relation: payload.relation || payload.new_relation || 'related_to', + weight: payload.old_weight, + description: payload.old_description, + evidence: payload.old_evidence, + chunk_id: payload.old_chunk_id, + notes: payload.old_notes, + tags: payload.old_tags + }); + } + if ( + change.operation === 'manual_expunge_concept' && + Array.isArray(payload.removed_relationships) + ) { + candidates.push( + ...payload.removed_relationships.map((relationship) => ({ + ...relationship, + expunged_by_concept: payload.concept + })) + ); + } + for (const [index, relationship] of candidates.entries()) { + const source = String(relationship.source || '').trim(); + const target = String(relationship.target || '').trim(); + const relation = String(relationship.relation || 'related_to').trim() || 'related_to'; + const key = relationshipIdentityKey(source, target, relation); + if (!source || !target || activeKeys.has(key) || seen.has(key)) continue; + seen.add(key); + rows.push({ + id: `expunged-relationship-${change.event_id || key}${change.operation === 'manual_expunge_concept' ? `-${index}` : ''}`, + type: 'RELATES_TO', + source: `concept:${source}`, + target: `concept:${target}`, + label: relation, + weight: valueOrUndefined(relationship.weight), + data: { + source, + target, + source_label: source, + target_label: target, + relation, + description: relationship.description || '', + evidence: relationship.evidence || '', + chunk_id: relationship.chunk_id || '', + notes: relationship.notes || '', + tags: Array.isArray(relationship.tags) ? relationship.tags : [], + verification_state: 'rejected', + expunged: true, + expunge_event_id: change.event_id, + expunged_at: change.timestamp, + expunged_by_concept: relationship.expunged_by_concept || '' + } + }); + } + } + return rows; +} + +/** + * @param {GraphChange[]} sourceChanges + * @param {GraphNode[]} activeNodes + * @returns {GraphNode[]} + */ +export function buildExpungedConcepts(sourceChanges, activeNodes) { + const activeNames = new Set((activeNodes || []).map((node) => String(node.data?.name || ''))); + const seen = new Set(); + const rows = []; + for (const change of sourceChanges || []) { + if (change.operation !== 'manual_expunge_concept') continue; + const payload = changePayload(change); + const concept = String(payload.concept || change.concepts?.[0] || '').trim(); + if (!concept || activeNames.has(concept) || seen.has(concept)) continue; + seen.add(concept); + const removedMentions = Array.isArray(payload.removed_chunk_mentions) + ? payload.removed_chunk_mentions + : []; + const removedRelationships = Array.isArray(payload.removed_relationships) + ? payload.removed_relationships + : []; + rows.push({ + id: `expunged-concept-${change.event_id || concept}`, + type: 'concept', + label: concept, + data: { + name: concept, + entity_type: 'concept', + chunk_count: removedMentions.length, + relationship_count: removedRelationships.length, + notes: payload.old_notes || '', + tags: Array.isArray(payload.old_tags) ? payload.old_tags : [], + verification_state: 'rejected', + expunged: true, + expunge_event_id: change.event_id, + expunged_at: change.timestamp + } + }); + } + return rows; +} + +export function relationshipIdentityKey(source, target, relation) { + return [source, target, relation] + .map((value) => + String(value || '') + .trim() + .toLowerCase() + ) + .join('|'); +} + +export function valueOrUndefined(value) { + return value === undefined || value === null || value === '' ? undefined : value; +} diff --git a/lamb-kb-server/backend/config.py b/lamb-kb-server/backend/config.py index 0f270561d..c44358597 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") @@ -52,6 +56,81 @@ 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 + ), + } + + +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/models.py b/lamb-kb-server/backend/database/models.py index 2e89fb9c5..f21ee3034 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,13 @@ 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) + # --- 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 e87659b4f..5ad3347ae 100644 --- a/lamb-kb-server/backend/main.py +++ b/lamb-kb-server/backend/main.py @@ -99,6 +99,17 @@ async def log_requests(request, call_next): 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: + 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() -> None: 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..b19c484f4 --- /dev/null +++ b/lamb-kb-server/backend/plugins/kg_rag_query.py @@ -0,0 +1,282 @@ +"""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)) + graph_limit_factor = int( + params.get("graph_limit_factor") or kg_config.get("limit_factor") or 4 + ) + graph_limit_factor = max(1, min(graph_limit_factor, 20)) + + 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": [], + "traversed_edges": [], + "expanded_chunk_ids": [], + "latest_changes": [], + "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) + + graph_start = time.perf_counter() + try: + expansion = graph_store.expand_from_chunks( + collection_id=str(collection.id), + org_id=str(collection.organization_id), + seed_chunk_ids=seed_chunk_ids, + depth=graph_depth, + limit=top_k * graph_limit_factor, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("KG-RAG graph expansion failed: %s", exc) + trace["warnings"].append(f"Graph expansion failed: {exc}") + trace["graph_latency_ms"] = (time.perf_counter() - graph_start) * 1000 + return self._attach_trace(baseline_results, trace, include_trace) + + trace.update( + { + "graph_expanded": bool(expansion.get("expanded_chunk_ids")), + "entry_concepts": expansion.get("entry_concepts", []), + "traversed_edges": expansion.get("traversed_edges", []), + "expanded_chunk_ids": expansion.get("expanded_chunk_ids", []), + "latest_changes": expansion.get("latest_changes", []), + "graph_latency_ms": expansion.get( + "graph_latency_ms", + (time.perf_counter() - graph_start) * 1000, + ), + } + ) + + expanded_ids = [ + cid + for cid in expansion.get("expanded_chunk_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._merge_results(baseline_results + expanded_results, top_k=top_k) + 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) + + # ------------------------------------------------------------------ + # 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 additional chunks discovered by graph expansion. + + The new vector backend abstraction (``VectorDBBackend``) exposes + ``query`` / ``add_chunks`` / ``delete_by_source`` but not a direct + ``get_by_id``. ChromaDB collections opened through the backend's + client cache support ``get`` natively, so we reach into that cache + for the ChromaDB backend. Other backends fall back to an empty + result, which the trace surfaces as a warning. + """ + if not expanded_ids: + return [] + + # ChromaDB backend has a module-level client cache we can reuse. + # Other backends do not have a stable lookup-by-id surface yet, so + # we return [] and let the trace warn. + try: + from plugins.vector_db import chromadb_backend as _chroma_mod + + client = _chroma_mod._get_client(collection.storage_path) + from plugins.vector_db.chromadb_backend import _to_chroma_ef + + chroma_collection = client.get_collection( + name=collection.backend_collection_id or str(collection.id), + embedding_function=_to_chroma_ef(embedding_function), + ) + rows = chroma_collection.get( + ids=expanded_ids, + include=["documents", "metadatas"], + ) + except Exception as exc: # noqa: BLE001 + logger.debug("Could not fetch expanded chunks: %s", exc) + return [] + + ids = rows.get("ids") or expanded_ids + documents = rows.get("documents") or [] + metadatas = rows.get("metadatas") or [] + results: list[dict[str, Any]] = [] + for index, chunk_id in enumerate(ids): + document = documents[index] if index < len(documents) else "" + metadata = dict(metadatas[index] or {}) if index < len(metadatas) else {} + metadata.setdefault("document_id", chunk_id) + metadata["kg_rag_origin"] = "graph_expansion" + data = metadata.get("parent_text") if return_parent_context else None + results.append( + { + "similarity": 0.72, + "data": data or document, + "metadata": metadata, + } + ) + return results + + def _merge_results( + self, results: list[dict[str, Any]], top_k: int + ) -> list[dict[str, Any]]: + by_chunk_id: dict[str, dict[str, Any]] = {} + for result in results: + chunk_id = self._result_chunk_id(result) or str(result.get("data", ""))[:120] + existing = by_chunk_id.get(chunk_id) + if existing is None or float(result.get("similarity", 0.0)) > float( + existing.get("similarity", 0.0) + ): + by_chunk_id[chunk_id] = result + + ordered = sorted( + by_chunk_id.values(), + key=lambda item: float(item.get("similarity", 0.0)), + reverse=True, + ) + return ordered[: max(top_k, 1) + 3] + + @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/routers/benchmarks.py b/lamb-kb-server/backend/routers/benchmarks.py new file mode 100644 index 000000000..e33c5d391 --- /dev/null +++ b/lamb-kb-server/backend/routers/benchmarks.py @@ -0,0 +1,93 @@ +"""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, Body, Depends +from schemas.benchmark import ( + BenchmarkDataset, + BenchmarkDatasetSummary, + BenchmarkRunAllRequest, + BenchmarkRunAllResponse, + BenchmarkRunRequest, + BenchmarkRunResponse, +) +from schemas.content import EmbeddingCredentials +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, + request: BenchmarkRunRequest, + embedding_credentials: EmbeddingCredentials = Body( + default_factory=EmbeddingCredentials, + embed=True, + ), + db: Session = Depends(get_session), +) -> BenchmarkRunResponse: + return BenchmarkService.run( + db=db, + collection_id=collection_id, + request=request, + embedding_credentials={ + "api_key": embedding_credentials.api_key, + "api_endpoint": embedding_credentials.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, + request: BenchmarkRunAllRequest, + embedding_credentials: EmbeddingCredentials = Body( + default_factory=EmbeddingCredentials, + embed=True, + ), + db: Session = Depends(get_session), +) -> BenchmarkRunAllResponse: + return BenchmarkService.run_all( + db=db, + collection_id=collection_id, + dataset_ids=request.dataset_ids, + threshold=request.threshold, + embedding_credentials={ + "api_key": embedding_credentials.api_key, + "api_endpoint": embedding_credentials.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..8d9b09e3e --- /dev/null +++ b/lamb-kb-server/backend/routers/graph.py @@ -0,0 +1,546 @@ +"""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 ( + GraphAuditRequest, + GraphAuditResponse, + 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, + } + + +def _collection_dict(collection: Collection) -> dict[str, Any]: + """Adapt a new Collection ORM row to the legacy ``{id, owner, ...}`` + dict shape that graph_store internals still expect.""" + return { + "id": collection.id, + "name": collection.name, + "organization_id": collection.organization_id, + "owner": collection.organization_id, + } + + +@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() + + # Pull chunks from ChromaDB (the only backend with a stable get-by-id + # path right now). Qdrant migration is a TODO. + if collection.vector_db_backend != "chromadb": + raise HTTPException( + status_code=400, + detail=( + "Graph migration currently only supports the chromadb vector " + f"backend (collection uses '{collection.vector_db_backend}')." + ), + ) + + from plugins.vector_db import chromadb_backend as _chroma_mod + + client = _chroma_mod._get_client(collection.storage_path) + try: + chroma_collection = client.get_collection( + name=collection.backend_collection_id or collection.id + ) + except Exception as exc: + raise HTTPException( + status_code=404, + detail=f"ChromaDB collection for {collection.id} not found", + ) from exc + + ids: list[str] = [] + texts: list[str] = [] + metadatas: list[dict[str, Any]] = [] + offset = 0 + batch_size = 500 + + while True: + result = chroma_collection.get( + include=["documents", "metadatas"], + limit=batch_size, + offset=offset, + ) + result_ids = result.get("ids") or [] + documents = result.get("documents") or [] + result_metadatas = result.get("metadatas") or [] + if not result_ids: + break + + for i, chunk_id in enumerate(result_ids): + text = documents[i] if i < len(documents) else None + if not text: + continue + metadata = ( + result_metadatas[i] + if i < len(result_metadatas) + and isinstance(result_metadatas[i], dict) + else {} + ) + ids.append(chunk_id) + texts.append(text) + metadatas.append(metadata) + + offset += len(result_ids) + + 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=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.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}/audit-trace", + response_model=GraphAuditResponse, + summary="Audit a KG-RAG graph retrieval trace", +) +async def audit_graph_trace( + collection_id: str, + request: GraphAuditRequest, + token: str = Depends(verify_token), + db: Session = Depends(get_session), +): + collection = _get_collection_or_404(db, collection_id) + graph_store = _graph_store_or_503() + trace = graph_store.expand_from_chunks( + collection_id=collection_id, + org_id=str(collection.organization_id), + seed_chunk_ids=request.seed_chunk_ids, + depth=request.graph_depth, + limit=request.limit, + ) + return { + "collection_id": collection_id, + "seed_chunk_ids": request.seed_chunk_ids, + "trace": trace, + } + + +@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/schemas/benchmark.py b/lamb-kb-server/backend/schemas/benchmark.py new file mode 100644 index 000000000..8f2bf1585 --- /dev/null +++ b/lamb-kb-server/backend/schemas/benchmark.py @@ -0,0 +1,109 @@ +"""Schemas for KG-RAG benchmark evaluation.""" + +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, 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 BenchmarkRunRequest(BaseModel): + 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) + + +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) + + +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 6c63cf0a7..27f9ab379 100644 --- a/lamb-kb-server/backend/schemas/collection.py +++ b/lamb-kb-server/backend/schemas/collection.py @@ -48,6 +48,14 @@ class CreateCollectionRequest(BaseModel): default="chromadb", description="Registered vector DB backend name.", ) + 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." + ), + ) class UpdateCollectionRequest(BaseModel): @@ -88,6 +96,7 @@ class CollectionResponse(BaseModel): chunking_params: dict[str, Any] embedding: EmbeddingConfig vector_db_backend: str + graph_enabled: bool = False status: str document_count: int chunk_count: int @@ -117,6 +126,7 @@ 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)), 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..f96e05254 --- /dev/null +++ b/lamb-kb-server/backend/schemas/graph.py @@ -0,0 +1,156 @@ +"""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 GraphAuditRequest(BaseModel): + seed_chunk_ids: List[str] = Field( + ..., description="Chroma/graph chunk IDs to use as graph entry points" + ) + graph_depth: int = Field(2, description="RELATES_TO traversal depth") + limit: int = Field(20, description="Maximum number of expanded chunks to return") + + +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 GraphAuditResponse(BaseModel): + collection_id: str = Field(..., description="Collection ID") + seed_chunk_ids: List[str] = Field(..., description="Requested seed chunk IDs") + trace: Dict[str, Any] = Field(..., description="Graph expansion trace") + + +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/services/benchmark.py b/lamb-kb-server/backend/services/benchmark.py new file mode 100644 index 000000000..abfbe00b6 --- /dev/null +++ b/lamb-kb-server/backend/services/benchmark.py @@ -0,0 +1,699 @@ +"""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_response = query_with_plugin( + db=db, + collection_id=collection_id, + query_text=question.question, + plugin_name="kg_rag_query", + plugin_params={ + "top_k": top_k, + "threshold": threshold, + "graph_depth": graph_depth, + "include_trace": True, + }, + 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 d866ca754..046317f11 100644 --- a/lamb-kb-server/backend/services/collection_service.py +++ b/lamb-kb-server/backend/services/collection_service.py @@ -143,6 +143,7 @@ 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)), status="ready", document_count=0, chunk_count=0, 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..89ef6d093 --- /dev/null +++ b/lamb-kb-server/backend/services/concept_extraction.py @@ -0,0 +1,328 @@ +"""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 + +try: + from openai import OpenAI +except Exception: # pragma: no cover - exercised when optional dependency is absent + OpenAI = None + + +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 + if not any(char.isalpha() or char.isdigit() for char in normalized): + 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, + client: Optional[Any] = None, + ): + 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 = 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)) + self.client = client + + api_key = self.config.get("openai_api_key") or "" + if self.client is None and api_key and OpenAI is not None: + self.client = OpenAI(api_key=api_key) + + 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.client 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", + } + ], + }, + "limits": {"max_entities": 5, "max_relationships": 6}, + "text": text[:6000], + } + try: + response = self._create_json_completion( + model=self.model, + messages=[ + {"role": "system", "content": system}, + {"role": "user", "content": json.dumps(user, ensure_ascii=False)}, + ], + ) + content = response.choices[0].message.content or "{}" + parsed = json.loads(content) + except Exception as exc: + logger.warning("KG-RAG concept extraction failed: %s", exc) + return {"entities": [], "relationships": []} + return ( + parsed + if isinstance(parsed, dict) + else {"entities": [], "relationships": []} + ) + + def _create_json_completion( + self, model: str, messages: List[Dict[str, str]] + ) -> Any: + if self.client is None: + raise RuntimeError("OpenAI client is not configured") + try: + return self.client.chat.completions.create( + model=model, + messages=messages, + response_format={"type": "json_object"}, + temperature=0, + ) + except Exception: + fallback_model = self.chat_model + if model == fallback_model: + raise + return self.client.chat.completions.create( + model=fallback_model, + messages=messages, + response_format={"type": "json_object"}, + temperature=0, + ) + + 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[:5]: + 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[:6]: + 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..5f8045ace --- /dev/null +++ b/lamb-kb-server/backend/services/graph_indexing.py @@ -0,0 +1,180 @@ +"""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"} + + api_key = (openai_api_key or kg_config.get("openai_api_key") or "").strip() + if 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) + + 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..568651588 --- /dev/null +++ b/lamb-kb-server/backend/services/graph_store.py @@ -0,0 +1,2472 @@ +"""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 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 coalesce(concept.verification_state, '') <> 'rejected' + 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, + concept.verification_state AS verification_state, + 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}, + } + + 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, + 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": row.get("verification_state") + 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 "", + "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"} + + 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}, + } + + 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, + concept.verification_state = CASE WHEN $verification_state IS NULL THEN concept.verification_state ELSE $verification_state END + """, + org_id=org_id, + concept=concept, + notes=notes, + tags=tags, + verification_state=verification_state, + 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 {} + 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 + 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, + filename=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), + 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_chunks( + self, + collection_id: str, + org_id: str, + seed_chunk_ids: List[str], + depth: int, + limit: int, + ) -> Dict[str, Any]: + depth = max(1, min(int(depth or 2), 4)) + limit = max(1, int(limit or 10)) + start = time.perf_counter() + if not seed_chunk_ids: + 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: + entry_rows = session.run( + """ + MATCH (chunk:Chunk {collection_id: $collection_id})-[:MENTIONS]->(concept:Concept {org_id: $org_id}) + WHERE chunk.chunk_id IN $seed_chunk_ids + AND coalesce(concept.verification_state, '') <> 'rejected' + RETURN concept.name AS name, count(*) AS mentions + ORDER BY mentions DESC, name ASC + LIMIT 8 + """, + collection_id=collection_id, + org_id=org_id, + seed_chunk_ids=seed_chunk_ids, + ).data() + entry_concepts = [row["name"] for row in entry_rows] + + if not entry_concepts: + return self._empty_expansion(start) + + relation_path_query = f""" + MATCH (entry:Concept {{org_id: $org_id}}) + WHERE entry.name IN $entry_concepts + 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 coalesce(node.verification_state, '') <> 'rejected') + AND all(rel IN relationships(path) WHERE coalesce(rel.verification_state, '') <> 'rejected') + WITH entry, related, path, + length(path) AS hops, + reduce(score = 0.0, rel IN relationships(path) | + score + 2.0 * coalesce(rel.weight, 1.0) + ) AS path_score + ORDER BY path_score DESC, hops ASC, related.name ASC + LIMIT $limit + OPTIONAL MATCH (related)<-[:MENTIONS]-(chunk:Chunk {{collection_id: $collection_id}}) + RETURN entry.name AS entry, + related.name AS related, + [rel IN relationships(path) | {{type: coalesce(rel.relation, type(rel)), raw_type: type(rel), source: startNode(rel).name, target: endNode(rel).name, weight: coalesce(rel.weight, 1), description: coalesce(rel.description, '')}}] AS edges, + collect(DISTINCT chunk.chunk_id)[0..6] AS chunk_ids, + hops AS hops, + path_score AS score + ORDER BY score DESC, hops ASC + """ + expanded_rows = session.run( + relation_path_query, + org_id=org_id, + collection_id=collection_id, + entry_concepts=entry_concepts, + limit=limit, + ).data() + + direct_rows = session.run( + """ + MATCH (concept:Concept {org_id: $org_id})<-[:MENTIONS]-(chunk:Chunk {collection_id: $collection_id}) + WHERE concept.name IN $entry_concepts + AND coalesce(concept.verification_state, '') <> 'rejected' + RETURN chunk.chunk_id AS chunk_id, + count(*) AS mentions, + min(chunk.source_label) AS source_label + ORDER BY mentions DESC, source_label ASC + LIMIT $limit + """, + org_id=org_id, + collection_id=collection_id, + entry_concepts=entry_concepts, + limit=limit, + ).data() + + chunk_ids: List[str] = [] + seen_chunk_ids: Set[str] = set() + + def add_chunk_id(chunk_id: str) -> None: + if chunk_id and chunk_id not in seen_chunk_ids: + seen_chunk_ids.add(chunk_id) + chunk_ids.append(chunk_id) + + for row in direct_rows: + add_chunk_id(row.get("chunk_id")) + + edges: List[Dict[str, Any]] = [] + related_concepts: Set[str] = set(entry_concepts) + for row in expanded_rows: + if row.get("related"): + related_concepts.add(row.get("related")) + for chunk_id in row.get("chunk_ids", []): + add_chunk_id(chunk_id) + edges.extend(row.get("edges") or []) + + changes = session.run( + """ + MATCH (event:ChangeEvent {collection_id: $collection_id, org_id: $org_id}) + WHERE any(concept IN coalesce(event.concepts, []) WHERE concept IN $concepts) + RETURN event.operation AS operation, + event.actor AS actor, + event.timestamp AS timestamp, + event.filename AS filename, + event.concepts AS concepts, + event.payload_json AS payload_json + ORDER BY event.timestamp DESC + LIMIT 10 + """, + collection_id=collection_id, + org_id=org_id, + concepts=[concept for concept in related_concepts if concept], + ).data() + + return { + "entry_concepts": entry_concepts, + "expanded_chunk_ids": chunk_ids[:limit], + "traversed_edges": self._dedupe_edges(edges), + "latest_changes": 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, + } + + @staticmethod + def _dedupe_edges(edges: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + seen = set() + deduped = [] + for edge in edges: + key = (edge.get("source"), edge.get("target"), edge.get("type")) + if key in seen: + continue + seen.add(key) + deduped.append(edge) + return deduped[:80] diff --git a/lamb-kb-server/backend/services/ingestion_service.py b/lamb-kb-server/backend/services/ingestion_service.py index f2ca94b68..3a05eeb28 100644 --- a/lamb-kb-server/backend/services/ingestion_service.py +++ b/lamb-kb-server/backend/services/ingestion_service.py @@ -275,6 +275,105 @@ 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: object, + embedding_function: object, + 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. + """ + 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 + + # The chunks were already embedded and stored; pull them back from the + # vector backend so we have stable IDs and the exact text that's + # searchable. + try: + from plugins.vector_db import chromadb_backend as _chroma_mod # noqa: PLC0415 + + if collection.vector_db_backend != "chromadb": + logger.warning( + "Graph indexing currently only supports chromadb backend; " + "skipping for collection %s (backend=%s).", + collection.id, + collection.vector_db_backend, + ) + return + + client = _chroma_mod._get_client(collection.storage_path) + from plugins.vector_db.chromadb_backend import _to_chroma_ef # noqa: PLC0415 + + chroma_collection = client.get_collection( + name=collection.backend_collection_id or collection.id, + embedding_function=_to_chroma_ef(embedding_function), # type: ignore[arg-type] + ) + except Exception as exc: # noqa: BLE001 + logger.warning( + "Graph indexing: could not open ChromaDB collection %s: %s", + collection.id, + exc, + ) + return + + source_ids = [doc["source_item_id"] for doc in docs_list] + try: + from services.graph_indexing import ( # noqa: PLC0415 + index_chunks_for_collection, + ) + + for source_id in source_ids: + rows = chroma_collection.get( + where={"source_item_id": source_id}, + include=["documents", "metadatas"], + ) + ids = rows.get("ids") or [] + texts = rows.get("documents") or [] + metadatas = rows.get("metadatas") or [] + if not ids: + continue + result = index_chunks_for_collection( + collection=collection, + ids=list(ids), + texts=list(texts), + metadatas=[dict(m or {}) for m in metadatas], + 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"], + ) + except Exception as exc: # noqa: BLE001 + logger.warning( + "Graph indexing failed for collection %s: %s", collection.id, exc + ) + def delete_vectors( db: Session, collection_id: str, source_item_id: str diff --git a/lamb-kb-server/backend/services/query_service.py b/lamb-kb-server/backend/services/query_service.py index ade1dbbf7..97c988ec4 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 @@ -69,3 +70,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/pyproject.toml b/lamb-kb-server/pyproject.toml index c3b4e9e11..2f57d53dc 100644 --- a/lamb-kb-server/pyproject.toml +++ b/lamb-kb-server/pyproject.toml @@ -47,8 +47,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 = [ 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..6ed213a5e --- /dev/null +++ b/lamb-kb-server/tests/unit/test_kg_rag.py @@ -0,0 +1,187 @@ +"""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 +from collections.abc import Iterator + +import pytest + + +@pytest.fixture() +def reload_config(monkeypatch) -> Iterator[None]: + """Local copy of the reload_config fixture from test_config.py.""" + import config # noqa: PLC0415 + + yield + importlib.reload(config) + + +# --------------------------------------------------------------------------- +# 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"]) From e5152243d8c3cfc1c79ff0eefd20637984b1b6e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A0=20Guilera?= Date: Sun, 17 May 2026 14:50:53 +0200 Subject: [PATCH 002/156] fix(kg-rag): address review feedback from self-review Blocking fixes: - Benchmark route now uses a single body model with embedded ``embedding_credentials``, so the LAMB proxy's flat JSON validates (previously FastAPI's Body(embed=True) required wrapping under ``request``). Adds an EmbeddingCredentialsBody schema. - init_db now runs lightweight ALTER TABLE migrations so existing installations get the new ``graph_enabled`` column without an ``OperationalError: no such column``. Uses a direct sqlite3 connection (not the SQLAlchemy engine) so the migration check doesn't leave WAL state that breaks fork-based lock tests. Only runs against pre-existing DBs to keep the hot path zero-cost. Significant: - Permalinks now flow through to Neo4j Chunk nodes (permalink_original / permalink_full_markdown / permalink_page). Snapshot endpoint exposes them on chunk node data for citations. Moderate: - Vector backend gets public ``get_chunks_by_id``, ``get_chunks_by_source``, and ``iter_all_chunks`` surfaces. KG-RAG plugin and migration endpoint now call those instead of reaching into ChromaDB's private client cache. Default ``[]`` / NotImplementedError lets non-chromadb backends degrade cleanly. - Graph migration endpoint surfaces unsupported-backend (e.g. qdrant) as a clear 400 with explanation; the frontend hides the migrate button and explains why instead of letting the user click into an error. - ConceptExtractor passes a configurable per-request OpenAI timeout (default 60s) so a hung vendor can't pin an ingestion worker indefinitely. New ``KG_RAG_OPENAI_TIMEOUT_SECONDS`` env knob. - Graph + Benchmark tabs are now gated on ``ks.graph_enabled``, with a graceful fallback to "show tabs if the store could be migrated". Documented ingest-time slowdown in .env.next.example. Nits: - Drop unused ``_collection_dict`` helper in routers/graph.py. - Lazy-import OrganizationConfigResolver in the LAMB proxy router to match existing convention. - Move the ``reload_config`` fixture into tests/unit/conftest.py so any unit test can pick it up. Tests: - New regression tests for the proxy-flat body shape and the schema migration. Idempotent re-run check is included. Full unit (308) + integration (179 + 1 skipped) suites pass. --- .env.next.example | 7 + .../knowledge_store_graph_router.py | 3 +- .../KnowledgeStoreDetail.svelte | 19 ++- .../KnowledgeStoreGraphView.svelte | 64 +++++--- lamb-kb-server/backend/config.py | 6 + lamb-kb-server/backend/database/connection.py | 61 ++++++++ lamb-kb-server/backend/plugins/base.py | 51 +++++++ .../backend/plugins/kg_rag_query.py | 50 +++--- .../plugins/vector_db/chromadb_backend.py | 121 +++++++++++++++ lamb-kb-server/backend/routers/benchmarks.py | 37 +++-- lamb-kb-server/backend/routers/graph.py | 106 ++++++------- lamb-kb-server/backend/schemas/benchmark.py | 27 ++++ .../backend/services/concept_extraction.py | 8 +- .../backend/services/graph_store.py | 32 +++- .../backend/services/ingestion_service.py | 114 +++++++------- lamb-kb-server/tests/unit/conftest.py | 16 ++ lamb-kb-server/tests/unit/test_kg_rag.py | 142 ++++++++++++++++-- 17 files changed, 654 insertions(+), 210 deletions(-) diff --git a/.env.next.example b/.env.next.example index ece16f648..6966994b7 100644 --- a/.env.next.example +++ b/.env.next.example @@ -123,6 +123,11 @@ OWI_ADMIN_PASSWORD=admin # 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 @@ -133,3 +138,5 @@ OWI_ADMIN_PASSWORD=admin # 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/backend/creator_interface/knowledge_store_graph_router.py b/backend/creator_interface/knowledge_store_graph_router.py index 969d59b72..bc1a2ee67 100644 --- a/backend/creator_interface/knowledge_store_graph_router.py +++ b/backend/creator_interface/knowledge_store_graph_router.py @@ -19,7 +19,6 @@ from pydantic import BaseModel, Field from lamb.auth_context import AuthContext, get_auth_context -from lamb.completions.org_config_resolver import OrganizationConfigResolver from lamb.database_manager import LambDatabaseManager from .knowledge_store_client import KnowledgeStoreClient @@ -91,6 +90,7 @@ async def migrate_knowledge_store_to_graph( 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 "" @@ -308,6 +308,7 @@ async def run_benchmark( auth: AuthContext = Depends(get_auth_context), ): _assert_ks_access(ks_id, auth) + from lamb.completions.org_config_resolver import OrganizationConfigResolver resolver = OrganizationConfigResolver(auth.user.get("email")) # Resolve embedding key the same way ingestion does, so the benchmark # baseline vector pass works without the caller threading creds. diff --git a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte index 14d47fa37..2b3fffc89 100644 --- a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte +++ b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte @@ -459,7 +459,16 @@
- + + {@const canEnableGraph = + graphStatus?.enabled && ks?.vector_db_backend === 'chromadb'} + {@const showGraphTabs = !!ks?.graph_enabled || canEnableGraph}
+ {#if migrationSupported} +

+ Run the migration below to extract concepts and relationships from + existing chunks. This calls the LLM extractor and writes results to + Neo4j; vector retrieval keeps working either way. +

+
+ + +
+ {:else} +

+ Graph migration is not yet supported for the + {vectorDbBackend} + vector backend. Migration currently requires + chromadb. + Create a new Knowledge Store with chromadb to use Graph RAG. +

+ {/if} {/if} diff --git a/lamb-kb-server/backend/config.py b/lamb-kb-server/backend/config.py index c44358597..2fce36068 100644 --- a/lamb-kb-server/backend/config.py +++ b/lamb-kb-server/backend/config.py @@ -125,6 +125,12 @@ def get_kg_rag_config() -> dict[str, Any]: "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 + ), } diff --git a/lamb-kb-server/backend/database/connection.py b/lamb-kb-server/backend/database/connection.py index 2c19ad44b..0bbd2a93e 100644 --- a/lamb-kb-server/backend/database/connection.py +++ b/lamb-kb-server/backend/database/connection.py @@ -59,6 +59,14 @@ def init_db() -> None: "Only one instance may run per data directory." ) from exc + # Remember whether we're opening a pre-existing DB. If yes, after + # ``create_all`` (which is a no-op on existing tables) we additionally + # run lightweight ALTER TABLE migrations so new columns get added to + # already-populated installations. On a fresh DB the model definition + # already includes everything, so we skip the migration round-trip — + # that keeps init_db's hot path minimal. + db_existed = DB_PATH.exists() + _engine = create_engine( f"sqlite:///{DB_PATH}", pool_pre_ping=True, @@ -68,12 +76,65 @@ def init_db() -> None: event.listen(_engine, "connect", _enable_sqlite_wal) Base.metadata.create_all(bind=_engine) + if db_existed: + _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", + ), + ] + # 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/plugins/base.py b/lamb-kb-server/backend/plugins/base.py index 2a1ddc937..49fc72009 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 [] diff --git a/lamb-kb-server/backend/plugins/kg_rag_query.py b/lamb-kb-server/backend/plugins/kg_rag_query.py index b19c484f4..c3aaf2f76 100644 --- a/lamb-kb-server/backend/plugins/kg_rag_query.py +++ b/lamb-kb-server/backend/plugins/kg_rag_query.py @@ -193,53 +193,37 @@ def _fetch_expanded_results( expanded_ids: list[str], return_parent_context: bool, ) -> list[dict[str, Any]]: - """Look up additional chunks discovered by graph expansion. - - The new vector backend abstraction (``VectorDBBackend``) exposes - ``query`` / ``add_chunks`` / ``delete_by_source`` but not a direct - ``get_by_id``. ChromaDB collections opened through the backend's - client cache support ``get`` natively, so we reach into that cache - for the ChromaDB backend. Other backends fall back to an empty - result, which the trace surfaces as a warning. + """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 [] - # ChromaDB backend has a module-level client cache we can reuse. - # Other backends do not have a stable lookup-by-id surface yet, so - # we return [] and let the trace warn. try: - from plugins.vector_db import chromadb_backend as _chroma_mod - - client = _chroma_mod._get_client(collection.storage_path) - from plugins.vector_db.chromadb_backend import _to_chroma_ef - - chroma_collection = client.get_collection( - name=collection.backend_collection_id or str(collection.id), - embedding_function=_to_chroma_ef(embedding_function), - ) - rows = chroma_collection.get( - ids=expanded_ids, - include=["documents", "metadatas"], + 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 [] - ids = rows.get("ids") or expanded_ids - documents = rows.get("documents") or [] - metadatas = rows.get("metadatas") or [] results: list[dict[str, Any]] = [] - for index, chunk_id in enumerate(ids): - document = documents[index] if index < len(documents) else "" - metadata = dict(metadatas[index] or {}) if index < len(metadatas) else {} - metadata.setdefault("document_id", chunk_id) + for item in fetched: + metadata = dict(item.metadata or {}) metadata["kg_rag_origin"] = "graph_expansion" - data = metadata.get("parent_text") if return_parent_context else None + data = metadata.pop("parent_text", None) if return_parent_context else None results.append( { - "similarity": 0.72, - "data": data or document, + "similarity": item.score or 0.72, + "data": data or item.text, "metadata": metadata, } ) 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 d517b2fce..57361a999 100644 --- a/lamb-kb-server/backend/plugins/vector_db/chromadb_backend.py +++ b/lamb-kb-server/backend/plugins/vector_db/chromadb_backend.py @@ -229,6 +229,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, *, diff --git a/lamb-kb-server/backend/routers/benchmarks.py b/lamb-kb-server/backend/routers/benchmarks.py index e33c5d391..ce341cfa1 100644 --- a/lamb-kb-server/backend/routers/benchmarks.py +++ b/lamb-kb-server/backend/routers/benchmarks.py @@ -4,7 +4,7 @@ from database.connection import get_session from dependencies import verify_token -from fastapi import APIRouter, Body, Depends +from fastapi import APIRouter, Depends from schemas.benchmark import ( BenchmarkDataset, BenchmarkDatasetSummary, @@ -13,7 +13,6 @@ BenchmarkRunRequest, BenchmarkRunResponse, ) -from schemas.content import EmbeddingCredentials from services.benchmark import BenchmarkService from sqlalchemy.orm import Session @@ -49,20 +48,23 @@ async def get_benchmark_dataset(dataset_id: str) -> BenchmarkDataset: ) async def run_collection_benchmark( collection_id: str, - request: BenchmarkRunRequest, - embedding_credentials: EmbeddingCredentials = Body( - default_factory=EmbeddingCredentials, - embed=True, - ), + 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=request, + request=body, embedding_credentials={ - "api_key": embedding_credentials.api_key, - "api_endpoint": embedding_credentials.api_endpoint, + "api_key": creds.api_key, + "api_endpoint": creds.api_endpoint, }, ) @@ -74,20 +76,17 @@ async def run_collection_benchmark( ) async def run_all_collection_benchmarks( collection_id: str, - request: BenchmarkRunAllRequest, - embedding_credentials: EmbeddingCredentials = Body( - default_factory=EmbeddingCredentials, - embed=True, - ), + body: BenchmarkRunAllRequest, db: Session = Depends(get_session), ) -> BenchmarkRunAllResponse: + creds = body.embedding_credentials return BenchmarkService.run_all( db=db, collection_id=collection_id, - dataset_ids=request.dataset_ids, - threshold=request.threshold, + dataset_ids=body.dataset_ids, + threshold=body.threshold, embedding_credentials={ - "api_key": embedding_credentials.api_key, - "api_endpoint": embedding_credentials.api_endpoint, + "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 index 8d9b09e3e..50f0c7517 100644 --- a/lamb-kb-server/backend/routers/graph.py +++ b/lamb-kb-server/backend/routers/graph.py @@ -78,17 +78,6 @@ def _graph_status_payload() -> dict[str, Any]: } -def _collection_dict(collection: Collection) -> dict[str, Any]: - """Adapt a new Collection ORM row to the legacy ``{id, owner, ...}`` - dict shape that graph_store internals still expect.""" - return { - "id": collection.id, - "name": collection.name, - "organization_id": collection.organization_id, - "owner": collection.organization_id, - } - - @router.get("/status", summary="Get Graph RAG feature availability") async def get_graph_status(token: str = Depends(verify_token)): return _graph_status_payload() @@ -117,64 +106,65 @@ async def migrate_collection_to_graph( collection = _get_collection_or_404(db, collection_id) _graph_store_or_503() - # Pull chunks from ChromaDB (the only backend with a stable get-by-id - # path right now). Qdrant migration is a TODO. - if collection.vector_db_backend != "chromadb": + from plugins.base import EmbeddingRegistry, VectorDBRegistry # noqa: PLC0415 + + backend = VectorDBRegistry.get(collection.vector_db_backend) + if backend is None: raise HTTPException( - status_code=400, + status_code=503, detail=( - "Graph migration currently only supports the chromadb vector " - f"backend (collection uses '{collection.vector_db_backend}')." + f"Vector DB backend '{collection.vector_db_backend}' is not " + "available." ), ) - from plugins.vector_db import chromadb_backend as _chroma_mod + # 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]] = [] - client = _chroma_mod._get_client(collection.storage_path) try: - chroma_collection = client.get_collection( - name=collection.backend_collection_id or collection.id - ) - except Exception as exc: + 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"ChromaDB collection for {collection.id} not found", + detail=f"Vector backend collection for {collection.id} not found", ) from exc - ids: list[str] = [] - texts: list[str] = [] - metadatas: list[dict[str, Any]] = [] - offset = 0 - batch_size = 500 - - while True: - result = chroma_collection.get( - include=["documents", "metadatas"], - limit=batch_size, - offset=offset, - ) - result_ids = result.get("ids") or [] - documents = result.get("documents") or [] - result_metadatas = result.get("metadatas") or [] - if not result_ids: - break - - for i, chunk_id in enumerate(result_ids): - text = documents[i] if i < len(documents) else None - if not text: - continue - metadata = ( - result_metadatas[i] - if i < len(result_metadatas) - and isinstance(result_metadatas[i], dict) - else {} - ) - ids.append(chunk_id) - texts.append(text) - metadatas.append(metadata) - - offset += len(result_ids) - if not ids: collection.graph_enabled = True db.commit() diff --git a/lamb-kb-server/backend/schemas/benchmark.py b/lamb-kb-server/backend/schemas/benchmark.py index 8f2bf1585..29abab010 100644 --- a/lamb-kb-server/backend/schemas/benchmark.py +++ b/lamb-kb-server/backend/schemas/benchmark.py @@ -72,6 +72,21 @@ class BenchmarkComparison(BaseModel): 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): dataset_id: Optional[str] = Field( "educational", description="Built-in dataset ID to use when questions are omitted" @@ -82,6 +97,14 @@ class BenchmarkRunRequest(BaseModel): 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): @@ -101,6 +124,10 @@ class BenchmarkRunAllRequest(BaseModel): 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): diff --git a/lamb-kb-server/backend/services/concept_extraction.py b/lamb-kb-server/backend/services/concept_extraction.py index 89ef6d093..80936b987 100644 --- a/lamb-kb-server/backend/services/concept_extraction.py +++ b/lamb-kb-server/backend/services/concept_extraction.py @@ -126,7 +126,13 @@ def __init__( api_key = self.config.get("openai_api_key") or "" if self.client is None and api_key and OpenAI is not None: - self.client = OpenAI(api_key=api_key) + # Bound the OpenAI call so a slow / hung vendor can't pin an + # ingestion worker thread indefinitely. Configurable via env + # so operators can stretch it for big chunks if needed. + timeout_seconds = float( + self.config.get("openai_timeout_seconds") or 60.0 + ) + self.client = OpenAI(api_key=api_key, timeout=timeout_seconds) def extract_for_chunks(self, chunks: List[TextChunk]) -> GraphExtraction: if not chunks: diff --git a/lamb-kb-server/backend/services/graph_store.py b/lamb-kb-server/backend/services/graph_store.py index 568651588..3d5b7825b 100644 --- a/lamb-kb-server/backend/services/graph_store.py +++ b/lamb-kb-server/backend/services/graph_store.py @@ -477,6 +477,9 @@ def get_collection_graph( 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 @@ -572,6 +575,12 @@ def get_collection_graph( "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 [], }, } @@ -2208,6 +2217,21 @@ def _ingest_tx( 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}) @@ -2220,7 +2244,10 @@ def _ingest_tx( chunk.text = $text, chunk.parent_text = $parent_text, chunk.section_title = $section_title, - chunk.source_label = $source_label + 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, @@ -2233,6 +2260,9 @@ def _ingest_tx( 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, []): diff --git a/lamb-kb-server/backend/services/ingestion_service.py b/lamb-kb-server/backend/services/ingestion_service.py index 3a05eeb28..8b0a29cf0 100644 --- a/lamb-kb-server/backend/services/ingestion_service.py +++ b/lamb-kb-server/backend/services/ingestion_service.py @@ -289,8 +289,8 @@ def _maybe_index_graph( *, collection: Collection, docs_list: list[dict], - backend: object, - embedding_function: object, + backend, + embedding_function, openai_api_key: str = "", ) -> None: """Run graph indexing for this batch if the collection opted in. @@ -298,6 +298,12 @@ def _maybe_index_graph( 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 @@ -308,71 +314,57 @@ def _maybe_index_graph( if not kg_config.get("enabled") or not kg_config.get("index_on_ingest", True): return - # The chunks were already embedded and stored; pull them back from the - # vector backend so we have stable IDs and the exact text that's - # searchable. - try: - from plugins.vector_db import chromadb_backend as _chroma_mod # noqa: PLC0415 + from services.graph_indexing import ( # noqa: PLC0415 + index_chunks_for_collection, + ) - if collection.vector_db_backend != "chromadb": + 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 currently only supports chromadb backend; " - "skipping for collection %s (backend=%s).", - collection.id, + "Graph indexing: backend %s failed to fetch source %s: %s", collection.vector_db_backend, + source_id, + exc, ) - return - - client = _chroma_mod._get_client(collection.storage_path) - from plugins.vector_db.chromadb_backend import _to_chroma_ef # noqa: PLC0415 - - chroma_collection = client.get_collection( - name=collection.backend_collection_id or collection.id, - embedding_function=_to_chroma_ef(embedding_function), # type: ignore[arg-type] - ) - except Exception as exc: # noqa: BLE001 - logger.warning( - "Graph indexing: could not open ChromaDB collection %s: %s", - collection.id, - exc, - ) - return - - source_ids = [doc["source_item_id"] for doc in docs_list] - try: - from services.graph_indexing import ( # noqa: PLC0415 - index_chunks_for_collection, - ) - - for source_id in source_ids: - rows = chroma_collection.get( - where={"source_item_id": source_id}, - include=["documents", "metadatas"], - ) - ids = rows.get("ids") or [] - texts = rows.get("documents") or [] - metadatas = rows.get("metadatas") or [] - if not ids: - continue - result = index_chunks_for_collection( - collection=collection, - ids=list(ids), - texts=list(texts), - metadatas=[dict(m or {}) for m in metadatas], - openai_api_key=openai_api_key, - filename=source_id, + 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, ) - if result.get("error"): - logger.warning( - "Graph indexing for source %s in collection %s reported: %s", - source_id, - collection.id, - result["error"], - ) - except Exception as exc: # noqa: BLE001 - logger.warning( - "Graph indexing failed for collection %s: %s", collection.id, exc + 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( 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_kg_rag.py b/lamb-kb-server/tests/unit/test_kg_rag.py index 6ed213a5e..3bd9dfa4f 100644 --- a/lamb-kb-server/tests/unit/test_kg_rag.py +++ b/lamb-kb-server/tests/unit/test_kg_rag.py @@ -19,20 +19,10 @@ from __future__ import annotations import importlib -from collections.abc import Iterator import pytest -@pytest.fixture() -def reload_config(monkeypatch) -> Iterator[None]: - """Local copy of the reload_config fixture from test_config.py.""" - import config # noqa: PLC0415 - - yield - importlib.reload(config) - - # --------------------------------------------------------------------------- # config # --------------------------------------------------------------------------- @@ -185,3 +175,135 @@ def test_plugin_returns_baseline_when_no_seed_ids(monkeypatch): ) 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() From 7ec425eaa0913061d097901a35610d98190f55b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A0=20Guilera?= Date: Mon, 18 May 2026 00:24:32 +0200 Subject: [PATCH 003/156] feat(kg-rag): wire Graph RAG into the Knowledge Store UI (locked-at-create toggle, vendor/model picker, verified-only retrieval, Sigma full-graph view) and drop the in-app benchmark surface fix(kb-v2): per-collection LLM extraction config + plugin system (OpenAI/Ollama), cascade-aware curation actions, and Cypher tightened to use only verified concepts + relationships --- .../creator_interface/kb_server_manager.py | 10 +- .../knowledge_store_client.py | 47 +- .../knowledge_store_graph_router.py | 67 +-- .../knowledge_store_router.py | 22 + backend/creator_interface/main.py | 2 +- frontend/svelte-app/package.json | 3 + .../KnowledgeStoreBenchmarkView.svelte | 207 ------- .../KnowledgeStoreDetail.svelte | 52 +- .../KnowledgeStoreGraphView.svelte | 565 ++++++++++++------ .../knowledgeStores/SigmaGraphModal.svelte | 461 ++++++++++++++ .../modals/CreateKnowledgeStoreModal.svelte | 175 +++++- .../src/lib/services/benchmarkService.js | 43 -- .../src/lib/services/knowledgeStoreService.js | 17 + lamb-kb-server/backend/database/connection.py | 15 + lamb-kb-server/backend/database/models.py | 10 + lamb-kb-server/backend/main.py | 3 + lamb-kb-server/backend/plugins/base.py | 82 +++ .../backend/plugins/kg_rag_query.py | 243 +++++++- .../plugins/llm_extraction/__init__.py | 6 + .../backend/plugins/llm_extraction/ollama.py | 144 +++++ .../backend/plugins/llm_extraction/openai.py | 148 +++++ .../plugins/vector_db/chromadb_backend.py | 10 +- lamb-kb-server/backend/routers/graph.py | 2 +- lamb-kb-server/backend/routers/system.py | 18 +- lamb-kb-server/backend/schemas/benchmark.py | 13 +- lamb-kb-server/backend/schemas/collection.py | 43 ++ lamb-kb-server/backend/services/benchmark.py | 27 +- .../backend/services/collection_service.py | 46 +- .../backend/services/concept_extraction.py | 127 ++-- .../backend/services/graph_indexing.py | 21 +- .../backend/services/graph_store.py | 149 ++++- 31 files changed, 2136 insertions(+), 642 deletions(-) delete mode 100644 frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreBenchmarkView.svelte create mode 100644 frontend/svelte-app/src/lib/components/knowledgeStores/SigmaGraphModal.svelte delete mode 100644 frontend/svelte-app/src/lib/services/benchmarkService.js create mode 100644 lamb-kb-server/backend/plugins/llm_extraction/__init__.py create mode 100644 lamb-kb-server/backend/plugins/llm_extraction/ollama.py create mode 100644 lamb-kb-server/backend/plugins/llm_extraction/openai.py diff --git a/backend/creator_interface/kb_server_manager.py b/backend/creator_interface/kb_server_manager.py index 0d1d88616..97c408f1f 100644 --- a/backend/creator_interface/kb_server_manager.py +++ b/backend/creator_interface/kb_server_manager.py @@ -21,9 +21,13 @@ # 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'} -LAMB_KB_SERVER = _KB_REDIRECTS.get(_raw_kb_server, _raw_kb_server) or 'http://172.17.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) LAMB_KB_SERVER_TOKEN = os.getenv('LAMB_KB_SERVER_TOKEN') if not LAMB_KB_SERVER_TOKEN: raise ValueError("LAMB_KB_SERVER_TOKEN environment variable is required") diff --git a/backend/creator_interface/knowledge_store_client.py b/backend/creator_interface/knowledge_store_client.py index a0c2955fd..81de4df93 100644 --- a/backend/creator_interface/knowledge_store_client.py +++ b/backend/creator_interface/knowledge_store_client.py @@ -254,6 +254,9 @@ async def create_collection( chunking_params: Dict[str, Any] = None, embedding_endpoint: str = "", 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. @@ -277,8 +280,22 @@ async def create_collection( "vector_db_backend": vector_db_backend, "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.""" @@ -480,36 +497,6 @@ async def list_graph_changes( params=params or {}, ) - async def run_benchmark( - self, - knowledge_store_id: str, - body: Dict[str, Any], - embedding_api_key: str = "", - embedding_api_endpoint: str = "", - creator_user: Dict[str, Any] = None, - ) -> Dict: - """Run a single benchmark dataset against a Knowledge Store.""" - config = self._get_ks_config(creator_user) - payload = { - **body, - "embedding_credentials": { - "api_key": embedding_api_key or "", - "api_endpoint": embedding_api_endpoint or "", - }, - } - return await self._request( - "POST", - f"/benchmarks/collections/{knowledge_store_id}/run", - config, - json=payload, - ) - - async def list_benchmark_datasets( - self, creator_user: Dict[str, Any] = None - ) -> Dict: - config = self._get_ks_config(creator_user) - return await self._request("GET", "/benchmarks/datasets", config) - async def graph_concept_rename( self, knowledge_store_id: str, diff --git a/backend/creator_interface/knowledge_store_graph_router.py b/backend/creator_interface/knowledge_store_graph_router.py index bc1a2ee67..a4a611bd9 100644 --- a/backend/creator_interface/knowledge_store_graph_router.py +++ b/backend/creator_interface/knowledge_store_graph_router.py @@ -1,13 +1,15 @@ """Creator-Interface proxy for KG-RAG / semantic-graph endpoints. -Routes mount at ``/creator/knowledge-stores/{ks_id}/graph/...`` and -``/creator/knowledge-stores/{ks_id}/benchmarks/...``. Each call resolves -the per-org KB Server URL/token through ``KnowledgeStoreClient`` and -proxies to the new KB Server's ``/graph`` and ``/benchmarks`` routers. - -The KB Server only exposes those routers when ``KG_RAG_ENABLED=true``; if -the flag is off the proxy will return 503 from the KB Server. The -frontend uses ``/graph/status`` to gate its UI accordingly. +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 @@ -117,7 +119,7 @@ async def get_graph_snapshot( 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=200), + limit: int = Query(default=60, ge=1, le=10000), auth: AuthContext = Depends(get_auth_context), ): _assert_ks_access(ks_id, auth) @@ -282,52 +284,5 @@ async def curate_relationship( ) -# ---------------------------------------------------------------------- -# Benchmarks -# ---------------------------------------------------------------------- - - -class BenchmarkRunBody(BaseModel): - dataset_id: Optional[str] = "educational" - top_k: Optional[int] = None - graph_depth: Optional[int] = None - threshold: float = 0.0 - - -@router.get("/benchmarks/datasets") -async def list_benchmark_datasets( - auth: AuthContext = Depends(get_auth_context), -): - return await _client.list_benchmark_datasets(creator_user=auth.user) -@router.post("/{ks_id}/benchmarks/run") -async def run_benchmark( - ks_id: str, - body: BenchmarkRunBody, - auth: AuthContext = Depends(get_auth_context), -): - _assert_ks_access(ks_id, auth) - from lamb.completions.org_config_resolver import OrganizationConfigResolver - resolver = OrganizationConfigResolver(auth.user.get("email")) - # Resolve embedding key the same way ingestion does, so the benchmark - # baseline vector pass works without the caller threading creds. - ks = _db.get_knowledge_store(ks_id) - embedding_api_key = "" - embedding_api_endpoint = "" - try: - embedding_api_key = ( - resolver.get_provider_api_key(ks.get("embedding_vendor")) or "" - ) - embedding_api_endpoint = ( - resolver.get_provider_endpoint(ks.get("embedding_vendor")) or "" - ) - except Exception: # noqa: BLE001 - embedding_api_key = "" - return await _client.run_benchmark( - knowledge_store_id=ks_id, - body=body.model_dump(exclude_none=True), - embedding_api_key=embedding_api_key, - embedding_api_endpoint=embedding_api_endpoint, - creator_user=auth.user, - ) diff --git a/backend/creator_interface/knowledge_store_router.py b/backend/creator_interface/knowledge_store_router.py index e13992d0c..c69f8088d 100644 --- a/backend/creator_interface/knowledge_store_router.py +++ b/backend/creator_interface/knowledge_store_router.py @@ -54,6 +54,13 @@ class KnowledgeStoreCreate(BaseModel): # 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): @@ -145,6 +152,14 @@ async def get_options(auth: AuthContext = Depends(get_auth_context)): return await _client.get_org_options(creator_user=auth.user) +@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 # ====================================================================== @@ -219,6 +234,9 @@ async def create_knowledge_store( chunking_params=body.chunking_params, embedding_endpoint=resolved_endpoint or "", 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: @@ -266,11 +284,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") diff --git a/backend/creator_interface/main.py b/backend/creator_interface/main.py index 9cbed5795..b65ec4e96 100644 --- a/backend/creator_interface/main.py +++ b/backend/creator_interface/main.py @@ -148,7 +148,7 @@ 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") -# Graph + benchmark proxy. Mounted on the same prefix so the frontend +# 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( knowledge_store_graph_router, prefix="/knowledge-stores" diff --git a/frontend/svelte-app/package.json b/frontend/svelte-app/package.json index f67fdf6db..f857a9d61 100644 --- a/frontend/svelte-app/package.json +++ b/frontend/svelte-app/package.json @@ -51,8 +51,11 @@ "dompurify": "^3.4.1", "flowbite-svelte": "^0.48.6", "flowbite-svelte-icons": "^2.1.1", + "graphology": "^0.26.0", + "graphology-layout-forceatlas2": "^0.10.1", "marked": "^15.0.12", "openai": "^4.53.3", + "sigma": "^3.0.3", "svelte-i18n": "^4.0.1" } } diff --git a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreBenchmarkView.svelte b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreBenchmarkView.svelte deleted file mode 100644 index 61b17b474..000000000 --- a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreBenchmarkView.svelte +++ /dev/null @@ -1,207 +0,0 @@ - - - -
- {#if !graphEnabled} -
- Graph RAG is not enabled on this Knowledge Store, so the - kg_rag_query arm of the benchmark will - degrade to the vector baseline. Enable the graph first for a meaningful - comparison. -
- {/if} - -
- - - - - -
- - {#if error} -
{error}
- {/if} - - {#if result} -
-

Summary

- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
P@kR@kMRRAvg vectorAvg graphAvg total
Baseline{fmt(result?.baseline?.precision_at_k)}{fmt(result?.baseline?.recall_at_k)}{fmt(result?.baseline?.mrr)}{fmtMs(result?.baseline?.avg_vector_ms)}{fmtMs(result?.baseline?.avg_graph_ms)}{fmtMs(result?.baseline?.avg_total_ms)}
KG-RAG{fmt(result?.kg_rag?.precision_at_k)}{fmt(result?.kg_rag?.recall_at_k)}{fmt(result?.kg_rag?.mrr)}{fmtMs(result?.kg_rag?.avg_vector_ms)}{fmtMs(result?.kg_rag?.avg_graph_ms)}{fmtMs(result?.kg_rag?.avg_total_ms)}
- {#if result?.comparison} -

- ΔP@k {fmt(result.comparison.delta_precision_at_k)} · ΔR@k {fmt( - result.comparison.delta_recall_at_k, - )} · ΔMRR {fmt(result.comparison.delta_mrr)} · graph overhead {fmtMs( - result.comparison.graph_overhead_ms, - )} -

-

{result.comparison.expected_behavior}

- {/if} -
- -
-

Per question

- - - - - - - - - - - - {#each result?.results || [] as row (row.question_id)} - - - - - - - - {/each} - -
QuestionKindBaseline P@kKG-RAG P@kΔMRR
{row.question}{row.kind}{fmt(row?.baseline?.precision_at_k)}{fmt(row?.kg_rag?.precision_at_k)} - {fmt( - (row?.kg_rag?.mrr ?? 0) - (row?.baseline?.mrr ?? 0), - )} -
-
- {/if} -
diff --git a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte index 2b3fffc89..063ef758c 100644 --- a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte +++ b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte @@ -22,7 +22,6 @@ import ConfirmationModal from '$lib/components/modals/ConfirmationModal.svelte'; import AddContentToKSModal from '$lib/components/knowledgeStores/AddContentToKSModal.svelte'; import KnowledgeStoreGraphView from '$lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte'; - import KnowledgeStoreBenchmarkView from '$lib/components/knowledgeStores/KnowledgeStoreBenchmarkView.svelte'; import { getGraphStatus } from '$lib/services/graphService'; /** @type {{ ksId: string }} */ @@ -448,27 +447,46 @@
{$_('knowledgeStores.vectorDb', { default: 'Vector DB' })}
-
{ks.vector_db_backend}
+
+ {ks.vector_db_backend} + {#if ks.graph_enabled} + + {$_('knowledgeStores.graphEnabledBadge', { default: 'Graph RAG' })} + + {/if} +
+ {#if ks.graph_enabled && (ks.extraction?.vendor || ks.extraction?.model)} +
+ {$_('knowledgeStores.extractionSummary', { + default: 'Extractor:' + })} + {ks.extraction?.vendor || '—'} + {#if ks.extraction?.model} · {ks.extraction.model}{/if} +
+ {/if}
{$_('knowledgeStores.lockedNotice', { default: - 'Chunking strategy, embedding vendor / model, and vector DB are locked at creation and cannot be changed.' + 'Chunking strategy, embedding vendor / model, vector DB, and Graph RAG are locked at creation and cannot be changed.' })}
- {@const canEnableGraph = - graphStatus?.enabled && ks?.vector_db_backend === 'chromadb'} - {@const showGraphTabs = !!ks?.graph_enabled || canEnableGraph} + {@const showGraphTabs = !!ks?.graph_enabled}
@@ -507,10 +516,7 @@ - {:else if activeTab === 'benchmark'} - {:else}
diff --git a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte index 9d5b64b38..aaeefb57c 100644 --- a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte +++ b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte @@ -17,58 +17,32 @@ import { onMount } from 'svelte'; import { getGraphSnapshot, - listGraphChanges, - migrateToGraph, renameConcept, - mergeConcepts, curateConcept, editRelationship, curateRelationship, } from '$lib/services/graphService'; - import { - changeOperationLabel, - changeMetaLine, - changeDetail, - } from '$lib/utils/graphCuration'; + import SigmaGraphModal from './SigmaGraphModal.svelte'; import { _ } from '$lib/i18n'; - /** @type {{ ksId: string, graphEnabled: boolean, vectorDbBackend?: string }} */ - let { ksId, graphEnabled, vectorDbBackend = '' } = $props(); - - // Graph migration currently only supports the chromadb backend - // (qdrant doesn't yet expose an iter-all-chunks surface). When the - // store uses a different backend we hide the migrate button entirely - // and explain why, rather than letting the user click and get a 400. - let migrationSupported = $derived( - !vectorDbBackend || vectorDbBackend === 'chromadb', - ); + /** @type {{ ksId: string, graphEnabled: boolean }} */ + let { ksId, graphEnabled } = $props(); let loading = $state(false); let error = $state(''); - let success = $state(''); let snapshot = $state(/** @type {any} */ (null)); - let changes = $state(/** @type {any[]} */ ([])); let filter = $state({ concept: '', filename: '', document_id: '' }); - let migrating = $state(false); - let migrateApiKey = $state(''); - - // Curation modals - let curationTarget = $state(/** @type {any} */ (null)); + let sigmaOpen = $state(false); async function loadAll() { loading = true; error = ''; try { - const [snap, hist] = await Promise.all([ - getGraphSnapshot(ksId, { - ...stripEmpty(filter), - limit: 80, - include_chunks: 'true', - }), - listGraphChanges(ksId, { ...stripEmpty(filter), limit: 25 }), - ]); - snapshot = snap; - changes = Array.isArray(hist) ? hist : []; + snapshot = await getGraphSnapshot(ksId, { + ...stripEmpty(filter), + limit: 80, + include_chunks: 'true', + }); } catch (/** @type {*} */ err) { error = err?.response?.data?.detail || err?.message || 'Failed to load graph'; } finally { @@ -86,82 +60,232 @@ return out; } - async function onMigrate() { - migrating = true; - error = ''; + let bulkBusy = $state(false); + let editingConcept = $state(''); + let editConceptName = $state(''); + let editingEdgeId = $state(''); + let editEdgeRelation = $state(''); + + /** + * Set a concept's verification_state. Used by per-item toggle and + * by the bulk operations. + * @param {string} name + * @param {'verified'|'unverified'|'rejected'} state + */ + async function setConceptState(name, state) { try { - const body = migrateApiKey ? { openai_api_key: migrateApiKey } : {}; - const result = await migrateToGraph(ksId, body); - success = `Migrated: ${result.chunks_seen ?? result.chunks ?? 0} chunks processed`; - await loadAll(); + await curateConcept(ksId, name, { + verification_state: state, + reason: `Set to ${state} via Knowledge Store UI`, + }); } catch (/** @type {*} */ err) { - error = - err?.response?.data?.detail || err?.message || 'Migration failed'; - } finally { - migrating = false; + error = err?.response?.data?.detail || err?.message || `Failed to ${state}`; + throw err; } } - /** @param {string} name */ - async function approveConcept(name) { + /** + * Set a relationship's verification_state. + * @param {{ source: string, target: string, relation: string }} rel + * @param {'verified'|'unverified'|'rejected'} state + */ + async function setRelationshipState(rel, state) { try { - await curateConcept(ksId, name, { - verification_state: 'verified', - reason: 'Approved via Knowledge Store UI', + await curateRelationship(ksId, { + source_concept: rel.source, + target_concept: rel.target, + relation: rel.relation, + verification_state: state, + reason: `Set to ${state} via Knowledge Store UI`, }); - await loadAll(); } catch (/** @type {*} */ err) { - error = err?.response?.data?.detail || err?.message || 'Approval failed'; + error = err?.response?.data?.detail || err?.message || `Failed to ${state}`; + throw err; } } - /** @param {string} name */ - async function rejectConcept(name) { + /** @param {any} node */ + async function toggleConceptVerification(node) { + const name = displayName(node.data?.name || node.label); + const current = String(node.data?.verification_state || 'unverified'); + const next = current === 'verified' ? 'unverified' : 'verified'; + await setConceptState(name, next); + await loadAll(); + } + + /** @param {any} edge */ + async function toggleRelationshipVerification(edge) { + const rel = edgeEndpoints(edge); + const current = String(edge.data?.verification_state || 'unverified'); + const next = current === 'verified' ? 'unverified' : 'verified'; + await setRelationshipState(rel, next); + await loadAll(); + } + + /** @param {'verified' | 'rejected'} state */ + async function bulkConcepts(state) { + if (!snapshot?.nodes?.length) return; + const reason = + state === 'verified' ? 'Bulk approval' : 'Bulk rejection'; + if ( + !confirm( + `${reason} of ALL concepts in this Knowledge Store. Continue?`, + ) + ) + return; + bulkBusy = true; try { - await curateConcept(ksId, name, { - verification_state: 'rejected', - reason: 'Rejected via Knowledge Store UI', - }); + const concepts = snapshot.nodes.filter( + (/** @type {any} */ n) => n.type === 'concept', + ); + for (const node of concepts) { + const cur = String(node.data?.verification_state || 'unverified'); + if (cur === state) continue; + const name = displayName(node.data?.name || node.label); + try { + await curateConcept(ksId, name, { + verification_state: state, + reason, + }); + } catch (/** @type {*} */ err) { + console.warn('Bulk concept curate failed for', name, err); + } + } await loadAll(); - } catch (/** @type {*} */ err) { - error = - err?.response?.data?.detail || err?.message || 'Rejection failed'; + } finally { + bulkBusy = false; } } - /** @param {{ source: string, target: string, relation: string }} rel */ - async function approveRelationship(rel) { + /** @param {'verified' | 'rejected'} state */ + async function bulkRelationships(state) { + if (!snapshot?.edges?.length) return; + const reason = + state === 'verified' ? 'Bulk approval' : 'Bulk rejection'; + if ( + !confirm( + `${reason} of ALL relationships in this Knowledge Store. Continue?`, + ) + ) + return; + bulkBusy = true; try { - await curateRelationship(ksId, { - source_concept: rel.source, - target_concept: rel.target, - relation: rel.relation, - verification_state: 'verified', - reason: 'Approved via Knowledge Store UI', + const rels = snapshot.edges.filter( + (/** @type {any} */ e) => e.type === 'RELATES_TO', + ); + for (const edge of rels) { + const cur = String(edge.data?.verification_state || 'unverified'); + if (cur === state) continue; + const rel = edgeEndpoints(edge); + try { + await curateRelationship(ksId, { + source_concept: rel.source, + target_concept: rel.target, + relation: rel.relation, + verification_state: state, + reason, + }); + } catch (/** @type {*} */ err) { + console.warn('Bulk relationship curate failed for', rel, err); + } + } + await loadAll(); + } finally { + bulkBusy = false; + } + } + + /** @param {any} node */ + function startEditConcept(node) { + editingConcept = displayName(node.data?.name || node.label); + editConceptName = editingConcept; + } + + function cancelEditConcept() { + editingConcept = ''; + editConceptName = ''; + } + + async function commitEditConcept() { + const oldName = editingConcept; + const newName = editConceptName.trim(); + if (!oldName || !newName || newName === oldName) { + cancelEditConcept(); + return; + } + try { + await renameConcept(ksId, oldName, { + new_name: newName, + reason: 'Renamed via Knowledge Store UI', }); + cancelEditConcept(); await loadAll(); } catch (/** @type {*} */ err) { - error = err?.response?.data?.detail || err?.message || 'Approval failed'; + error = err?.response?.data?.detail || err?.message || 'Rename failed'; } } - /** @param {{ source: string, target: string, relation: string }} rel */ - async function rejectRelationship(rel) { + /** @param {any} edge */ + function startEditEdge(edge) { + editingEdgeId = edge.id; + editEdgeRelation = String(edge.data?.relation || edge.label || ''); + } + + function cancelEditEdge() { + editingEdgeId = ''; + editEdgeRelation = ''; + } + + /** @param {any} edge */ + async function commitEditEdge(edge) { + const newRelation = editEdgeRelation.trim(); + const { source, target, relation: oldRelation } = edgeEndpoints(edge); + if (!newRelation || newRelation === oldRelation) { + cancelEditEdge(); + return; + } try { await editRelationship(ksId, { - source_concept: rel.source, - target_concept: rel.target, - relation: rel.relation, - verification_state: 'rejected', - reason: 'Rejected via Knowledge Store UI', + source_concept: source, + target_concept: target, + relation: oldRelation, + new_relation: newRelation, + reason: 'Edited via Knowledge Store UI', }); + cancelEditEdge(); await loadAll(); } catch (/** @type {*} */ err) { - error = - err?.response?.data?.detail || err?.message || 'Rejection failed'; + error = err?.response?.data?.detail || err?.message || 'Edit failed'; } } + /** + * Strip the ``concept:`` node-ID prefix when no friendlier label is + * available. Node IDs in Neo4j are stored as ``concept:`` but + * the user-facing surface should show just the entity name. + * @param {string | undefined | null} value + * @returns {string} + */ + function displayName(value) { + if (!value) return ''; + return String(value).replace(/^concept:/i, ''); + } + + /** + * Extract the clean (prefix-less) source and target concept names + * from a snapshot edge. The snapshot returns ``edge.data.source`` and + * ``edge.data.target`` as canonical names, but we fall back to + * stripping the ``concept:`` prefix from the raw ID for robustness. + * @param {any} edge + */ + function edgeEndpoints(edge) { + return { + source: String(edge.data?.source || displayName(edge.source)), + target: String(edge.data?.target || displayName(edge.target)), + relation: String(edge.data?.relation || edge.label || ''), + }; + } + onMount(() => { if (graphEnabled) loadAll(); }); @@ -169,49 +293,33 @@
{#if !graphEnabled} -
+

Graph RAG is not enabled on this Knowledge Store.

- {#if migrationSupported} -

- Run the migration below to extract concepts and relationships from - existing chunks. This calls the LLM extractor and writes results to - Neo4j; vector retrieval keeps working either way. -

-
- - -
- {:else} -

- Graph migration is not yet supported for the - {vectorDbBackend} - vector backend. Migration currently requires - chromadb. - Create a new Knowledge Store with chromadb to use Graph RAG. -

- {/if} +

+ Graph RAG is locked at creation time alongside chunking, embedding, and + vector DB. To use Graph RAG, create a new Knowledge Store with the + Enable Graph RAG toggle. +

{/if} {#if error}
{error}
{/if} - {#if success} -
{success}
- {/if} + {#if graphEnabled} +
+ Retrieval policy: + only concepts and relationships you mark as + verified + are used to enhance LLM retrieval. The LLM receives the + retrieved chunks as context — the graph drives + which chunks are returned (via question-entity expansion + RRF fusion + with the vector baseline) but the concept/relation triples themselves are not + added to the prompt. Approve the items you trust to make them + contribute to the retrieval. +
+ {/if}
{#if loading} @@ -270,28 +389,89 @@
-

Concepts

+
+

Concepts

+
+ + +
+
{#if !snapshot?.nodes?.length}

No concept nodes yet.

{:else}
    {#each snapshot.nodes.filter((/** @type {any} */ n) => n.type === 'concept') as node (node.id)} + {@const conceptName = displayName(node.data?.name || node.label)} + {@const state = String(node.data?.verification_state || 'unverified')} + {@const isEditing = editingConcept === conceptName}
  • -
    -
    {node.label}
    -
    {node.data?.verification_state || 'unverified'}
    +
    + {#if isEditing} + { + if (e.key === 'Enter') commitEditConcept(); + if (e.key === 'Escape') cancelEditConcept(); + }} + autofocus + /> + {:else} +
    {displayName(node.label)}
    +
    + {state} +
    + {/if}
    - - + {#if isEditing} + + + {:else} + {#if state === 'verified'} + + {:else} + + {/if} + + {/if}
    • {/each} @@ -300,42 +480,95 @@
-

Relationships

+
+

Relationships

+
+ + +
+
{#if !snapshot?.edges?.length}

No relationships yet.

{:else}
    {#each snapshot.edges.filter((/** @type {any} */ e) => e.type === 'RELATES_TO') as edge (edge.id)} + {@const state = String(edge.data?.verification_state || 'unverified')} + {@const isEditing = editingEdgeId === edge.id}
  • -
    - {edge.data?.source_label || edge.source} - → - {edge.data?.target_label || edge.target} - - {edge.label || edge.data?.relation} - +
    +
    + {displayName(edge.data?.source_label || edge.source)} + → + {displayName(edge.data?.target_label || edge.target)} +
    +
    + {#if isEditing} + { + if (e.key === 'Enter') commitEditEdge(edge); + if (e.key === 'Escape') cancelEditEdge(); + }} + autofocus + /> + {:else} + + {edge.label || edge.data?.relation} + + {state} + {/if} +
    - - + {#if isEditing} + + + {:else} + {#if state === 'verified'} + + {:else} + + {/if} + + {/if}
    • {/each} @@ -344,20 +577,6 @@
{/if} -
-

Recent changes

- {#if !changes.length} -

No change events yet.

- {:else} -
    - {#each changes as change (change.event_id)} -
  • -
    {changeOperationLabel(change)}
    -
    {changeMetaLine(change)}
    -
    {changeDetail(change)}
    -
    • - {/each} -
- {/if} -
+ + (sigmaOpen = false)} /> diff --git a/frontend/svelte-app/src/lib/components/knowledgeStores/SigmaGraphModal.svelte b/frontend/svelte-app/src/lib/components/knowledgeStores/SigmaGraphModal.svelte new file mode 100644 index 000000000..22695765a --- /dev/null +++ b/frontend/svelte-app/src/lib/components/knowledgeStores/SigmaGraphModal.svelte @@ -0,0 +1,461 @@ + + + +{#if open} + +{/if} diff --git a/frontend/svelte-app/src/lib/components/modals/CreateKnowledgeStoreModal.svelte b/frontend/svelte-app/src/lib/components/modals/CreateKnowledgeStoreModal.svelte index d0bd973c2..9ef256314 100644 --- a/frontend/svelte-app/src/lib/components/modals/CreateKnowledgeStoreModal.svelte +++ b/frontend/svelte-app/src/lib/components/modals/CreateKnowledgeStoreModal.svelte @@ -16,6 +16,7 @@ import axios from 'axios'; import { getOptions, + getLlmVendors, createKnowledgeStore, toggleSharing } from '$lib/services/knowledgeStoreService'; @@ -30,6 +31,7 @@ let name = $state(''); let description = $state(''); let isShared = $state(false); + let graphEnabled = $state(false); let advancedOpen = $state(false); // ── Options + config ────────────────────────────────────────────── @@ -53,6 +55,13 @@ let embeddingEndpoint = $state(''); let vectorDb = $state(''); + // LLM extraction picker — only relevant when graph_enabled === true. + // Loaded lazily once the modal opens. + let llmVendors = $state(/** @type {Array} */ ([])); + let extractionVendor = $state(''); + let extractionModel = $state(''); + let extractionEndpoint = $state(''); + let currentStrategyParams = $derived.by(() => { const s = (options.chunking_strategies ?? []).find( (/** @type {any} */ s) => s.name === chunkingStrategy @@ -73,6 +82,16 @@ return options.embedding_models?.[embeddingVendor] ?? []; }); + // Extraction vendor params (schema declared by the kb-server plugin). + let currentExtractionVendor = $derived.by(() => + llmVendors.find((/** @type {any} */ v) => v.name === extractionVendor) + ); + let extractionModelChoices = $derived.by(() => { + const params = currentExtractionVendor?.parameters || []; + const modelParam = params.find((/** @type {any} */ p) => p.name === 'model'); + return Array.isArray(modelParam?.choices) ? modelParam.choices : []; + }); + // When the user switches strategy, fill any missing param keys with the // declared defaults — preserves user edits, fills blanks. Uses Svelte 5 // dependencies on `currentStrategyParams` only; reading chunkingParams @@ -162,6 +181,27 @@ if (!vectorDb && options.vector_db_backends?.length) { vectorDb = options.vector_db_backends[0].name; } + // LLM extraction vendors (best-effort — older kb-servers may not + // expose /llm-vendors yet; we tolerate that gracefully). + try { + const llmRes = await getLlmVendors(); + llmVendors = Array.isArray(llmRes?.vendors) ? llmRes.vendors : []; + if (!extractionVendor && llmVendors.length > 0) { + extractionVendor = llmVendors[0].name; + } + if (!extractionModel) { + const params = llmVendors.find( + (/** @type {any} */ v) => v.name === extractionVendor + )?.parameters || []; + const modelParam = params.find( + (/** @type {any} */ p) => p.name === 'model' + ); + extractionModel = modelParam?.default || ''; + } + } catch (/** @type {unknown} */ llmErr) { + console.warn('getLlmVendors failed (graph picker unavailable)', llmErr); + llmVendors = []; + } } catch (/** @type {unknown} */ err) { optionsError = readableError(err, 'Failed to load options'); console.error('loadOptions failed', err); @@ -198,6 +238,21 @@ name = ''; description = ''; isShared = false; + graphEnabled = false; + // Keep llmVendors loaded across reopen so we don't re-fetch every + // time; just reset the user's picks back to the loaded defaults. + if (llmVendors.length > 0) { + extractionVendor = llmVendors[0].name; + const params = llmVendors[0].parameters || []; + const modelParam = params.find( + (/** @type {any} */ p) => p.name === 'model' + ); + extractionModel = modelParam?.default || ''; + } else { + extractionVendor = ''; + extractionModel = ''; + } + extractionEndpoint = ''; error = ''; nameError = ''; isSubmitting = false; @@ -245,6 +300,7 @@ error = ''; try { + const wantsGraph = graphEnabled && vectorDb === 'chromadb'; const ks = await createKnowledgeStore({ name: name.trim(), description: description.trim() || '', @@ -253,7 +309,13 @@ embedding_vendor: embeddingVendor, embedding_model: embeddingModel, embedding_endpoint: embeddingEndpoint.trim() || undefined, - vector_db_backend: vectorDb + vector_db_backend: vectorDb, + graph_enabled: wantsGraph, + extraction_vendor: wantsGraph ? extractionVendor || undefined : undefined, + extraction_model: wantsGraph ? extractionModel.trim() || undefined : undefined, + extraction_endpoint: wantsGraph + ? extractionEndpoint.trim() || undefined + : undefined }); if (isShared) { // Sharing is a separate endpoint; failure here shouldn't @@ -374,6 +436,35 @@ + +
@@ -392,7 +483,7 @@ > {$_('knowledgeStores.createModal.lockedNotice', { default: - 'Chunking strategy, embedding vendor/model, and vector DB are locked once the Knowledge Store is created. Chunking parameters can be edited later but only apply to newly ingested content.' + 'Chunking strategy, embedding vendor/model, vector DB, and Graph RAG are locked once the Knowledge Store is created. Chunking parameters can be edited later but only apply to newly ingested content.' })}
@@ -561,6 +652,86 @@ disabled={isSubmitting} /> + + {#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} diff --git a/frontend/svelte-app/src/lib/services/benchmarkService.js b/frontend/svelte-app/src/lib/services/benchmarkService.js deleted file mode 100644 index c48a27cac..000000000 --- a/frontend/svelte-app/src/lib/services/benchmarkService.js +++ /dev/null @@ -1,43 +0,0 @@ -/** - * @module benchmarkService - * KG-RAG benchmark API client for Knowledge Stores. - * - * Targets the LAMB backend proxy at - * ``/creator/knowledge-stores/{ksId}/benchmarks/...``, which forwards to the - * new KB Server's ``/benchmarks`` router. Like the graph endpoints, these - * are only available when ``KG_RAG_ENABLED=true`` on the KB Server. - */ - -import axios from 'axios'; -import { browser } from '$app/environment'; -import { getApiUrl } from '$lib/config'; - -function authHeaders() { - const token = localStorage.getItem('userToken'); - if (!token) throw new Error('User not authenticated.'); - return { Authorization: `Bearer ${token}` }; -} - -/** - * List built-in benchmark datasets. - */ -export async function listDatasets() { - if (!browser) throw new Error('Browser only.'); - const url = getApiUrl('/knowledge-stores/benchmarks/datasets'); - const response = await axios.get(url, { headers: authHeaders() }); - return response.data; -} - -/** - * Run a benchmark dataset against a Knowledge Store, comparing - * ``simple_query`` (vector baseline) against ``kg_rag_query`` (vector + - * graph expansion). - * @param {string} ksId - * @param {{ dataset_id?: string, top_k?: number, graph_depth?: number, threshold?: number }} body - */ -export async function runBenchmark(ksId, body) { - if (!browser) throw new Error('Browser only.'); - const url = getApiUrl(`/knowledge-stores/${ksId}/benchmarks/run`); - const response = await axios.post(url, body, { headers: authHeaders() }); - return response.data; -} diff --git a/frontend/svelte-app/src/lib/services/knowledgeStoreService.js b/frontend/svelte-app/src/lib/services/knowledgeStoreService.js index 0c8f824f4..261196ed4 100644 --- a/frontend/svelte-app/src/lib/services/knowledgeStoreService.js +++ b/frontend/svelte-app/src/lib/services/knowledgeStoreService.js @@ -97,6 +97,19 @@ export async function getOptions() { return response.data; } +/** + * Fetch registered LLM extraction vendors (KG-RAG concept extraction) + * along with their parameter schemas so the create form can render a + * vendor + model picker. + * @returns {Promise<{ vendors: Array<{ name: string, description: string, parameters: Array<{ name: string, type: string, default: any, choices: any }> }> }>} + */ +export async function getLlmVendors() { + if (!browser) throw new Error('Browser only.'); + const url = getApiUrl('/knowledge-stores/llm-vendors'); + const response = await axios.get(url, { headers: authHeaders() }); + return response.data; +} + // --------------------------------------------------------------------------- // CRUD // --------------------------------------------------------------------------- @@ -136,6 +149,10 @@ export async function getKnowledgeStore(ksId) { * embedding_model: string, * embedding_endpoint?: string, * vector_db_backend: string, + * graph_enabled?: boolean, + * extraction_vendor?: string, + * extraction_model?: string, + * extraction_endpoint?: string, * }} data * @returns {Promise} */ diff --git a/lamb-kb-server/backend/database/connection.py b/lamb-kb-server/backend/database/connection.py index 0bbd2a93e..d3582e267 100644 --- a/lamb-kb-server/backend/database/connection.py +++ b/lamb-kb-server/backend/database/connection.py @@ -103,6 +103,21 @@ def _run_lightweight_migrations(engine: Engine) -> None: "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 diff --git a/lamb-kb-server/backend/database/models.py b/lamb-kb-server/backend/database/models.py index f21ee3034..03a24d8cf 100644 --- a/lamb-kb-server/backend/database/models.py +++ b/lamb-kb-server/backend/database/models.py @@ -78,6 +78,16 @@ class Collection(Base): # 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 5ad3347ae..ddd1c7575 100644 --- a/lamb-kb-server/backend/main.py +++ b/lamb-kb-server/backend/main.py @@ -135,6 +135,9 @@ def _discover_plugins() -> None: "plugins.embedding.openai", "plugins.embedding.ollama", "plugins.embedding.local", + # LLM extraction backends (KG-RAG concept extraction) + "plugins.llm_extraction.openai", + "plugins.llm_extraction.ollama", ] import importlib # noqa: PLC0415 diff --git a/lamb-kb-server/backend/plugins/base.py b/lamb-kb-server/backend/plugins/base.py index 49fc72009..bf52a9dc4 100644 --- a/lamb-kb-server/backend/plugins/base.py +++ b/lamb-kb-server/backend/plugins/base.py @@ -299,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 # --------------------------------------------------------------------------- @@ -443,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/kg_rag_query.py b/lamb-kb-server/backend/plugins/kg_rag_query.py index c3aaf2f76..58aef05bf 100644 --- a/lamb-kb-server/backend/plugins/kg_rag_query.py +++ b/lamb-kb-server/backend/plugins/kg_rag_query.py @@ -124,24 +124,74 @@ def augment( trace["graph_latency_ms"] = (time.perf_counter() - graph_start) * 1000 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. + # The expansion is intentionally limited (a few chunks per + # entity, only specific entities with ≤25 mentions) to avoid + # flooding the candidate set with chunks that mention common + # nouns. + question_entities = self._extract_question_entities(query_text) + 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) + + # Decide which expansion arms contribute. Chunk-seeded + # expansion can be useful when the question is short and + # ambiguous, but it tends to inflate noise on + # encyclopedic / Wikipedia-style corpora because shared + # high-frequency entities (``World War II``, ``United + # States``) connect unrelated chunks. Toggleable via + # ``params.use_chunk_expansion``. + use_chunk_expansion = self._as_bool( + params.get("use_chunk_expansion", False) + ) + + merged_entry_concepts = list( + dict.fromkeys( + question_expansion.get("entry_concepts", []) + + (expansion.get("entry_concepts", []) if use_chunk_expansion else []) + ) + ) + # Question-entity expansion ranks first: those chunks were + # reached via question-derived concepts, so their precision is + # higher than the chunk-seeded expansion which can drift via + # generic shared entities. + merged_expanded_ids = list( + dict.fromkeys( + question_expansion.get("expanded_chunk_ids", []) + + (expansion.get("expanded_chunk_ids", []) if use_chunk_expansion else []) + ) + ) + trace.update( { - "graph_expanded": bool(expansion.get("expanded_chunk_ids")), - "entry_concepts": expansion.get("entry_concepts", []), + "graph_expanded": bool(merged_expanded_ids), + "entry_concepts": merged_entry_concepts, "traversed_edges": expansion.get("traversed_edges", []), - "expanded_chunk_ids": expansion.get("expanded_chunk_ids", []), + "expanded_chunk_ids": merged_expanded_ids, "latest_changes": expansion.get("latest_changes", []), + "question_entities": question_entities, "graph_latency_ms": expansion.get( "graph_latency_ms", (time.perf_counter() - graph_start) * 1000, - ), + ) + + question_expansion.get("graph_latency_ms", 0.0), } ) expanded_ids = [ - cid - for cid in expansion.get("expanded_chunk_ids", []) - if cid not in seed_chunk_ids + cid for cid in merged_expanded_ids if cid not in seed_chunk_ids ] expanded_results = self._fetch_expanded_results( backend=backend, @@ -151,11 +201,114 @@ def augment( return_parent_context=return_parent_context, ) - merged = self._merge_results(baseline_results + expanded_results, top_k=top_k) + merged = self._rrf_merge( + baseline_results=baseline_results, + expanded_results=expanded_results, + top_k=top_k, + rrf_k=int(params.get("rrf_k", 40) or 40), + graph_weight=float(params.get("graph_weight", 0.5) or 0.5), + ) 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) -> list[str]: + """Use a small LLM to extract named entities from the question. + + Returns an empty list if the OpenAI client isn't configured 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 [] + if question in cls._question_cache: + return cls._question_cache[question] + + try: + import openai # noqa: PLC0415 + except Exception: # noqa: BLE001 + return [] + + import config as config_module # noqa: PLC0415 + + kg = config_module.get_kg_rag_config() + api_key = (kg.get("openai_api_key") or "").strip() + if not api_key: + return [] + + # Use a small/fast model. ``KG_RAG_QUESTION_EXTRACTION_MODEL`` + # overrides; otherwise fall back to the chat model already used + # by the extractor (typically gpt-4o-mini). + import os # noqa: PLC0415 + + model = os.getenv("KG_RAG_QUESTION_EXTRACTION_MODEL") or ( + kg.get("chat_model") or "gpt-4o-mini" + ) + + client = openai.OpenAI(api_key=api_key, timeout=15.0) + prompt = ( + "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." + ) + try: + resp = client.chat.completions.create( + model=model, + messages=[ + {"role": "system", "content": prompt}, + {"role": "user", "content": question}, + ], + response_format={"type": "json_object"}, + temperature=0, + ) + raw = resp.choices[0].message.content or "{}" + import json as _json # noqa: PLC0415 + + data = _json.loads(raw) + entities = data.get("entities") or [] + if not isinstance(entities, list): + entities = [] + entities = [str(e).strip() for e in entities if str(e).strip()] + except Exception as exc: # noqa: BLE001 + logger.debug("Question-entity LLM extraction failed: %s", exc) + entities = [] + + # FIFO cache eviction + if len(cls._question_cache) >= cls._QUESTION_CACHE_SIZE: + try: + cls._question_cache.pop(next(iter(cls._question_cache))) + except StopIteration: + pass + cls._question_cache[question] = entities + return entities + # ------------------------------------------------------------------ # Internals # ------------------------------------------------------------------ @@ -229,24 +382,64 @@ def _fetch_expanded_results( ) return results - def _merge_results( - self, results: list[dict[str, Any]], top_k: int + def _rrf_merge( + self, + *, + baseline_results: list[dict[str, Any]], + expanded_results: list[dict[str, Any]], + top_k: int, + rrf_k: int = 60, + graph_weight: float = 0.6, ) -> list[dict[str, Any]]: - by_chunk_id: dict[str, dict[str, Any]] = {} - for result in results: - chunk_id = self._result_chunk_id(result) or str(result.get("data", ""))[:120] - existing = by_chunk_id.get(chunk_id) - if existing is None or float(result.get("similarity", 0.0)) > float( - existing.get("similarity", 0.0) - ): - by_chunk_id[chunk_id] = result - - ordered = sorted( - by_chunk_id.values(), - key=lambda item: float(item.get("similarity", 0.0)), - reverse=True, - ) - return ordered[: max(top_k, 1) + 3] + """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( 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 57361a999..47c871d34 100644 --- a/lamb-kb-server/backend/plugins/vector_db/chromadb_backend.py +++ b/lamb-kb-server/backend/plugins/vector_db/chromadb_backend.py @@ -383,12 +383,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]: + 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/routers/graph.py b/lamb-kb-server/backend/routers/graph.py index 50f0c7517..8b77cb419 100644 --- a/lamb-kb-server/backend/routers/graph.py +++ b/lamb-kb-server/backend/routers/graph.py @@ -211,7 +211,7 @@ async def get_graph_snapshot( chunk_id: str | None = Query(None), filename: str | None = Query(None), include_chunks: bool = Query(True), - limit: int = Query(60, ge=1, le=200), + limit: int = Query(60, ge=1, le=10000), token: str = Depends(verify_token), db: Session = Depends(get_session), ): diff --git a/lamb-kb-server/backend/routers/system.py b/lamb-kb-server/backend/routers/system.py index ab984ac96..3723a0d81 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 @@ -79,3 +84,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 index 29abab010..bef93ce43 100644 --- a/lamb-kb-server/backend/schemas/benchmark.py +++ b/lamb-kb-server/backend/schemas/benchmark.py @@ -2,7 +2,7 @@ from typing import Any, Dict, List, Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field class BenchmarkQuestion(BaseModel): @@ -88,6 +88,17 @@ class EmbeddingCredentialsBody(BaseModel): 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" ) diff --git a/lamb-kb-server/backend/schemas/collection.py b/lamb-kb-server/backend/schemas/collection.py index 27f9ab379..0d8174bae 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 --- @@ -56,6 +80,14 @@ class CreateCollectionRequest(BaseModel): "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): @@ -97,6 +129,7 @@ class CollectionResponse(BaseModel): embedding: EmbeddingConfig vector_db_backend: str graph_enabled: bool = False + extraction: ExtractionConfig | None = None status: str document_count: int chunk_count: int @@ -127,6 +160,16 @@ def from_orm_row(cls, row: Any) -> "CollectionResponse": ), 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/services/benchmark.py b/lamb-kb-server/backend/services/benchmark.py index abfbe00b6..206a4bb1d 100644 --- a/lamb-kb-server/backend/services/benchmark.py +++ b/lamb-kb-server/backend/services/benchmark.py @@ -446,17 +446,32 @@ def run( 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={ - "top_k": top_k, - "threshold": threshold, - "graph_depth": graph_depth, - "include_trace": True, - }, + plugin_params=kg_params, embedding_credentials=creds, ) diff --git a/lamb-kb-server/backend/services/collection_service.py b/lamb-kb-server/backend/services/collection_service.py index 046317f11..880b87971 100644 --- a/lamb-kb-server/backend/services/collection_service.py +++ b/lamb-kb-server/backend/services/collection_service.py @@ -8,7 +8,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 @@ -35,6 +40,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, @@ -130,6 +146,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, @@ -144,6 +164,9 @@ def create_collection(db: Session, req: CreateCollectionRequest) -> Collection: 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, @@ -309,6 +332,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 @@ -327,6 +351,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 index 80936b987..6e94e2409 100644 --- a/lamb-kb-server/backend/services/concept_extraction.py +++ b/lamb-kb-server/backend/services/concept_extraction.py @@ -11,11 +11,7 @@ from typing import Any, Dict, List, Optional, Set, Tuple import config as config_module - -try: - from openai import OpenAI -except Exception: # pragma: no cover - exercised when optional dependency is absent - OpenAI = None +from plugins.base import LLMExtractionFunction, LLMExtractionRegistry logger = logging.getLogger("lamb-kb") @@ -115,24 +111,59 @@ class ConceptExtractor: def __init__( self, kg_config: Optional[Dict[str, Any]] = None, - client: Optional[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 = self.config.get("extraction_model") or self.chat_model + 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)) - self.client = client - - api_key = self.config.get("openai_api_key") or "" - if self.client is None and api_key and OpenAI is not None: - # Bound the OpenAI call so a slow / hung vendor can't pin an - # ingestion worker thread indefinitely. Configurable via env - # so operators can stretch it for big chunks if needed. - timeout_seconds = float( - self.config.get("openai_timeout_seconds") or 60.0 + + 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 "") ) - self.client = OpenAI(api_key=api_key, timeout=timeout_seconds) + 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: @@ -141,7 +172,7 @@ def extract_for_chunks(self, chunks: List[TextChunk]) -> GraphExtraction: extraction = GraphExtraction( concepts_by_chunk={chunk.chunk_id: [] for chunk in chunks} ) - if self.client is None: + if self.backend is None: return extraction parent_groups: Dict[str, List[TextChunk]] = {} @@ -227,50 +258,26 @@ def _extract_parent_text( } ], }, - "limits": {"max_entities": 5, "max_relationships": 6}, + # 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], } - try: - response = self._create_json_completion( - model=self.model, - messages=[ - {"role": "system", "content": system}, - {"role": "user", "content": json.dumps(user, ensure_ascii=False)}, - ], - ) - content = response.choices[0].message.content or "{}" - parsed = json.loads(content) - except Exception as exc: - logger.warning("KG-RAG concept extraction failed: %s", exc) + if self.backend is None: return {"entities": [], "relationships": []} - return ( - parsed - if isinstance(parsed, dict) - else {"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, ) - - def _create_json_completion( - self, model: str, messages: List[Dict[str, str]] - ) -> Any: - if self.client is None: - raise RuntimeError("OpenAI client is not configured") - try: - return self.client.chat.completions.create( - model=model, - messages=messages, - response_format={"type": "json_object"}, - temperature=0, - ) - except Exception: - fallback_model = self.chat_model - if model == fallback_model: - raise - return self.client.chat.completions.create( - model=fallback_model, - messages=messages, - response_format={"type": "json_object"}, - temperature=0, - ) + 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] = {} @@ -278,7 +285,7 @@ def _parse_payload(self, payload: Dict[str, Any], chunk_id: str) -> GraphExtract if not isinstance(raw_entities, list): raw_entities = [] - for item in raw_entities[:5]: + for item in raw_entities[:12]: if not isinstance(item, dict): continue display_name = _clean_text(item.get("name"), limit=120) @@ -299,7 +306,7 @@ def _parse_payload(self, payload: Dict[str, Any], chunk_id: str) -> GraphExtract raw_relationships = [] related_names: Set[str] = set() - for item in raw_relationships[:6]: + for item in raw_relationships[:12]: if not isinstance(item, dict): continue source = normalize_concept(_clean_text(item.get("source"), limit=120)) diff --git a/lamb-kb-server/backend/services/graph_indexing.py b/lamb-kb-server/backend/services/graph_indexing.py index 5f8045ace..21b58d184 100644 --- a/lamb-kb-server/backend/services/graph_indexing.py +++ b/lamb-kb-server/backend/services/graph_indexing.py @@ -92,8 +92,19 @@ def index_chunks_for_collection( 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 not api_key: + if resolved_vendor == "openai" and not api_key: return { "indexed": False, "chunks": len(chunks), @@ -106,7 +117,13 @@ def index_chunks_for_collection( extractor_config = dict(kg_config) extractor_config["openai_api_key"] = api_key - extractor = ConceptExtractor(kg_config=extractor_config) + 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: diff --git a/lamb-kb-server/backend/services/graph_store.py b/lamb-kb-server/backend/services/graph_store.py index 3d5b7825b..71b6b9679 100644 --- a/lamb-kb-server/backend/services/graph_store.py +++ b/lamb-kb-server/backend/services/graph_store.py @@ -2255,7 +2255,11 @@ def _ingest_tx( collection_id=collection_id, org_id=org_id, file_id=file_id, - filename=filename, + # 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"), @@ -2362,19 +2366,34 @@ def expand_from_chunks( warning="Neo4j is not configured or available; KG expansion skipped", ) + # ``MAX_ENTRY_MENTIONS`` filters out highly-cited concepts: a + # concept that appears in N+ chunks (e.g. "World War II", + # "United States") behaves as a stop-word in the graph — it + # would expand to dozens of unrelated chunks that share only a + # generic theme. The threshold is conservative because the + # extractor caps each chunk at 12 entities, so a truly specific + # bridging concept rarely exceeds ~10 mentions in a + # benchmark-sized corpus. + MAX_ENTRY_MENTIONS = 20 + with self.driver.session() as session: entry_rows = session.run( """ MATCH (chunk:Chunk {collection_id: $collection_id})-[:MENTIONS]->(concept:Concept {org_id: $org_id}) WHERE chunk.chunk_id IN $seed_chunk_ids - AND coalesce(concept.verification_state, '') <> 'rejected' - RETURN concept.name AS name, count(*) AS mentions + AND concept.verification_state = 'verified' + WITH concept, count(*) AS local_mentions + OPTIONAL MATCH (concept)<-[:MENTIONS]-(global:Chunk {collection_id: $collection_id}) + WITH concept, local_mentions, count(global) AS total_mentions + WHERE total_mentions <= $max_total_mentions + RETURN concept.name AS name, local_mentions AS mentions ORDER BY mentions DESC, name ASC LIMIT 8 """, collection_id=collection_id, org_id=org_id, seed_chunk_ids=seed_chunk_ids, + max_total_mentions=MAX_ENTRY_MENTIONS, ).data() entry_concepts = [row["name"] for row in entry_rows] @@ -2387,8 +2406,8 @@ def expand_from_chunks( 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 coalesce(node.verification_state, '') <> 'rejected') - AND all(rel IN relationships(path) WHERE coalesce(rel.verification_state, '') <> 'rejected') + AND all(node IN nodes(path) WHERE node.verification_state = 'verified') + AND all(rel IN relationships(path) WHERE rel.verification_state = 'verified') WITH entry, related, path, length(path) AS hops, reduce(score = 0.0, rel IN relationships(path) | @@ -2417,7 +2436,7 @@ def expand_from_chunks( """ MATCH (concept:Concept {org_id: $org_id})<-[:MENTIONS]-(chunk:Chunk {collection_id: $collection_id}) WHERE concept.name IN $entry_concepts - AND coalesce(concept.verification_state, '') <> 'rejected' + AND concept.verification_state = 'verified' RETURN chunk.chunk_id AS chunk_id, count(*) AS mentions, min(chunk.source_label) AS source_label @@ -2476,6 +2495,124 @@ def add_chunk_id(chunk_id: str) -> None: "graph_latency_ms": (time.perf_counter() - start) * 1000, } + 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. + + Unlike :meth:`expand_from_chunks`, which discovers entry concepts + via the seed chunks' MENTIONS edges, 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}) + WHERE concept.name IN $names + AND concept.verification_state = 'verified' + OPTIONAL MATCH (concept)<-[:MENTIONS]-(chunk:Chunk {collection_id: $collection_id}) + WITH concept, count(chunk) AS mentions + WHERE mentions > 0 AND 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]] = [] From 229bd0bad841b321c8c123ad94f9f82506ccd02a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A0=20Guilera?= Date: Mon, 18 May 2026 11:38:01 +0200 Subject: [PATCH 004/156] fix(kg-rag): question-entity extractor uses per-collection model The question-time entity extractor was reading KG_RAG_QUESTION_EXTRACTION_MODEL env and falling back to the server-level chat_model, ignoring the per-collection extraction config picked by the user at KS creation. As a result the entity names produced at query time could be normalised differently from those stored at build time, silently zeroing the question-entity seeding path. It now resolves vendor/model/endpoint from the collection's stored extraction_* fields via LLMExtractionRegistry.build(...), the same path used by the build-time concept extractor. Cache key is now (vendor, model, question) so cross-collection lookups don't collide. --- .../backend/plugins/kg_rag_query.py | 110 ++++++++++-------- 1 file changed, 62 insertions(+), 48 deletions(-) diff --git a/lamb-kb-server/backend/plugins/kg_rag_query.py b/lamb-kb-server/backend/plugins/kg_rag_query.py index 58aef05bf..15d04e73d 100644 --- a/lamb-kb-server/backend/plugins/kg_rag_query.py +++ b/lamb-kb-server/backend/plugins/kg_rag_query.py @@ -132,7 +132,7 @@ def augment( # entity, only specific entities with ≤25 mentions) to avoid # flooding the candidate set with chunks that mention common # nouns. - question_entities = self._extract_question_entities(query_text) + question_entities = self._extract_question_entities(query_text, collection) question_expansion: dict[str, Any] = {} if question_entities: try: @@ -234,42 +234,71 @@ def augment( _question_cache: dict[str, list[str]] = {} @classmethod - def _extract_question_entities(cls, question: str) -> list[str]: - """Use a small LLM to extract named entities from the question. - - Returns an empty list if the OpenAI client isn't configured or - the call fails — the graph expansion silently degrades to the + 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 [] - if question in cls._question_cache: - return cls._question_cache[question] - - try: - import openai # noqa: PLC0415 - except Exception: # noqa: BLE001 - return [] import config as config_module # noqa: PLC0415 kg = config_module.get_kg_rag_config() - api_key = (kg.get("openai_api_key") or "").strip() - if not api_key: - return [] - - # Use a small/fast model. ``KG_RAG_QUESTION_EXTRACTION_MODEL`` - # overrides; otherwise fall back to the chat model already used - # by the extractor (typically gpt-4o-mini). - import os # noqa: PLC0415 + 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 - model = os.getenv("KG_RAG_QUESTION_EXTRACTION_MODEL") or ( - kg.get("chat_model") or "gpt-4o-mini" - ) + 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 [] - client = openai.OpenAI(api_key=api_key, timeout=15.0) - prompt = ( + 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 " @@ -278,27 +307,12 @@ def _extract_question_entities(cls, question: str) -> list[str]: "Do not output entity types, descriptions or explanations — only " "the surface strings." ) - try: - resp = client.chat.completions.create( - model=model, - messages=[ - {"role": "system", "content": prompt}, - {"role": "user", "content": question}, - ], - response_format={"type": "json_object"}, - temperature=0, - ) - raw = resp.choices[0].message.content or "{}" - import json as _json # noqa: PLC0415 - - data = _json.loads(raw) - entities = data.get("entities") or [] - if not isinstance(entities, list): - entities = [] - entities = [str(e).strip() for e in entities if str(e).strip()] - except Exception as exc: # noqa: BLE001 - logger.debug("Question-entity LLM extraction failed: %s", exc) - entities = [] + + 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: @@ -306,7 +320,7 @@ def _extract_question_entities(cls, question: str) -> list[str]: cls._question_cache.pop(next(iter(cls._question_cache))) except StopIteration: pass - cls._question_cache[question] = entities + cls._question_cache[cache_key] = entities return entities # ------------------------------------------------------------------ From 029568e6ad51a560c32367f6b42afefac7f91b85 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A0=20Guilera?= Date: Sun, 24 May 2026 20:54:55 +0200 Subject: [PATCH 005/156] feat(kg-rag): auto-route /query through KG-RAG plugin for graph-enabled KSs and surface extracted entities in Test Query UI --- .../KnowledgeStoreDetail.svelte | 33 +++++++++++++ .../src/lib/services/knowledgeStoreService.js | 2 +- lamb-kb-server/backend/routers/query.py | 12 +++++ lamb-kb-server/backend/schemas/query.py | 8 ++++ .../backend/services/query_service.py | 48 +++++++++++++++++++ 5 files changed, 102 insertions(+), 1 deletion(-) diff --git a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte index 063ef758c..19e2d2486 100644 --- a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte +++ b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreDetail.svelte @@ -49,6 +49,8 @@ let querying = $state(false); let queryResults = $state([]); let queryError = $state(''); + /** @type {string[] | null} */ + let queryEntities = $state(null); // Modals let showAddContent = $state(false); @@ -269,12 +271,16 @@ querying = true; queryError = ''; queryResults = []; + queryEntities = null; try { const data = await queryKnowledgeStore(ksId, { queryText, topK: queryTopK }); queryResults = data?.results ?? []; + // Present only when the KS has graph_enabled and the backend + // routed through the KG-RAG plugin. + queryEntities = Array.isArray(data?.entities) ? data.entities : null; } catch (/** @type {unknown} */ err) { queryError = err instanceof Error ? err.message : 'Query failed'; } finally { @@ -727,6 +733,33 @@ {/if} + {#if queryEntities} +
+
+ {$_('knowledgeStores.extractedEntities', { + default: 'Extracted entities (KG-RAG)' + })} +
+ {#if queryEntities.length === 0} +
+ {$_('knowledgeStores.extractedEntitiesEmpty', { + default: 'No named entities found in the question.' + })} +
+ {:else} +
+ {#each queryEntities as entity (entity)} + + {entity} + + {/each} +
+ {/if} +
+ {/if} + {#if queryResults.length > 0}
{#each queryResults as r, i (i)} diff --git a/frontend/svelte-app/src/lib/services/knowledgeStoreService.js b/frontend/svelte-app/src/lib/services/knowledgeStoreService.js index 261196ed4..33468b5cf 100644 --- a/frontend/svelte-app/src/lib/services/knowledgeStoreService.js +++ b/frontend/svelte-app/src/lib/services/knowledgeStoreService.js @@ -281,7 +281,7 @@ export async function removeContent(ksId, libraryItemId) { * builder's "test query" affordance and by the KS detail panel. * @param {string} ksId * @param {{ queryText: string, topK?: number }} data - * @returns {Promise<{ results: KSQueryResult[], query: string, top_k: number }>} + * @returns {Promise<{ results: KSQueryResult[], query: string, top_k: number, entities?: string[] | null }>} */ export async function queryKnowledgeStore(ksId, data) { if (!browser) throw new Error('Browser only.'); 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/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/query_service.py b/lamb-kb-server/backend/services/query_service.py index 97c988ec4..e12d4f424 100644 --- a/lamb-kb-server/backend/services/query_service.py +++ b/lamb-kb-server/backend/services/query_service.py @@ -63,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, From 23687fe991ec2e8c443a7dd13a409113315bb0f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A0=20Guilera?= Date: Wed, 27 May 2026 23:22:53 +0200 Subject: [PATCH 006/156] feat(graph): per-collection concept verification, drag/zoom graph, curation filters and pagination - Fix concept verification state to be scoped per Knowledge Store via MENTIONS.verification_state instead of the org-level Concept node, so approving concepts in one KS does not affect other KSs - Fix change-detection in curation tx to compare against MENTIONS.verification_state (the per-collection value) rather than concept.verification_state, which could be 'verified' from another KS and silently suppress the write - Delete associated graph data (Document, Chunk nodes, orphaned Concepts) when a linked library item is removed from a Knowledge Store - Add node drag support and zoom-adaptive labels to the Sigma.js graph view - Add verification status filter (All/Unverified/Verified/Rejected) to Concepts and Relationships sections in the curation panel - Add client-side pagination (20 items/page) to both curation lists - Show percentage verified in stats cards and section headers - Optimistic in-place mutation on verify/unverify toggle to eliminate flicker --- .../KnowledgeStoreGraphView.svelte | 226 ++++++++++++++---- .../knowledgeStores/SigmaGraphModal.svelte | 70 +++++- .../backend/services/graph_store.py | 132 +++++++++- .../backend/services/ingestion_service.py | 31 +++ 4 files changed, 396 insertions(+), 63 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte index aaeefb57c..303dfe484 100644 --- a/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte +++ b/frontend/svelte-app/src/lib/components/knowledgeStores/KnowledgeStoreGraphView.svelte @@ -34,13 +34,95 @@ let filter = $state({ concept: '', filename: '', document_id: '' }); let sigmaOpen = $state(false); + // Status filters (client-side, applied to already-loaded snapshot) + let conceptStatusFilter = $state('all'); + let relStatusFilter = $state('all'); + + // Pagination + const PAGE_SIZE = 20; + let conceptPage = $state(1); + let relPage = $state(1); + + // --- Derived lists --- + + let allConcepts = $derived( + (snapshot?.nodes || []).filter((/** @type {any} */ n) => n.type === 'concept'), + ); + let allRels = $derived( + (snapshot?.edges || []).filter((/** @type {any} */ e) => e.type === 'RELATES_TO'), + ); + + let filteredConcepts = $derived( + conceptStatusFilter === 'all' + ? allConcepts + : allConcepts.filter( + (/** @type {any} */ n) => + (n.data?.verification_state || 'unverified') === conceptStatusFilter, + ), + ); + let filteredRels = $derived( + relStatusFilter === 'all' + ? allRels + : allRels.filter( + (/** @type {any} */ e) => + (e.data?.verification_state || 'unverified') === relStatusFilter, + ), + ); + + let conceptTotalPages = $derived(Math.max(1, Math.ceil(filteredConcepts.length / PAGE_SIZE))); + let relTotalPages = $derived(Math.max(1, Math.ceil(filteredRels.length / PAGE_SIZE))); + + let pagedConcepts = $derived( + filteredConcepts.slice((conceptPage - 1) * PAGE_SIZE, conceptPage * PAGE_SIZE), + ); + let pagedRels = $derived( + filteredRels.slice((relPage - 1) * PAGE_SIZE, relPage * PAGE_SIZE), + ); + + // --- Stats --- + + let conceptVerifiedCount = $derived( + allConcepts.filter( + (/** @type {any} */ n) => (n.data?.verification_state || 'unverified') === 'verified', + ).length, + ); + let relVerifiedCount = $derived( + allRels.filter( + (/** @type {any} */ e) => (e.data?.verification_state || 'unverified') === 'verified', + ).length, + ); + + /** @param {number} num @param {number} den */ + function pct(num, den) { + if (!den) return '—'; + return Math.round((num / den) * 100) + '%'; + } + + // Reset pages to 1 when filter or snapshot changes + $effect(() => { + // eslint-disable-next-line no-unused-expressions + conceptStatusFilter; + conceptPage = 1; + }); + $effect(() => { + // eslint-disable-next-line no-unused-expressions + relStatusFilter; + relPage = 1; + }); + $effect(() => { + // eslint-disable-next-line no-unused-expressions + snapshot; + conceptPage = 1; + relPage = 1; + }); + async function loadAll() { loading = true; error = ''; try { snapshot = await getGraphSnapshot(ksId, { ...stripEmpty(filter), - limit: 80, + limit: 200, include_chunks: 'true', }); } catch (/** @type {*} */ err) { @@ -110,7 +192,9 @@ const current = String(node.data?.verification_state || 'unverified'); const next = current === 'verified' ? 'unverified' : 'verified'; await setConceptState(name, next); - await loadAll(); + // Mutate in place — avoids a full re-fetch and the resulting flicker. + // $state deep-proxy picks up the change immediately. + node.data.verification_state = next; } /** @param {any} edge */ @@ -119,34 +203,23 @@ const current = String(edge.data?.verification_state || 'unverified'); const next = current === 'verified' ? 'unverified' : 'verified'; await setRelationshipState(rel, next); - await loadAll(); + edge.data.verification_state = next; } /** @param {'verified' | 'rejected'} state */ async function bulkConcepts(state) { if (!snapshot?.nodes?.length) return; - const reason = - state === 'verified' ? 'Bulk approval' : 'Bulk rejection'; - if ( - !confirm( - `${reason} of ALL concepts in this Knowledge Store. Continue?`, - ) - ) - return; + const reason = state === 'verified' ? 'Bulk approval' : 'Bulk rejection'; + if (!confirm(`${reason} of ALL concepts in this Knowledge Store. Continue?`)) return; bulkBusy = true; try { - const concepts = snapshot.nodes.filter( - (/** @type {any} */ n) => n.type === 'concept', - ); + const concepts = snapshot.nodes.filter((/** @type {any} */ n) => n.type === 'concept'); for (const node of concepts) { const cur = String(node.data?.verification_state || 'unverified'); if (cur === state) continue; const name = displayName(node.data?.name || node.label); try { - await curateConcept(ksId, name, { - verification_state: state, - reason, - }); + await curateConcept(ksId, name, { verification_state: state, reason }); } catch (/** @type {*} */ err) { console.warn('Bulk concept curate failed for', name, err); } @@ -160,19 +233,11 @@ /** @param {'verified' | 'rejected'} state */ async function bulkRelationships(state) { if (!snapshot?.edges?.length) return; - const reason = - state === 'verified' ? 'Bulk approval' : 'Bulk rejection'; - if ( - !confirm( - `${reason} of ALL relationships in this Knowledge Store. Continue?`, - ) - ) - return; + const reason = state === 'verified' ? 'Bulk approval' : 'Bulk rejection'; + if (!confirm(`${reason} of ALL relationships in this Knowledge Store. Continue?`)) return; bulkBusy = true; try { - const rels = snapshot.edges.filter( - (/** @type {any} */ e) => e.type === 'RELATES_TO', - ); + const rels = snapshot.edges.filter((/** @type {any} */ e) => e.type === 'RELATES_TO'); for (const edge of rels) { const cur = String(edge.data?.verification_state || 'unverified'); if (cur === state) continue; @@ -320,6 +385,7 @@ contribute to the retrieval.
{/if} +
diff --git a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte index 9d0190847..7bfae87ba 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte @@ -28,11 +28,14 @@ selectedKnowledgeBases = $bindable([]), loadingKnowledgeBases = false, knowledgeBaseError = '', - userFiles = [], - selectedFilePath = $bindable(''), - loadingFiles = false, - fileError = '', - onFilesChanged, + libraries = [], + selectedLibraryId = $bindable(''), + loadingLibraries = false, + libraryError = '', + libraryItems = [], + selectedItemId = $bindable(''), + loadingItems = false, + itemsError = '', onchange } = $props(); @@ -202,12 +205,15 @@ bind:selectedKnowledgeBases loadingKnowledgeBases={loadingKnowledgeBases} knowledgeBaseError={knowledgeBaseError} - {userFiles} - bind:selectedFilePath - loadingFiles={loadingFiles} - fileError={fileError} + {libraries} + bind:selectedLibraryId + {loadingLibraries} + {libraryError} + {libraryItems} + bind:selectedItemId + {loadingItems} + {itemsError} {formState} - {onFilesChanged} /> {/if} diff --git a/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte b/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte index 45c4d4a15..109847536 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte @@ -3,7 +3,7 @@ import { _ } from '$lib/i18n'; import { isKbBasedRag, isSingleFileRag, isRubricRag } from '$lib/utils/ragProcessorHelpers.js'; import KnowledgeBaseSelector from './KnowledgeBaseSelector.svelte'; - import SingleFileSelector from './SingleFileSelector.svelte'; + import LibraryItemSelector from './LibraryItemSelector.svelte'; let { selectedRagProcessor = '', @@ -13,12 +13,15 @@ selectedKnowledgeBases = $bindable([]), loadingKnowledgeBases = false, knowledgeBaseError = '', - userFiles = [], - selectedFilePath = $bindable(''), - loadingFiles = false, - fileError = '', - formState, - onFilesChanged + libraries = [], + selectedLibraryId = $bindable(''), + loadingLibraries = false, + libraryError = '', + libraryItems = [], + selectedItemId = $bindable(''), + loadingItems = false, + itemsError = '', + formState } = $props(); let showKbSelector = $derived(isKbBasedRag(selectedRagProcessor)); @@ -57,13 +60,16 @@ /> {/if} {#if showFileSelector} - {/if} From 59eb89e69ce866d6bbdd753915e3cd1e8363cc10 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Thu, 28 May 2026 17:23:40 +0200 Subject: [PATCH 014/156] feat: assistant payload uses library_id+item_id instead of file_path --- .../assistants/assistantFormSubmit.test.js | 40 +++++++++++++++++-- .../assistants/logic/assistantFormSubmit.js | 6 ++- 2 files changed, 41 insertions(+), 5 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js b/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js index dd96206e3..980e0a8b1 100644 --- a/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js +++ b/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js @@ -37,7 +37,8 @@ describe('buildAssistantPayload', () => { selectedConnector: 'openai', selectedLlm: 'gpt-4', selectedRagProcessor: 'no_rag', - selectedFilePath: '', + selectedLibraryId: '', + selectedItemId: '', visionEnabled: false, imageGenerationEnabled: false, selectedKnowledgeBases: [], @@ -46,7 +47,10 @@ describe('buildAssistantPayload', () => { }; const payload = buildAssistantPayload(form); expect(payload.name).toBe('test'); - expect(JSON.parse(payload.metadata).connector).toBe('openai'); + const metadata = JSON.parse(payload.metadata); + expect(metadata.connector).toBe('openai'); + expect(metadata.library_id).toBeUndefined(); + expect(metadata.item_id).toBeUndefined(); }); test('includes rubric fields when rubric_rag is selected', () => { @@ -60,7 +64,8 @@ describe('buildAssistantPayload', () => { selectedConnector: 'openai', selectedLlm: 'gpt-4', selectedRagProcessor: 'rubric_rag', - selectedFilePath: '', + selectedLibraryId: '', + selectedItemId: '', visionEnabled: false, imageGenerationEnabled: false, selectedKnowledgeBases: [], @@ -84,7 +89,8 @@ describe('buildAssistantPayload', () => { selectedConnector: 'openai', selectedLlm: 'gpt-4', selectedRagProcessor: 'simple_rag', - selectedFilePath: '', + selectedLibraryId: '', + selectedItemId: '', visionEnabled: true, imageGenerationEnabled: false, selectedKnowledgeBases: ['kb1', 'kb2'], @@ -97,4 +103,30 @@ describe('buildAssistantPayload', () => { const metadata = JSON.parse(payload.metadata); expect(metadata.capabilities.vision).toBe(true); }); + + test('includes library_id and item_id when single_file_rag is selected', () => { + const form = { + name: 'test', + description: '', + system_prompt: '', + prompt_template: '', + RAG_Top_k: 3, + selectedPromptProcessor: 'simple_augment', + selectedConnector: 'openai', + selectedLlm: 'gpt-4o-mini', + selectedRagProcessor: 'single_file_rag', + selectedLibraryId: 'lib-123', + selectedItemId: 'item-456', + visionEnabled: false, + imageGenerationEnabled: false, + selectedKnowledgeBases: [], + selectedRubricId: '', + rubricFormat: 'markdown' + }; + const payload = buildAssistantPayload(form); + const metadata = JSON.parse(payload.metadata); + expect(metadata.library_id).toBe('lib-123'); + expect(metadata.item_id).toBe('item-456'); + expect(metadata.file_path).toBeUndefined(); + }); }); diff --git a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js index 12765fca0..6dc322f9f 100644 --- a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js +++ b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js @@ -30,13 +30,17 @@ export function buildAssistantPayload(form) { connector: form.selectedConnector, llm: form.selectedLlm, rag_processor: form.selectedRagProcessor, - file_path: isSingleFileRag(form.selectedRagProcessor) ? form.selectedFilePath : '', capabilities: { vision: form.visionEnabled, image_generation: form.imageGenerationEnabled } }; + if (isSingleFileRag(form.selectedRagProcessor)) { + metadataObj.library_id = form.selectedLibraryId || ''; + metadataObj.item_id = form.selectedItemId || ''; + } + if (isRubricRag(form.selectedRagProcessor)) { metadataObj.rubric_id = form.selectedRubricId; metadataObj.rubric_format = form.rubricFormat; From e2276aa3f7b49f4601fd2fdff626337b2389f997 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Thu, 28 May 2026 17:24:33 +0200 Subject: [PATCH 015/156] feat: populate library_id+item_id in edit mode for single_file_rag --- .../assistants/logic/assistantFormState.svelte.js | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js index 3d03cf67a..927d4fa3c 100644 --- a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js +++ b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js @@ -203,6 +203,12 @@ export function populateFormFields(form, data, getAvailableModels, preserveDescr } } + // Library fields (single_file_rag with Library Manager) + if (isSingleFileRag(form.selectedRagProcessor)) { + form.selectedLibraryId = metadata?.library_id || ''; + form.selectedItemId = metadata?.item_id || ''; + } + // Vision capability try { form.visionEnabled = metadata?.capabilities?.vision || false; From 74833d5d63d9f57879a830e677d7eebbbb2f4a56 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Thu, 28 May 2026 17:45:04 +0200 Subject: [PATCH 016/156] fix: library items not loading when selecting library in single_file_rag --- .../assistants/AssistantForm.svelte | 1 - .../assistants/assistantFormLibrary.test.js | 22 ++++++++++++++++--- .../assistants/logic/assistantFormFetchers.js | 4 +--- 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte index 0faf9711b..5e5b00678 100644 --- a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte @@ -177,7 +177,6 @@ // Effect to fetch library items when selected library changes $effect(() => { if (isSingleFileRag(form.selectedRagProcessor) && form.selectedLibraryId && form.configInitialized) { - form.itemsFetchAttempted = false; doFetchLibraryItems(form.selectedLibraryId); } }); diff --git a/frontend/svelte-app/src/lib/components/assistants/assistantFormLibrary.test.js b/frontend/svelte-app/src/lib/components/assistants/assistantFormLibrary.test.js index ded2bbf68..f7bded5d2 100644 --- a/frontend/svelte-app/src/lib/components/assistants/assistantFormLibrary.test.js +++ b/frontend/svelte-app/src/lib/components/assistants/assistantFormLibrary.test.js @@ -84,14 +84,13 @@ describe('fetchLibraryItems', () => { const form = { libraryItems: [], loadingItems: false, - itemsError: '', - itemsFetchAttempted: false + itemsError: '' }; await fetchLibraryItems(form, 'lib-1'); + expect(getItems).toHaveBeenCalledWith('lib-1', { limit: 100 }); expect(form.libraryItems).toEqual(mockResponse.items); - expect(form.itemsFetchAttempted).toBe(true); expect(form.loadingItems).toBe(false); }); @@ -110,4 +109,21 @@ describe('fetchLibraryItems', () => { expect(form.libraryItems).toEqual([]); expect(form.itemsError).toBeTruthy(); }); + + it('re-fetches when called again without force flag', async () => { + const mockResponse = { items: [{ id: 'item-1', title: 'Doc A', status: 'ready' }], total: 1 }; + getItems.mockResolvedValue(mockResponse); + + const form = { + libraryItems: [], + loadingItems: false, + itemsError: '' + }; + + await fetchLibraryItems(form, 'lib-1'); + expect(getItems).toHaveBeenCalledTimes(1); + + await fetchLibraryItems(form, 'lib-1'); + expect(getItems).toHaveBeenCalledTimes(2); + }); }); diff --git a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormFetchers.js b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormFetchers.js index 59387cee3..7ce0eedf0 100644 --- a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormFetchers.js +++ b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormFetchers.js @@ -124,11 +124,10 @@ export async function fetchLibraries(form, force = false) { */ export async function fetchLibraryItems(form, libraryId, force = false) { if (!libraryId) return; - if (form.itemsFetchAttempted && !force) return; form.loadingItems = true; form.itemsError = ''; try { - const response = await getItems(libraryId, { status: 'ready', limit: 100 }); + const response = await getItems(libraryId, { limit: 100 }); form.libraryItems = response.items || []; } catch (err) { if (err instanceof Error && err.message.startsWith('Session expired')) return; @@ -136,6 +135,5 @@ export async function fetchLibraryItems(form, libraryId, force = false) { form.libraryItems = []; } finally { form.loadingItems = false; - form.itemsFetchAttempted = true; } } From 28e308e814b9affcb3cc9d82cbf63c0336943fe9 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Thu, 28 May 2026 18:56:51 +0200 Subject: [PATCH 017/156] chore: stop tracking changelog (already in .gitignore) --- ...ngle-file-rag-library-manager-changelog.md | 26 ------------------- 1 file changed, 26 deletions(-) delete mode 100644 docs/superpowers/plans/2026-05-28-single-file-rag-library-manager-changelog.md diff --git a/docs/superpowers/plans/2026-05-28-single-file-rag-library-manager-changelog.md b/docs/superpowers/plans/2026-05-28-single-file-rag-library-manager-changelog.md deleted file mode 100644 index 3477d708d..000000000 --- a/docs/superpowers/plans/2026-05-28-single-file-rag-library-manager-changelog.md +++ /dev/null @@ -1,26 +0,0 @@ -# Single File RAG → Library Manager: Implementation Changelog - -## Task 1: Tests for single_file_rag (RED) -- Created: `testing/unit-tests/completions/test_single_file_rag.py` -- 11 test cases covering: LM success, connection error, 404, 500, missing env vars, file_path backward compat, file not found, missing metadata, partial metadata, invalid JSON, library_id precedence over file_path -- Results: 7 failed, 4 passed (correct for RED phase) -- Added `id` attribute to MockAssistant (current single_file_rag.py accesses assistant.id in debug log) - -## Task 2: Implement single_file_rag with LM support (GREEN) -- Modified: `backend/lamb/completions/rag/single_file_rag.py` (full rewrite, ~95 lines) -- New functions: `_fetch_from_library_manager()`, `_read_from_static_file()` -- Priority: library_id+item_id > file_path -- Sources format aligned with simple_rag: {title, url, similarity} - - LM path: title="Library Document", url="/docs/{lib_id}/{item_id}", similarity=1.0 - - file_path: title=basename, url="/static/public/{path}", similarity=1.0 -- Error handling: LM errors return empty context; file_path errors preserve old behavior (error message in context) -- Uses httpx.get() with LAMB_LIBRARY_SERVER and LAMB_LIBRARY_TOKEN env vars -- Backward compat: file_path still reads from static/public/ - -## Task 3: Metadata validation for single_file_rag -- Created: `testing/unit-tests/assistant_router/test_metadata_validation.py` (7 tests) -- Created: `backend/creator_interface/metadata_validators.py` (extracted from assistant_router.py, zero heavy deps) -- Modified: `backend/creator_interface/assistant_router.py` (delegates to metadata_validators) -- Validates: library_id+item_id pair, or file_path, for single_file_rag -- Other RAG processors are unaffected -- Extraction avoids heavy import chain in tests From 943faa4e8358fca80548c05eb2fe10cafad93306 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Thu, 28 May 2026 19:00:44 +0200 Subject: [PATCH 018/156] chore: fix .gitignore formatting (remove blank line, add final newline) --- .gitignore | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 6676206d8..169a0ec46 100644 --- a/.gitignore +++ b/.gitignore @@ -186,7 +186,6 @@ docs/ frontend-architecture-diagrams.md opencode.json package-lock.json - .playwright-mcp aac_logs/ @@ -195,4 +194,4 @@ aac_logs/ # Worktrees .worktrees/ .agents/ -skills-lock.json \ No newline at end of file +skills-lock.json From 6fa49637f3990436c6008f7349c9ada3c4602364 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 12:32:56 +0200 Subject: [PATCH 019/156] test: add metadata validation tests for document_rag field --- .../test_metadata_validation.py | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/testing/unit-tests/assistant_router/test_metadata_validation.py b/testing/unit-tests/assistant_router/test_metadata_validation.py index 3b29a9def..b65a82514 100644 --- a/testing/unit-tests/assistant_router/test_metadata_validation.py +++ b/testing/unit-tests/assistant_router/test_metadata_validation.py @@ -102,3 +102,53 @@ def test_simple_rag_not_affected(): result, error = validate_update_plugin_metadata(body) assert error is None assert result is not None + + +class TestDocumentRagValidation: + def _make_body(self, **overrides): + base = { + "metadata": { + "prompt_processor": "simple_augment", + "connector": "openai", + "llm": "gpt-4o-mini", + "rag_processor": "no_rag", + "document_rag": "", + } + } + base["metadata"].update(overrides) + return base + + def test_document_rag_single_file_requires_library_or_file(self): + body = self._make_body(document_rag="single_file_rag") + _, error = validate_update_plugin_metadata(body) + assert error is not None + assert "document_rag" in error + + def test_document_rag_with_library_id_and_item_id_is_valid(self): + body = self._make_body( + document_rag="single_file_rag", + library_id="lib-1", + item_id="item-1", + ) + metadata, error = validate_update_plugin_metadata(body) + assert error is None + assert metadata is not None + + def test_document_rag_library_id_without_item_id(self): + body = self._make_body( + document_rag="single_file_rag", + library_id="lib-1", + ) + _, error = validate_update_plugin_metadata(body) + assert error is not None + assert "item_id" in error + + def test_document_rag_empty_is_valid(self): + body = self._make_body(document_rag="") + metadata, error = validate_update_plugin_metadata(body) + assert error is None + + def test_non_document_rag_unaffected(self): + body = self._make_body(rag_processor="simple_rag", document_rag="") + metadata, error = validate_update_plugin_metadata(body) + assert error is None From dc60013aa0b03c4d3cb76771638792c9fe9d9908 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 12:32:56 +0200 Subject: [PATCH 020/156] feat: metadata_validators validates document_rag field --- backend/creator_interface/metadata_validators.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/backend/creator_interface/metadata_validators.py b/backend/creator_interface/metadata_validators.py index 09991f6ee..fae6c0a3d 100644 --- a/backend/creator_interface/metadata_validators.py +++ b/backend/creator_interface/metadata_validators.py @@ -73,5 +73,19 @@ def validate_update_plugin_metadata( 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") == "single_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=single_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 From 2ac57739fc0ab6934fd69725fb21b20e82714c53 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 12:35:15 +0200 Subject: [PATCH 021/156] test: add simple_augment tests for document_context in system prompt --- .../completions/test_simple_augment.py | 116 ++++++++++++++++++ 1 file changed, 116 insertions(+) create mode 100644 testing/unit-tests/completions/test_simple_augment.py diff --git a/testing/unit-tests/completions/test_simple_augment.py b/testing/unit-tests/completions/test_simple_augment.py new file mode 100644 index 000000000..e096f9bbe --- /dev/null +++ b/testing/unit-tests/completions/test_simple_augment.py @@ -0,0 +1,116 @@ +import pytest +from lamb.completions.pps.simple_augment import prompt_processor + + +class MockAssistant: + def __init__(self, system_prompt="", prompt_template="", metadata=None): + self.system_prompt = system_prompt + self.prompt_template = prompt_template + self.metadata = metadata + self.api_callback = metadata + + +def _make_request(messages): + return {"messages": messages} + + +class TestDocumentContextInSystemPrompt: + def test_document_context_appended_to_system_prompt(self): + assistant = MockAssistant( + system_prompt="You are a tutor.", + prompt_template="Answer: {user_input}\nContext: {context}", + ) + request = _make_request([{"role": "user", "content": "What is 2+2?"}]) + doc_ctx = {"context": "# Reference Document\nSome content here.", "sources": []} + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + system_msg = result[0] + assert system_msg["role"] == "system" + assert "You are a tutor." in system_msg["content"] + assert "# Reference Document\nSome content here." in system_msg["content"] + + def test_document_context_as_system_prompt_when_none_exists(self): + assistant = MockAssistant( + system_prompt="", + prompt_template="Answer: {user_input}\nContext: {context}", + ) + request = _make_request([{"role": "user", "content": "Hello"}]) + doc_ctx = {"context": "Document content.", "sources": []} + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + system_msg = result[0] + assert system_msg["role"] == "system" + assert system_msg["content"] == "Document content." + + def test_no_document_context_leaves_system_prompt_unchanged(self): + assistant = MockAssistant( + system_prompt="Original prompt.", + prompt_template="{user_input}", + ) + request = _make_request([{"role": "user", "content": "Hi"}]) + + result = prompt_processor(request, assistant=assistant) + + system_msg = result[0] + assert system_msg["content"] == "Original prompt." + + def test_empty_document_context_ignored(self): + assistant = MockAssistant( + system_prompt="Original.", + prompt_template="{user_input}", + ) + request = _make_request([{"role": "user", "content": "Hi"}]) + doc_ctx = {"context": "", "sources": []} + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + system_msg = result[0] + assert system_msg["content"] == "Original." + + def test_document_sources_not_in_user_message(self): + assistant = MockAssistant( + system_prompt="Sys.", + prompt_template="Q: {user_input}\nCtx: {context}", + ) + request = _make_request([{"role": "user", "content": "Question?"}]) + doc_ctx = { + "context": "Doc text.", + "sources": [{"title": "Doc", "url": "/docs/1/2", "similarity": 1.0}], + } + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + last_msg = result[-1] + assert "Doc text." not in last_msg["content"] + assert "## Available Sources" not in last_msg["content"] + + def test_document_context_and_rag_context_coexist(self): + assistant = MockAssistant( + system_prompt="Sys.", + prompt_template="Q: {user_input}\nCtx: {context}", + ) + request = _make_request([{"role": "user", "content": "Question?"}]) + doc_ctx = {"context": "Reference doc.", "sources": []} + rag_ctx = {"context": "RAG chunks.", "sources": [{"title": "KB", "url": "/kb/1", "similarity": 0.9}]} + + result = prompt_processor( + request, assistant=assistant, rag_context=rag_ctx, document_context=doc_ctx + ) + + system_msg = result[0] + assert "Reference doc." in system_msg["content"] + + last_msg = result[-1] + assert "RAG chunks." in last_msg["content"] + assert "Reference doc." not in last_msg["content"] + + def test_no_assistant_document_context_ignored(self): + request = _make_request([{"role": "user", "content": "Hi"}]) + doc_ctx = {"context": "Doc.", "sources": []} + + result = prompt_processor(request, document_context=doc_ctx) + + assert len(result) == 1 + assert result[0]["role"] == "user" From ea6bf494626024515c0c28db8fcd1a405bfc7bb0 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 12:35:21 +0200 Subject: [PATCH 022/156] feat: simple_augment injects document_context into system prompt --- backend/lamb/completions/pps/simple_augment.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/backend/lamb/completions/pps/simple_augment.py b/backend/lamb/completions/pps/simple_augment.py index 414dd8654..47fb7c055 100644 --- a/backend/lamb/completions/pps/simple_augment.py +++ b/backend/lamb/completions/pps/simple_augment.py @@ -58,7 +58,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: @@ -78,11 +79,16 @@ def prompt_processor( processed_messages = [] if assistant: - # Add system message from assistant if available - if assistant.system_prompt: + # Build system prompt with optional document context + system_content = assistant.system_prompt or "" + if document_context and isinstance(document_context, dict): + doc_text = document_context.get("context", "") + if doc_text: + system_content = (system_content + "\n\n" + doc_text) if system_content else doc_text + if system_content: processed_messages.append({ "role": "system", - "content": assistant.system_prompt + "content": system_content }) # Add previous messages except the last one From 41b124cf1a8082148f426c73edea7dbf513aead1 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 12:51:04 +0200 Subject: [PATCH 023/156] test: add pipeline integration tests for document_rag --- .../completions/test_document_rag_pipeline.py | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 testing/unit-tests/completions/test_document_rag_pipeline.py diff --git a/testing/unit-tests/completions/test_document_rag_pipeline.py b/testing/unit-tests/completions/test_document_rag_pipeline.py new file mode 100644 index 000000000..08149cd07 --- /dev/null +++ b/testing/unit-tests/completions/test_document_rag_pipeline.py @@ -0,0 +1,76 @@ +import json +import pytest +from lamb.completions.main import parse_plugin_config, process_completion_request + + +class MockAssistant: + def __init__(self, metadata="", system_prompt="", prompt_template=""): + self.metadata = metadata + self.api_callback = metadata + self.system_prompt = system_prompt + self.prompt_template = prompt_template + self.id = 1 + self.name = "test" + self.owner = "test_owner" + + +class TestParsePluginConfigDocumentRag: + def test_reads_document_rag_from_metadata(self): + metadata = json.dumps({ + "prompt_processor": "simple_augment", + "connector": "openai", + "llm": "gpt-4o-mini", + "rag_processor": "context_aware_rag", + "document_rag": "single_file_rag", + }) + assistant = MockAssistant(metadata=metadata) + config = parse_plugin_config(assistant) + assert config["document_rag"] == "single_file_rag" + + def test_document_rag_defaults_to_empty(self): + metadata = json.dumps({ + "prompt_processor": "simple_augment", + "connector": "openai", + "llm": "gpt-4o-mini", + "rag_processor": "no_rag", + }) + assistant = MockAssistant(metadata=metadata) + config = parse_plugin_config(assistant) + assert config["document_rag"] == "" + + def test_legacy_single_file_rag_stays_in_rag_processor(self): + metadata = json.dumps({ + "prompt_processor": "simple_augment", + "connector": "openai", + "llm": "gpt-4o-mini", + "rag_processor": "single_file_rag", + }) + assistant = MockAssistant(metadata=metadata) + config = parse_plugin_config(assistant) + assert config["rag_processor"] == "single_file_rag" + assert config["document_rag"] == "" + + +class TestProcessCompletionRequestDocumentContext: + def test_passes_document_context_to_pps(self): + captured_kwargs = {} + + def mock_pps(request, assistant=None, rag_context=None, **kwargs): + captured_kwargs["document_context"] = kwargs.get("document_context") + captured_kwargs["rag_context"] = rag_context + return [{"role": "user", "content": "test"}] + + pps = {"simple_augment": mock_pps} + plugin_config = {"prompt_processor": "simple_augment"} + doc_ctx = {"context": "Doc.", "sources": []} + + process_completion_request( + request={"messages": [{"role": "user", "content": "Hi"}]}, + assistant_details=None, + plugin_config=plugin_config, + rag_context=None, + pps=pps, + document_context=doc_ctx, + ) + + assert captured_kwargs.get("document_context") == doc_ctx From ea429a80bb0e0c3452b0d669ed6aa080c7c39938 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 12:51:07 +0200 Subject: [PATCH 024/156] feat: pipeline reads document_rag and passes document_context to PPS --- backend/lamb/completions/main.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/backend/lamb/completions/main.py b/backend/lamb/completions/main.py index 8cd712d31..0c49e2258 100644 --- a/backend/lamb/completions/main.py +++ b/backend/lamb/completions/main.py @@ -168,8 +168,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}") @@ -324,7 +325,8 @@ def parse_plugin_config(assistant_details) -> Dict[str, str]: "prompt_processor": "default", "connector": "openai", "llm": "gpt-4", - "rag_processor": "" + "rag_processor": "", + "document_rag": "" } # Apply defaults for missing keys @@ -352,6 +354,9 @@ 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") return pps, connectors, rag_processors @@ -378,12 +383,12 @@ 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: +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) + 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 @@ -494,7 +499,8 @@ 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) stream = request.get("stream", False) llm = plugin_config.get("llm") # Get LLM from config From f481093ce0a091737fa0a223e06d1e280f45892d Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 12:57:24 +0200 Subject: [PATCH 025/156] fix: add env vars to pipeline test for import chain --- .../completions/test_document_rag_pipeline.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/testing/unit-tests/completions/test_document_rag_pipeline.py b/testing/unit-tests/completions/test_document_rag_pipeline.py index 08149cd07..b0c5ab8a7 100644 --- a/testing/unit-tests/completions/test_document_rag_pipeline.py +++ b/testing/unit-tests/completions/test_document_rag_pipeline.py @@ -1,5 +1,24 @@ import json +import os import pytest + +_dummy_env = { + "LAMB_WEB_HOST": "http://localhost:9099", + "LAMB_BACKEND_HOST": "http://localhost:9099", + "LAMB_BEARER_TOKEN": "test-token", + "OPENAI_BASE_URL": "http://localhost:11434", + "OPENAI_MODEL": "gpt-4", + "OWI_BASE_URL": "http://localhost:8080", + "OWI_PATH": "/tmp/owi", + "SIGNUP_SECRET_KEY": "test-secret", + "OWI_ADMIN_NAME": "admin", + "OWI_ADMIN_EMAIL": "admin@test.com", + "OWI_ADMIN_PASSWORD": "test-password", + "LAMB_DB_PATH": "/tmp/test_lamb.db", +} +for _k, _v in _dummy_env.items(): + os.environ.setdefault(_k, _v) + from lamb.completions.main import parse_plugin_config, process_completion_request From 9f59babe049931f004f09bd5010c8651dd2a0057 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 13:00:00 +0200 Subject: [PATCH 026/156] test: add frontend tests for document_rag submit and state --- .../assistantFormDocumentRag.test.js | 156 ++++++++++++++++++ .../assistants/assistantFormSubmit.test.js | 91 +++++++++- 2 files changed, 243 insertions(+), 4 deletions(-) create mode 100644 frontend/svelte-app/src/lib/components/assistants/assistantFormDocumentRag.test.js diff --git a/frontend/svelte-app/src/lib/components/assistants/assistantFormDocumentRag.test.js b/frontend/svelte-app/src/lib/components/assistants/assistantFormDocumentRag.test.js new file mode 100644 index 000000000..fc1b5b4ff --- /dev/null +++ b/frontend/svelte-app/src/lib/components/assistants/assistantFormDocumentRag.test.js @@ -0,0 +1,156 @@ +import { describe, it, expect, vi } from 'vitest'; + +vi.mock('$lib/utils/ragProcessorHelpers.js', () => ({ + isKbBasedRag: (p) => ['simple_rag', 'context_aware_rag', 'hierarchical_rag'].includes(p), + isSingleFileRag: (p) => p === 'single_file_rag', + isRubricRag: (p) => p === 'rubric_rag', + normalizeRagProcessor: (p) => p +})); + +vi.mock('$lib/stores/assistantConfigStore', () => ({ + assistantConfigStore: { + subscribe: vi.fn(), + configDefaults: { config: {} } + } +})); + +vi.mock('$lib/utils/assistantData', () => ({ + getAssistantMetadataObject: (data) => { + if (!data.metadata) return {}; + return typeof data.metadata === 'string' ? JSON.parse(data.metadata) : data.metadata; + } +})); + +vi.mock('./logic/assistantFormUtils.svelte.js', () => ({ + loadRagPlaceholders: () => [], + selectModel: (model) => model +})); + +vi.mock('svelte/store', () => ({ + get: (store) => store.configDefaults || { config: {} } +})); + +import { populateFormFields, resetFormFieldsToDefaults } from './logic/assistantFormState.svelte.js'; + +function makeForm(overrides = {}) { + return { + formState: 'create', + initialAssistantData: null, + name: '', + description: '', + system_prompt: '', + prompt_template: '', + RAG_Top_k: 3, + isAdvancedMode: false, + promptProcessors: ['simple_augment'], + connectorsList: ['openai'], + ragProcessors: ['no_rag', 'simple_rag', 'context_aware_rag', 'single_file_rag'], + selectedPromptProcessor: 'simple_augment', + selectedConnector: 'openai', + selectedLlm: 'gpt-4o-mini', + selectedRagProcessor: 'no_rag', + documentRagEnabled: false, + visionEnabled: false, + imageGenerationEnabled: false, + ownedKnowledgeBases: [], + sharedKnowledgeBases: [], + selectedKnowledgeBases: [], + loadingKnowledgeBases: false, + knowledgeBaseError: '', + kbFetchAttempted: false, + pendingKBSelections: null, + userFiles: [], + selectedFilePath: '', + loadingFiles: false, + fileError: '', + filesFetchAttempted: false, + libraries: [], + selectedLibraryId: '', + loadingLibraries: false, + libraryError: '', + librariesFetchAttempted: false, + libraryItems: [], + selectedItemId: '', + loadingItems: false, + itemsError: '', + itemsFetchAttempted: false, + accessibleRubrics: [], + selectedRubricId: '', + rubricFormat: 'markdown', + loadingRubrics: false, + rubricError: '', + rubricsFetchAttempted: false, + formError: '', + formLoading: false, + configInitialized: true, + successMessage: '', + importError: '', + formDirty: false, + previousAssistantId: null, + ragPlaceholders: [], + get accessibleKnowledgeBases() { + return [...this.ownedKnowledgeBases, ...this.sharedKnowledgeBases]; + }, + ...overrides + }; +} + +describe('documentRagEnabled in form state', () => { + it('populateFormFields sets documentRagEnabled from metadata', () => { + const form = makeForm({ formState: 'edit' }); + const data = { + name: 'Test', + description: '', + system_prompt: '', + prompt_template: '', + RAG_Top_k: 3, + metadata: JSON.stringify({ + prompt_processor: 'simple_augment', + connector: 'openai', + llm: 'gpt-4o-mini', + rag_processor: 'no_rag', + document_rag: 'single_file_rag', + library_id: 'lib-1', + item_id: 'item-1' + }) + }; + + populateFormFields(form, data, () => ['gpt-4o-mini']); + + expect(form.documentRagEnabled).toBe(true); + expect(form.selectedLibraryId).toBe('lib-1'); + expect(form.selectedItemId).toBe('item-1'); + }); + + it('populateFormFields reads file_path for legacy antiguo', () => { + const form = makeForm({ formState: 'edit' }); + const data = { + name: 'Legacy Old', + description: '', + system_prompt: '', + prompt_template: '', + RAG_Top_k: 3, + metadata: JSON.stringify({ + prompt_processor: 'simple_augment', + connector: 'openai', + llm: 'gpt-4o-mini', + rag_processor: 'single_file_rag', + file_path: 'docs/mi_documento.md' + }) + }; + + populateFormFields(form, data, () => ['gpt-4o-mini']); + + expect(form.documentRagEnabled).toBe(false); + expect(form.selectedRagProcessor).toBe('single_file_rag'); + expect(form.selectedFilePath).toBe('docs/mi_documento.md'); + expect(form.selectedLibraryId).toBe(''); + expect(form.selectedItemId).toBe(''); + }); + + it('resetFormFieldsToDefaults clears documentRagEnabled', () => { + const form = makeForm({ documentRagEnabled: true, selectedLibraryId: 'lib-1' }); + resetFormFieldsToDefaults(form, () => ['gpt-4o-mini']); + expect(form.documentRagEnabled).toBe(false); + }); +}); diff --git a/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js b/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js index 980e0a8b1..3c590ca38 100644 --- a/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js +++ b/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js @@ -1,4 +1,4 @@ -import { describe, test, expect, vi } from 'vitest'; +import { describe, test, it, expect, vi } from 'vitest'; vi.mock('$lib/utils/ragProcessorHelpers.js', () => ({ isKbBasedRag: (p) => ['simple_rag', 'context_aware_rag', 'hierarchical_rag'].includes(p), @@ -104,7 +104,7 @@ describe('buildAssistantPayload', () => { expect(metadata.capabilities.vision).toBe(true); }); - test('includes library_id and item_id when single_file_rag is selected', () => { + test('legacy single_file_rag without documentRagEnabled does not include library refs', () => { const form = { name: 'test', description: '', @@ -117,6 +117,8 @@ describe('buildAssistantPayload', () => { selectedRagProcessor: 'single_file_rag', selectedLibraryId: 'lib-123', selectedItemId: 'item-456', + selectedFilePath: '', + documentRagEnabled: false, visionEnabled: false, imageGenerationEnabled: false, selectedKnowledgeBases: [], @@ -125,8 +127,89 @@ describe('buildAssistantPayload', () => { }; const payload = buildAssistantPayload(form); const metadata = JSON.parse(payload.metadata); - expect(metadata.library_id).toBe('lib-123'); - expect(metadata.item_id).toBe('item-456'); + expect(metadata.library_id).toBeUndefined(); + expect(metadata.item_id).toBeUndefined(); expect(metadata.file_path).toBeUndefined(); + expect(metadata.document_rag).toBeUndefined(); + }); +}); + +describe('buildAssistantPayload with document_rag', () => { + it('includes document_rag and library refs when documentRagEnabled', () => { + const form = { + name: 'Test', + description: '', + system_prompt: '', + prompt_template: '', + selectedPromptProcessor: 'simple_augment', + selectedConnector: 'openai', + selectedLlm: 'gpt-4o-mini', + selectedRagProcessor: 'no_rag', + documentRagEnabled: true, + selectedLibraryId: 'lib-1', + selectedItemId: 'item-1', + selectedFilePath: '', + visionEnabled: false, + imageGenerationEnabled: false, + RAG_Top_k: 3, + selectedKnowledgeBases: [] + }; + const payload = buildAssistantPayload(form); + const metadata = JSON.parse(payload.metadata); + expect(metadata.document_rag).toBe('single_file_rag'); + expect(metadata.library_id).toBe('lib-1'); + expect(metadata.item_id).toBe('item-1'); + expect(metadata.rag_processor).toBe('no_rag'); + }); + + it('omits document_rag when documentRagEnabled is false', () => { + const form = { + name: 'Test', + description: '', + system_prompt: '', + prompt_template: '', + selectedPromptProcessor: 'simple_augment', + selectedConnector: 'openai', + selectedLlm: 'gpt-4o-mini', + selectedRagProcessor: 'context_aware_rag', + documentRagEnabled: false, + selectedLibraryId: '', + selectedItemId: '', + selectedFilePath: '', + visionEnabled: false, + imageGenerationEnabled: false, + RAG_Top_k: 3, + selectedKnowledgeBases: ['kb-1'] + }; + const payload = buildAssistantPayload(form); + const metadata = JSON.parse(payload.metadata); + expect(metadata.document_rag).toBeUndefined(); + expect(metadata.library_id).toBeUndefined(); + }); + + it('legacy single_file_rag with file_path preserves file_path', () => { + const form = { + name: 'Legacy Old', + description: '', + system_prompt: '', + prompt_template: '', + selectedPromptProcessor: 'simple_augment', + selectedConnector: 'openai', + selectedLlm: 'gpt-4o-mini', + selectedRagProcessor: 'single_file_rag', + documentRagEnabled: false, + selectedLibraryId: '', + selectedItemId: '', + selectedFilePath: 'docs/mi_documento.md', + visionEnabled: false, + imageGenerationEnabled: false, + RAG_Top_k: 3, + selectedKnowledgeBases: [] + }; + const payload = buildAssistantPayload(form); + const metadata = JSON.parse(payload.metadata); + expect(metadata.rag_processor).toBe('single_file_rag'); + expect(metadata.file_path).toBe('docs/mi_documento.md'); + expect(metadata.document_rag).toBeUndefined(); }); }); From 599c2cff28c9691047e72eb72af973c460c167b6 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 13:00:03 +0200 Subject: [PATCH 027/156] feat: frontend submit and state support document_rag field --- .../logic/assistantFormState.svelte.js | 16 +++++++++++++--- .../assistants/logic/assistantFormSubmit.js | 5 ++++- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js index 927d4fa3c..1eeb387a0 100644 --- a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js +++ b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js @@ -89,6 +89,8 @@ export function createAssistantFormState() { itemsError: '', itemsFetchAttempted: false, + documentRagEnabled: false, + // --- Rubric state --- /** @type {Array<{rubric_id: string, title: string, description: string, is_mine: boolean, is_showcase: boolean, is_public: boolean}>} */ accessibleRubrics: [], @@ -148,6 +150,7 @@ export function resetFormFieldsToDefaults(form, getAvailableModels) { form.selectedFilePath = ''; form.selectedLibraryId = ''; form.selectedItemId = ''; + form.documentRagEnabled = false; form.visionEnabled = false; form.imageGenerationEnabled = false; } @@ -203,8 +206,14 @@ export function populateFormFields(form, data, getAvailableModels, preserveDescr } } - // Library fields (single_file_rag with Library Manager) + // Library fields: legacy path (rag_processor=single_file_rag) if (isSingleFileRag(form.selectedRagProcessor)) { + form.selectedFilePath = metadata?.file_path || ''; + } + + // Document RAG: new path (document_rag=single_file_rag) + form.documentRagEnabled = metadata?.document_rag === 'single_file_rag'; + if (form.documentRagEnabled && !isSingleFileRag(form.selectedRagProcessor)) { form.selectedLibraryId = metadata?.library_id || ''; form.selectedItemId = metadata?.item_id || ''; } @@ -258,8 +267,9 @@ export function clearRagDependentState(form) { form.rubricFormat = 'markdown'; } - // Library state - if (form.selectedLibraryId || form.libraries.length > 0 || form.selectedItemId || form.libraryItems.length > 0) { + // Library state — clear when neither document_rag nor legacy single_file_rag is active + const needsLibraryState = form.documentRagEnabled || isSingleFileRag(form.selectedRagProcessor); + if (!needsLibraryState && (form.selectedLibraryId || form.libraries.length > 0 || form.selectedItemId || form.libraryItems.length > 0)) { form.libraries = []; form.selectedLibraryId = ''; form.libraryItems = []; diff --git a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js index 6dc322f9f..bb98f4ee7 100644 --- a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js +++ b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js @@ -36,9 +36,12 @@ export function buildAssistantPayload(form) { } }; - if (isSingleFileRag(form.selectedRagProcessor)) { + if (form.documentRagEnabled) { + metadataObj.document_rag = 'single_file_rag'; metadataObj.library_id = form.selectedLibraryId || ''; metadataObj.item_id = form.selectedItemId || ''; + } else if (isSingleFileRag(form.selectedRagProcessor) && form.selectedFilePath) { + metadataObj.file_path = form.selectedFilePath; } if (isRubricRag(form.selectedRagProcessor)) { From 32d2890af10128bb2293b904431ee859b26d4591 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 13:05:17 +0200 Subject: [PATCH 028/156] feat: add independent Document Reference section, legacy single_file_rag read-only --- .../assistants/AssistantForm.svelte | 24 +++++++-- .../components/ConfigurationPanel.svelte | 49 ++++++++++++++++--- .../components/RagOptionsPanel.svelte | 24 ++++----- 3 files changed, 75 insertions(+), 22 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte index 5e5b00678..cd7751056 100644 --- a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte @@ -181,6 +181,22 @@ } }); + // Effect to fetch libraries when document_rag is enabled (independent of RAG processor) + $effect(() => { + if (form.documentRagEnabled && form.configInitialized) { + if (!form.librariesFetchAttempted && !form.loadingLibraries) { + doFetchLibraries(); + } + } + }); + + // Effect to fetch library items when selected library changes (for document_rag) + $effect(() => { + if (form.documentRagEnabled && form.selectedLibraryId && form.configInitialized) { + doFetchLibraryItems(form.selectedLibraryId); + } + }); + // --- Mode Switching Functions --- function switchToViewMode() { // Revert fields to initial state @@ -469,9 +485,11 @@ libraryError={form.libraryError} libraryItems={form.libraryItems} bind:selectedItemId={form.selectedItemId} - loadingItems={form.loadingItems} - itemsError={form.itemsError} - onchange={() => handleFieldChange(form)} + loadingItems={form.loadingItems} + itemsError={form.itemsError} + bind:documentRagEnabled={form.documentRagEnabled} + selectedFilePath={form.selectedFilePath} + onchange={() => handleFieldChange(form)} /> diff --git a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte index 7bfae87ba..db7708c6c 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte @@ -8,6 +8,7 @@ extractModelsMetadata } from '../logic/assistantFormUtils.svelte.js'; import RagOptionsPanel from './RagOptionsPanel.svelte'; + import LibraryItemSelector from './LibraryItemSelector.svelte'; let { formState, @@ -36,6 +37,8 @@ selectedItemId = $bindable(''), loadingItems = false, itemsError = '', + documentRagEnabled = $bindable(false), + selectedFilePath = '', onchange } = $props(); @@ -52,6 +55,13 @@ let currentModelMetadata = $derived(currentModelsMetadata.find((m) => m.id === selectedLlm) || null); let imageGenerationForced = $derived(currentModelMetadata?.forced_capabilities?.image_generation === true); let showRagOptions = $derived(hasRagOptions(selectedRagProcessor)); + let isLegacySingleFileRag = $derived(selectedRagProcessor === 'single_file_rag'); + let showDocumentSection = $derived(!isLegacySingleFileRag); + let filteredRAGProcessors = $derived( + formState === 'edit' + ? ragProcessors + : ragProcessors.filter((p) => p !== 'single_file_rag') + ); async function handleConnectorChange() { await tick(); @@ -190,12 +200,36 @@ + {#if showDocumentSection} +
+ +

{$_('assistants.form.documentRag.description', { default: 'Attach a reference document that will be available in the system context for all messages.' })}

+ {#if documentRagEnabled} + + {/if} +
+ {/if} + {#if showRagOptions} {/if} diff --git a/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte b/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte index 109847536..319f3da9b 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte @@ -21,12 +21,14 @@ selectedItemId = $bindable(''), loadingItems = false, itemsError = '', + selectedFilePath = '', formState } = $props(); let showKbSelector = $derived(isKbBasedRag(selectedRagProcessor)); let showFileSelector = $derived(isSingleFileRag(selectedRagProcessor)); let showTopK = $derived(isKbBasedRag(selectedRagProcessor)); + let isLegacyWithFilePath = $derived(isSingleFileRag(selectedRagProcessor) && !!selectedFilePath);
@@ -59,17 +61,15 @@ error={knowledgeBaseError} /> {/if} - {#if showFileSelector} - + {#if showFileSelector && isLegacyWithFilePath} +
+

+ {$_('assistants.form.ragOptions.legacySingleFileNotice', { default: 'This assistant uses the legacy single-file RAG. The document cannot be changed. To use a different document, create a new assistant with the Reference Document option.' })} +

+
+
+

{$_('assistants.form.ragOptions.document', { default: 'Document' })}

+

{selectedFilePath}

+
{/if}
From daeee0b49710cfbaa9f127c839ffd24873febeb8 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 13:39:17 +0200 Subject: [PATCH 029/156] refactor: extract parse_plugin_config and process_completion_request to lightweight module --- backend/lamb/completions/main.py | 43 +--------------- backend/lamb/completions/plugin_config.py | 51 +++++++++++++++++++ .../completions/test_document_rag_pipeline.py | 20 +------- 3 files changed, 53 insertions(+), 61 deletions(-) create mode 100644 backend/lamb/completions/plugin_config.py diff --git a/backend/lamb/completions/main.py b/backend/lamb/completions/main.py index 0c49e2258..1c48d5198 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 @@ -304,39 +305,6 @@ 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": "", - "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 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. @@ -383,15 +351,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], 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 - def load_plugins(plugin_type: str) -> Dict[str, Any]: """ Dynamically load plugins from the specified directory 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/testing/unit-tests/completions/test_document_rag_pipeline.py b/testing/unit-tests/completions/test_document_rag_pipeline.py index b0c5ab8a7..7b36a719a 100644 --- a/testing/unit-tests/completions/test_document_rag_pipeline.py +++ b/testing/unit-tests/completions/test_document_rag_pipeline.py @@ -1,25 +1,7 @@ import json -import os import pytest -_dummy_env = { - "LAMB_WEB_HOST": "http://localhost:9099", - "LAMB_BACKEND_HOST": "http://localhost:9099", - "LAMB_BEARER_TOKEN": "test-token", - "OPENAI_BASE_URL": "http://localhost:11434", - "OPENAI_MODEL": "gpt-4", - "OWI_BASE_URL": "http://localhost:8080", - "OWI_PATH": "/tmp/owi", - "SIGNUP_SECRET_KEY": "test-secret", - "OWI_ADMIN_NAME": "admin", - "OWI_ADMIN_EMAIL": "admin@test.com", - "OWI_ADMIN_PASSWORD": "test-password", - "LAMB_DB_PATH": "/tmp/test_lamb.db", -} -for _k, _v in _dummy_env.items(): - os.environ.setdefault(_k, _v) - -from lamb.completions.main import parse_plugin_config, process_completion_request +from lamb.completions.plugin_config import parse_plugin_config, process_completion_request class MockAssistant: From da8998d9a5a29d4167ea71a8ae60c901605cca8f Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 13:51:42 +0200 Subject: [PATCH 030/156] fix: disable document_rag toggle in edit mode --- .../components/assistants/components/ConfigurationPanel.svelte | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte index db7708c6c..6181211dc 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte @@ -209,7 +209,7 @@ {#if showDocumentSection}
From b043133873f9b2d7a66a69e66e1083d4121e4ea2 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Fri, 29 May 2026 13:51:43 +0200 Subject: [PATCH 031/156] fix: remove empty options from LibraryItemSelector dropdowns --- .../assistants/components/LibraryItemSelector.svelte | 6 ------ 1 file changed, 6 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/components/LibraryItemSelector.svelte b/frontend/svelte-app/src/lib/components/assistants/components/LibraryItemSelector.svelte index 08a57c0e9..c3c59a1c0 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/LibraryItemSelector.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/LibraryItemSelector.svelte @@ -56,9 +56,6 @@ onchange={handleLibraryChange} class="block w-full px-3 py-2 border border-gray-300 rounded-md shadow-sm focus:outline-none focus:ring-brand focus:border-brand sm:text-sm bg-white text-gray-900" > - {#each libraries as lib (lib.id)} {/each} @@ -92,9 +89,6 @@ onchange={handleItemChange} class="block w-full px-3 py-2 border border-gray-300 rounded-md shadow-sm focus:outline-none focus:ring-brand focus:border-brand sm:text-sm bg-white text-gray-900" > - {#each readyItems as item (item.id)} {/each} From deb37fd03cfcd0f2203b83ac4a8bdfaec9c7a23b Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 01:16:07 +0200 Subject: [PATCH 032/156] fix:assistantform bug from merge --- .../src/lib/components/assistants/AssistantForm.svelte | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte index a9b9e4d79..eb3d47471 100644 --- a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte @@ -29,7 +29,7 @@ // --- Props --- // --- Props --- // Use $props for Svelte 5 runes mode - let { + let { assistant = null, onFormSuccess = /** @type {(e: { assistantId: number }) => void} */ (() => {}), From 119231a879b027d01d5745d1e74bb18ca0c98894 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 01:47:35 +0200 Subject: [PATCH 033/156] sintax fix --- .../src/lib/components/assistants/AssistantForm.svelte | 2 -- 1 file changed, 2 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte index eb3d47471..c25421b36 100644 --- a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte @@ -507,7 +507,5 @@ {/if} - {/if} -
From bbc611feceb76d7a53fc3d113a08016a2efca450 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 18:57:19 +0200 Subject: [PATCH 034/156] fix: KB server legacy auth and connectivity + KB server v2 HuggingFace startup - Add LAMB_API_KEY to legacy KB service (docker-compose) to fix 401 errors - Remove _KB_REDIRECTS from kb_server_manager.py that broke connectivity when running inside Docker (redirect pointed to unreachable 172.18.0.1) - Add HuggingFace env vars and cache setup to KB server v2 service - Disable HF telemetry in local embedding plugin to prevent token read errors when running as non-root user --- backend/creator_interface/kb_server_manager.py | 10 +++------- docker-compose-example.yaml | 11 +++++++++-- lamb-kb-server/backend/plugins/embedding/local.py | 7 ++++++- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/backend/creator_interface/kb_server_manager.py b/backend/creator_interface/kb_server_manager.py index 0d1d88616..a68315a3f 100644 --- a/backend/creator_interface/kb_server_manager.py +++ b/backend/creator_interface/kb_server_manager.py @@ -21,9 +21,7 @@ # 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'} -LAMB_KB_SERVER = _KB_REDIRECTS.get(_raw_kb_server, _raw_kb_server) or 'http://172.17.0.1:9090' +LAMB_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: raise ValueError("LAMB_KB_SERVER_TOKEN environment variable is required") @@ -76,9 +74,8 @@ 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, + 'url': org_url, 'token': api_token } else: @@ -90,9 +87,8 @@ def _get_kb_config_for_user(self, creator_user: Dict[str, Any]) -> Dict[str, str # Fallback to global environment variables if not self.global_kb_server_url: raise ValueError("LAMB_KB_SERVER environment variable is required") - resolved_url = _KB_REDIRECTS.get(self.global_kb_server_url, self.global_kb_server_url) return { - 'url': resolved_url, + 'url': self.global_kb_server_url, 'token': self.global_kb_server_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/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: From eb0a9ea79d9cace2233c0899c088dd6070bc7649 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:14:46 +0200 Subject: [PATCH 035/156] feat: extract shared KS query helpers into _ks_query_helpers.py --- .../lamb/completions/rag/_ks_query_helpers.py | 112 ++++++++++++++++++ backend/tests/test_ks_query_helpers.py | 62 ++++++++++ 2 files changed, 174 insertions(+) create mode 100644 backend/lamb/completions/rag/_ks_query_helpers.py create mode 100644 backend/tests/test_ks_query_helpers.py 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..0d8ef637c --- /dev/null +++ b/backend/lamb/completions/rag/_ks_query_helpers.py @@ -0,0 +1,112 @@ +"""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 _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.""" + 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"), + } + 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 + sources.append(source) + return sources diff --git a/backend/tests/test_ks_query_helpers.py b/backend/tests/test_ks_query_helpers.py new file mode 100644 index 000000000..b5bc910f6 --- /dev/null +++ b/backend/tests/test_ks_query_helpers.py @@ -0,0 +1,62 @@ +import pytest +from lamb.completions.rag._ks_query_helpers import _extract_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" From deeff41a2f24893b3c1aea9c2e80776ba54b58ef Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:16:16 +0200 Subject: [PATCH 036/156] feat: extract shared query rewriting helper into _query_rewriting_helper.py --- .../rag/_query_rewriting_helper.py | 116 ++++++++++++++++++ backend/tests/test_query_rewriting_helper.py | 47 +++++++ 2 files changed, 163 insertions(+) create mode 100644 backend/lamb/completions/rag/_query_rewriting_helper.py create mode 100644 backend/tests/test_query_rewriting_helper.py 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/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) == "" From 1f69aeca17822b3b6fd9db383ee5f264a9684364 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:19:28 +0200 Subject: [PATCH 037/156] refactor: knowledge_store_rag uses shared _ks_query_helpers --- .../completions/rag/knowledge_store_rag.py | 141 ++---------------- 1 file changed, 9 insertions(+), 132 deletions(-) diff --git a/backend/lamb/completions/rag/knowledge_store_rag.py b/backend/lamb/completions/rag/knowledge_store_rag.py index 4ec066d51..f49741878 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, + _extract_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,21 +87,18 @@ 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 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") @@ -242,11 +125,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) From afc1b2faf72a6d44dc4100a0e1d09f6e20bfc951 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:22:28 +0200 Subject: [PATCH 038/156] refactor: context_aware_rag uses shared _query_rewriting_helper --- .../lamb/completions/rag/context_aware_rag.py | 147 +----------------- 1 file changed, 4 insertions(+), 143 deletions(-) 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 = [] From 3c63612df9135e176edd1544435388bc11cb7b1a Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:27:45 +0200 Subject: [PATCH 039/156] feat: add query_rewriting_ks_rag processor with small-fast-model query optimization --- .../completions/rag/query_rewriting_ks_rag.py | 122 +++++++++++++ backend/pytest.ini | 1 + backend/tests/test_query_rewriting_ks_rag.py | 161 ++++++++++++++++++ 3 files changed, 284 insertions(+) create mode 100644 backend/lamb/completions/rag/query_rewriting_ks_rag.py create mode 100644 backend/tests/test_query_rewriting_ks_rag.py 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..fb1c10236 --- /dev/null +++ b/backend/lamb/completions/rag/query_rewriting_ks_rag.py @@ -0,0 +1,122 @@ +"""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, + _extract_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] = {} + sources: List[Dict[str, Any]] = [] + contexts: List[str] = [] + + 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") + sources.extend(_extract_sources(ks_id, chunks)) + for chunk in chunks: + text = chunk.get("text") or "" + if text: + contexts.append(text) + else: + logger.warning(f"KS {ks_id}: {result.get('error', 'unknown error')}") + + combined_context = "\n\n".join(contexts) if contexts else "" + + 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/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_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") From 5c7839a2b40c2ea241e910057c284885ce8bc010 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:28:41 +0200 Subject: [PATCH 040/156] feat: add KS_BASED RAG type classification for knowledge store processors --- .../svelte-app/src/lib/utils/ragProcessorHelpers.js | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js index f2c9499e8..d022e4dbc 100644 --- a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js +++ b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js @@ -15,6 +15,8 @@ export const RAG_TYPES = Object.freeze({ /** RAG processors that use knowledge base collections */ KB_BASED: ['simple_rag', 'context_aware_rag', 'hierarchical_rag'], + /** RAG processors that use knowledge stores (new KB Server v2) */ + KS_BASED: ['query_rewriting_ks_rag', 'knowledge_store_rag'], /** RAG processor that uses a single file */ SINGLE_FILE: ['single_file_rag'], /** RAG processor that uses rubrics */ @@ -32,6 +34,15 @@ export function isKbBasedRag(processor) { return RAG_TYPES.KB_BASED.includes(processor); } +/** + * Returns true if the processor uses knowledge stores. + * @param {string} processor + * @returns {boolean} + */ +export function isKsBasedRag(processor) { + return RAG_TYPES.KS_BASED.includes(processor); +} + /** * Returns true if the processor uses a single file. * @param {string} processor From 07faad86b32c7b229652ba1bd96570df02c1269d Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:30:11 +0200 Subject: [PATCH 041/156] feat: add Knowledge Store state fields to assistant form state --- .../logic/assistantFormState.svelte.js | 36 ++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js index 1eeb387a0..25607f42a 100644 --- a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js +++ b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js @@ -16,7 +16,7 @@ import { get } from 'svelte/store'; import { assistantConfigStore } from '$lib/stores/assistantConfigStore'; -import { isKbBasedRag, isSingleFileRag, isRubricRag, normalizeRagProcessor } from '$lib/utils/ragProcessorHelpers.js'; +import { isKbBasedRag, isKsBasedRag, isSingleFileRag, isRubricRag, normalizeRagProcessor } from '$lib/utils/ragProcessorHelpers.js'; import { loadRagPlaceholders, selectModel } from './assistantFormUtils.svelte.js'; import { getAssistantMetadataObject } from '$lib/utils/assistantData'; @@ -67,6 +67,19 @@ export function createAssistantFormState() { /** @type {string[] | null} */ pendingKBSelections: null, + // --- Knowledge Store state --- + /** @type {Array<{id: string, name: string, is_shared: boolean, owner_email: string, embedding_vendor: string, embedding_model: string}>} */ + ownedKnowledgeStores: [], + /** @type {Array<{id: string, name: string, is_shared: boolean, owner_email: string, embedding_vendor: string, embedding_model: string}>} */ + sharedKnowledgeStores: [], + /** @type {string[]} */ + selectedKnowledgeStores: [], + loadingKnowledgeStores: false, + knowledgeStoreError: '', + ksFetchAttempted: false, + /** @type {string[] | null} */ + pendingKSSelections: null, + // --- File state --- /** @type {Array<{name: string, path: string}>} */ userFiles: [], @@ -114,6 +127,9 @@ export function createAssistantFormState() { // --- Derived --- get accessibleKnowledgeBases() { return [...this.ownedKnowledgeBases, ...this.sharedKnowledgeBases]; + }, + get accessibleKnowledgeStores() { + return [...this.ownedKnowledgeStores, ...this.sharedKnowledgeStores]; } }); @@ -147,6 +163,9 @@ export function resetFormFieldsToDefaults(form, getAvailableModels) { form.selectedLlm = selectModel(defaults.llm || '', getAvailableModels()); form.selectedKnowledgeBases = []; + form.ownedKnowledgeStores = []; + form.sharedKnowledgeStores = []; + form.selectedKnowledgeStores = []; form.selectedFilePath = ''; form.selectedLibraryId = ''; form.selectedItemId = ''; @@ -194,6 +213,13 @@ export function populateFormFields(form, data, getAvailableModels, preserveDescr form.pendingKBSelections = null; } + // Deferred selection for KSs + if (isKsBasedRag(form.selectedRagProcessor)) { + form.pendingKSSelections = data.RAG_collections?.split(',').filter(Boolean) || []; + } else { + form.pendingKSSelections = null; + } + // Rubric fields if (isRubricRag(form.selectedRagProcessor)) { try { @@ -258,6 +284,14 @@ export function clearRagDependentState(form) { form.kbFetchAttempted = false; } + if (!isKsBasedRag(form.selectedRagProcessor) && (form.selectedKnowledgeStores.length > 0 || form.ownedKnowledgeStores.length > 0 || form.sharedKnowledgeStores.length > 0 || form.ksFetchAttempted)) { + form.ownedKnowledgeStores = []; + form.sharedKnowledgeStores = []; + form.selectedKnowledgeStores = []; + form.knowledgeStoreError = ''; + form.ksFetchAttempted = false; + } + if (!isSingleFileRag(form.selectedRagProcessor) && (form.selectedFilePath || form.userFiles.length > 0)) { form.selectedFilePath = ''; } From b2b01593131613b856b857071744ed53b89420a2 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:30:52 +0200 Subject: [PATCH 042/156] feat: add fetchKnowledgeStores to assistant form fetchers --- .../assistants/logic/assistantFormFetchers.js | 39 ++++++++++++++++++- 1 file changed, 38 insertions(+), 1 deletion(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormFetchers.js b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormFetchers.js index 7ce0eedf0..5b8546628 100644 --- a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormFetchers.js +++ b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormFetchers.js @@ -5,9 +5,10 @@ */ import { getUserKnowledgeBases, getSharedKnowledgeBases } from '$lib/services/knowledgeBaseService'; +import { getKnowledgeStores } from '$lib/services/knowledgeStoreService'; import { fetchAccessibleRubrics } from '$lib/services/rubricService'; import { apiJson } from '$lib/services/apiClient'; -import { isKbBasedRag, isRubricRag } from '$lib/utils/ragProcessorHelpers.js'; +import { isKbBasedRag, isKsBasedRag, isRubricRag } from '$lib/utils/ragProcessorHelpers.js'; import { getAssistantMetadataObject } from '$lib/utils/assistantData'; import { getLibraries, getItems } from '$lib/services/libraryService'; @@ -43,6 +44,42 @@ export async function fetchKnowledgeBases(form) { } } +/** + * Fetches accessible knowledge stores, separating owned vs shared + * by comparing owner_email against the logged-in user's email. + * @param {import('./assistantFormState.svelte.js').createAssistantFormState} form + */ +export async function fetchKnowledgeStores(form) { + if (form.loadingKnowledgeStores || form.ksFetchAttempted) return; + if (!isKsBasedRag(form.selectedRagProcessor)) return; + + form.loadingKnowledgeStores = true; + form.knowledgeStoreError = ''; + + try { + const stores = await getKnowledgeStores(); + const currentEmail = localStorage.getItem('userEmail') || ''; + + const owned = stores.filter((ks) => ks.owner_email === currentEmail); + const shared = stores.filter((ks) => ks.owner_email !== currentEmail); + + owned.sort((a, b) => a.name.localeCompare(b.name)); + shared.sort((a, b) => a.name.localeCompare(b.name)); + + form.ownedKnowledgeStores = owned; + form.sharedKnowledgeStores = shared; + } catch (err) { + if (err instanceof Error && err.message.startsWith('Session expired')) return; + console.error('Error fetching knowledge stores:', err); + form.knowledgeStoreError = err instanceof Error ? err.message : 'Failed to load knowledge stores'; + form.ownedKnowledgeStores = []; + form.sharedKnowledgeStores = []; + } finally { + form.loadingKnowledgeStores = false; + form.ksFetchAttempted = true; + } +} + /** * Fetches accessible rubrics. * @param {import('./assistantFormState.svelte.js').createAssistantFormState} form From 7761dce1c6ed214df721565520302bf288b679cb Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:31:09 +0200 Subject: [PATCH 043/156] feat: include Knowledge Store IDs in assistant submission payload --- .../components/assistants/logic/assistantFormSubmit.js | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js index bb98f4ee7..de308a78b 100644 --- a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js +++ b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js @@ -4,7 +4,7 @@ * Extracted from AssistantForm.svelte to enable isolated testing. */ -import { isKbBasedRag, isSingleFileRag, isRubricRag } from '$lib/utils/ragProcessorHelpers.js'; +import { isKbBasedRag, isKsBasedRag, isSingleFileRag, isRubricRag } from '$lib/utils/ragProcessorHelpers.js'; /** * Validates form data before submission. @@ -55,7 +55,11 @@ export function buildAssistantPayload(form) { system_prompt: form.system_prompt, prompt_template: form.prompt_template, RAG_Top_k: Number(form.RAG_Top_k) || 3, - RAG_collections: isKbBasedRag(form.selectedRagProcessor) ? form.selectedKnowledgeBases.join(',') : '', + RAG_collections: isKbBasedRag(form.selectedRagProcessor) + ? form.selectedKnowledgeBases.join(',') + : isKsBasedRag(form.selectedRagProcessor) + ? form.selectedKnowledgeStores.join(',') + : '', metadata: JSON.stringify(metadataObj), pre_retrieval_endpoint: '', post_retrieval_endpoint: '', From c18c2e89727449be8409bf0588de52a86087b661 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:31:59 +0200 Subject: [PATCH 044/156] feat: add KnowledgeStoreSelector component for assistant form --- .../components/KnowledgeStoreSelector.svelte | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 frontend/svelte-app/src/lib/components/assistants/components/KnowledgeStoreSelector.svelte diff --git a/frontend/svelte-app/src/lib/components/assistants/components/KnowledgeStoreSelector.svelte b/frontend/svelte-app/src/lib/components/assistants/components/KnowledgeStoreSelector.svelte new file mode 100644 index 000000000..9ece37b9d --- /dev/null +++ b/frontend/svelte-app/src/lib/components/assistants/components/KnowledgeStoreSelector.svelte @@ -0,0 +1,72 @@ + + + +
+

+ {$_('assistants.form.knowledgeStores.label', { default: 'Knowledge Stores' })} +

+ {#if loading} +

+ {$_('assistants.form.knowledgeStores.loading', { default: 'Loading knowledge stores...' })} +

+ {:else if error} +

+ {$_('assistants.form.knowledgeStores.error', { default: 'Error loading knowledge stores:' })} {error} +

+ {:else if ownedKnowledgeStores.length === 0 && sharedKnowledgeStores.length === 0} +

+ {$_('assistants.form.knowledgeStores.noneFound', { default: 'No accessible knowledge stores found.' })} +

+ {:else} +
+ {#if ownedKnowledgeStores.length > 0} +
+
+ {$_('assistants.form.knowledgeStores.myKS', { default: 'My Knowledge Stores' })} +
+
+ {#each ownedKnowledgeStores as ks (ks.id)} + + {/each} +
+
+ {/if} + {#if sharedKnowledgeStores.length > 0} +
+
+ {$_('assistants.form.knowledgeStores.sharedKS', { default: 'Shared Knowledge Stores' })} +
+
+ {#each sharedKnowledgeStores as ks (ks.id)} + + {/each} +
+
+ {/if} +
+ {/if} +
From 7be4649484a09876bffb0c1045525b5d75ff08b3 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:32:32 +0200 Subject: [PATCH 045/156] feat: wire KnowledgeStoreSelector into RagOptionsPanel with Top K support --- .../components/RagOptionsPanel.svelte | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte b/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte index 319f3da9b..92e645ccc 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/RagOptionsPanel.svelte @@ -1,8 +1,9 @@ @@ -61,6 +68,15 @@ error={knowledgeBaseError} /> {/if} + {#if showKsSelector} + + {/if} {#if showFileSelector && isLegacyWithFilePath}

From e0bc7f39006a34d82d869124676f9387a6e9ac4b Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:33:06 +0200 Subject: [PATCH 046/156] feat: pass Knowledge Store props through ConfigurationPanel --- .../assistants/components/ConfigurationPanel.svelte | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte index 6181211dc..1aa2798bf 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte @@ -29,6 +29,11 @@ selectedKnowledgeBases = $bindable([]), loadingKnowledgeBases = false, knowledgeBaseError = '', + ownedKnowledgeStores = [], + sharedKnowledgeStores = [], + selectedKnowledgeStores = $bindable([]), + loadingKnowledgeStores = false, + knowledgeStoreError = '', libraries = [], selectedLibraryId = $bindable(''), loadingLibraries = false, @@ -239,6 +244,11 @@ bind:selectedKnowledgeBases loadingKnowledgeBases={loadingKnowledgeBases} knowledgeBaseError={knowledgeBaseError} + {ownedKnowledgeStores} + {sharedKnowledgeStores} + bind:selectedKnowledgeStores + loadingKnowledgeStores={loadingKnowledgeStores} + knowledgeStoreError={knowledgeStoreError} {libraries} bind:selectedLibraryId loadingLibraries={loadingLibraries} From 1a83b60fe2deefd9d7248d2e9f283ce4052f5fbc Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:33:58 +0200 Subject: [PATCH 047/156] feat: wire Knowledge Store fetcher and effects into AssistantForm orchestrator --- .../assistants/AssistantForm.svelte | 38 +++++++++++++++---- 1 file changed, 31 insertions(+), 7 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte index c25421b36..eb47d1dc3 100644 --- a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte @@ -9,7 +9,7 @@ import { isKbBasedRag, isSingleFileRag, isRubricRag } from '$lib/utils/ragProcessorHelpers.js'; import { validateImportedAssistant } from './logic/importAssistantValidator.js'; import { createAssistantFormState, resetFormFieldsToDefaults, populateFormFields, revertToInitial, clearRagDependentState, handleFieldChange } from './logic/assistantFormState.svelte.js'; - import { fetchKnowledgeBases, fetchRubricsList, fetchLibraries, fetchLibraryItems } from './logic/assistantFormFetchers.js'; + import { fetchKnowledgeBases, fetchRubricsList, fetchLibraries, fetchLibraryItems, fetchKnowledgeStores } from './logic/assistantFormFetchers.js'; import { validateSubmission, buildAssistantPayload } from './logic/assistantFormSubmit.js'; import AssistantFormHeader from './components/AssistantFormHeader.svelte'; import AssistantNameField from './components/AssistantNameField.svelte'; @@ -54,6 +54,11 @@ await fetchKnowledgeBases(form); } + async function doFetchKnowledgeStores() { + if (!isMounted) return; + await fetchKnowledgeStores(form); + } + async function doFetchLibraries(force = false) { if (!isMounted) return; await fetchLibraries(form, force); @@ -152,6 +157,16 @@ } }); + // Effect to apply pending KS selections when list becomes available + $effect(() => { + if (form.pendingKSSelections && form.accessibleKnowledgeStores.length > 0) { + form.selectedKnowledgeStores = form.pendingKSSelections.filter((id) => + form.accessibleKnowledgeStores.some((ks) => ks.id === id) + ); + form.pendingKSSelections = null; + } + }); + // Effect to fetch KBs/Files when RAG processor changes $effect(() => { if ((isKbBasedRag(form.selectedRagProcessor)) && form.configInitialized) { @@ -161,6 +176,10 @@ } else { // Already attempted or loading } + } else if (isKsBasedRag(form.selectedRagProcessor) && form.configInitialized) { + if (!form.ksFetchAttempted && !form.loadingKnowledgeStores) { + doFetchKnowledgeStores(); + } } else if (isSingleFileRag(form.selectedRagProcessor) && form.configInitialized) { // Fetch libraries when switching to single_file_rag if (!form.librariesFetchAttempted && !form.loadingLibraries) { @@ -477,12 +496,17 @@ bind:visionEnabled={form.visionEnabled} bind:imageGenerationEnabled={form.imageGenerationEnabled} bind:RAG_Top_k={form.RAG_Top_k} - ownedKnowledgeBases={form.ownedKnowledgeBases} - sharedKnowledgeBases={form.sharedKnowledgeBases} - bind:selectedKnowledgeBases={form.selectedKnowledgeBases} - loadingKnowledgeBases={form.loadingKnowledgeBases} - knowledgeBaseError={form.knowledgeBaseError} - libraries={form.libraries} + ownedKnowledgeBases={form.ownedKnowledgeBases} + sharedKnowledgeBases={form.sharedKnowledgeBases} + bind:selectedKnowledgeBases={form.selectedKnowledgeBases} + loadingKnowledgeBases={form.loadingKnowledgeBases} + knowledgeBaseError={form.knowledgeBaseError} + ownedKnowledgeStores={form.ownedKnowledgeStores} + sharedKnowledgeStores={form.sharedKnowledgeStores} + bind:selectedKnowledgeStores={form.selectedKnowledgeStores} + loadingKnowledgeStores={form.loadingKnowledgeStores} + knowledgeStoreError={form.knowledgeStoreError} + libraries={form.libraries} bind:selectedLibraryId={form.selectedLibraryId} loadingLibraries={form.loadingLibraries} libraryError={form.libraryError} From 42b7ea42c0a24e322e2b63b21290bfdddc8330fd Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 20:34:31 +0200 Subject: [PATCH 048/156] feat: add query_rewriting_ks_rag to fallback capabilities and i18n keys --- frontend/svelte-app/src/lib/locales/en.json | 8 ++++++++ frontend/svelte-app/src/lib/locales/es.json | 8 ++++++++ .../svelte-app/src/lib/stores/assistantConfigStore.js | 2 +- 3 files changed, 17 insertions(+), 1 deletion(-) diff --git a/frontend/svelte-app/src/lib/locales/en.json b/frontend/svelte-app/src/lib/locales/en.json index d2c0d565e..6988a7223 100644 --- a/frontend/svelte-app/src/lib/locales/en.json +++ b/frontend/svelte-app/src/lib/locales/en.json @@ -328,6 +328,14 @@ "sharedKB": "Shared Knowledge Bases", "shared": "Shared by {shared_by}" }, + "knowledgeStores": { + "label": "Knowledge Stores", + "loading": "Loading knowledge stores...", + "error": "Error loading knowledge stores:", + "noneFound": "No accessible knowledge stores found.", + "myKS": "My Knowledge Stores", + "sharedKS": "Shared Knowledge Stores" + }, "singleFile": { "label": "Select File", "selectedLabel": "Selected File", diff --git a/frontend/svelte-app/src/lib/locales/es.json b/frontend/svelte-app/src/lib/locales/es.json index 74fd8c00b..271d48691 100644 --- a/frontend/svelte-app/src/lib/locales/es.json +++ b/frontend/svelte-app/src/lib/locales/es.json @@ -190,6 +190,14 @@ "sharedKB": "Bases de Conocimiento Compartidas", "shared": "Compartida por {shared_by}" }, + "knowledgeStores": { + "label": "Knowledge Stores", + "loading": "Cargando knowledge stores...", + "error": "Error al cargar knowledge stores:", + "noneFound": "No se encontraron knowledge stores accesibles.", + "myKS": "Mis Knowledge Stores", + "sharedKS": "Knowledge Stores compartidos" + }, "singleFile": { "label": "Seleccionar Archivo", "selectedLabel": "Archivo Seleccionado", diff --git a/frontend/svelte-app/src/lib/stores/assistantConfigStore.js b/frontend/svelte-app/src/lib/stores/assistantConfigStore.js index 30890f04b..922621feb 100644 --- a/frontend/svelte-app/src/lib/stores/assistantConfigStore.js +++ b/frontend/svelte-app/src/lib/stores/assistantConfigStore.js @@ -84,7 +84,7 @@ function getFallbackCapabilities() { const capabilities = { prompt_processors: ['simple_augment'], connectors: {}, - rag_processors: ['no_rag', 'simple_rag', 'context_aware_rag', 'single_file_rag'] + rag_processors: ['no_rag', 'simple_rag', 'context_aware_rag', 'single_file_rag', 'query_rewriting_ks_rag'] }; return capabilities; } From e954a444544dbef6b2058dd02c1ceb0ed210bd5d Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 1 Jun 2026 22:51:06 +0200 Subject: [PATCH 049/156] fix: add missing isKsBasedRag import and fix ConfigurationPanel prop indentation --- .../assistants/AssistantForm.svelte | 36 +++++++++---------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte index eb47d1dc3..8f39f146b 100644 --- a/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/AssistantForm.svelte @@ -6,7 +6,7 @@ import { get } from 'svelte/store'; import { createAssistant, updateAssistant } from '$lib/services/assistantService'; import { extractModelsFromConnectorData, selectModel } from './logic/assistantFormUtils.svelte.js'; - import { isKbBasedRag, isSingleFileRag, isRubricRag } from '$lib/utils/ragProcessorHelpers.js'; + import { isKbBasedRag, isKsBasedRag, isSingleFileRag, isRubricRag } from '$lib/utils/ragProcessorHelpers.js'; import { validateImportedAssistant } from './logic/importAssistantValidator.js'; import { createAssistantFormState, resetFormFieldsToDefaults, populateFormFields, revertToInitial, clearRagDependentState, handleFieldChange } from './logic/assistantFormState.svelte.js'; import { fetchKnowledgeBases, fetchRubricsList, fetchLibraries, fetchLibraryItems, fetchKnowledgeStores } from './logic/assistantFormFetchers.js'; @@ -496,28 +496,28 @@ bind:visionEnabled={form.visionEnabled} bind:imageGenerationEnabled={form.imageGenerationEnabled} bind:RAG_Top_k={form.RAG_Top_k} - ownedKnowledgeBases={form.ownedKnowledgeBases} - sharedKnowledgeBases={form.sharedKnowledgeBases} - bind:selectedKnowledgeBases={form.selectedKnowledgeBases} - loadingKnowledgeBases={form.loadingKnowledgeBases} - knowledgeBaseError={form.knowledgeBaseError} - ownedKnowledgeStores={form.ownedKnowledgeStores} - sharedKnowledgeStores={form.sharedKnowledgeStores} - bind:selectedKnowledgeStores={form.selectedKnowledgeStores} - loadingKnowledgeStores={form.loadingKnowledgeStores} - knowledgeStoreError={form.knowledgeStoreError} - libraries={form.libraries} + ownedKnowledgeBases={form.ownedKnowledgeBases} + sharedKnowledgeBases={form.sharedKnowledgeBases} + bind:selectedKnowledgeBases={form.selectedKnowledgeBases} + loadingKnowledgeBases={form.loadingKnowledgeBases} + knowledgeBaseError={form.knowledgeBaseError} + ownedKnowledgeStores={form.ownedKnowledgeStores} + sharedKnowledgeStores={form.sharedKnowledgeStores} + bind:selectedKnowledgeStores={form.selectedKnowledgeStores} + loadingKnowledgeStores={form.loadingKnowledgeStores} + knowledgeStoreError={form.knowledgeStoreError} + libraries={form.libraries} bind:selectedLibraryId={form.selectedLibraryId} loadingLibraries={form.loadingLibraries} libraryError={form.libraryError} libraryItems={form.libraryItems} bind:selectedItemId={form.selectedItemId} - loadingItems={form.loadingItems} - itemsError={form.itemsError} - bind:documentRagEnabled={form.documentRagEnabled} - selectedFilePath={form.selectedFilePath} - onchange={() => handleFieldChange(form)} - /> + loadingItems={form.loadingItems} + itemsError={form.itemsError} + bind:documentRagEnabled={form.documentRagEnabled} + selectedFilePath={form.selectedFilePath} + onchange={() => handleFieldChange(form)} + />

From b5cc2ae3ba478111d7da0369c8fa074a850676f5 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Tue, 2 Jun 2026 17:27:07 +0200 Subject: [PATCH 050/156] fix: remove remaining _KB_REDIRECTS reference in create_knowledgebase --- backend/creator_interface/kb_server_manager.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/backend/creator_interface/kb_server_manager.py b/backend/creator_interface/kb_server_manager.py index a68315a3f..e8c7653d7 100644 --- a/backend/creator_interface/kb_server_manager.py +++ b/backend/creator_interface/kb_server_manager.py @@ -493,10 +493,6 @@ async def create_knowledge_base(self, kb_data: KnowledgeBaseCreate, creator_user 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) - if _alt_url: - kb_server_url = _alt_url async with httpx.AsyncClient() as client: kb_server_collections_url = f"{kb_server_url}/collections" logger.info(f"Creating collection in KB server at {kb_server_collections_url}: {sanitized_name}") From 25034560d475b92ee74b20a308ad847210416b2a Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Tue, 2 Jun 2026 19:33:40 +0200 Subject: [PATCH 051/156] feat: wrap Document RAG content with descriptive REFERENCE DOCUMENT label in system prompt --- .../lamb/completions/pps/simple_augment.py | 9 ++++++- backend/tests/test_simple_augment.py | 25 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/backend/lamb/completions/pps/simple_augment.py b/backend/lamb/completions/pps/simple_augment.py index a9e3bef50..6a8d0a7f9 100644 --- a/backend/lamb/completions/pps/simple_augment.py +++ b/backend/lamb/completions/pps/simple_augment.py @@ -98,7 +98,14 @@ def prompt_processor( if document_context and isinstance(document_context, dict): doc_text = document_context.get("context", "") if doc_text: - system_content = (system_content + "\n\n" + doc_text) if system_content else 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}" + ) + system_content = (system_content + labeled_doc) if system_content else labeled_doc if system_content: processed_messages.append({ "role": "system", diff --git a/backend/tests/test_simple_augment.py b/backend/tests/test_simple_augment.py index d60a62f3d..d4da9c31b 100644 --- a/backend/tests/test_simple_augment.py +++ b/backend/tests/test_simple_augment.py @@ -89,6 +89,31 @@ def test_explicit_template_with_context_placeholder_unchanged(): assert "Some chunk." in augmented +def test_document_context_gets_descriptive_label(): + """Document RAG content should be wrapped with a descriptive REFERENCE DOCUMENT label.""" + assistant = _make_assistant(system_prompt="You are helpful.", prompt_template="Answer: {user_input}") + request = { + "messages": [ + {"role": "user", "content": "What is the answer?"} + ] + } + document_context = {"context": "This is the reference document content.", "sources": []} + + result = prompt_processor(request, assistant=assistant, document_context=document_context) + + system_msg = result[0] + assert system_msg["role"] == "system" + assert "REFERENCE DOCUMENT" in system_msg["content"] + assert "This is the reference document content." in system_msg["content"] + # Should contain the explanation for the LLM + assert "selected by the assistant creator" in system_msg["content"] + # Should come after the original system prompt + assert "You are helpful." in system_msg["content"] + doc_index = system_msg["content"].index("REFERENCE DOCUMENT") + sys_index = system_msg["content"].index("You are helpful.") + assert doc_index > sys_index + + 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 From c90addfef0c718afb7dd35725903183e2e10dfb7 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Tue, 2 Jun 2026 19:35:15 +0200 Subject: [PATCH 052/156] feat: rename query_rewriting_ks_rag display to 'Context Aware Rag', hide legacy processors from create dropdown --- .../components/ConfigurationPanel.svelte | 6 +-- .../src/lib/stores/assistantConfigStore.js | 2 +- .../src/lib/utils/ragProcessorHelpers.js | 44 ++++++++++++++++--- 3 files changed, 41 insertions(+), 11 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte index 1aa2798bf..8715a9611 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte @@ -3,7 +3,7 @@ import { _ } from '$lib/i18n'; import { tick } from 'svelte'; import { assistantConfigStore } from '$lib/stores/assistantConfigStore'; - import { hasRagOptions } from '$lib/utils/ragProcessorHelpers.js'; + import { hasRagOptions, isHiddenInCreate, getRagProcessorDisplayName } from '$lib/utils/ragProcessorHelpers.js'; import { extractModelsMetadata } from '../logic/assistantFormUtils.svelte.js'; @@ -65,7 +65,7 @@ let filteredRAGProcessors = $derived( formState === 'edit' ? ragProcessors - : ragProcessors.filter((p) => p !== 'single_file_rag') + : ragProcessors.filter((p) => !isHiddenInCreate(p)) ); async function handleConnectorChange() { @@ -206,7 +206,7 @@ disabled={formState === 'edit'} class="mt-1 block w-full pl-3 pr-10 py-2 text-base text-gray-900 border border-gray-300 focus:outline-none focus:ring-brand focus:border-brand sm:text-sm rounded-md bg-white disabled:bg-gray-100 disabled:cursor-not-allowed"> {#each filteredRAGProcessors as processor (processor)} - + {/each} diff --git a/frontend/svelte-app/src/lib/stores/assistantConfigStore.js b/frontend/svelte-app/src/lib/stores/assistantConfigStore.js index 922621feb..1647757b6 100644 --- a/frontend/svelte-app/src/lib/stores/assistantConfigStore.js +++ b/frontend/svelte-app/src/lib/stores/assistantConfigStore.js @@ -84,7 +84,7 @@ function getFallbackCapabilities() { const capabilities = { prompt_processors: ['simple_augment'], connectors: {}, - rag_processors: ['no_rag', 'simple_rag', 'context_aware_rag', 'single_file_rag', 'query_rewriting_ks_rag'] + rag_processors: ['no_rag', 'simple_rag', 'context_aware_rag', 'hierarchical_rag', 'single_file_rag', 'query_rewriting_ks_rag', 'knowledge_store_rag', 'rubric_rag'] }; return capabilities; } diff --git a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js index d022e4dbc..813015c17 100644 --- a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js +++ b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js @@ -6,17 +6,16 @@ * specific RAG processor strings, we centralize the classification * logic here. Each function encapsulates a "strategy" for determining * which UI/behavior to apply based on the RAG processor type. - * - * To add a new RAG processor type, add it to the appropriate constant - * array and all conditionals across the app update automatically. */ /** @readonly */ export const RAG_TYPES = Object.freeze({ - /** RAG processors that use knowledge base collections */ + /** RAG processors that use knowledge base collections (legacy, port 9090) */ KB_BASED: ['simple_rag', 'context_aware_rag', 'hierarchical_rag'], - /** RAG processors that use knowledge stores (new KB Server v2) */ + /** RAG processors that use knowledge stores (new KB Server v2, port 9092) */ KS_BASED: ['query_rewriting_ks_rag', 'knowledge_store_rag'], + /** Processors hidden from the create dropdown (kept for edit/backward-compat only) */ + HIDDEN_IN_CREATE: ['simple_rag', 'context_aware_rag', 'hierarchical_rag', 'single_file_rag', 'knowledge_store_rag'], /** RAG processor that uses a single file */ SINGLE_FILE: ['single_file_rag'], /** RAG processor that uses rubrics */ @@ -26,7 +25,7 @@ export const RAG_TYPES = Object.freeze({ }); /** - * Returns true if the processor uses knowledge base collections. + * Returns true if the processor uses knowledge base collections (legacy). * @param {string} processor * @returns {boolean} */ @@ -35,7 +34,7 @@ export function isKbBasedRag(processor) { } /** - * Returns true if the processor uses knowledge stores. + * Returns true if the processor uses knowledge stores (new). * @param {string} processor * @returns {boolean} */ @@ -43,6 +42,15 @@ export function isKsBasedRag(processor) { return RAG_TYPES.KS_BASED.includes(processor); } +/** + * Returns true if the processor should be hidden in the create dropdown. + * @param {string} processor + * @returns {boolean} + */ +export function isHiddenInCreate(processor) { + return RAG_TYPES.HIDDEN_IN_CREATE.includes(processor); +} + /** * Returns true if the processor uses a single file. * @param {string} processor @@ -90,3 +98,25 @@ export function normalizeRagProcessor(processor) { if (normalized === 'no rag') return 'no_rag'; return normalized; } + +/** + * Returns a human-readable display name for a RAG processor. + * Maps internal names to user-friendly labels: + * query_rewriting_ks_rag → "Context Aware Rag" + * context_aware_rag → "Context Aware Rag (Old)" + * knowledge_store_rag → "Knowledge Store Rag (Legacy)" + * @param {string} processor + * @returns {string} + */ +export function getRagProcessorDisplayName(processor) { + if (!processor) return ''; + const displayNames = { + 'query_rewriting_ks_rag': 'Context Aware Rag', + 'context_aware_rag': 'Context Aware Rag (Old)', + 'knowledge_store_rag': 'Knowledge Store Rag (Legacy)', + }; + if (displayNames[processor]) { + return displayNames[processor]; + } + return processor.replace(/_/g, ' ').replace(/\b\w/g, (l) => l.toUpperCase()); +} From d707fa33d7ec036fbcaee0bbf401185541785193 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Tue, 2 Jun 2026 19:39:21 +0200 Subject: [PATCH 053/156] feat: add missing i18n keys for Document RAG, Knowledge Stores, and singleFile across all 4 locales --- frontend/svelte-app/src/lib/locales/ca.json | 24 +++++++++++++++++++-- frontend/svelte-app/src/lib/locales/en.json | 16 ++++++++++++-- frontend/svelte-app/src/lib/locales/es.json | 16 ++++++++++++-- frontend/svelte-app/src/lib/locales/eu.json | 24 +++++++++++++++++++-- 4 files changed, 72 insertions(+), 8 deletions(-) diff --git a/frontend/svelte-app/src/lib/locales/ca.json b/frontend/svelte-app/src/lib/locales/ca.json index f4f6ae254..02b9e4ba5 100644 --- a/frontend/svelte-app/src/lib/locales/ca.json +++ b/frontend/svelte-app/src/lib/locales/ca.json @@ -175,7 +175,9 @@ "fileReadError": "Error llegint l'arxiu seleccionat." }, "ragOptions": { - "title": "Opcions RAG" + "title": "Opcions RAG", + "legacySingleFileNotice": "Aquest assistent utilitza el RAG de fitxer únic legacy. El document no es pot canviar. Per utilitzar un document diferent, crea un nou assistent amb l'opció Document de Referència.", + "document": "Document" }, "ragTopK": { "label": "RAG Top K", @@ -190,6 +192,18 @@ "sharedKB": "Bases de Coneixement Compartides", "shared": "Compartida per {shared_by}" }, + "knowledgeStores": { + "label": "Knowledge Stores", + "loading": "Carregant knowledge stores...", + "error": "Error en carregar knowledge stores:", + "noneFound": "No s'han trobat knowledge stores accessibles.", + "myKS": "Els meus Knowledge Stores", + "sharedKS": "Knowledge Stores compartits" + }, + "documentRag": { + "label": "Document de Referència", + "description": "Adjunta un document de referència que estarà disponible en el context del sistema per a tots els missatges." + }, "singleFile": { "label": "Seleccionar Arxiu", "selectedLabel": "Arxiu Seleccionat", @@ -198,7 +212,13 @@ "loading": "Carregant arxius...", "error": "Error en carregar arxius:", "noneFound": "No s'han trobat arxius. Si us plau puja un arxiu.", - "required": "Si us plau selecciona un arxiu" + "required": "Si us plau selecciona un arxiu", + "libraryLabel": "Biblioteca", + "loadingLibraries": "Carregant biblioteques...", + "noLibraries": "No hi ha biblioteques disponibles.", + "manageLibraries": "Gestionar biblioteques", + "itemLabel": "Document", + "loadingItems": "Carregant documents..." } }, "noDescription": "Sense descripció", diff --git a/frontend/svelte-app/src/lib/locales/en.json b/frontend/svelte-app/src/lib/locales/en.json index 6988a7223..a9bf60e37 100644 --- a/frontend/svelte-app/src/lib/locales/en.json +++ b/frontend/svelte-app/src/lib/locales/en.json @@ -313,7 +313,9 @@ "fileReadError": "Error reading the selected file." }, "ragOptions": { - "title": "RAG Options" + "title": "RAG Options", + "legacySingleFileNotice": "This assistant uses the legacy single-file RAG. The document cannot be changed. To use a different document, create a new assistant with the Reference Document option.", + "document": "Document" }, "ragTopK": { "label": "RAG Top K", @@ -336,6 +338,10 @@ "myKS": "My Knowledge Stores", "sharedKS": "Shared Knowledge Stores" }, + "documentRag": { + "label": "Reference Document", + "description": "Attach a reference document that will be available in the system context for all messages." + }, "singleFile": { "label": "Select File", "selectedLabel": "Selected File", @@ -344,7 +350,13 @@ "loading": "Loading files...", "error": "Error loading files:", "noneFound": "No files found. Please upload a file.", - "required": "Please select a file" + "required": "Please select a file", + "libraryLabel": "Library", + "loadingLibraries": "Loading libraries...", + "noLibraries": "No libraries available.", + "manageLibraries": "Manage libraries", + "itemLabel": "Document", + "loadingItems": "Loading documents..." } } }, diff --git a/frontend/svelte-app/src/lib/locales/es.json b/frontend/svelte-app/src/lib/locales/es.json index 271d48691..4057ea618 100644 --- a/frontend/svelte-app/src/lib/locales/es.json +++ b/frontend/svelte-app/src/lib/locales/es.json @@ -175,7 +175,9 @@ "fileReadError": "Error leyendo el archivo seleccionado." }, "ragOptions": { - "title": "Opciones RAG" + "title": "Opciones RAG", + "legacySingleFileNotice": "Este asistente usa el RAG de archivo único legacy. El documento no se puede cambiar. Para usar un documento diferente, crea un nuevo asistente con la opción Documento de Referencia.", + "document": "Documento" }, "ragTopK": { "label": "RAG Top K", @@ -198,6 +200,10 @@ "myKS": "Mis Knowledge Stores", "sharedKS": "Knowledge Stores compartidos" }, + "documentRag": { + "label": "Documento de Referencia", + "description": "Adjunta un documento de referencia que estará disponible en el contexto del sistema para todos los mensajes." + }, "singleFile": { "label": "Seleccionar Archivo", "selectedLabel": "Archivo Seleccionado", @@ -206,7 +212,13 @@ "loading": "Cargando archivos...", "error": "Error al cargar archivos:", "noneFound": "No se encontraron archivos. Por favor sube un archivo.", - "required": "Por favor selecciona un archivo" + "required": "Por favor selecciona un archivo", + "libraryLabel": "Biblioteca", + "loadingLibraries": "Cargando bibliotecas...", + "noLibraries": "No hay bibliotecas disponibles.", + "manageLibraries": "Gestionar bibliotecas", + "itemLabel": "Documento", + "loadingItems": "Cargando documentos..." } }, "noDescription": "Sin descripción", diff --git a/frontend/svelte-app/src/lib/locales/eu.json b/frontend/svelte-app/src/lib/locales/eu.json index 928f8da7e..f795bcd42 100644 --- a/frontend/svelte-app/src/lib/locales/eu.json +++ b/frontend/svelte-app/src/lib/locales/eu.json @@ -313,7 +313,9 @@ "fileReadError": "Errorea hautatutako fitxategia irakurtzean." }, "ragOptions": { - "title": "RAG Aukerak" + "title": "RAG Aukerak", + "legacySingleFileNotice": "Laguntzaile honek fitxategi bakarreko RAG legacy-a erabiltzen du. Dokumentua ezin da aldatu. Dokumentu desberdin bat erabiltzeko, sortu laguntzaile berri bat Erreferentzia Dokumentua aukerarekin.", + "document": "Dokumentua" }, "ragTopK": { "label": "RAG Top K", @@ -328,6 +330,18 @@ "sharedKB": "Partekatutako Ezagutza Baseak", "shared": "{shared_by}(e)k partekatua" }, + "knowledgeStores": { + "label": "Knowledge Stores", + "loading": "Knowledge Stores kargatzen...", + "error": "Errorea knowledge Stores kargatzean:", + "noneFound": "Ez da knowledge Stores sarbiderik aurkitu.", + "myKS": "Nire Knowledge Stores", + "sharedKS": "Partekatutako Knowledge Stores" + }, + "documentRag": { + "label": "Erreferentzia Dokumentua", + "description": "Erreferentzia dokumentu bat erantsi, sistemaren testuan erabilgarri egongo dena mezu guztietarako." + }, "singleFile": { "label": "Fitxategia Hautatu", "selectedLabel": "Hautatutako Fitxategia", @@ -336,7 +350,13 @@ "loading": "Fitxategiak kargatzen...", "error": "Errorea fitxategiak kargatzean:", "noneFound": "Ez da fitxategirik aurkitu. Mesedez igo fitxategi bat.", - "required": "Mesedez hautatu fitxategi bat" + "required": "Mesedez hautatu fitxategi bat", + "libraryLabel": "Liburutegia", + "loadingLibraries": "Liburutegiak kargatzen...", + "noLibraries": "Ez dago liburutegi erabilgarririk.", + "manageLibraries": "Liburutegiak kudeatu", + "itemLabel": "Dokumentua", + "loadingItems": "Dokumentuak kargatzen..." } } }, From d5eaa54199a24355a716251b78cc46a69051535b Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Tue, 2 Jun 2026 22:34:54 +0200 Subject: [PATCH 054/156] feat: translate Knowledge Bases and Knowledge Stores labels to Fuentes de Conocimiento / Fonts de Coneixement / Ezagutza Iturriak in ES, CA, EU --- frontend/svelte-app/src/lib/locales/ca.json | 30 ++++++++++----------- frontend/svelte-app/src/lib/locales/es.json | 30 ++++++++++----------- frontend/svelte-app/src/lib/locales/eu.json | 30 ++++++++++----------- 3 files changed, 45 insertions(+), 45 deletions(-) diff --git a/frontend/svelte-app/src/lib/locales/ca.json b/frontend/svelte-app/src/lib/locales/ca.json index 02b9e4ba5..5cd71819b 100644 --- a/frontend/svelte-app/src/lib/locales/ca.json +++ b/frontend/svelte-app/src/lib/locales/ca.json @@ -184,21 +184,21 @@ "help": "Nombre de documents rellevants a recuperar (1-10)." }, "knowledgeBases": { - "label": "Bases de Coneixement", - "loading": "Carregant bases de coneixement...", - "error": "Error en carregar bases de coneixement:", - "noneFound": "No s'han trobat bases de coneixement accessibles.", - "myKB": "Les Meves Bases de Coneixement", + "label": "Fonts de Coneixement", + "loading": "Carregant fonts de coneixement...", + "error": "Error en carregar fonts de coneixement:", + "noneFound": "No s'han trobat fonts de coneixement accessibles.", + "myKB": "Les Meves Fonts de Coneixement", "sharedKB": "Bases de Coneixement Compartides", "shared": "Compartida per {shared_by}" }, "knowledgeStores": { - "label": "Knowledge Stores", - "loading": "Carregant knowledge stores...", - "error": "Error en carregar knowledge stores:", - "noneFound": "No s'han trobat knowledge stores accessibles.", - "myKS": "Els meus Knowledge Stores", - "sharedKS": "Knowledge Stores compartits" + "label": "Fonts de Coneixement", + "loading": "Carregant fonts de coneixement...", + "error": "Error en carregar fonts de coneixement:", + "noneFound": "No s'han trobat fonts de coneixement accessibles.", + "myKS": "Les meves Fonts de Coneixement", + "sharedKS": "Fonts de Coneixement compartides" }, "documentRag": { "label": "Document de Referència", @@ -424,10 +424,10 @@ "resetColumns": "Restablir columnes" }, "knowledgeBases": { - "title": "Bases de Coneixement", - "pageTitle": "Bases de Coneixement", - "pageDescription": "Gestiona les teves bases de coneixement per utilitzar amb assistents d'aprenentatge.", - "myKnowledgeBases": "Les Meves Bases de Coneixement", + "title": "Fonts de Coneixement", + "pageTitle": "Fonts de Coneixement", + "pageDescription": "Gestiona les teves fonts de coneixement per utilitzar amb assistents d'aprenentatge.", + "myKnowledgeBases": "Les Meves Fonts de Coneixement", "createNew": "Crear Nova Base de Coneixement", "createFirst": "Crea la Teva Primera Base de Coneixement", "loading": "Carregant bases de coneixement...", diff --git a/frontend/svelte-app/src/lib/locales/es.json b/frontend/svelte-app/src/lib/locales/es.json index 4057ea618..7b15cd9c9 100644 --- a/frontend/svelte-app/src/lib/locales/es.json +++ b/frontend/svelte-app/src/lib/locales/es.json @@ -184,21 +184,21 @@ "help": "Número de documentos relevantes a recuperar (1-10)." }, "knowledgeBases": { - "label": "Bases de Conocimiento", - "loading": "Cargando bases de conocimiento...", - "error": "Error al cargar bases de conocimiento:", - "noneFound": "No se encontraron bases de conocimiento accesibles.", - "myKB": "Mis Bases de Conocimiento", + "label": "Fuentes de Conocimiento", + "loading": "Cargando fuentes de conocimiento...", + "error": "Error al cargar fuentes de conocimiento:", + "noneFound": "No se encontraron fuentes de conocimiento accesibles.", + "myKB": "Mis Fuentes de Conocimiento", "sharedKB": "Bases de Conocimiento Compartidas", "shared": "Compartida por {shared_by}" }, "knowledgeStores": { - "label": "Knowledge Stores", - "loading": "Cargando knowledge stores...", - "error": "Error al cargar knowledge stores:", - "noneFound": "No se encontraron knowledge stores accesibles.", - "myKS": "Mis Knowledge Stores", - "sharedKS": "Knowledge Stores compartidos" + "label": "Fuentes de Conocimiento", + "loading": "Cargando fuentes de conocimiento...", + "error": "Error al cargar fuentes de conocimiento:", + "noneFound": "No se encontraron fuentes de conocimiento accesibles.", + "myKS": "Mis Fuentes de Conocimiento", + "sharedKS": "Fuentes de Conocimiento compartidas" }, "documentRag": { "label": "Documento de Referencia", @@ -424,10 +424,10 @@ "resetColumns": "Restablecer columnas" }, "knowledgeBases": { - "title": "Bases de Conocimiento", - "pageTitle": "Bases de Conocimiento", - "pageDescription": "Gestiona tus bases de conocimiento para usar con asistentes de aprendizaje.", - "myKnowledgeBases": "Mis Bases de Conocimiento", + "title": "Fuentes de Conocimiento", + "pageTitle": "Fuentes de Conocimiento", + "pageDescription": "Gestiona tus fuentes de conocimiento para usar con asistentes de aprendizaje.", + "myKnowledgeBases": "Mis Fuentes de Conocimiento", "createNew": "Crear Nueva Base de Conocimiento", "createFirst": "Crea Tu Primera Base de Conocimiento", "loading": "Cargando bases de conocimiento...", diff --git a/frontend/svelte-app/src/lib/locales/eu.json b/frontend/svelte-app/src/lib/locales/eu.json index f795bcd42..206b82fc9 100644 --- a/frontend/svelte-app/src/lib/locales/eu.json +++ b/frontend/svelte-app/src/lib/locales/eu.json @@ -322,21 +322,21 @@ "help": "Berreskuratzeko dokumentu garrantzitsuen kopurua (1-10)." }, "knowledgeBases": { - "label": "Ezagutza Baseak", - "loading": "Ezagutza baseak kargatzen...", - "error": "Errorea ezagutza baseak kargatzean:", - "noneFound": "Ez da aurkitu eskuragarri dagoen ezagutza baserik.", - "myKB": "Nire Ezagutza Baseak", + "label": "Ezagutza Iturriak", + "loading": "Ezagutza iturriak kargatzen...", + "error": "Errorea ezagutza iturriak kargatzean:", + "noneFound": "Ez da aurkitu eskuragarri dagoen ezagutza iturririk.", + "myKB": "Nire Ezagutza Iturriak", "sharedKB": "Partekatutako Ezagutza Baseak", "shared": "{shared_by}(e)k partekatua" }, "knowledgeStores": { - "label": "Knowledge Stores", - "loading": "Knowledge Stores kargatzen...", - "error": "Errorea knowledge Stores kargatzean:", - "noneFound": "Ez da knowledge Stores sarbiderik aurkitu.", - "myKS": "Nire Knowledge Stores", - "sharedKS": "Partekatutako Knowledge Stores" + "label": "Ezagutza Iturriak", + "loading": "Ezagutza iturriak kargatzen...", + "error": "Errorea ezagutza iturriak kargatzean:", + "noneFound": "Ez da ezagutza iturri sarbiderik aurkitu.", + "myKS": "Nire Ezagutza Iturriak", + "sharedKS": "Partekatutako Ezagutza Iturriak" }, "documentRag": { "label": "Erreferentzia Dokumentua", @@ -424,10 +424,10 @@ "resetColumns": "Zutabeak berrezarri" }, "knowledgeBases": { - "title": "Ezagutza Baseak", - "pageTitle": "Ezagutza Baseak", - "pageDescription": "Kudeatu zure ezagutza baseak ikaskuntza laguntzaileekin erabiltzeko.", - "myKnowledgeBases": "Nire Ezagutza Baseak", + "title": "Ezagutza Iturriak", + "pageTitle": "Ezagutza Iturriak", + "pageDescription": "Kudeatu zure ezagutza iturriak ikaskuntza laguntzaileekin erabiltzeko.", + "myKnowledgeBases": "Nire Ezagutza Iturriak", "createNew": "Ezagutza Base Berria Sortu", "createFirst": "Sortu Zure Lehen Ezagutza Basea", "loading": "Ezagutza baseak kargatzen...", From 9c1b51519223ffb380fc7ae2b0fabf72b7686c2e Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Wed, 3 Jun 2026 00:29:49 +0200 Subject: [PATCH 055/156] Add E2E Playwright test for query_rewriting_ks_rag with dual document/KS RAG channels Introduces context_aware_new.spec.js: two-library-item setup, KS ingest, assistant UI with query_rewriting_ks_rag and document RAG, chat verification via bypass echo or real LLM dual-token probe, and full cleanup. --- .../tests/context_aware_new.spec.js | 720 ++++++++++++++++++ 1 file changed, 720 insertions(+) create mode 100644 testing/playwright/tests/context_aware_new.spec.js diff --git a/testing/playwright/tests/context_aware_new.spec.js b/testing/playwright/tests/context_aware_new.spec.js new file mode 100644 index 000000000..289cd8c03 --- /dev/null +++ b/testing/playwright/tests/context_aware_new.spec.js @@ -0,0 +1,720 @@ +const { test, expect } = require("@playwright/test"); +const path = require("path"); +require("dotenv").config({ path: path.join(__dirname, "..", ".env"), quiet: true }); + +/** + * E2E: Library → Knowledge Store → Context-Aware KS RAG + Document RAG → UI chat. + * + * Validates the query_rewriting_ks_rag + single_file_rag (document reference) stack + * described in docs/superpowers/plans/2026-06-01-query-rewriting-ks-rag-implementation-summary.md. + * + * Flow: + * 1. Create library + upload two distinct markdown files (API): + * - one indexed in the Knowledge Store (KS token), + * - one used only as the Document RAG reference (doc-RAG token). + * 2. Create Knowledge Store + link the KS-only item + poll until ready (API). + * 3. Create assistant via UI with query_rewriting_ks_rag, document RAG, and KS binding. + * 4. Chat via the assistant detail UI; verify each token appears in its own channel. + * 5. Cleanup assistant, KS content link, KS, library (API). + */ + +/** Indexed in the Knowledge Store — must appear only in the user Context: block. */ +const KS_RAG_TOKEN = "LAMBKSCTX9001"; +const KS_DOC_CONTENT = `# E2E KS Retrieval Document + +The KS retrieval verification token is ${KS_RAG_TOKEN}. +This file is indexed in the Knowledge Store for dynamic RAG retrieval. +`; + +/** Attached as Document RAG reference — must appear only in the system REFERENCE DOCUMENT body. */ +const DOC_RAG_TOKEN = "LAMBDOCRAG9002"; +const DOC_RAG_CONTENT = `# E2E Reference Document + +The document RAG reference token is ${DOC_RAG_TOKEN}. +This file is attached as a static reference document in the system prompt. +`; + +const SYSTEM_PROMPT = + "You are an assistant for an automated E2E test. " + + "Answer using only information present in your context (the reference document and any retrieved knowledge store content). " + + "When asked about verification tokens, quote each token exactly as it appears in the context " + + "and indicate whether it came from the reference document or from retrieved knowledge store content."; + +/** User message for real LLM chat — probes both tokens without naming them. */ +const LLM_TOKEN_PROBE_MESSAGE = + "What are the two verification tokens described in your context? " + + "One should appear in the static reference document attached to this assistant. " + + "The other should appear in the knowledge store retrieval context about dynamic RAG. " + + "Reply with both exact token strings and label which source each came from."; + +/** Shorter probe for bypass echo (triggers KS retrieval query). */ +const BYPASS_CHAT_MESSAGE = + "Do you have reference context from both channels? KS retrieval verification probe."; + +// Matches simple_augment.DEFAULT_RAG_PROMPT_TEMPLATE — predictable {context} placement for KS RAG. +const 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}"; + +const REFERENCE_DOCUMENT_BOILERPLATE = + /## REFERENCE DOCUMENT\s+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\.\s+/i; + +/** Collapse whitespace for stable substring / equality checks. */ +function normalizeText(text) { + return String(text || "") + .replace(/\r\n/g, "\n") + .replace(/\s+/g, " ") + .trim(); +} + +/** Split full-conversation-bypass echo into system and user sections. */ +function splitBypassConversation(responseText) { + const userMatch = responseText.match(/\suser:\s/i); + if (!userMatch || userMatch.index === undefined) { + const systemOnly = responseText.replace(/^system:\s/i, "").trim(); + return { system: systemOnly, user: "" }; + } + const system = responseText.slice(0, userMatch.index).replace(/^system:\s/i, "").trim(); + let user = responseText.slice(userMatch.index).replace(/^\s*user:\s/i, "").trim(); + const assistantMatch = user.match(/\sassistant:\s/i); + if (assistantMatch && assistantMatch.index !== undefined) { + user = user.slice(0, assistantMatch.index).trim(); + } + return { system, user }; +} + +/** Document body injected after the REFERENCE DOCUMENT boilerplate in system prompt. */ +function extractReferenceDocumentBody(systemText) { + const markerIdx = systemText.search(/## REFERENCE DOCUMENT/i); + if (markerIdx === -1) return null; + const afterMarker = systemText.slice(markerIdx); + const bodyMatch = afterMarker.replace(REFERENCE_DOCUMENT_BOILERPLATE, ""); + if (bodyMatch === afterMarker) return null; + return bodyMatch.trim(); +} + +/** RAG chunks substituted into the last user message ({context} placeholder). */ +function extractRagContextBlock(userText) { + if (!userText) return null; + const patterns = [ + /Context:\s*([\s\S]*?)\sQuestion:\s/i, + /Context:\s*([\s\S]*?)\n\nQuestion:/i, + /This is the context:\s*([\s\S]*?)\n(?:This is the user input:|Now answer|$)/i, + ]; + for (const re of patterns) { + const match = userText.match(re); + if (match?.[1]?.trim()) return match[1].trim(); + } + return null; +} + +/** True when the reply is a full-conversation-bypass echo (system:/user: roles). */ +function isBypassEcho(responseText) { + return /^system:\s/i.test(responseText) || (/\suser:\s/i.test(responseText) && /REFERENCE DOCUMENT/i.test(responseText)); +} + +/** Real LLM answered with both channel tokens in natural language. */ +function llmResponseHasBothTokens(responseText) { + return ( + !isBypassEcho(responseText) && + responseText.includes(KS_RAG_TOKEN) && + responseText.includes(DOC_RAG_TOKEN) && + !/thinking/i.test(responseText) + ); +} + +function chatBackendUnavailable(responseText) { + const noOpenAiKey = /incorrect api key|invalid_api_key|your-openai-api-key/i.test(responseText); + const noOllamaChat = /ollama.*not|connection refused|model not found|\[ollama error\]/i.test( + responseText, + ); + const hasError = + /error|failed|unavailable/i.test(responseText) && responseText.length < 300; + return noOpenAiKey || noOllamaChat || hasError; +} + +test.describe.serial("Context-aware KS RAG + document RAG (UI chat)", () => { + const ts = Date.now(); + const libraryName = `pw_ctx_lib_${ts}`; + const ksName = `pw_ctx_ks_${ts}`; + const assistantName = `pw_ctx_asst_${ts}`; + + let token; + let libraryId; + let ksItemId; + let docRagItemId; + let knowledgeStoreId; + let assistantId; + let assistantUsesBypass = true; + let pipelineSkipReason = null; + + let chunkingStrategy = "simple"; + let embeddingVendor = "openai"; + let embeddingModel = "text-embedding-3-small"; + let vectorDbBackend = "chromadb"; + + async function apiCall(page, method, urlPath, options = {}) { + return page.evaluate( + async ({ method, urlPath, token, body }) => { + const headers = { Authorization: `Bearer ${token}` }; + const init = { method, headers }; + if (body) { + headers["Content-Type"] = "application/json"; + init.body = JSON.stringify(body); + } + const res = await fetch(urlPath, init); + const text = await res.text(); + let data; + try { + data = JSON.parse(text); + } catch { + data = text; + } + return { status: res.status, data }; + }, + { method, urlPath, token, body: options.body }, + ); + } + + async function pollUntil(page, fn, predicate, { timeoutMs = 90000, baseDelay = 1000 } = {}) { + let delay = baseDelay; + const deadline = Date.now() + timeoutMs; + let last = null; + while (Date.now() < deadline) { + last = await fn(); + if (predicate(last)) return last; + await page.waitForTimeout(delay); + delay = Math.min(delay * 2, 8000); + } + return last; + } + + function pickAllowed(plugins, fallback) { + if (Array.isArray(plugins) && plugins.length > 0) { + const first = plugins[0]; + return (first && first.name) || fallback; + } + return fallback; + } + + async function uploadLibraryMarkdown(page, { filename, title, docContent }) { + return page.evaluate( + async ({ libraryId, token, filename, title, docContent }) => { + const blob = new Blob([docContent], { type: "text/markdown" }); + const form = new FormData(); + form.append("file", blob, filename); + form.append("title", title); + + const r = await fetch(`/creator/libraries/${libraryId}/upload`, { + method: "POST", + headers: { Authorization: `Bearer ${token}` }, + body: form, + }); + const text = await r.text(); + let data; + try { + data = JSON.parse(text); + } catch { + data = text; + } + return { status: r.status, data }; + }, + { libraryId, token, filename, title, docContent }, + ); + } + + async function pollItemReady(page, itemId) { + const final = await pollUntil( + page, + () => apiCall(page, "GET", `/creator/libraries/${libraryId}/items/${itemId}/status`), + (r) => r && r.status === 200 && (r.data.status === "ready" || r.data.status === "failed"), + { timeoutMs: 60000 }, + ); + expect(final?.data?.status).toBe("ready"); + } + + test.beforeAll(async ({ browser }) => { + const context = await browser.newContext({ + storageState: path.join(__dirname, "..", ".auth", "state.json"), + }); + const page = await context.newPage(); + await page.goto("/"); + await page.waitForLoadState("domcontentloaded"); + token = await page.evaluate(() => localStorage.getItem("userToken")); + expect(token).toBeTruthy(); + await context.close(); + }); + + test.beforeEach(async ({ page }) => { + await page.goto("/"); + await page.waitForLoadState("domcontentloaded"); + }); + + test("step 1: create library", async ({ page }) => { + const res = await apiCall(page, "POST", "/creator/libraries", { + body: { name: libraryName, description: "Context-aware RAG E2E" }, + }); + expect(res.status).toBe(200); + expect(res.data).toHaveProperty("id"); + libraryId = res.data.id; + }); + + test("step 2: upload KS document (indexed in Knowledge Store)", async ({ page }) => { + expect(libraryId).toBeTruthy(); + + const res = await uploadLibraryMarkdown(page, { + filename: "ks-retrieval-doc.md", + title: "KS Retrieval Doc", + docContent: KS_DOC_CONTENT, + }); + + expect(res.status).toBe(200); + expect(res.data).toHaveProperty("item_id"); + ksItemId = res.data.item_id; + }); + + test("step 3: poll KS library item until ready", async ({ page }) => { + await pollItemReady(page, ksItemId); + }); + + test("step 4: upload Document RAG reference (not linked to KS)", async ({ page }) => { + expect(libraryId).toBeTruthy(); + + const res = await uploadLibraryMarkdown(page, { + filename: "doc-rag-reference.md", + title: "Doc RAG Reference", + docContent: DOC_RAG_CONTENT, + }); + + expect(res.status).toBe(200); + expect(res.data).toHaveProperty("item_id"); + docRagItemId = res.data.item_id; + }); + + test("step 5: poll Document RAG library item until ready", async ({ page }) => { + await pollItemReady(page, docRagItemId); + }); + + test("step 6: create knowledge store", async ({ page }) => { + const opts = await apiCall(page, "GET", "/creator/knowledge-stores/options"); + if (opts.status === 200 && opts.data) { + chunkingStrategy = pickAllowed(opts.data.chunking_strategies, chunkingStrategy); + vectorDbBackend = pickAllowed(opts.data.vector_db_backends, vectorDbBackend); + embeddingVendor = pickAllowed(opts.data.embedding_vendors, embeddingVendor); + const modelsForVendor = (opts.data.embedding_models || {})[embeddingVendor]; + if (Array.isArray(modelsForVendor) && modelsForVendor.length > 0) { + embeddingModel = modelsForVendor[0]; + } + } + + const res = await apiCall(page, "POST", "/creator/knowledge-stores", { + body: { + name: ksName, + description: "Context-aware RAG E2E", + chunking_strategy: chunkingStrategy, + embedding_vendor: embeddingVendor, + embedding_model: embeddingModel, + vector_db_backend: vectorDbBackend, + }, + }); + + if (res.status === 503) { + pipelineSkipReason = "KB Server v2 not reachable (503)"; + test.skip(true, pipelineSkipReason); + return; + } + expect(res.status).toBe(200); + knowledgeStoreId = res.data.id; + }); + + test("step 7: link KS document to knowledge store and wait until ready", async ({ page }) => { + test.skip(!!pipelineSkipReason, pipelineSkipReason || ""); + expect(knowledgeStoreId).toBeTruthy(); + + const linkRes = await apiCall( + page, + "POST", + `/creator/knowledge-stores/${knowledgeStoreId}/content`, + { body: { library_id: libraryId, item_ids: [ksItemId] } }, + ); + + if (linkRes.status === 503) { + pipelineSkipReason = "KB Server unreachable on content link"; + test.skip(true, pipelineSkipReason); + return; + } + expect(linkRes.status).toBe(200); + + const final = await pollUntil( + page, + () => + apiCall( + page, + "GET", + `/creator/knowledge-stores/${knowledgeStoreId}/content/${ksItemId}`, + ), + (r) => + r && + r.status === 200 && + (r.data.status === "ready" || r.data.status === "failed"), + { timeoutMs: 120000 }, + ); + + if (!final || final.data.status !== "ready") { + pipelineSkipReason = + "KS embedding ingestion did not reach ready (likely missing embedding API key)"; + test.skip(true, pipelineSkipReason); + return; + } + }); + + test("step 8-9: create assistant with context-aware RAG + document reference", async ({ + page, + }) => { + test.skip(!!pipelineSkipReason, pipelineSkipReason || ""); + + await page.goto("/assistants?view=create"); + await page.waitForLoadState("domcontentloaded"); + + const createBtn = page.getByRole("button", { + name: /\+\s*Create|Create Assistant|Crear/i, + }).first(); + await expect(createBtn).toBeVisible({ timeout: 10000 }); + await createBtn.click(); + + const form = page.locator("#assistant-form-main"); + await expect(form).toBeVisible({ timeout: 30000 }); + + await page.fill("#assistant-name", assistantName); + await page.fill("#assistant-description", `E2E context-aware RAG test ${ts}`); + await page.fill("#system-prompt", SYSTEM_PROMPT); + await page.fill("#prompt_template", RAG_PROMPT_TEMPLATE); + + // Advanced mode exposes connector / LLM pickers. + const advancedToggle = page.getByText(/Advanced Mode/i).first(); + if (await advancedToggle.count()) { + await advancedToggle.click(); + } + + await page.waitForLoadState("networkidle").catch(() => {}); + + // Connector: bypass by default (deterministic CI). Set PLAYWRIGHT_PREFER_REAL_LLM=1 + // to exercise OpenAI/Ollama and the natural-language dual-token chat assertions. + const preferRealLlm = process.env.PLAYWRIGHT_PREFER_REAL_LLM === "1"; + const connectorSel = page.locator("#connector"); + if (await connectorSel.count()) { + const connectorValues = await connectorSel.evaluate((sel) => + Array.from(sel.options).map((o) => o.value), + ); + const bypassVal = connectorValues.find((v) => v === "bypass"); + const ollamaVal = connectorValues.find((v) => /ollama/i.test(v)); + const openaiVal = connectorValues.find((v) => v === "openai"); + + if (preferRealLlm && ollamaVal) { + assistantUsesBypass = false; + await connectorSel.selectOption(ollamaVal); + await page.waitForTimeout(500); + const modelSel = page.locator("#llm"); + const modelValues = await modelSel.evaluate((sel) => + Array.from(sel.options).map((o) => o.value), + ); + const chatModel = modelValues.find((v) => + /qwen|llama|phi|mistral|gemma/i.test(v), + ); + if (chatModel) { + await modelSel.selectOption(chatModel); + } + } else if (preferRealLlm && openaiVal) { + assistantUsesBypass = false; + await connectorSel.selectOption(openaiVal); + } else if (bypassVal) { + assistantUsesBypass = true; + await connectorSel.selectOption(bypassVal); + await page.waitForTimeout(300); + const modelSel = page.locator("#llm"); + const modelValues = await modelSel.evaluate((sel) => + Array.from(sel.options).map((o) => o.value), + ); + const fullConv = modelValues.find((v) => v === "full-conversation-bypass"); + if (fullConv) { + await modelSel.selectOption(fullConv); + } + } else if (ollamaVal) { + assistantUsesBypass = false; + await connectorSel.selectOption(ollamaVal); + await page.waitForTimeout(500); + const modelSel = page.locator("#llm"); + const modelValues = await modelSel.evaluate((sel) => + Array.from(sel.options).map((o) => o.value), + ); + const chatModel = modelValues.find((v) => + /qwen|llama|phi|mistral|gemma/i.test(v), + ); + if (chatModel) { + await modelSel.selectOption(chatModel); + } + } else if (openaiVal) { + assistantUsesBypass = false; + await connectorSel.selectOption(openaiVal); + } + } + + // Context-aware KS RAG (query rewriting + Knowledge Store retrieval). + const ragSelect = page.locator("#rag-processor"); + await expect(ragSelect).toBeVisible({ timeout: 10000 }); + await ragSelect.selectOption("query_rewriting_ks_rag"); + + // Document RAG: attach the library item as reference document. + const docRagToggle = page.locator("label").filter({ hasText: /Reference Document/i }).first(); + await expect(docRagToggle).toBeVisible({ timeout: 10000 }); + await docRagToggle.click(); + + // Wait for libraries to load, then pick our library + item. + const librarySel = page.locator("#library-selector"); + await expect(librarySel).toBeVisible({ timeout: 30000 }); + await expect(librarySel.locator(`option[value="${libraryId}"]`)).toHaveCount(1, { + timeout: 30000, + }); + await librarySel.selectOption(libraryId); + + const itemSel = page.locator("#item-selector"); + await expect(itemSel).toBeVisible({ timeout: 30000 }); + await expect(itemSel.locator(`option[value="${docRagItemId}"]`)).toHaveCount(1, { + timeout: 30000, + }); + await itemSel.selectOption(docRagItemId); + + // Bind the Knowledge Store created earlier. + const ksPicker = page.locator('[data-testid="ks-picker"]'); + await expect(ksPicker).toBeVisible({ timeout: 30000 }); + const ksCheckbox = ksPicker.getByRole("checkbox", { name: new RegExp(ksName) }); + await expect(ksCheckbox).toBeVisible({ timeout: 30000 }); + await ksCheckbox.check(); + + const saveButton = page.locator('button[type="submit"][form="assistant-form-main"]'); + await expect(saveButton).toBeEnabled({ timeout: 60000 }); + + const createRequest = page.waitForResponse((response) => { + if (response.request().method() !== "POST") return false; + try { + const url = new URL(response.url()); + return ( + url.pathname.endsWith("/assistant/create_assistant") && + response.status() >= 200 && + response.status() < 300 + ); + } catch { + return false; + } + }); + + await Promise.all([createRequest, saveButton.click()]); + await page.waitForURL(/\/assistants(\?|$)/, { timeout: 30000 }).catch(() => {}); + + const list = await apiCall(page, "GET", "/creator/assistant/get_assistants"); + const all = list.data?.assistants || list.data?.data || list.data || []; + const slug = assistantName.toLowerCase(); + const found = (Array.isArray(all) ? all : []).find( + (a) => a && typeof a.name === "string" && a.name.toLowerCase().includes(slug), + ); + expect(found, `assistant ${assistantName} must exist after save`).toBeTruthy(); + assistantId = found.id; + + const detail = await apiCall(page, "GET", `/creator/assistant/get_assistant/${assistantId}`); + const meta = JSON.parse(detail.data?.metadata || "{}"); + expect(meta.rag_processor).toBe("query_rewriting_ks_rag"); + expect(meta.document_rag).toBe("single_file_rag"); + expect(meta.library_id).toBe(libraryId); + expect(meta.item_id).toBe(docRagItemId); + expect(meta.item_id).not.toBe(ksItemId); + expect(String(detail.data?.RAG_collections || "")).toContain(knowledgeStoreId); + if (meta.connector && meta.connector !== "bypass") { + assistantUsesBypass = false; + } + }); + + test("step 10-12: chat via UI and verify both RAG channels independently", async ({ + page, + }) => { + test.setTimeout(120_000); + test.skip(!!pipelineSkipReason, pipelineSkipReason || ""); + expect(assistantId, "assistant must exist from previous step").toBeTruthy(); + + await page.goto("/assistants"); + await page.waitForLoadState("networkidle").catch(() => {}); + + const searchBox = page.locator('input[placeholder*="Search" i]'); + if (await searchBox.count()) { + await searchBox.fill(assistantName); + await page.waitForTimeout(500); + } + + await page.getByText(assistantName, { exact: false }).first().click(); + + const chatTab = page.getByRole("button", { name: /Chat/i }).first(); + await expect(chatTab).toBeVisible({ timeout: 10000 }); + await chatTab.click(); + + const chatInput = page.getByPlaceholder("Type your message..."); + await expect(chatInput).toBeVisible({ timeout: 30000 }); + + const chatMessage = assistantUsesBypass ? BYPASS_CHAT_MESSAGE : LLM_TOKEN_PROBE_MESSAGE; + await chatInput.fill(chatMessage); + await page.getByRole("button", { name: /^Send$/i }).click(); + + const assistantBubble = page.locator(".bg-gray-200.text-gray-800").last(); + await expect(assistantBubble).toBeVisible({ timeout: 120000 }); + + let responseText = ""; + await expect + .poll( + async () => { + responseText = (await assistantBubble.innerText()).trim(); + if (chatBackendUnavailable(responseText)) { + return responseText; + } + if (isBypassEcho(responseText)) { + const { system, user } = splitBypassConversation(responseText); + const docBody = extractReferenceDocumentBody(system); + const ragBlock = extractRagContextBlock(user); + const bypassReady = + docBody && + normalizeText(docBody) === normalizeText(DOC_RAG_CONTENT) && + ragBlock && + ragBlock.includes(KS_RAG_TOKEN) && + !ragBlock.includes(DOC_RAG_TOKEN); + return bypassReady ? responseText : ""; + } + if (llmResponseHasBothTokens(responseText)) { + return responseText; + } + return ""; + }, + { timeout: 120000, intervals: [500, 1000, 2000] }, + ) + .not.toBe(""); + + console.log( + `[${assistantUsesBypass ? "bypass" : "real-llm"}] Assistant UI response: "${responseText.slice(0, 1600)}"`, + ); + + if (chatBackendUnavailable(responseText)) { + test.skip( + true, + `Chat backend not available — setup verified, chat skipped. Got: ${responseText.slice(0, 300)}`, + ); + } + + if (isBypassEcho(responseText)) { + const { system, user } = splitBypassConversation(responseText); + + expect(system, "bypass echo must include a system message").toBeTruthy(); + expect(user, "bypass echo must include an augmented user message").toBeTruthy(); + + // 1) Document RAG: system body must match the reference doc only (not the KS doc). + const injectedDoc = extractReferenceDocumentBody(system); + expect(injectedDoc, "system message must contain document body after REFERENCE DOCUMENT boilerplate").toBeTruthy(); + expect(normalizeText(injectedDoc)).toBe(normalizeText(DOC_RAG_CONTENT)); + expect(injectedDoc, "document RAG body must carry the doc-RAG token").toContain(DOC_RAG_TOKEN); + expect(injectedDoc, "document RAG body must not leak the KS-only token").not.toContain( + KS_RAG_TOKEN, + ); + + // 2) KS RAG: user Context block must carry the indexed doc only (not the reference doc). + const injectedRag = extractRagContextBlock(user); + expect(injectedRag, "user message must contain a Context: block with KS retrieval").toBeTruthy(); + expect(injectedRag.length, "KS context block must be non-empty").toBeGreaterThan(20); + expect(injectedRag, "KS context must include the KS-only token").toContain(KS_RAG_TOKEN); + expect(injectedRag, "KS context should include the KS document heading").toMatch( + /E2E KS Retrieval Document/i, + ); + expect(injectedRag, "KS context must not leak the document-RAG token").not.toContain( + DOC_RAG_TOKEN, + ); + + // Cross-channel isolation: each token stays in its pipeline. + expect(system, "doc-RAG token must appear in system (document RAG)").toContain(DOC_RAG_TOKEN); + expect(system, "KS token must not appear in system document body").not.toContain(KS_RAG_TOKEN); + expect(user, "KS token must appear in augmented user message (KS RAG)").toContain(KS_RAG_TOKEN); + expect(user, "doc-RAG token must not appear in user Context block").not.toContain(DOC_RAG_TOKEN); + return; + } + + // Real LLM: both tokens must appear in the assistant's natural-language answer, + // proving it read document RAG (system) and KS RAG (retrieved context). + expect( + responseText, + "LLM response must include the knowledge-store verification token", + ).toContain(KS_RAG_TOKEN); + expect( + responseText, + "LLM response must include the document-RAG reference token", + ).toContain(DOC_RAG_TOKEN); + }); + + test("step 13: delete assistant", async ({ page }) => { + if (!assistantId) return; + + await apiCall(page, "DELETE", `/creator/assistant/delete_assistant/${assistantId}`); + + const list = await apiCall(page, "GET", "/creator/assistant/get_assistants"); + const all = list.data?.assistants || list.data?.data || list.data || []; + const stillThere = (Array.isArray(all) ? all : []).some((a) => String(a.id) === String(assistantId)); + expect(stillThere).toBe(false); + assistantId = null; + }); + + test("step 14: remove KS content link and delete knowledge store", async ({ page }) => { + if (!knowledgeStoreId) return; + + if (ksItemId) { + await apiCall( + page, + "DELETE", + `/creator/knowledge-stores/${knowledgeStoreId}/content/${ksItemId}`, + ).catch(() => {}); + } + + await apiCall(page, "DELETE", `/creator/knowledge-stores/${knowledgeStoreId}`); + knowledgeStoreId = null; + }); + + test("step 15: delete library", async ({ page }) => { + if (!libraryId) return; + + await apiCall(page, "DELETE", `/creator/libraries/${libraryId}`); + libraryId = null; + }); + + test.afterAll(async ({ browser }) => { + const context = await browser.newContext({ + storageState: path.join(__dirname, "..", ".auth", "state.json"), + }); + const page = await context.newPage(); + await page.goto("/"); + await page.waitForLoadState("domcontentloaded"); + + if (assistantId) { + await apiCall(page, "DELETE", `/creator/assistant/delete_assistant/${assistantId}`).catch( + () => {}, + ); + } + if (knowledgeStoreId && ksItemId) { + await apiCall( + page, + "DELETE", + `/creator/knowledge-stores/${knowledgeStoreId}/content/${ksItemId}`, + ).catch(() => {}); + } + if (knowledgeStoreId) { + await apiCall(page, "DELETE", `/creator/knowledge-stores/${knowledgeStoreId}`).catch( + () => {}, + ); + } + if (libraryId) { + await apiCall(page, "DELETE", `/creator/libraries/${libraryId}`).catch(() => {}); + } + await context.close(); + }); +}); From c382020cde1f3df4fefbb40d1a3f0693d239c0b3 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Wed, 3 Jun 2026 00:42:37 +0200 Subject: [PATCH 056/156] fix: add recency-bias reminder at end of Document RAG system prompt --- backend/lamb/completions/pps/simple_augment.py | 4 ++++ backend/tests/test_simple_augment.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/backend/lamb/completions/pps/simple_augment.py b/backend/lamb/completions/pps/simple_augment.py index 6a8d0a7f9..aacb06dea 100644 --- a/backend/lamb/completions/pps/simple_augment.py +++ b/backend/lamb/completions/pps/simple_augment.py @@ -104,6 +104,10 @@ def prompt_processor( "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 = (system_content + labeled_doc) if system_content else labeled_doc if system_content: diff --git a/backend/tests/test_simple_augment.py b/backend/tests/test_simple_augment.py index d4da9c31b..3ff91aa38 100644 --- a/backend/tests/test_simple_augment.py +++ b/backend/tests/test_simple_augment.py @@ -107,6 +107,10 @@ def test_document_context_gets_descriptive_label(): assert "This is the reference document content." in system_msg["content"] # Should contain the explanation for the LLM assert "selected by the assistant creator" in system_msg["content"] + # Should contain the recency-bias reminder at the end + assert "available for this entire conversation" in system_msg["content"] + # Reminder should be at the very end of the system content + assert system_msg["content"].rstrip().endswith("use it.") # Should come after the original system prompt assert "You are helpful." in system_msg["content"] doc_index = system_msg["content"].index("REFERENCE DOCUMENT") From 2e0ad267423544225a539f62da49537f44b831c5 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 12:30:44 +0200 Subject: [PATCH 057/156] refactor: revert simple_augment.py to dev state, remove document_context tests --- .../lamb/completions/pps/simple_augment.py | 14 +-- .../completions/test_simple_augment.py | 118 +----------------- 2 files changed, 6 insertions(+), 126 deletions(-) diff --git a/backend/lamb/completions/pps/simple_augment.py b/backend/lamb/completions/pps/simple_augment.py index 47fb7c055..414dd8654 100644 --- a/backend/lamb/completions/pps/simple_augment.py +++ b/backend/lamb/completions/pps/simple_augment.py @@ -58,8 +58,7 @@ 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, - document_context: Optional[Dict[str, Any]] = None, + rag_context: Optional[Dict[str, Any]] = None ) -> List[Dict[str, str]]: """ Simple augment prompt processor that: @@ -79,16 +78,11 @@ def prompt_processor( processed_messages = [] if assistant: - # Build system prompt with optional document context - system_content = assistant.system_prompt or "" - if document_context and isinstance(document_context, dict): - doc_text = document_context.get("context", "") - if doc_text: - system_content = (system_content + "\n\n" + doc_text) if system_content else doc_text - if system_content: + # Add system message from assistant if available + if assistant.system_prompt: processed_messages.append({ "role": "system", - "content": system_content + "content": assistant.system_prompt }) # Add previous messages except the last one diff --git a/testing/unit-tests/completions/test_simple_augment.py b/testing/unit-tests/completions/test_simple_augment.py index e096f9bbe..4c9718ec9 100644 --- a/testing/unit-tests/completions/test_simple_augment.py +++ b/testing/unit-tests/completions/test_simple_augment.py @@ -1,116 +1,2 @@ -import pytest -from lamb.completions.pps.simple_augment import prompt_processor - - -class MockAssistant: - def __init__(self, system_prompt="", prompt_template="", metadata=None): - self.system_prompt = system_prompt - self.prompt_template = prompt_template - self.metadata = metadata - self.api_callback = metadata - - -def _make_request(messages): - return {"messages": messages} - - -class TestDocumentContextInSystemPrompt: - def test_document_context_appended_to_system_prompt(self): - assistant = MockAssistant( - system_prompt="You are a tutor.", - prompt_template="Answer: {user_input}\nContext: {context}", - ) - request = _make_request([{"role": "user", "content": "What is 2+2?"}]) - doc_ctx = {"context": "# Reference Document\nSome content here.", "sources": []} - - result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) - - system_msg = result[0] - assert system_msg["role"] == "system" - assert "You are a tutor." in system_msg["content"] - assert "# Reference Document\nSome content here." in system_msg["content"] - - def test_document_context_as_system_prompt_when_none_exists(self): - assistant = MockAssistant( - system_prompt="", - prompt_template="Answer: {user_input}\nContext: {context}", - ) - request = _make_request([{"role": "user", "content": "Hello"}]) - doc_ctx = {"context": "Document content.", "sources": []} - - result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) - - system_msg = result[0] - assert system_msg["role"] == "system" - assert system_msg["content"] == "Document content." - - def test_no_document_context_leaves_system_prompt_unchanged(self): - assistant = MockAssistant( - system_prompt="Original prompt.", - prompt_template="{user_input}", - ) - request = _make_request([{"role": "user", "content": "Hi"}]) - - result = prompt_processor(request, assistant=assistant) - - system_msg = result[0] - assert system_msg["content"] == "Original prompt." - - def test_empty_document_context_ignored(self): - assistant = MockAssistant( - system_prompt="Original.", - prompt_template="{user_input}", - ) - request = _make_request([{"role": "user", "content": "Hi"}]) - doc_ctx = {"context": "", "sources": []} - - result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) - - system_msg = result[0] - assert system_msg["content"] == "Original." - - def test_document_sources_not_in_user_message(self): - assistant = MockAssistant( - system_prompt="Sys.", - prompt_template="Q: {user_input}\nCtx: {context}", - ) - request = _make_request([{"role": "user", "content": "Question?"}]) - doc_ctx = { - "context": "Doc text.", - "sources": [{"title": "Doc", "url": "/docs/1/2", "similarity": 1.0}], - } - - result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) - - last_msg = result[-1] - assert "Doc text." not in last_msg["content"] - assert "## Available Sources" not in last_msg["content"] - - def test_document_context_and_rag_context_coexist(self): - assistant = MockAssistant( - system_prompt="Sys.", - prompt_template="Q: {user_input}\nCtx: {context}", - ) - request = _make_request([{"role": "user", "content": "Question?"}]) - doc_ctx = {"context": "Reference doc.", "sources": []} - rag_ctx = {"context": "RAG chunks.", "sources": [{"title": "KB", "url": "/kb/1", "similarity": 0.9}]} - - result = prompt_processor( - request, assistant=assistant, rag_context=rag_ctx, document_context=doc_ctx - ) - - system_msg = result[0] - assert "Reference doc." in system_msg["content"] - - last_msg = result[-1] - assert "RAG chunks." in last_msg["content"] - assert "Reference doc." not in last_msg["content"] - - def test_no_assistant_document_context_ignored(self): - request = _make_request([{"role": "user", "content": "Hi"}]) - doc_ctx = {"context": "Doc.", "sources": []} - - result = prompt_processor(request, document_context=doc_ctx) - - assert len(result) == 1 - assert result[0]["role"] == "user" +# Archivo vacío — todos los tests eran de document_context +# y se mueven a test_kvcache_augment.py From 5f4bdd72e839040afb6e987a670bfc02f0861320 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 12:32:24 +0200 Subject: [PATCH 058/156] refactor: revert single_file_rag.py to dev state, keep only file_path tests --- .../lamb/completions/rag/single_file_rag.py | 152 ++++++++-------- .../completions/test_single_file_rag.py | 168 ++---------------- 2 files changed, 93 insertions(+), 227 deletions(-) diff --git a/backend/lamb/completions/rag/single_file_rag.py b/backend/lamb/completions/rag/single_file_rag.py index c3babd5e1..d0769fced 100644 --- a/backend/lamb/completions/rag/single_file_rag.py +++ b/backend/lamb/completions/rag/single_file_rag.py @@ -3,104 +3,102 @@ from lamb.lamb_classes import Assistant import json import logging -import httpx +# Set up logger logger = logging.getLogger('lamb.completions.rag.single_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: + """ + A RAG processor that returns the content of a single file as context. + The file path should be specified in the assistant's metadata field under 'file_path'. + The file path is relative to the project's static/public folder. + """ + logger.debug(f"Starting single_file_rag processor with assistant: {assistant.id if assistant else 'None'}") + + if not assistant: 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") - file_path = metadata.get("file_path") - - if library_id and item_id: - return _fetch_from_library_manager(library_id, item_id) - elif file_path: - return _read_from_static_file(file_path) - else: - logger.warning("No library_id+item_id or file_path in metadata") - return empty_result - - -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}"} + return { + "context": "", + "sources": [] + } try: - response = httpx.get(url, params={"format": "markdown"}, headers=headers, timeout=30.0) - if response.status_code == 200: - content = response.text + # Parse the metadata to get the file path + logger.debug(f"Full metadata content: {assistant.metadata}") + + # Handle empty metadata + if not assistant.metadata or assistant.metadata.strip() == '': + logger.warning(f"Empty metadata for assistant {assistant.id if assistant else 'unknown'}") + config = {} + else: + config = json.loads(assistant.metadata) + + logger.debug(f"Parsed metadata config: {config}") + file_path = config.get('file_path') + logger.debug(f"Extracted file_path from metadata: {file_path}") + + if not file_path: + logger.warning("No file_path found in metadata") return { - "context": content, - "sources": [{ - "title": "Library Document", - "url": f"/docs/{library_id}/{item_id}", - "similarity": 1.0, - }], + "context": "", + "sources": [] } - 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 + # Construct absolute path from project's static/public folder + base_path = os.path.join('static', 'public') + full_path = os.path.join(base_path, file_path) + logger.debug(f"Base path: {base_path}") + logger.debug(f"Full path: {full_path}") -def _read_from_static_file(file_path: str) -> Dict[str, Any]: - base_path = os.path.join('static', 'public') - full_path = os.path.join(base_path, file_path) + # Ensure the path doesn't escape the static/public directory + if '..' in file_path or not os.path.abspath(full_path).startswith(os.path.abspath(base_path)): + logger.error(f"Security check failed - path attempts to escape base directory: {full_path}") + return { + "context": "Error: Invalid file path", + "sources": [] + } - if '..' in file_path or not os.path.abspath(full_path).startswith(os.path.abspath(base_path)): - error_msg = f"Error: Invalid file path: {file_path}" - logger.warning(f"Path traversal attempt detected: {file_path}") - return {"context": error_msg, "sources": []} + # Check if file exists + if not os.path.exists(full_path): + logger.error(f"File not found: {full_path}") + return { + "context": f"Error: File not found: {file_path}", + "sources": [] + } - if not os.path.exists(full_path): - error_msg = f"Error: File not found: {file_path}" - logger.warning(f"File not found: {full_path}") - return {"context": error_msg, "sources": []} + # Log file stats + file_stats = os.stat(full_path) + logger.debug(f"File stats - size: {file_stats.st_size} bytes, last modified: {file_stats.st_mtime}") - try: - with open(full_path, 'r', encoding='utf-8') as f: - content = f.read() + # Read the file content if it exists + logger.debug(f"Reading file content from: {full_path}") + with open(full_path, 'r', encoding='utf-8') as file: + content = file.read() + logger.debug(f"Successfully read {len(content)} characters from file") + return { "context": content, "sources": [{ - "title": os.path.basename(file_path), - "url": f"/static/public/{file_path}", - "similarity": 1.0, - }], + "source": file_path, + "content": content, + "score": 1.0 + }] + } + + except json.JSONDecodeError as e: + logger.error(f"Failed to parse metadata JSON: {e}") + return { + "context": f"Error processing file: Invalid metadata format", + "sources": [] } except Exception as e: - error_msg = f"Error reading file {file_path}: {str(e)}" - logger.warning(f"Error reading file {full_path}: {e}") - return {"context": error_msg, "sources": []} + logger.error(f"Error processing file: {str(e)}", exc_info=True) + return { + "context": f"Error processing file: {str(e)}", + "sources": [] + } \ No newline at end of file diff --git a/testing/unit-tests/completions/test_single_file_rag.py b/testing/unit-tests/completions/test_single_file_rag.py index e7530c52f..07e09cf1d 100644 --- a/testing/unit-tests/completions/test_single_file_rag.py +++ b/testing/unit-tests/completions/test_single_file_rag.py @@ -1,6 +1,6 @@ import json import os -from unittest.mock import patch, MagicMock, mock_open +from unittest.mock import patch, mock_open import pytest from backend.lamb.completions.rag.single_file_rag import rag_processor @@ -14,127 +14,16 @@ def __init__(self, metadata_dict, owner="test@example.com"): self.owner = owner -MOCK_LM_URL = "http://localhost:9091" -MOCK_LM_TOKEN = "test-token" -MOCK_ENV = { - "LAMB_LIBRARY_SERVER": MOCK_LM_URL, - "LAMB_LIBRARY_TOKEN": MOCK_LM_TOKEN, -} - - def _make_messages(user_message="What does the document say?"): return [{"role": "user", "content": user_message}] -# --- Library Manager path (new behavior) --- - -@patch.dict(os.environ, MOCK_ENV) -@patch("backend.lamb.completions.rag.single_file_rag.httpx.get") -def test_library_item_success(mock_get): - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.text = "# Document Title\nThis is the content." - mock_get.return_value = mock_response - - assistant = MockAssistant({ - "rag_processor": "single_file_rag", - "library_id": "lib-123", - "item_id": "item-456", - }) - - result = rag_processor(messages=_make_messages(), assistant=assistant) - - assert "# Document Title" in result["context"] - assert "This is the content." in result["context"] - assert len(result["sources"]) == 1 - assert result["sources"][0]["title"] == "Library Document" - assert result["sources"][0]["url"] == "/docs/lib-123/item-456" - assert result["sources"][0]["similarity"] == 1.0 - mock_get.assert_called_once_with( - f"{MOCK_LM_URL}/libraries/lib-123/items/item-456/content", - params={"format": "markdown"}, - headers={"Authorization": f"Bearer {MOCK_LM_TOKEN}"}, - timeout=30.0, - ) - - -@patch.dict(os.environ, MOCK_ENV) -@patch("backend.lamb.completions.rag.single_file_rag.httpx.get") -def test_library_item_connection_error_returns_empty_context(mock_get): - import httpx - mock_get.side_effect = httpx.ConnectError("Connection refused") - - assistant = MockAssistant({ - "rag_processor": "single_file_rag", - "library_id": "lib-123", - "item_id": "item-456", - }) - - result = rag_processor(messages=_make_messages(), assistant=assistant) - - assert result["context"] == "" - assert result["sources"] == [] - - -@patch.dict(os.environ, MOCK_ENV) -@patch("backend.lamb.completions.rag.single_file_rag.httpx.get") -def test_library_item_not_found_returns_empty_context(mock_get): - mock_response = MagicMock() - mock_response.status_code = 404 - mock_response.text = "Not found" - mock_get.return_value = mock_response - - assistant = MockAssistant({ - "rag_processor": "single_file_rag", - "library_id": "lib-123", - "item_id": "item-nonexistent", - }) - - result = rag_processor(messages=_make_messages(), assistant=assistant) - - assert result["context"] == "" - assert result["sources"] == [] - - -@patch.dict(os.environ, MOCK_ENV) -@patch("backend.lamb.completions.rag.single_file_rag.httpx.get") -def test_library_item_server_error_returns_empty_context(mock_get): - mock_response = MagicMock() - mock_response.status_code = 500 - mock_response.text = "Internal Server Error" - mock_get.return_value = mock_response - - assistant = MockAssistant({ - "rag_processor": "single_file_rag", - "library_id": "lib-123", - "item_id": "item-456", - }) - - result = rag_processor(messages=_make_messages(), assistant=assistant) - - assert result["context"] == "" - assert result["sources"] == [] - - -@patch.dict(os.environ, {}, clear=True) -def test_library_item_missing_env_vars_returns_empty_context(): - assistant = MockAssistant({ - "rag_processor": "single_file_rag", - "library_id": "lib-123", - "item_id": "item-456", - }) - - result = rag_processor(messages=_make_messages(), assistant=assistant) - - assert result["context"] == "" - assert result["sources"] == [] - - -# --- Backward compatibility: file_path (existing behavior) --- - @patch("backend.lamb.completions.rag.single_file_rag.os.path.exists", return_value=True) +@patch("backend.lamb.completions.rag.single_file_rag.os.stat") @patch("builtins.open", new_callable=mock_open, read_data="File content here") -def test_file_path_backward_compat_success(mock_file, mock_exists): +def test_file_path_success(mock_file, mock_stat, mock_exists): + mock_stat.return_value.st_size = 100 + mock_stat.return_value.st_mtime = 0 assistant = MockAssistant({ "rag_processor": "single_file_rag", "file_path": "1/notes.md", @@ -144,13 +33,12 @@ def test_file_path_backward_compat_success(mock_file, mock_exists): assert "File content here" in result["context"] assert len(result["sources"]) == 1 - assert result["sources"][0]["title"] == "notes.md" - assert result["sources"][0]["url"] == "/static/public/1/notes.md" - assert result["sources"][0]["similarity"] == 1.0 + assert result["sources"][0]["source"] == "1/notes.md" + assert result["sources"][0]["score"] == 1.0 @patch("backend.lamb.completions.rag.single_file_rag.os.path.exists", return_value=False) -def test_file_path_not_found_returns_error_in_context(mock_exists): +def test_file_path_not_found(mock_exists): assistant = MockAssistant({ "rag_processor": "single_file_rag", "file_path": "1/missing.md", @@ -158,13 +46,11 @@ def test_file_path_not_found_returns_error_in_context(mock_exists): result = rag_processor(messages=_make_messages(), assistant=assistant) - assert "Error" in result["context"] or "not found" in result["context"].lower() + assert "not found" in result["context"].lower() or "Error" in result["context"] assert result["sources"] == [] -# --- Error cases --- - -def test_no_file_path_no_library_id_returns_empty_context(): +def test_no_file_path_returns_empty_context(): assistant = MockAssistant({ "rag_processor": "single_file_rag", }) @@ -175,19 +61,14 @@ def test_no_file_path_no_library_id_returns_empty_context(): assert result["sources"] == [] -def test_library_id_without_item_id_returns_empty_context(): - assistant = MockAssistant({ - "rag_processor": "single_file_rag", - "library_id": "lib-123", - }) - - result = rag_processor(messages=_make_messages(), assistant=assistant) +def test_no_assistant_returns_empty_context(): + result = rag_processor(messages=_make_messages(), assistant=None) assert result["context"] == "" assert result["sources"] == [] -def test_invalid_metadata_json_returns_empty_context(): +def test_invalid_metadata_json(): assistant = MockAssistant.__new__(MockAssistant) assistant.id = "mock-assistant-id" assistant.api_callback = "not valid json" @@ -196,30 +77,17 @@ def test_invalid_metadata_json_returns_empty_context(): result = rag_processor(messages=_make_messages(), assistant=assistant) - assert result["context"] == "" + assert "Error" in result["context"] or result["context"] == "" assert result["sources"] == [] -# --- Priority: library_id+item_id takes precedence over file_path --- - -@patch.dict(os.environ, MOCK_ENV) -@patch("backend.lamb.completions.rag.single_file_rag.httpx.get") -def test_library_id_takes_precedence_over_file_path(mock_get): - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.text = "LM content" - mock_get.return_value = mock_response - +def test_path_traversal_blocked(): assistant = MockAssistant({ "rag_processor": "single_file_rag", - "library_id": "lib-123", - "item_id": "item-456", - "file_path": "1/old_file.md", + "file_path": "../../etc/passwd", }) result = rag_processor(messages=_make_messages(), assistant=assistant) - assert "LM content" in result["context"] - assert result["sources"][0]["title"] == "Library Document" - assert result["sources"][0]["url"] == "/docs/lib-123/item-456" - mock_get.assert_called_once() + assert "Error" in result["context"] or "Invalid" in result["context"] + assert result["sources"] == [] From 40964b7ec99bba50a5b99fb0aef830c33a736c33 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 12:33:47 +0200 Subject: [PATCH 059/156] feat: add kvcache_augment PPS with COMPATIBLE_RAG and document_context injection --- .../lamb/completions/pps/kvcache_augment.py | 151 ++++++++++++++++++ .../completions/test_kvcache_augment.py | 127 +++++++++++++++ 2 files changed, 278 insertions(+) create mode 100644 backend/lamb/completions/pps/kvcache_augment.py create mode 100644 testing/unit-tests/completions/test_kvcache_augment.py diff --git a/backend/lamb/completions/pps/kvcache_augment.py b/backend/lamb/completions/pps/kvcache_augment.py new file mode 100644 index 000000000..2180af2ae --- /dev/null +++ b/backend/lamb/completions/pps/kvcache_augment.py @@ -0,0 +1,151 @@ +from typing import Dict, Any, List, Optional +from lamb.lamb_classes import Assistant +import json +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", +] + + +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 + + 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: + system_content = (system_content + "\n\n" + doc_text) if system_content else doc_text + 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: + context = rag_context.get("context", "") if isinstance(rag_context, dict) else str(rag_context) + sources_text = "" + if isinstance(rag_context, dict) and "sources" in rag_context: + sources = rag_context["sources"] + if sources: + sources_text = "\n\n## Available Sources\n\n" + for i, source in enumerate(sources, 1): + title = source.get("title", "Unknown") + url = source.get("url", "") + similarity = source.get("similarity", 0) + sources_text += f"{i}. [{title}]({url}) (similarity: {similarity:.3f})\n" + full_context = context + sources_text + 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: + context = rag_context.get("context", "") if isinstance(rag_context, dict) else str(rag_context) + sources_text = "" + if isinstance(rag_context, dict) and "sources" in rag_context: + sources = rag_context["sources"] + if sources: + sources_text = "\n\n## Available Sources\n\n" + for i, source in enumerate(sources, 1): + title = source.get("title", "Unknown") + url = source.get("url", "") + similarity = source.get("similarity", 0) + sources_text += f"{i}. [{title}]({url}) (similarity: {similarity:.3f})\n" + full_context = context + sources_text + 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: + processed_messages.append(messages[-1]) + + return processed_messages + + return messages diff --git a/testing/unit-tests/completions/test_kvcache_augment.py b/testing/unit-tests/completions/test_kvcache_augment.py new file mode 100644 index 000000000..3aedad2af --- /dev/null +++ b/testing/unit-tests/completions/test_kvcache_augment.py @@ -0,0 +1,127 @@ +import pytest +from lamb.completions.pps.kvcache_augment import prompt_processor, COMPATIBLE_RAG + + +class MockAssistant: + def __init__(self, system_prompt="", prompt_template="", metadata=None): + self.system_prompt = system_prompt + self.prompt_template = prompt_template + self.metadata = metadata + self.api_callback = metadata + + +def _make_request(messages): + return {"messages": messages} + + +class TestCompatibleRag: + def test_declares_compatible_rag_list(self): + assert isinstance(COMPATIBLE_RAG, list) + assert "library_file_rag" in COMPATIBLE_RAG + assert "no_rag" in COMPATIBLE_RAG + + def test_does_not_include_legacy_rags(self): + assert "simple_rag" not in COMPATIBLE_RAG + assert "single_file_rag" not in COMPATIBLE_RAG + + +class TestDocumentContextInSystemPrompt: + def test_document_context_appended_to_system_prompt(self): + assistant = MockAssistant( + system_prompt="You are a tutor.", + prompt_template="Answer: {user_input}\nContext: {context}", + ) + request = _make_request([{"role": "user", "content": "What is 2+2?"}]) + doc_ctx = {"context": "# Reference Document\nSome content here.", "sources": []} + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + system_msg = result[0] + assert system_msg["role"] == "system" + assert "You are a tutor." in system_msg["content"] + assert "# Reference Document\nSome content here." in system_msg["content"] + + def test_document_context_as_system_prompt_when_none_exists(self): + assistant = MockAssistant( + system_prompt="", + prompt_template="Answer: {user_input}\nContext: {context}", + ) + request = _make_request([{"role": "user", "content": "Hello"}]) + doc_ctx = {"context": "Document content.", "sources": []} + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + system_msg = result[0] + assert system_msg["role"] == "system" + assert system_msg["content"] == "Document content." + + def test_no_document_context_leaves_system_prompt_unchanged(self): + assistant = MockAssistant( + system_prompt="Original prompt.", + prompt_template="{user_input}", + ) + request = _make_request([{"role": "user", "content": "Hi"}]) + + result = prompt_processor(request, assistant=assistant) + + system_msg = result[0] + assert system_msg["content"] == "Original prompt." + + def test_empty_document_context_ignored(self): + assistant = MockAssistant( + system_prompt="Original.", + prompt_template="{user_input}", + ) + request = _make_request([{"role": "user", "content": "Hi"}]) + doc_ctx = {"context": "", "sources": []} + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + system_msg = result[0] + assert system_msg["content"] == "Original." + + def test_document_sources_not_in_user_message(self): + assistant = MockAssistant( + system_prompt="Sys.", + prompt_template="Q: {user_input}\nCtx: {context}", + ) + request = _make_request([{"role": "user", "content": "Question?"}]) + doc_ctx = { + "context": "Doc text.", + "sources": [{"title": "Doc", "url": "/docs/1/2", "similarity": 1.0}], + } + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + last_msg = result[-1] + assert "Doc text." not in last_msg["content"] + assert "## Available Sources" not in last_msg["content"] + + def test_document_context_and_rag_context_coexist(self): + assistant = MockAssistant( + system_prompt="Sys.", + prompt_template="Q: {user_input}\nCtx: {context}", + ) + request = _make_request([{"role": "user", "content": "Question?"}]) + doc_ctx = {"context": "Reference doc.", "sources": []} + rag_ctx = {"context": "RAG chunks.", "sources": [{"title": "KB", "url": "/kb/1", "similarity": 0.9}]} + + result = prompt_processor( + request, assistant=assistant, rag_context=rag_ctx, document_context=doc_ctx + ) + + system_msg = result[0] + assert "Reference doc." in system_msg["content"] + + last_msg = result[-1] + assert "RAG chunks." in last_msg["content"] + assert "Reference doc." not in last_msg["content"] + + def test_no_assistant_document_context_ignored(self): + request = _make_request([{"role": "user", "content": "Hi"}]) + doc_ctx = {"context": "Doc.", "sources": []} + + result = prompt_processor(request, document_context=doc_ctx) + + assert len(result) == 1 + assert result[0]["role"] == "user" From 4ee164853751adf1a345f7fb735b23b0c7d00b80 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 12:35:01 +0200 Subject: [PATCH 060/156] feat: add library_file_rag RAG processor with Library Manager HTTP integration --- .../lamb/completions/rag/library_file_rag.py | 72 ++++++++ .../completions/test_library_file_rag.py | 171 ++++++++++++++++++ 2 files changed, 243 insertions(+) create mode 100644 backend/lamb/completions/rag/library_file_rag.py create mode 100644 testing/unit-tests/completions/test_library_file_rag.py 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/testing/unit-tests/completions/test_library_file_rag.py b/testing/unit-tests/completions/test_library_file_rag.py new file mode 100644 index 000000000..84458dfb2 --- /dev/null +++ b/testing/unit-tests/completions/test_library_file_rag.py @@ -0,0 +1,171 @@ +import json +import os +from unittest.mock import patch, MagicMock +import pytest + +from backend.lamb.completions.rag.library_file_rag import rag_processor + + +class MockAssistant: + def __init__(self, metadata_dict, owner="test@example.com"): + self.id = "mock-assistant-id" + self.api_callback = json.dumps(metadata_dict) + self.metadata = self.api_callback + self.owner = owner + + +MOCK_LM_URL = "http://localhost:9091" +MOCK_LM_TOKEN = "test-token" +MOCK_ENV = { + "LAMB_LIBRARY_SERVER": MOCK_LM_URL, + "LAMB_LIBRARY_TOKEN": MOCK_LM_TOKEN, +} + + +def _make_messages(user_message="What does the document say?"): + return [{"role": "user", "content": user_message}] + + +@patch.dict(os.environ, MOCK_ENV) +@patch("backend.lamb.completions.rag.library_file_rag.httpx.get") +def test_library_item_success(mock_get): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.text = "# Document Title\nThis is the content." + mock_get.return_value = mock_response + + assistant = MockAssistant({ + "rag_processor": "library_file_rag", + "library_id": "lib-123", + "item_id": "item-456", + }) + + result = rag_processor(messages=_make_messages(), assistant=assistant) + + assert "# Document Title" in result["context"] + assert "This is the content." in result["context"] + assert len(result["sources"]) == 1 + assert result["sources"][0]["title"] == "Library Document" + assert result["sources"][0]["url"] == "/docs/lib-123/item-456" + assert result["sources"][0]["similarity"] == 1.0 + mock_get.assert_called_once_with( + f"{MOCK_LM_URL}/libraries/lib-123/items/item-456/content", + params={"format": "markdown"}, + headers={"Authorization": f"Bearer {MOCK_LM_TOKEN}"}, + timeout=30.0, + ) + + +@patch.dict(os.environ, MOCK_ENV) +@patch("backend.lamb.completions.rag.library_file_rag.httpx.get") +def test_library_item_connection_error_returns_empty_context(mock_get): + import httpx + mock_get.side_effect = httpx.ConnectError("Connection refused") + + assistant = MockAssistant({ + "rag_processor": "library_file_rag", + "library_id": "lib-123", + "item_id": "item-456", + }) + + result = rag_processor(messages=_make_messages(), assistant=assistant) + + assert result["context"] == "" + assert result["sources"] == [] + + +@patch.dict(os.environ, MOCK_ENV) +@patch("backend.lamb.completions.rag.library_file_rag.httpx.get") +def test_library_item_not_found_returns_empty_context(mock_get): + mock_response = MagicMock() + mock_response.status_code = 404 + mock_response.text = "Not found" + mock_get.return_value = mock_response + + assistant = MockAssistant({ + "rag_processor": "library_file_rag", + "library_id": "lib-123", + "item_id": "item-nonexistent", + }) + + result = rag_processor(messages=_make_messages(), assistant=assistant) + + assert result["context"] == "" + assert result["sources"] == [] + + +@patch.dict(os.environ, MOCK_ENV) +@patch("backend.lamb.completions.rag.library_file_rag.httpx.get") +def test_library_item_server_error_returns_empty_context(mock_get): + mock_response = MagicMock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + mock_get.return_value = mock_response + + assistant = MockAssistant({ + "rag_processor": "library_file_rag", + "library_id": "lib-123", + "item_id": "item-456", + }) + + result = rag_processor(messages=_make_messages(), assistant=assistant) + + assert result["context"] == "" + assert result["sources"] == [] + + +@patch.dict(os.environ, {}, clear=True) +def test_library_item_missing_env_vars_returns_empty_context(): + assistant = MockAssistant({ + "rag_processor": "library_file_rag", + "library_id": "lib-123", + "item_id": "item-456", + }) + + result = rag_processor(messages=_make_messages(), assistant=assistant) + + assert result["context"] == "" + assert result["sources"] == [] + + +def test_no_library_id_returns_empty_context(): + assistant = MockAssistant({ + "rag_processor": "library_file_rag", + }) + + result = rag_processor(messages=_make_messages(), assistant=assistant) + + assert result["context"] == "" + assert result["sources"] == [] + + +def test_library_id_without_item_id_returns_empty_context(): + assistant = MockAssistant({ + "rag_processor": "library_file_rag", + "library_id": "lib-123", + }) + + result = rag_processor(messages=_make_messages(), assistant=assistant) + + assert result["context"] == "" + assert result["sources"] == [] + + +def test_no_assistant_returns_empty_context(): + result = rag_processor(messages=_make_messages(), assistant=None) + + assert result["context"] == "" + assert result["sources"] == [] + + +def test_invalid_metadata_json_returns_empty_context(): + assistant = MockAssistant.__new__(MockAssistant) + assistant.id = "mock-assistant-id" + assistant.api_callback = "not valid json" + assistant.metadata = "not valid json" + assistant.owner = "test@example.com" + + result = rag_processor(messages=_make_messages(), assistant=assistant) + + assert result["context"] == "" + assert result["sources"] == [] From e4bab0367699a99437172a4ca166429354c016d5 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 12:35:47 +0200 Subject: [PATCH 061/156] feat: add COMPATIBLE_RAG declaration to simple_augment (legacy PPS) --- backend/lamb/completions/pps/simple_augment.py | 9 +++++++++ .../unit-tests/completions/test_simple_augment.py | 15 +++++++++++++-- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/backend/lamb/completions/pps/simple_augment.py b/backend/lamb/completions/pps/simple_augment.py index 414dd8654..36979a29c 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 _has_vision_capability(assistant: Assistant) -> bool: """ diff --git a/testing/unit-tests/completions/test_simple_augment.py b/testing/unit-tests/completions/test_simple_augment.py index 4c9718ec9..e888e6305 100644 --- a/testing/unit-tests/completions/test_simple_augment.py +++ b/testing/unit-tests/completions/test_simple_augment.py @@ -1,2 +1,13 @@ -# Archivo vacío — todos los tests eran de document_context -# y se mueven a test_kvcache_augment.py +from lamb.completions.pps.simple_augment import COMPATIBLE_RAG + + +class TestSimpleAugmentCompatibleRag: + def test_declares_compatible_rag_list(self): + assert isinstance(COMPATIBLE_RAG, list) + assert "simple_rag" in COMPATIBLE_RAG + assert "context_aware_rag" in COMPATIBLE_RAG + assert "single_file_rag" in COMPATIBLE_RAG + assert "no_rag" in COMPATIBLE_RAG + + def test_does_not_include_new_rags(self): + assert "library_file_rag" not in COMPATIBLE_RAG From 44ea1e1050a7681a3f144467d860ed337973f2a3 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 12:54:17 +0200 Subject: [PATCH 062/156] feat: add COMPATIBLE_RAG validation in load_and_validate_plugins --- backend/lamb/completions/main.py | 23 +++++ .../test_compatible_rag_validation.py | 94 +++++++++++++++++++ testing/unit-tests/conftest.py | 17 ++++ 3 files changed, 134 insertions(+) create mode 100644 testing/unit-tests/completions/test_compatible_rag_validation.py create mode 100644 testing/unit-tests/conftest.py diff --git a/backend/lamb/completions/main.py b/backend/lamb/completions/main.py index 1c48d5198..218aa95ed 100644 --- a/backend/lamb/completions/main.py +++ b/backend/lamb/completions/main.py @@ -308,6 +308,7 @@ def _check_quota(assistant_id: int, assistant_details) -> None: 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') @@ -326,6 +327,28 @@ def load_and_validate_plugins(plugin_config: Dict[str, str]) -> Tuple[Dict[str, 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 async def get_rag_context(request: Dict[str, Any], rag_processors: Dict[str, Any], rag_processor: str, assistant_details: Any) -> Any: diff --git a/testing/unit-tests/completions/test_compatible_rag_validation.py b/testing/unit-tests/completions/test_compatible_rag_validation.py new file mode 100644 index 000000000..87a93eec7 --- /dev/null +++ b/testing/unit-tests/completions/test_compatible_rag_validation.py @@ -0,0 +1,94 @@ +import pytest +from unittest.mock import patch, MagicMock +from fastapi import HTTPException + +from lamb.completions import main as main_module +from lamb.completions.main import load_and_validate_plugins + + +class TestCompatibleRagValidation: + def _make_plugin_config(self, pps="simple_augment", rag="no_rag", doc_rag=""): + return { + "prompt_processor": pps, + "connector": "openai", + "llm": "gpt-4", + "rag_processor": rag, + "document_rag": doc_rag, + } + + @patch.object(main_module, "load_plugins") + @patch("importlib.import_module") + def test_compatible_rag_passes_when_rag_in_list(self, mock_import, mock_load): + mock_pps_module = MagicMock() + mock_pps_module.COMPATIBLE_RAG = ["simple_rag", "no_rag"] + mock_import.return_value = mock_pps_module + + mock_load.side_effect = lambda t: ( + {"simple_augment": MagicMock()} if t == "pps" + else {"simple_rag": MagicMock(), "no_rag": MagicMock()} if t == "rag" + else {"openai": MagicMock()} + ) + + config = self._make_plugin_config(pps="simple_augment", rag="simple_rag") + pps, connectors, rag_processors = load_and_validate_plugins(config) + + assert "simple_augment" in pps + + @patch.object(main_module, "load_plugins") + @patch("importlib.import_module") + def test_compatible_rag_rejects_incompatible_rag(self, mock_import, mock_load): + mock_pps_module = MagicMock() + mock_pps_module.COMPATIBLE_RAG = ["library_file_rag", "no_rag"] + mock_import.return_value = mock_pps_module + + mock_load.side_effect = lambda t: ( + {"kvcache_augment": MagicMock()} if t == "pps" + else {"simple_rag": MagicMock(), "no_rag": MagicMock(), "library_file_rag": MagicMock()} if t == "rag" + else {"openai": MagicMock()} + ) + + config = self._make_plugin_config(pps="kvcache_augment", rag="simple_rag") + + with pytest.raises(HTTPException) as exc_info: + load_and_validate_plugins(config) + assert exc_info.value.status_code == 400 + assert "not compatible" in str(exc_info.value.detail) + + @patch.object(main_module, "load_plugins") + @patch("importlib.import_module") + def test_compatible_rag_validates_document_rag(self, mock_import, mock_load): + mock_pps_module = MagicMock() + mock_pps_module.COMPATIBLE_RAG = ["library_file_rag", "no_rag"] + mock_import.return_value = mock_pps_module + + mock_load.side_effect = lambda t: ( + {"kvcache_augment": MagicMock()} if t == "pps" + else {"no_rag": MagicMock(), "library_file_rag": MagicMock(), "single_file_rag": MagicMock()} if t == "rag" + else {"openai": MagicMock()} + ) + + config = self._make_plugin_config( + pps="kvcache_augment", rag="no_rag", doc_rag="single_file_rag" + ) + + with pytest.raises(HTTPException) as exc_info: + load_and_validate_plugins(config) + assert exc_info.value.status_code == 400 + assert "document_rag" in str(exc_info.value.detail) + + @patch.object(main_module, "load_plugins") + @patch("importlib.import_module") + def test_no_compatible_rag_skips_validation(self, mock_import, mock_load): + mock_pps_module = MagicMock(spec=[]) + mock_import.return_value = mock_pps_module + + mock_load.side_effect = lambda t: ( + {"some_pps": MagicMock()} if t == "pps" + else {"no_rag": MagicMock()} if t == "rag" + else {"openai": MagicMock()} + ) + + config = self._make_plugin_config(pps="some_pps", rag="no_rag") + pps, connectors, rag_processors = load_and_validate_plugins(config) + + assert "some_pps" in pps diff --git a/testing/unit-tests/conftest.py b/testing/unit-tests/conftest.py new file mode 100644 index 000000000..9de80a8b9 --- /dev/null +++ b/testing/unit-tests/conftest.py @@ -0,0 +1,17 @@ +import os + +os.environ.setdefault("LAMB_WEB_HOST", "localhost") +os.environ.setdefault("PIPELINES_HOST", "localhost") +os.environ.setdefault("LAMB_BACKEND_HOST", "localhost") +os.environ.setdefault("LAMB_BEARER_TOKEN", "test-bearer-token") +os.environ.setdefault("PIPELINES_BEARER_TOKEN", "test-bearer-token") +os.environ.setdefault("OPENAI_BASE_URL", "http://localhost:9099") +os.environ.setdefault("OPENAI_MODEL", "gpt-4") +os.environ.setdefault("OWI_BASE_URL", "http://localhost:8080") +os.environ.setdefault("OWI_PATH", "/tmp/owi") +os.environ.setdefault("SIGNUP_SECRET_KEY", "test-secret-key") +os.environ.setdefault("OWI_ADMIN_NAME", "admin") +os.environ.setdefault("OWI_ADMIN_EMAIL", "admin@test.com") +os.environ.setdefault("OWI_ADMIN_PASSWORD", "admin-password") +os.environ.setdefault("LAMB_PROJECT_PATH", "/tmp/lamb") +os.environ.setdefault("LAMB_DB_PATH", "/tmp/lamb_test.db") From 688d4accd2ba78f002aa1113c584dcabfc668911 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 12:55:18 +0200 Subject: [PATCH 063/156] refactor: metadata_validators accepts library_file_rag for document_rag --- backend/creator_interface/metadata_validators.py | 4 ++-- .../assistant_router/test_metadata_validation.py | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/backend/creator_interface/metadata_validators.py b/backend/creator_interface/metadata_validators.py index fae6c0a3d..4d95a0a59 100644 --- a/backend/creator_interface/metadata_validators.py +++ b/backend/creator_interface/metadata_validators.py @@ -74,13 +74,13 @@ def validate_update_plugin_metadata( return None, "single_file_rag: item_id requires library_id" # document_rag specific validation - if metadata_dict.get("document_rag") == "single_file_rag": + 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=single_file_rag requires library_id + item_id in metadata" + "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" diff --git a/testing/unit-tests/assistant_router/test_metadata_validation.py b/testing/unit-tests/assistant_router/test_metadata_validation.py index b65a82514..91299ff8f 100644 --- a/testing/unit-tests/assistant_router/test_metadata_validation.py +++ b/testing/unit-tests/assistant_router/test_metadata_validation.py @@ -108,7 +108,7 @@ class TestDocumentRagValidation: def _make_body(self, **overrides): base = { "metadata": { - "prompt_processor": "simple_augment", + "prompt_processor": "kvcache_augment", "connector": "openai", "llm": "gpt-4o-mini", "rag_processor": "no_rag", @@ -118,15 +118,15 @@ def _make_body(self, **overrides): base["metadata"].update(overrides) return base - def test_document_rag_single_file_requires_library_or_file(self): - body = self._make_body(document_rag="single_file_rag") + def test_document_rag_library_file_requires_library_or_file(self): + body = self._make_body(document_rag="library_file_rag") _, error = validate_update_plugin_metadata(body) assert error is not None assert "document_rag" in error def test_document_rag_with_library_id_and_item_id_is_valid(self): body = self._make_body( - document_rag="single_file_rag", + document_rag="library_file_rag", library_id="lib-1", item_id="item-1", ) @@ -136,7 +136,7 @@ def test_document_rag_with_library_id_and_item_id_is_valid(self): def test_document_rag_library_id_without_item_id(self): body = self._make_body( - document_rag="single_file_rag", + document_rag="library_file_rag", library_id="lib-1", ) _, error = validate_update_plugin_metadata(body) From 25d83ec6f238a9c2f85827ec4cbb85b55b721f0a Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 12:56:03 +0200 Subject: [PATCH 064/156] test: adapt document_rag pipeline tests to library_file_rag + kvcache_augment --- .../completions/test_document_rag_pipeline.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/testing/unit-tests/completions/test_document_rag_pipeline.py b/testing/unit-tests/completions/test_document_rag_pipeline.py index 7b36a719a..0abc25af5 100644 --- a/testing/unit-tests/completions/test_document_rag_pipeline.py +++ b/testing/unit-tests/completions/test_document_rag_pipeline.py @@ -18,19 +18,19 @@ def __init__(self, metadata="", system_prompt="", prompt_template=""): class TestParsePluginConfigDocumentRag: def test_reads_document_rag_from_metadata(self): metadata = json.dumps({ - "prompt_processor": "simple_augment", + "prompt_processor": "kvcache_augment", "connector": "openai", "llm": "gpt-4o-mini", "rag_processor": "context_aware_rag", - "document_rag": "single_file_rag", + "document_rag": "library_file_rag", }) assistant = MockAssistant(metadata=metadata) config = parse_plugin_config(assistant) - assert config["document_rag"] == "single_file_rag" + assert config["document_rag"] == "library_file_rag" def test_document_rag_defaults_to_empty(self): metadata = json.dumps({ - "prompt_processor": "simple_augment", + "prompt_processor": "kvcache_augment", "connector": "openai", "llm": "gpt-4o-mini", "rag_processor": "no_rag", @@ -61,8 +61,8 @@ def mock_pps(request, assistant=None, rag_context=None, **kwargs): captured_kwargs["rag_context"] = rag_context return [{"role": "user", "content": "test"}] - pps = {"simple_augment": mock_pps} - plugin_config = {"prompt_processor": "simple_augment"} + pps = {"kvcache_augment": mock_pps} + plugin_config = {"prompt_processor": "kvcache_augment"} doc_ctx = {"context": "Doc.", "sources": []} process_completion_request( From dcf2933a84884de2460d52f1cad0c145d47c0947 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 12:57:47 +0200 Subject: [PATCH 065/156] refactor: frontend sends document_rag=library_file_rag instead of single_file_rag --- .../components/assistants/assistantFormDocumentRag.test.js | 2 +- .../src/lib/components/assistants/assistantFormSubmit.test.js | 2 +- .../components/assistants/logic/assistantFormState.svelte.js | 4 ++-- .../lib/components/assistants/logic/assistantFormSubmit.js | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/assistantFormDocumentRag.test.js b/frontend/svelte-app/src/lib/components/assistants/assistantFormDocumentRag.test.js index fc1b5b4ff..166c61980 100644 --- a/frontend/svelte-app/src/lib/components/assistants/assistantFormDocumentRag.test.js +++ b/frontend/svelte-app/src/lib/components/assistants/assistantFormDocumentRag.test.js @@ -109,7 +109,7 @@ describe('documentRagEnabled in form state', () => { connector: 'openai', llm: 'gpt-4o-mini', rag_processor: 'no_rag', - document_rag: 'single_file_rag', + document_rag: 'library_file_rag', library_id: 'lib-1', item_id: 'item-1' }) diff --git a/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js b/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js index 3c590ca38..3a16a88ba 100644 --- a/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js +++ b/frontend/svelte-app/src/lib/components/assistants/assistantFormSubmit.test.js @@ -156,7 +156,7 @@ describe('buildAssistantPayload with document_rag', () => { }; const payload = buildAssistantPayload(form); const metadata = JSON.parse(payload.metadata); - expect(metadata.document_rag).toBe('single_file_rag'); + expect(metadata.document_rag).toBe('library_file_rag'); expect(metadata.library_id).toBe('lib-1'); expect(metadata.item_id).toBe('item-1'); expect(metadata.rag_processor).toBe('no_rag'); diff --git a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js index 1eeb387a0..4fd908bc8 100644 --- a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js +++ b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormState.svelte.js @@ -211,8 +211,8 @@ export function populateFormFields(form, data, getAvailableModels, preserveDescr form.selectedFilePath = metadata?.file_path || ''; } - // Document RAG: new path (document_rag=single_file_rag) - form.documentRagEnabled = metadata?.document_rag === 'single_file_rag'; + // Document RAG: new path (document_rag=library_file_rag) + form.documentRagEnabled = metadata?.document_rag === 'library_file_rag'; if (form.documentRagEnabled && !isSingleFileRag(form.selectedRagProcessor)) { form.selectedLibraryId = metadata?.library_id || ''; form.selectedItemId = metadata?.item_id || ''; diff --git a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js index bb98f4ee7..5a4ce5e73 100644 --- a/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js +++ b/frontend/svelte-app/src/lib/components/assistants/logic/assistantFormSubmit.js @@ -37,7 +37,7 @@ export function buildAssistantPayload(form) { }; if (form.documentRagEnabled) { - metadataObj.document_rag = 'single_file_rag'; + metadataObj.document_rag = 'library_file_rag'; metadataObj.library_id = form.selectedLibraryId || ''; metadataObj.item_id = form.selectedItemId || ''; } else if (isSingleFileRag(form.selectedRagProcessor) && form.selectedFilePath) { From 2a4000b43abe2a964c7467c63bcd0acce3abd8cf Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 15:26:03 +0200 Subject: [PATCH 066/156] gitignore --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 169a0ec46..972d013f5 100644 --- a/.gitignore +++ b/.gitignore @@ -195,3 +195,6 @@ aac_logs/ .worktrees/ .agents/ skills-lock.json +test-results/ +testing/load/document-cache-test/ +testing/load/kv-cache-classroom/ From 4cc2a8ebfd557f9d299abfef40f1183f8a8191c2 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 16:49:42 +0200 Subject: [PATCH 067/156] feat: add DEFAULT_RAG_PROMPT_TEMPLATE to kvcache_augment #2 --- backend/lamb/completions/pps/kvcache_augment.py | 6 ++++++ .../unit-tests/completions/test_kvcache_augment.py | 12 +++++++++++- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/backend/lamb/completions/pps/kvcache_augment.py b/backend/lamb/completions/pps/kvcache_augment.py index 2180af2ae..a488daf23 100644 --- a/backend/lamb/completions/pps/kvcache_augment.py +++ b/backend/lamb/completions/pps/kvcache_augment.py @@ -13,6 +13,12 @@ "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}" +) + def _has_vision_capability(assistant: Assistant) -> bool: if not assistant: diff --git a/testing/unit-tests/completions/test_kvcache_augment.py b/testing/unit-tests/completions/test_kvcache_augment.py index 3aedad2af..9c090890b 100644 --- a/testing/unit-tests/completions/test_kvcache_augment.py +++ b/testing/unit-tests/completions/test_kvcache_augment.py @@ -1,5 +1,5 @@ import pytest -from lamb.completions.pps.kvcache_augment import prompt_processor, COMPATIBLE_RAG +from lamb.completions.pps.kvcache_augment import prompt_processor, COMPATIBLE_RAG, DEFAULT_RAG_PROMPT_TEMPLATE class MockAssistant: @@ -125,3 +125,13 @@ def test_no_assistant_document_context_ignored(self): assert len(result) == 1 assert result[0]["role"] == "user" + + +class TestDefaultRagPromptTemplate: + def test_constant_matches_expected_value(self): + 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 From 148cfe652aac25af9d02cbb74dd9b90e1098923e Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 16:50:55 +0200 Subject: [PATCH 068/156] feat: add D3 fallback with DEFAULT_RAG_PROMPT_TEMPLATE in kvcache_augment #2 --- .../lamb/completions/pps/kvcache_augment.py | 31 ++++++++++++- .../completions/test_kvcache_augment.py | 43 +++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/backend/lamb/completions/pps/kvcache_augment.py b/backend/lamb/completions/pps/kvcache_augment.py index a488daf23..bbd1fe6b7 100644 --- a/backend/lamb/completions/pps/kvcache_augment.py +++ b/backend/lamb/completions/pps/kvcache_augment.py @@ -150,7 +150,36 @@ def prompt_processor( "content": prompt }) else: - processed_messages.append(messages[-1]) + 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") + context = rag_context.get("context", "") if isinstance(rag_context, dict) else str(rag_context) + prompt = prompt.replace("{context}", "\n\n" + context + "\n\n") + + processed_messages.append({ + "role": messages[-1]['role'], + "content": prompt + }) + else: + processed_messages.append(messages[-1]) return processed_messages diff --git a/testing/unit-tests/completions/test_kvcache_augment.py b/testing/unit-tests/completions/test_kvcache_augment.py index 9c090890b..20b643c2e 100644 --- a/testing/unit-tests/completions/test_kvcache_augment.py +++ b/testing/unit-tests/completions/test_kvcache_augment.py @@ -135,3 +135,46 @@ def test_constant_matches_expected_value(self): "Context:\n{context}\n\nQuestion: {user_input}" ) assert DEFAULT_RAG_PROMPT_TEMPLATE == expected + + +class TestD3Fallback: + def test_default_template_used_when_no_prompt_template_but_rag_context(self): + assistant = MockAssistant( + system_prompt="You are a tutor.", + prompt_template="", + ) + request = _make_request([{"role": "user", "content": "What is 2+2?"}]) + rag_ctx = {"context": "RAG chunks here.", "sources": []} + + result = prompt_processor(request, assistant=assistant, rag_context=rag_ctx) + + last_msg = result[-1] + assert "RAG chunks here." in last_msg["content"] + assert "Use the following context" in last_msg["content"] + + def test_no_fallback_when_no_rag_context(self): + assistant = MockAssistant( + system_prompt="You are a tutor.", + prompt_template="", + ) + request = _make_request([{"role": "user", "content": "Hi"}]) + + result = prompt_processor(request, assistant=assistant) + + last_msg = result[-1] + assert last_msg["content"] == "Hi" + assert "Use the following context" not in last_msg["content"] + + def test_explicit_template_wins_over_default(self): + assistant = MockAssistant( + system_prompt="Sys.", + prompt_template="Custom: {context}\nQ: {user_input}", + ) + request = _make_request([{"role": "user", "content": "Q?"}]) + rag_ctx = {"context": "Some chunk.", "sources": []} + + result = prompt_processor(request, assistant=assistant, rag_context=rag_ctx) + + last_msg = result[-1] + assert "Custom:" in last_msg["content"] + assert "Some chunk." in last_msg["content"] From 2c130b9b1d6bc7ca036b913c88cfd5c121b9490b Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 16:52:21 +0200 Subject: [PATCH 069/156] feat: add labeled_doc wrapper with recency-bias in kvcache_augment #2 --- .../lamb/completions/pps/kvcache_augment.py | 13 +++- .../completions/test_kvcache_augment.py | 67 ++++++++++++++++++- 2 files changed, 78 insertions(+), 2 deletions(-) diff --git a/backend/lamb/completions/pps/kvcache_augment.py b/backend/lamb/completions/pps/kvcache_augment.py index bbd1fe6b7..dd2541ebe 100644 --- a/backend/lamb/completions/pps/kvcache_augment.py +++ b/backend/lamb/completions/pps/kvcache_augment.py @@ -66,7 +66,18 @@ def prompt_processor( if document_context and isinstance(document_context, dict): doc_text = document_context.get("context", "") if doc_text: - system_content = (system_content + "\n\n" + doc_text) if system_content else 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", diff --git a/testing/unit-tests/completions/test_kvcache_augment.py b/testing/unit-tests/completions/test_kvcache_augment.py index 20b643c2e..53ffe146b 100644 --- a/testing/unit-tests/completions/test_kvcache_augment.py +++ b/testing/unit-tests/completions/test_kvcache_augment.py @@ -53,7 +53,8 @@ def test_document_context_as_system_prompt_when_none_exists(self): system_msg = result[0] assert system_msg["role"] == "system" - assert system_msg["content"] == "Document content." + assert "## REFERENCE DOCUMENT" in system_msg["content"] + assert "Document content." in system_msg["content"] def test_no_document_context_leaves_system_prompt_unchanged(self): assistant = MockAssistant( @@ -178,3 +179,67 @@ def test_explicit_template_wins_over_default(self): last_msg = result[-1] assert "Custom:" in last_msg["content"] assert "Some chunk." in last_msg["content"] + + +class TestLabeledDocWrapper: + def test_labeled_doc_wrapper_in_system_prompt(self): + assistant = MockAssistant( + system_prompt="You are a tutor.", + prompt_template="Answer: {user_input}", + ) + request = _make_request([{"role": "user", "content": "Hello"}]) + doc_ctx = {"context": "# Reference Document\nSome content.", "sources": []} + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + system_msg = result[0] + assert "## REFERENCE DOCUMENT" in system_msg["content"] + assert "This document has been selected by the assistant creator" in system_msg["content"] + assert "# Reference Document\nSome content." in system_msg["content"] + assert "IMPORTANT: The reference document above is available" in system_msg["content"] + + def test_recency_bias_reminder_after_doc_content(self): + assistant = MockAssistant( + system_prompt="You are a tutor.", + prompt_template="Answer: {user_input}", + ) + request = _make_request([{"role": "user", "content": "Hello"}]) + doc_ctx = {"context": "Doc content.", "sources": []} + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + system_msg = result[0] + doc_index = system_msg["content"].index("Doc content.") + reminder_index = system_msg["content"].index("IMPORTANT: The reference document above") + assert doc_index < reminder_index + + def test_labeled_doc_prepended_before_system_prompt(self): + assistant = MockAssistant( + system_prompt="You are a tutor.", + prompt_template="Answer: {user_input}", + ) + request = _make_request([{"role": "user", "content": "Hello"}]) + doc_ctx = {"context": "Doc content.", "sources": []} + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + system_msg = result[0] + doc_index = system_msg["content"].index("## REFERENCE DOCUMENT") + sys_index = system_msg["content"].index("You are a tutor.") + assert doc_index < sys_index + + def test_labeled_doc_becomes_system_prompt_when_none_exists(self): + assistant = MockAssistant( + system_prompt="", + prompt_template="Answer: {user_input}", + ) + request = _make_request([{"role": "user", "content": "Hello"}]) + doc_ctx = {"context": "Document content.", "sources": []} + + result = prompt_processor(request, assistant=assistant, document_context=doc_ctx) + + system_msg = result[0] + assert system_msg["role"] == "system" + assert "## REFERENCE DOCUMENT" in system_msg["content"] + assert "Document content." in system_msg["content"] + assert system_msg["content"].rstrip().endswith("use it.") From 62011f742b59e26702221507555646bcfe02c3cd Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 16:54:14 +0200 Subject: [PATCH 070/156] refactor: remove DEFAULT_RAG_PROMPT_TEMPLATE and document_context from simple_augment #2 --- .../lamb/completions/pps/simple_augment.py | 25 +----- backend/tests/test_simple_augment.py | 82 +------------------ 2 files changed, 5 insertions(+), 102 deletions(-) diff --git a/backend/lamb/completions/pps/simple_augment.py b/backend/lamb/completions/pps/simple_augment.py index 202e58fa8..0b1f654f4 100644 --- a/backend/lamb/completions/pps/simple_augment.py +++ b/backend/lamb/completions/pps/simple_augment.py @@ -97,27 +97,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) @@ -136,7 +117,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") # Add RAG context if available if rag_context: @@ -192,7 +173,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") # Add RAG context if available if rag_context: diff --git a/backend/tests/test_simple_augment.py b/backend/tests/test_simple_augment.py index 3ff91aa38..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,44 +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_document_context_gets_descriptive_label(): - """Document RAG content should be wrapped with a descriptive REFERENCE DOCUMENT label.""" - assistant = _make_assistant(system_prompt="You are helpful.", prompt_template="Answer: {user_input}") - request = { - "messages": [ - {"role": "user", "content": "What is the answer?"} - ] - } - document_context = {"context": "This is the reference document content.", "sources": []} - - result = prompt_processor(request, assistant=assistant, document_context=document_context) - - system_msg = result[0] - assert system_msg["role"] == "system" - assert "REFERENCE DOCUMENT" in system_msg["content"] - assert "This is the reference document content." in system_msg["content"] - # Should contain the explanation for the LLM - assert "selected by the assistant creator" in system_msg["content"] - # Should contain the recency-bias reminder at the end - assert "available for this entire conversation" in system_msg["content"] - # Reminder should be at the very end of the system content - assert system_msg["content"].rstrip().endswith("use it.") - # Should come after the original system prompt - assert "You are helpful." in system_msg["content"] - doc_index = system_msg["content"].index("REFERENCE DOCUMENT") - sys_index = system_msg["content"].index("You are helpful.") - assert doc_index > sys_index - - -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 From d4f4390d6326232d7b872929b93cdbc743c24f2c Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 16:56:45 +0200 Subject: [PATCH 071/156] feat: add PPS_COMPATIBLE_RAG and helpers to ragProcessorHelpers #2 --- .../src/lib/utils/ragProcessorHelpers.js | 49 ++++++++++++++ .../src/lib/utils/ragProcessorHelpers.test.js | 64 ++++++++++++++++++- 2 files changed, 112 insertions(+), 1 deletion(-) diff --git a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js index 813015c17..bf0d7a692 100644 --- a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js +++ b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js @@ -120,3 +120,52 @@ export function getRagProcessorDisplayName(processor) { } return processor.replace(/_/g, ' ').replace(/\b\w/g, (l) => l.toUpperCase()); } + +/** + * Declaracio de compatibilitat PPS ↔ RAG (mirall del backend). + * Cada PPS declara quins RAGs son compatibles. + * Si un PPS no esta a la llista, s'assumeix que accepta qualsevol RAG. + */ +export const PPS_COMPATIBLE_RAG = Object.freeze({ + simple_augment: [ + 'simple_rag', + 'context_aware_rag', + 'hierarchical_rag', + 'single_file_rag', + 'rubric_rag', + 'no_rag' + ], + kvcache_augment: [ + 'library_file_rag', + 'knowledge_store_rag', + 'query_rewriting_ks_rag', + 'rubric_rag', + 'no_rag' + ] +}); + +/** + * Retorna els RAGs compatibles amb un PPS donat. + * Si el PPS no te declaracio, retorna tots els RAGs disponibles. + * @param {string} pps - Nom del prompt processor + * @param {string[]} allRagProcessors - Llista de tots els RAGs disponibles + * @returns {string[]} RAGs compatibles + */ +export function getCompatibleRagForPps(pps, allRagProcessors) { + const compatible = PPS_COMPATIBLE_RAG[pps]; + if (!compatible) { + return allRagProcessors; + } + return allRagProcessors.filter((rag) => compatible.includes(rag)); +} + +/** + * Retorna true si el PPS suporta document RAG. + * @param {string} pps - Nom del prompt processor + * @returns {boolean} + */ +export function ppsSupportsDocumentRag(pps) { + const compatible = PPS_COMPATIBLE_RAG[pps]; + if (!compatible) return false; + return compatible.includes('library_file_rag'); +} diff --git a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.test.js b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.test.js index cf742f003..ab76e83ef 100644 --- a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.test.js +++ b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.test.js @@ -6,7 +6,10 @@ import { isRubricRag, isNoRag, hasRagOptions, - normalizeRagProcessor + normalizeRagProcessor, + PPS_COMPATIBLE_RAG, + getCompatibleRagForPps, + ppsSupportsDocumentRag } from './ragProcessorHelpers.js'; describe('isKbBasedRag', () => { @@ -137,3 +140,62 @@ describe('normalizeRagProcessor', () => { expect(normalizeRagProcessor('')).toBe(''); }); }); + +describe('PPS_COMPATIBLE_RAG', () => { + test('declares simple_augment compatible RAGs', () => { + expect(PPS_COMPATIBLE_RAG.simple_augment).toContain('simple_rag'); + expect(PPS_COMPATIBLE_RAG.simple_augment).toContain('context_aware_rag'); + expect(PPS_COMPATIBLE_RAG.simple_augment).toContain('single_file_rag'); + expect(PPS_COMPATIBLE_RAG.simple_augment).toContain('no_rag'); + expect(PPS_COMPATIBLE_RAG.simple_augment).not.toContain('knowledge_store_rag'); + }); + + test('declares kvcache_augment compatible RAGs', () => { + expect(PPS_COMPATIBLE_RAG.kvcache_augment).toContain('knowledge_store_rag'); + expect(PPS_COMPATIBLE_RAG.kvcache_augment).toContain('query_rewriting_ks_rag'); + expect(PPS_COMPATIBLE_RAG.kvcache_augment).toContain('library_file_rag'); + expect(PPS_COMPATIBLE_RAG.kvcache_augment).toContain('no_rag'); + expect(PPS_COMPATIBLE_RAG.kvcache_augment).not.toContain('simple_rag'); + }); + + test('is frozen (immutable)', () => { + expect(Object.isFrozen(PPS_COMPATIBLE_RAG)).toBe(true); + }); +}); + +describe('getCompatibleRagForPps', () => { + const allRags = ['simple_rag', 'knowledge_store_rag', 'no_rag', 'library_file_rag']; + + test('filters RAGs to only compatible ones for simple_augment', () => { + const result = getCompatibleRagForPps('simple_augment', allRags); + expect(result).toContain('simple_rag'); + expect(result).toContain('no_rag'); + expect(result).not.toContain('knowledge_store_rag'); + }); + + test('filters RAGs to only compatible ones for kvcache_augment', () => { + const result = getCompatibleRagForPps('kvcache_augment', allRags); + expect(result).toContain('knowledge_store_rag'); + expect(result).toContain('library_file_rag'); + expect(result).not.toContain('simple_rag'); + }); + + test('returns all RAGs for unknown PPS', () => { + const result = getCompatibleRagForPps('unknown_pps', allRags); + expect(result).toEqual(allRags); + }); +}); + +describe('ppsSupportsDocumentRag', () => { + test('returns true for kvcache_augment', () => { + expect(ppsSupportsDocumentRag('kvcache_augment')).toBe(true); + }); + + test('returns false for simple_augment', () => { + expect(ppsSupportsDocumentRag('simple_augment')).toBe(false); + }); + + test('returns false for unknown PPS', () => { + expect(ppsSupportsDocumentRag('unknown_pps')).toBe(false); + }); +}); From 7e733df78dd217bb850dede92373906927514f98 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 16:58:28 +0200 Subject: [PATCH 072/156] feat: filter RAG processors by PPS compatibility in ConfigurationPanel #2 --- .../components/ConfigurationPanel.svelte | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte index 8715a9611..c76edb4a2 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte @@ -3,7 +3,7 @@ import { _ } from '$lib/i18n'; import { tick } from 'svelte'; import { assistantConfigStore } from '$lib/stores/assistantConfigStore'; - import { hasRagOptions, isHiddenInCreate, getRagProcessorDisplayName } from '$lib/utils/ragProcessorHelpers.js'; + import { hasRagOptions, isHiddenInCreate, getRagProcessorDisplayName, getCompatibleRagForPps, ppsSupportsDocumentRag } from '$lib/utils/ragProcessorHelpers.js'; import { extractModelsMetadata } from '../logic/assistantFormUtils.svelte.js'; @@ -61,12 +61,19 @@ let imageGenerationForced = $derived(currentModelMetadata?.forced_capabilities?.image_generation === true); let showRagOptions = $derived(hasRagOptions(selectedRagProcessor)); let isLegacySingleFileRag = $derived(selectedRagProcessor === 'single_file_rag'); - let showDocumentSection = $derived(!isLegacySingleFileRag); - let filteredRAGProcessors = $derived( - formState === 'edit' - ? ragProcessors - : ragProcessors.filter((p) => !isHiddenInCreate(p)) - ); + let showDocumentSection = $derived.by(() => { + if (!ppsSupportsDocumentRag(selectedPromptProcessor)) { + return false; + } + return !isLegacySingleFileRag; + }); + let filteredRAGProcessors = $derived.by(() => { + const compatibleWithPps = getCompatibleRagForPps(selectedPromptProcessor, ragProcessors); + if (formState === 'edit') { + return compatibleWithPps; + } + return compatibleWithPps.filter((p) => !isHiddenInCreate(p)); + }); async function handleConnectorChange() { await tick(); From e672d2d392c3d3b0625cd2d7ca568ba11007b5ff Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 17:43:21 +0200 Subject: [PATCH 073/156] fix: replace isHiddenInCreate with isDocumentRag filter in RAG dropdown #2 --- .../components/ConfigurationPanel.svelte | 12 +++---- .../src/lib/utils/ragProcessorHelpers.js | 14 +++++++- .../src/lib/utils/ragProcessorHelpers.test.js | 33 ++++++++++++++++++- 3 files changed, 49 insertions(+), 10 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte index c76edb4a2..62106a95f 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte @@ -3,7 +3,7 @@ import { _ } from '$lib/i18n'; import { tick } from 'svelte'; import { assistantConfigStore } from '$lib/stores/assistantConfigStore'; - import { hasRagOptions, isHiddenInCreate, getRagProcessorDisplayName, getCompatibleRagForPps, ppsSupportsDocumentRag } from '$lib/utils/ragProcessorHelpers.js'; + import { hasRagOptions, getRagProcessorDisplayName, getCompatibleRagForPps, ppsSupportsDocumentRag, isDocumentRag } from '$lib/utils/ragProcessorHelpers.js'; import { extractModelsMetadata } from '../logic/assistantFormUtils.svelte.js'; @@ -67,13 +67,9 @@ } return !isLegacySingleFileRag; }); - let filteredRAGProcessors = $derived.by(() => { - const compatibleWithPps = getCompatibleRagForPps(selectedPromptProcessor, ragProcessors); - if (formState === 'edit') { - return compatibleWithPps; - } - return compatibleWithPps.filter((p) => !isHiddenInCreate(p)); - }); + let filteredRAGProcessors = $derived( + getCompatibleRagForPps(selectedPromptProcessor, ragProcessors).filter((p) => !isDocumentRag(p)) + ); async function handleConnectorChange() { await tick(); diff --git a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js index bf0d7a692..19168bd90 100644 --- a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js +++ b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js @@ -21,7 +21,9 @@ export const RAG_TYPES = Object.freeze({ /** RAG processor that uses rubrics */ RUBRIC: ['rubric_rag'], /** No RAG processing */ - NONE: ['no_rag'] + NONE: ['no_rag'], + /** RAG processors that are document-level (selected via document section toggle, not RAG dropdown) */ + DOCUMENT_RAG: ['library_file_rag'] }); /** @@ -51,6 +53,16 @@ export function isHiddenInCreate(processor) { return RAG_TYPES.HIDDEN_IN_CREATE.includes(processor); } +/** + * Returns true if the processor is a document-level RAG (selected via document section toggle, + * not the RAG processor dropdown). These should never appear in the RAG dropdown. + * @param {string} processor + * @returns {boolean} + */ +export function isDocumentRag(processor) { + return RAG_TYPES.DOCUMENT_RAG.includes(processor); +} + /** * Returns true if the processor uses a single file. * @param {string} processor diff --git a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.test.js b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.test.js index ab76e83ef..c4d7dd9b2 100644 --- a/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.test.js +++ b/frontend/svelte-app/src/lib/utils/ragProcessorHelpers.test.js @@ -9,7 +9,8 @@ import { normalizeRagProcessor, PPS_COMPATIBLE_RAG, getCompatibleRagForPps, - ppsSupportsDocumentRag + ppsSupportsDocumentRag, + isDocumentRag } from './ragProcessorHelpers.js'; describe('isKbBasedRag', () => { @@ -199,3 +200,33 @@ describe('ppsSupportsDocumentRag', () => { expect(ppsSupportsDocumentRag('unknown_pps')).toBe(false); }); }); + +describe('isDocumentRag', () => { + test('returns true for library_file_rag', () => { + expect(isDocumentRag('library_file_rag')).toBe(true); + }); + + test('returns false for simple_rag', () => { + expect(isDocumentRag('simple_rag')).toBe(false); + }); + + test('returns false for knowledge_store_rag', () => { + expect(isDocumentRag('knowledge_store_rag')).toBe(false); + }); + + test('returns false for no_rag', () => { + expect(isDocumentRag('no_rag')).toBe(false); + }); + + test('returns false for undefined', () => { + expect(isDocumentRag(undefined)).toBe(false); + }); + + test('returns false for null', () => { + expect(isDocumentRag(null)).toBe(false); + }); + + test('returns false for empty string', () => { + expect(isDocumentRag('')).toBe(false); + }); +}); From 4f7be1f06f5b9f173da639d0b7459d270e8a6409 Mon Sep 17 00:00:00 2001 From: Fr4n9 Date: Mon, 8 Jun 2026 18:00:13 +0200 Subject: [PATCH 074/156] fix: disable prompt processor in edit mode + clean RAG display names #2 --- .../assistants/components/ConfigurationPanel.svelte | 3 ++- frontend/svelte-app/src/lib/utils/ragProcessorHelpers.js | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte index 62106a95f..39639cca9 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/ConfigurationPanel.svelte @@ -117,7 +117,8 @@
{#each filteredRAGProcessors as processor (processor)} diff --git a/frontend/svelte-app/src/lib/components/assistants/components/KnowledgeStoreSelector.svelte b/frontend/svelte-app/src/lib/components/assistants/components/KnowledgeStoreSelector.svelte index 9ece37b9d..59c49d232 100644 --- a/frontend/svelte-app/src/lib/components/assistants/components/KnowledgeStoreSelector.svelte +++ b/frontend/svelte-app/src/lib/components/assistants/components/KnowledgeStoreSelector.svelte @@ -36,7 +36,7 @@
{#each ownedKnowledgeStores as ks (ks.id)} -
@@ -573,14 +587,22 @@ class="rounded border border-gray-300 px-2 py-1 hover:bg-gray-50 disabled:opacity-40" onclick={() => conceptPage--} disabled={conceptPage <= 1} - >← Prev - Page {conceptPage} of {conceptTotalPages} · {filteredConcepts.length} items + >{$_('knowledgeStores.graph.prev')} + {$_('knowledgeStores.graph.pageInfo', { + values: { + page: conceptPage, + total: conceptTotalPages, + items: filteredConcepts.length, + }, + })} + >{$_('knowledgeStores.graph.next')}
{/if} {/if} @@ -590,38 +612,44 @@
-

Relationships

- {relVerifiedCount}/{allRels.length} verified +

+ {$_('knowledgeStores.graph.relationships')} +

+ {$_('knowledgeStores.graph.verifiedRatio', { + values: { count: relVerifiedCount, total: allRels.length }, + })}
+ >{$_('knowledgeStores.graph.approveAll')} + >{$_('knowledgeStores.graph.rejectAll')}
{#if !allRels.length} -

No relationships yet.

+

{$_('knowledgeStores.graph.noRelationships')}

{:else if !filteredRels.length} -

No relationships match the selected filter.

+

{$_('knowledgeStores.graph.noRelationshipsFilter')}

{:else}
    {#each pagedRels as edge (edge.id)} @@ -656,7 +684,7 @@ : state === 'rejected' ? 'bg-red-100 text-red-700' : 'bg-gray-100 text-gray-600'}" - >{state} + >{$_('knowledgeStores.graph.state.' + state)} {/if} @@ -666,31 +694,31 @@ type="button" class="rounded bg-[#2271b3] px-2 py-1 text-xs text-white hover:bg-[#1a5a90]" onclick={() => commitEditEdge(edge)} - >Save + >{$_('knowledgeStores.graph.save')} + >{$_('knowledgeStores.graph.cancel')} {:else} {#if state === 'verified'} + >{$_('knowledgeStores.graph.unverify')} {:else} + >{$_('knowledgeStores.graph.verify')} {/if} + >{$_('knowledgeStores.graph.edit')} {/if} @@ -704,14 +732,18 @@ class="rounded border border-gray-300 px-2 py-1 hover:bg-gray-50 disabled:opacity-40" onclick={() => relPage--} disabled={relPage <= 1} - >← Prev - Page {relPage} of {relTotalPages} · {filteredRels.length} items + >{$_('knowledgeStores.graph.prev')} + {$_('knowledgeStores.graph.pageInfo', { + values: { page: relPage, total: relTotalPages, items: filteredRels.length }, + })} + >{$_('knowledgeStores.graph.next')} {/if} {/if} diff --git a/frontend/svelte-app/src/lib/components/knowledgeStores/SigmaGraphModal.svelte b/frontend/svelte-app/src/lib/components/knowledgeStores/SigmaGraphModal.svelte index 9468fc52b..27926f024 100644 --- a/frontend/svelte-app/src/lib/components/knowledgeStores/SigmaGraphModal.svelte +++ b/frontend/svelte-app/src/lib/components/knowledgeStores/SigmaGraphModal.svelte @@ -15,6 +15,7 @@