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 01/39] 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 02/39] 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 03/39] 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 04/39] 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 05/39] 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 06/39] 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} +
@@ -535,31 +549,31 @@ type="button" class="rounded bg-[#2271b3] px-2 py-1 text-xs text-white hover:bg-[#1a5a90]" onclick={commitEditConcept} - >Save + >{$_('knowledgeStores.graph.save')} + >{$_('knowledgeStores.graph.cancel')} {:else} {#if state === 'verified'} + >{$_('knowledgeStores.graph.unverify')} {:else} + >{$_('knowledgeStores.graph.verify')} {/if} + >{$_('knowledgeStores.graph.edit')} {/if} @@ -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 @@