From ee1d301e96735315079cfd98974c19a14d7b6084 Mon Sep 17 00:00:00 2001 From: prakartii Date: Tue, 18 Aug 2026 01:07:06 +0530 Subject: [PATCH] Implement Prompt 8 and LLM Wiki service --- .env.example | 18 +- .gitignore | 3 + docker-compose.yml | 19 + packages/sdk/src/resources.ts | 75 ++- packages/types/src/index.ts | 68 ++- services/agents/app/core/config.py | 5 + services/agents/app/core/security.py | 95 ++++ services/agents/app/main.py | 9 + services/agents/app/memory/__init__.py | 3 + services/agents/app/memory/store.py | 234 +++++++++ services/agents/app/routers/conversations.py | 61 +++ services/agents/app/routers/memory.py | 69 +++ services/agents/requirements.txt | 1 + services/agents/tests/test_prompt8_memory.py | 245 ++++++++++ services/literature/app/core/config.py | 18 +- services/literature/app/core/security.py | 20 +- .../literature/app/integrations/kg_client.py | 100 ++-- .../app/integrations/wiki_client.py | 148 +++++- services/literature/app/knowledge/__init__.py | 15 + services/literature/app/knowledge/models.py | 190 ++++++++ .../literature/app/nlp/embedding_service.py | 4 + .../literature/app/orchestrator/stages.py | 124 ++++- services/literature/app/routers/ingestion.py | 7 +- services/literature/app/services/chunking.py | 129 +++++ .../app/services/literature_service.py | 19 +- .../tests/test_prompt8_knowledge_layer.py | 257 ++++++++++ services/llm-wiki/Dockerfile | 33 ++ services/llm-wiki/app/__init__.py | 0 services/llm-wiki/app/core/__init__.py | 0 services/llm-wiki/app/core/auth.py | 34 ++ services/llm-wiki/app/core/config.py | 30 ++ services/llm-wiki/app/db/__init__.py | 0 services/llm-wiki/app/db/pool.py | 93 ++++ services/llm-wiki/app/deps.py | 14 + services/llm-wiki/app/main.py | 44 ++ services/llm-wiki/app/models.py | 98 ++++ services/llm-wiki/app/repository.py | 456 ++++++++++++++++++ services/llm-wiki/app/routers/__init__.py | 0 services/llm-wiki/app/routers/health.py | 24 + services/llm-wiki/app/routers/wiki.py | 194 ++++++++ .../llm-wiki/migrations/001_wiki_schema.sql | 104 ++++ services/llm-wiki/pytest.ini | 2 + services/llm-wiki/requirements.txt | 8 + services/llm-wiki/tests/__init__.py | 0 services/llm-wiki/tests/conftest.py | 39 ++ services/llm-wiki/tests/test_auth.py | 62 +++ .../tests/test_compile_and_retrieve.py | 113 +++++ .../tests/test_e2e_literature_roundtrip.py | 195 ++++++++ services/llm-wiki/tests/test_health.py | 12 + .../tests/test_postgres_integration.py | 174 +++++++ .../llm-wiki/tests/test_query_endpoint.py | 38 ++ .../llm-wiki/tests/test_tenant_isolation.py | 64 +++ services/llm-wiki/tests/test_validation.py | 51 ++ services/llm-wiki/tests/test_versions.py | 39 ++ services/search/cmd/search/main.go | 1 + services/search/internal/handlers/context.go | 181 +++++++ services/search/internal/handlers/index.go | 26 +- services/search/internal/handlers/search.go | 14 +- services/search/internal/handlers/stream.go | 2 +- services/search/internal/search/okf_qmd.go | 89 +++- services/search/internal/search/provider.go | 7 + .../internal/search/providers_placeholder.go | 41 +- .../internal/search/tenant_isolation_test.go | 146 ++++++ 63 files changed, 4285 insertions(+), 79 deletions(-) create mode 100644 services/agents/app/core/security.py create mode 100644 services/agents/app/memory/__init__.py create mode 100644 services/agents/app/memory/store.py create mode 100644 services/agents/app/routers/conversations.py create mode 100644 services/agents/app/routers/memory.py create mode 100644 services/agents/tests/test_prompt8_memory.py create mode 100644 services/literature/app/knowledge/__init__.py create mode 100644 services/literature/app/knowledge/models.py create mode 100644 services/literature/app/services/chunking.py create mode 100644 services/literature/tests/test_prompt8_knowledge_layer.py create mode 100644 services/llm-wiki/Dockerfile create mode 100644 services/llm-wiki/app/__init__.py create mode 100644 services/llm-wiki/app/core/__init__.py create mode 100644 services/llm-wiki/app/core/auth.py create mode 100644 services/llm-wiki/app/core/config.py create mode 100644 services/llm-wiki/app/db/__init__.py create mode 100644 services/llm-wiki/app/db/pool.py create mode 100644 services/llm-wiki/app/deps.py create mode 100644 services/llm-wiki/app/main.py create mode 100644 services/llm-wiki/app/models.py create mode 100644 services/llm-wiki/app/repository.py create mode 100644 services/llm-wiki/app/routers/__init__.py create mode 100644 services/llm-wiki/app/routers/health.py create mode 100644 services/llm-wiki/app/routers/wiki.py create mode 100644 services/llm-wiki/migrations/001_wiki_schema.sql create mode 100644 services/llm-wiki/pytest.ini create mode 100644 services/llm-wiki/requirements.txt create mode 100644 services/llm-wiki/tests/__init__.py create mode 100644 services/llm-wiki/tests/conftest.py create mode 100644 services/llm-wiki/tests/test_auth.py create mode 100644 services/llm-wiki/tests/test_compile_and_retrieve.py create mode 100644 services/llm-wiki/tests/test_e2e_literature_roundtrip.py create mode 100644 services/llm-wiki/tests/test_health.py create mode 100644 services/llm-wiki/tests/test_postgres_integration.py create mode 100644 services/llm-wiki/tests/test_query_endpoint.py create mode 100644 services/llm-wiki/tests/test_tenant_isolation.py create mode 100644 services/llm-wiki/tests/test_validation.py create mode 100644 services/llm-wiki/tests/test_versions.py create mode 100644 services/search/internal/handlers/context.go create mode 100644 services/search/internal/search/tenant_isolation_test.go diff --git a/.env.example b/.env.example index 7741fa0..1ee09ff 100644 --- a/.env.example +++ b/.env.example @@ -39,10 +39,21 @@ BETTER_AUTH_SECRET=change_this_dev_secret_before_deploying BETTER_AUTH_URL=http://localhost:8089 # ── Search retrieval provider (services/search) ───────────────────────── -# pgvector is the default and only implemented backend; llm_wiki/ -# google_okf are placeholders (see services/search/README.md). -SEARCH_RETRIEVAL_PROVIDER=pgvector +# llm_wiki is the implemented default (services/search/internal/search +# has no pgvector code path). LLM_WIKI_URL is the canonical env var for +# the LLM Wiki backend (services/llm-wiki, a persistent FastAPI+Postgres +# service implementing POST /api/v1/wiki/compile for literature's write +# path and POST /llmwiki/query for search's read path). In +# docker-compose.yml this is already pointed at the llm-wiki container by +# default; leave it blank for a non-Docker/local run and literature falls +# back to a markdown volume while search falls back to its in-memory QMD +# index instead of calling a remote service. +SEARCH_RETRIEVAL_PROVIDER=llm_wiki LLM_WIKI_URL= +# Shared bearer token literature/search send to llm-wiki and llm-wiki +# checks incoming requests against. Leave blank for local dev (llm-wiki +# then accepts unauthenticated requests and logs a startup warning) -- +# set a real value before deploying anywhere reachable by others. LLM_WIKI_API_KEY= GOOGLE_OKF_URL= GOOGLE_OKF_API_KEY= @@ -61,6 +72,7 @@ REPORTS_SERVICE_URL=http://reports:8087 DOCKING_SERVICE_URL=http://docking:8088 AI_SERVICES_URL=http://ai-services:8090 KNOWLEDGE_SERVICE_URL=http://knowledge-service:8091 +LLM_WIKI_SERVICE_URL=http://llm-wiki:8092 # ── Frontend ───────────────────────────────────────────────────────────── NEXT_PUBLIC_API_BASE_URL=http://localhost:8080 diff --git a/.gitignore b/.gitignore index 6d95589..02bcf90 100644 --- a/.gitignore +++ b/.gitignore @@ -39,6 +39,9 @@ pnpm-debug.log* .DS_Store Thumbs.db +# Claude Code local session settings +.claude/ + # coverage coverage/ .nyc_output/ diff --git a/docker-compose.yml b/docker-compose.yml index de3a895..3c39fa8 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -109,10 +109,13 @@ services: env_file: .env environment: PORT: "8084" + LLM_WIKI_URL: http://llm-wiki:8092 + LLM_WIKI_API_KEY: ${LLM_WIKI_API_KEY:-} ports: ["8084:8084"] depends_on: postgres: { condition: service_healthy } opensearch: { condition: service_healthy } + llm-wiki: { condition: service_started } networks: [ai-rxos] restart: unless-stopped @@ -147,7 +150,15 @@ services: build: { context: ., dockerfile: services/literature/Dockerfile } environment: PORT: "8082" + LLM_WIKI_URL: http://llm-wiki:8092 + LLM_WIKI_API_KEY: ${LLM_WIKI_API_KEY:-} ports: ["8082:8082"] + # Overrides (not merges with) the x-py-service anchor's depends_on, so + # postgres/redis are repeated here alongside llm-wiki. + depends_on: + postgres: { condition: service_healthy } + redis: { condition: service_healthy } + llm-wiki: { condition: service_started } kg: <<: *py-service @@ -202,6 +213,14 @@ services: depends_on: neo4j: { condition: service_healthy } + llm-wiki: + <<: *py-service + build: { context: ., dockerfile: services/llm-wiki/Dockerfile } + environment: + PORT: "8092" + LLM_WIKI_API_KEY: ${LLM_WIKI_API_KEY:-} + ports: ["8092:8092"] + # ── Frontends ──────────────────────────────────────────────────────── web: build: { context: ., dockerfile: apps/web/Dockerfile } diff --git a/packages/sdk/src/resources.ts b/packages/sdk/src/resources.ts index 148941c..10a95e7 100644 --- a/packages/sdk/src/resources.ts +++ b/packages/sdk/src/resources.ts @@ -1,19 +1,56 @@ import type { + AgentMemoryEntry, AgentTask, + ConversationMessage, DockingResult, Molecule, Paginated, Paper, Report, SearchResult, + TenantScope, } from "@ai-rxos/types"; import type { AiRxOsClient } from "./client"; /** Thin typed wrappers over api-gateway routes. See architecture/03-api-contracts.md. */ +export interface ContextItem { + id: string; + title: string; + snippet: string; + source: string; + score: number; + citation: string; +} + +export interface ContextRequest extends TenantScope { + query?: string; + embedding?: number[]; + topK?: number; + sourceFilters?: string[]; + entityFilters?: string[]; + conversationId?: string; + agentId?: string; +} + export const search = (client: AiRxOsClient) => ({ query: (q: string, limit = 20) => client.get>(`/api/v1/search?q=${encodeURIComponent(q)}&limit=${limit}`), + /** Prompt 8 context-optimization endpoint: compact, cited context instead + * of full documents. See services/search/internal/handlers/context.go. */ + context: (req: ContextRequest) => + client.post<{ items: ContextItem[]; total: number }>(`/api/v1/search/context`, { + query: req.query, + embedding: req.embedding, + top_k: req.topK, + source_filters: req.sourceFilters, + entity_filters: req.entityFilters, + organization_id: req.organizationId, + workspace_id: req.workspaceId, + project_id: req.projectId, + conversation_id: req.conversationId, + agent_id: req.agentId, + }), }); export const literature = (client: AiRxOsClient) => ({ @@ -28,11 +65,47 @@ export const molecules = (client: AiRxOsClient) => ({ }); export const agents = (client: AiRxOsClient) => ({ + // services/agents exposes POST /api/v1/agents/invoke (see + // services/agents/app/main.py) — not /run, which was a pre-existing SDK/ + // service mismatch fixed alongside the memory/conversation additions below. run: (agentType: string, input: Record) => - client.post(`/api/v1/agents/run`, { agentType, input }), + client.post(`/api/v1/agents/invoke`, { agentType, input }), get: (id: string) => client.get(`/api/v1/agents/tasks/${id}`), }); +/** Agent memory (Prompt 8): store/retrieve/search agent-scoped memory. + * See services/agents/app/routers/memory.py. Tenant scoping is derived + * server-side from the caller's auth token, not from these arguments. */ +export const agentMemory = (client: AiRxOsClient) => ({ + store: (agentId: string, key: string, value: unknown, opts?: { provenance?: Record; persistLongTerm?: boolean }) => + client.post(`/api/v1/agents/memory`, { + agent_id: agentId, + key, + value, + provenance: opts?.provenance ?? {}, + persist_long_term: opts?.persistLongTerm ?? false, + }), + retrieve: (agentId: string, key: string) => + client.get(`/api/v1/agents/memory/${agentId}/${encodeURIComponent(key)}`), + search: (agentId: string, query?: string, limit = 10) => + client.get<{ items: AgentMemoryEntry[]; total: number }>( + `/api/v1/agents/memory/${agentId}?${query ? `query=${encodeURIComponent(query)}&` : ""}limit=${limit}`, + ), +}); + +/** Conversation memory (Prompt 8). See services/agents/app/routers/conversations.py. */ +export const conversations = (client: AiRxOsClient) => ({ + addMessage: (conversationId: string, message: Omit) => + client.post<{ messages: ConversationMessage[] }>( + `/api/v1/agents/conversations/${conversationId}/messages`, + message, + ), + getMessages: (conversationId: string, limit?: number) => + client.get<{ conversationId: string; messages: ConversationMessage[]; total: number }>( + `/api/v1/agents/conversations/${conversationId}/messages${limit ? `?limit=${limit}` : ""}`, + ), +}); + export const reports = (client: AiRxOsClient) => ({ list: (page = 1) => client.get>(`/api/v1/reports?page=${page}`), generate: (title: string, type: Report["type"]) => diff --git a/packages/types/src/index.ts b/packages/types/src/index.ts index 8c8a74e..cdffc2c 100644 --- a/packages/types/src/index.ts +++ b/packages/types/src/index.ts @@ -70,12 +70,78 @@ export type AgentTask = z.infer; export const SearchResultSchema = z.object({ id: z.string(), score: z.number(), - source: z.enum(["opensearch", "pgvector", "graph"]), + // "pgvector" is retained for backward compatibility with existing + // callers; services/search's actual retrieval providers are llm_wiki + // (default) and google_okf — see services/search/internal/search/provider.go. + source: z.enum(["opensearch", "pgvector", "llm_wiki", "google_okf", "graph"]), title: z.string(), snippet: z.string().optional(), }); export type SearchResult = z.infer; +/** Organization/workspace/project scope, mirroring packages/tenancy's + * TENANT_ID_CLAIM convention and services/auth's AuthPayload shape. */ +export const TenantScopeSchema = z.object({ + organizationId: z.string().optional(), + workspaceId: z.string().optional(), + projectId: z.string().optional(), +}); +export type TenantScope = z.infer; + +/** Canonical metadata attached to every LLM Wiki-indexed document/chunk. */ +export const KnowledgeMetadataSchema = z.object({ + documentId: z.string(), + sourceType: z.enum([ + "paper", + "conference_abstract", + "clinical_trial", + "patent", + "company", + "drug_pipeline", + "kg_derived", + "agent_memory", + "conversation", + ]), + sourceId: z.string(), + title: z.string(), + entityIds: z.array(z.string()).default([]), + entityTypes: z.array(z.string()).default([]), + version: z.number().int().positive().default(1), + createdAt: z.string().datetime().optional(), + updatedAt: z.string().datetime().optional(), + provenance: z.record(z.string(), z.unknown()).default({}), + citation: z.record(z.string(), z.unknown()).default({}), +}).merge(TenantScopeSchema); +export type KnowledgeMetadata = z.infer; + +/** A chunk of an indexed document, traceable back to its source and tenant. */ +export const ChunkSchema = z.object({ + chunkId: z.string(), + chunkIndex: z.number().int().nonnegative(), + text: z.string(), + metadata: KnowledgeMetadataSchema, +}); +export type Chunk = z.infer; + +/** A single entry in an agent's scoped memory (services/agents). */ +export const AgentMemoryEntrySchema = z.object({ + agentId: z.string(), + key: z.string(), + value: z.unknown(), + provenance: z.record(z.string(), z.unknown()).default({}), + storedAt: z.number().optional(), +}).merge(TenantScopeSchema); +export type AgentMemoryEntry = z.infer; + +/** A single conversation turn (services/agents conversation memory). */ +export const ConversationMessageSchema = z.object({ + role: z.enum(["user", "assistant", "system", "tool"]), + content: z.string(), + metadata: z.record(z.string(), z.unknown()).default({}), + createdAt: z.number().optional(), +}); +export type ConversationMessage = z.infer; + export const ReportSchema = z.object({ id: z.string().uuid(), title: z.string(), diff --git a/services/agents/app/core/config.py b/services/agents/app/core/config.py index fb4f071..d74136f 100644 --- a/services/agents/app/core/config.py +++ b/services/agents/app/core/config.py @@ -17,6 +17,11 @@ class Settings(BaseSettings): opensearch_url: str = "http://opensearch:9200" jwt_secret: str = "change_this_dev_secret_before_deploying" + # Canonical LLM Wiki URL (see root .env.example / services/search). When + # unset, long-term agent memory persistence is skipped rather than + # pointed at a service that doesn't exist in this repo's docker-compose. + llm_wiki_url: str | None = None + @lru_cache def get_settings() -> Settings: diff --git a/services/agents/app/core/security.py b/services/agents/app/core/security.py new file mode 100644 index 0000000..2e412fc --- /dev/null +++ b/services/agents/app/core/security.py @@ -0,0 +1,95 @@ +"""JWT verification and tenant-scope extraction for the agents service. + +services/agents had no auth wiring at all prior to this (every endpoint was +unauthenticated). This mirrors services/literature's app/core/security.py +pattern (same JWT secret convention, same claim names) rather than inventing +a third auth shape — see services/auth/src/tenantContext.ts::AuthPayload for +the canonical {sub, organizationId, roles} claim shape this aligns with. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import jwt +from fastapi import Depends, HTTPException, Security, status +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer + +from app.core.config import get_settings + +security = HTTPBearer(auto_error=False) +security_dependency = Security(security) + + +@dataclass(frozen=True) +class TenantContext: + organization_id: str | None = None + workspace_id: str | None = None + project_id: str | None = None + user_id: str | None = None + + def as_dict(self) -> dict[str, str]: + return { + k: v + for k, v in { + "organization_id": self.organization_id, + "workspace_id": self.workspace_id, + "project_id": self.project_id, + "user_id": self.user_id, + }.items() + if v + } + + +def get_current_user( + credentials: HTTPAuthorizationCredentials | None = security_dependency, +) -> dict[str, str]: + settings = get_settings() + if settings.environment == "test" and credentials is None: + return {"sub": "test-user-fallback"} + + if credentials is None or credentials.scheme.lower() != "bearer": + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="authentication required", + ) + + try: + payload = jwt.decode( + credentials.credentials, settings.jwt_secret, algorithms=["HS256"] + ) + except jwt.PyJWTError as exc: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="invalid authentication token", + ) from exc + + if not isinstance(payload, dict): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="invalid authentication payload", + ) + return payload + + +current_user_dependency = Depends(get_current_user) + + +def get_tenant_context( + auth_payload: dict[str, str] = current_user_dependency, +) -> TenantContext: + """Derive the caller's organization/workspace/project scope from the JWT. + + This is the only source of truth for scoping memory/conversation reads + and writes — callers must never be able to pass an organization_id/ + workspace_id in a request body and have it override this, or one + tenant could read another's agent/conversation memory. + """ + return TenantContext( + organization_id=auth_payload.get("organization_id") + or auth_payload.get("organizationId"), + workspace_id=auth_payload.get("workspace_id") + or auth_payload.get("workspaceId"), + project_id=auth_payload.get("project_id") or auth_payload.get("projectId"), + user_id=auth_payload.get("user_id") or auth_payload.get("sub"), + ) diff --git a/services/agents/app/main.py b/services/agents/app/main.py index 2dd7dd4..a1bf2b2 100644 --- a/services/agents/app/main.py +++ b/services/agents/app/main.py @@ -7,6 +7,9 @@ from pydantic import BaseModel from app.core.config import get_settings +from app.memory.store import AgentMemoryStore, ConversationMemoryStore +from app.routers import conversations as conversations_router +from app.routers import memory as memory_router settings = get_settings() app = FastAPI( @@ -19,6 +22,12 @@ _redis = redis.from_url(settings.redis_url, decode_responses=True) TASK_KEY = "agents:task:{id}" +agent_memory_store = AgentMemoryStore(_redis) +conversation_memory_store = ConversationMemoryStore(_redis) + +app.include_router(memory_router.router) +app.include_router(conversations_router.router) + class ToolInvocation(BaseModel): tool: str diff --git a/services/agents/app/memory/__init__.py b/services/agents/app/memory/__init__.py new file mode 100644 index 0000000..c8fa7a8 --- /dev/null +++ b/services/agents/app/memory/__init__.py @@ -0,0 +1,3 @@ +from app.memory.store import AgentMemoryStore, ConversationMemoryStore + +__all__ = ["AgentMemoryStore", "ConversationMemoryStore"] diff --git a/services/agents/app/memory/store.py b/services/agents/app/memory/store.py new file mode 100644 index 0000000..54ea9f4 --- /dev/null +++ b/services/agents/app/memory/store.py @@ -0,0 +1,234 @@ +"""Agent memory and conversation memory backends. + +No memory or conversation implementation existed anywhere in the repo +before this (services/agents was a ~65-line stub that only enqueued +AgentTask records to Redis — see app/main.py's TASK_KEY convention, which +this module follows). This deliberately does not implement the full +Postgres agents/conversations/messages schema sketched in +architecture/04-database-schemas.md — only the reusable Redis-backed +memory infrastructure future agents need, matching the existing +TASK_KEY = "agents:task:{id}" pattern and the architecture doc's own +`agent:context:{conversation_id}` Redis hash convention for conversation +state. + +Working/short-term memory lives in Redis with a TTL. "Long-term" durability +is delegated to LLM Wiki (see LLMWikiMemoryClient) exactly as literature's +wiki_client.py does: if LLM_WIKI_URL isn't configured, writes are recorded +as skipped rather than pretending a production LLM Wiki service exists. +""" + +from __future__ import annotations + +import json +import time +from typing import Any + +import httpx +import redis.asyncio as redis + +from app.core.config import get_settings +from app.core.security import TenantContext + +AGENT_MEMORY_KEY = "agent:memory:{org}:{workspace}:{agent_id}:{key}" +AGENT_MEMORY_INDEX_KEY = "agent:memory:index:{org}:{workspace}:{agent_id}" +CONVERSATION_KEY = "agent:context:{conversation_id}" + +_NO_ORG = "_none" +_NO_WORKSPACE = "_shared" + + +def _scope(tenant: TenantContext) -> tuple[str, str]: + return (tenant.organization_id or _NO_ORG, tenant.workspace_id or _NO_WORKSPACE) + + +class LLMWikiMemoryClient: + """Optional long-term persistence of memory entries through LLM Wiki. + + Mirrors services/literature/app/integrations/wiki_client.py's HTTP + contract (POST {url}/api/v1/wiki/compile) rather than inventing a new + one. When LLM_WIKI_URL is not configured — the default, since no real + LLM Wiki service is deployed anywhere in this repo — writes are + recorded as skipped instead of silently pretending to succeed. + """ + + def __init__(self, base_url: str | None = None, timeout_seconds: float = 5.0): + self.base_url = base_url + self.timeout_seconds = timeout_seconds + + def persist( + self, + *, + source_type: str, + source_id: str, + title: str, + text: str, + tenant: TenantContext, + ) -> dict[str, Any]: + if not self.base_url: + return {"success": True, "status": "skipped", "reason": "LLM_WIKI_URL not configured"} + try: + with httpx.Client(timeout=self.timeout_seconds) as client: + res = client.post( + f"{self.base_url.rstrip('/')}/api/v1/wiki/compile", + json={ + "document": {"source": source_type, "source_id": source_id, "title": title}, + "entities": [{"text": title, "category": source_type}], + "summary": {"concise_summary": text}, + "tenant": tenant.as_dict(), + }, + ) + if res.status_code in (200, 201): + return {"success": True, "status": "completed"} + return {"success": False, "status": "failed", "error": res.text} + except httpx.HTTPError as exc: + return {"success": False, "status": "failed", "error": str(exc)} + + +class AgentMemoryStore: + """Redis-backed agent memory scoped by organization/workspace/agent. + + Scoping comes exclusively from the caller-supplied TenantContext + (derived from the verified JWT — see app.core.security.get_tenant_context), + never from client-suppliable request fields, so one organization can + never read another's agent memory by guessing keys. + """ + + def __init__(self, redis_client: redis.Redis, ttl_seconds: int = 60 * 60 * 24 * 30): + self.redis = redis_client + self.ttl_seconds = ttl_seconds + settings = get_settings() + self.long_term = LLMWikiMemoryClient(getattr(settings, "llm_wiki_url", None)) + + async def store( + self, + *, + tenant: TenantContext, + agent_id: str, + key: str, + value: Any, + provenance: dict[str, Any] | None = None, + persist_long_term: bool = False, + ) -> dict[str, Any]: + org, workspace = _scope(tenant) + record = { + "organization_id": tenant.organization_id, + "workspace_id": tenant.workspace_id, + "project_id": tenant.project_id, + "agent_id": agent_id, + "key": key, + "value": value, + "provenance": provenance or {}, + "stored_at": time.time(), + } + scoped_key = AGENT_MEMORY_KEY.format(org=org, workspace=workspace, agent_id=agent_id, key=key) + await self.redis.set(scoped_key, json.dumps(record), ex=self.ttl_seconds) + + index_key = AGENT_MEMORY_INDEX_KEY.format(org=org, workspace=workspace, agent_id=agent_id) + await self.redis.sadd(index_key, key) + await self.redis.expire(index_key, self.ttl_seconds) + + if persist_long_term: + record["long_term"] = self.long_term.persist( + source_type="agent_memory", + source_id=f"{agent_id}:{key}", + title=key, + text=json.dumps(value) if not isinstance(value, str) else value, + tenant=tenant, + ) + return record + + async def retrieve(self, *, tenant: TenantContext, agent_id: str, key: str) -> dict[str, Any] | None: + org, workspace = _scope(tenant) + scoped_key = AGENT_MEMORY_KEY.format(org=org, workspace=workspace, agent_id=agent_id, key=key) + raw = await self.redis.get(scoped_key) + return json.loads(raw) if raw else None + + async def search( + self, *, tenant: TenantContext, agent_id: str, query: str | None = None, limit: int = 10 + ) -> list[dict[str, Any]]: + org, workspace = _scope(tenant) + index_key = AGENT_MEMORY_INDEX_KEY.format(org=org, workspace=workspace, agent_id=agent_id) + keys = await self.redis.smembers(index_key) + records: list[dict[str, Any]] = [] + for raw_key in keys: + scoped_key = AGENT_MEMORY_KEY.format(org=org, workspace=workspace, agent_id=agent_id, key=raw_key) + raw = await self.redis.get(scoped_key) + if not raw: + continue + record = json.loads(raw) + if query and query.lower() not in json.dumps(record.get("value", "")).lower(): + continue + records.append(record) + records.sort(key=lambda r: r.get("stored_at", 0), reverse=True) + return records[:limit] + + +class ConversationMemoryStore: + """Redis-backed conversation memory, keyed agent:context:{conversation_id} + per architecture/04-database-schemas.md's Redis convention. + + The tenant that first creates a conversation "owns" it; every + subsequent read/write must present a matching TenantContext or the + operation is treated as not-found (never as "forbidden", so a caller + probing conversation ids can't distinguish "wrong tenant" from + "doesn't exist"). + """ + + def __init__(self, redis_client: redis.Redis, ttl_seconds: int = 60 * 60 * 24, max_messages: int = 100): + self.redis = redis_client + self.ttl_seconds = ttl_seconds + self.max_messages = max_messages + + @staticmethod + def _key(conversation_id: str) -> str: + return CONVERSATION_KEY.format(conversation_id=conversation_id) + + async def _load(self, conversation_id: str) -> dict[str, Any] | None: + raw = await self.redis.get(self._key(conversation_id)) + return json.loads(raw) if raw else None + + @staticmethod + def _owns(record: dict[str, Any], tenant: TenantContext) -> bool: + owner = record.get("tenant") or {} + return owner.get("organization_id") == tenant.organization_id and owner.get( + "workspace_id" + ) == tenant.workspace_id + + async def add_message( + self, + *, + tenant: TenantContext, + conversation_id: str, + role: str, + content: str, + metadata: dict[str, Any] | None = None, + ) -> dict[str, Any] | None: + record = await self._load(conversation_id) + if record is None: + record = {"tenant": tenant.as_dict(), "messages": []} + elif not self._owns(record, tenant): + return None + + record["messages"].append( + { + "role": role, + "content": content, + "metadata": metadata or {}, + "created_at": time.time(), + "user_id": tenant.user_id, + } + ) + if len(record["messages"]) > self.max_messages: + record["messages"] = record["messages"][-self.max_messages :] + + await self.redis.set(self._key(conversation_id), json.dumps(record), ex=self.ttl_seconds) + return record + + async def get_messages( + self, *, tenant: TenantContext, conversation_id: str, limit: int | None = None + ) -> list[dict[str, Any]] | None: + record = await self._load(conversation_id) + if record is None or not self._owns(record, tenant): + return None + messages = record.get("messages", []) + return messages[-limit:] if limit else messages diff --git a/services/agents/app/routers/conversations.py b/services/agents/app/routers/conversations.py new file mode 100644 index 0000000..b950b0a --- /dev/null +++ b/services/agents/app/routers/conversations.py @@ -0,0 +1,61 @@ +from typing import Any + +from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel + +from app.core.security import TenantContext, get_tenant_context +from app.memory.store import ConversationMemoryStore + +router = APIRouter(prefix="/api/v1/agents/conversations", tags=["Conversation Memory"]) + +tenant_dependency = Depends(get_tenant_context) + + +class AddMessageRequest(BaseModel): + role: str + content: str + metadata: dict[str, Any] = {} + + +def get_conversation_store() -> ConversationMemoryStore: + from app.main import conversation_memory_store + + return conversation_memory_store + + +conversation_store_dependency = Depends(get_conversation_store) + + +@router.post("/{conversation_id}/messages", status_code=201) +async def add_message( + conversation_id: str, + req: AddMessageRequest, + tenant: TenantContext = tenant_dependency, + store: ConversationMemoryStore = conversation_store_dependency, +) -> dict[str, Any]: + record = await store.add_message( + tenant=tenant, + conversation_id=conversation_id, + role=req.role, + content=req.content, + metadata=req.metadata, + ) + if record is None: + # Conversation exists but belongs to a different tenant — reported + # as not-found rather than forbidden, so probing conversation ids + # can't be used to distinguish "wrong tenant" from "doesn't exist". + raise HTTPException(status_code=404, detail="conversation not found") + return record + + +@router.get("/{conversation_id}/messages") +async def get_messages( + conversation_id: str, + limit: int | None = None, + tenant: TenantContext = tenant_dependency, + store: ConversationMemoryStore = conversation_store_dependency, +) -> dict[str, Any]: + messages = await store.get_messages(tenant=tenant, conversation_id=conversation_id, limit=limit) + if messages is None: + raise HTTPException(status_code=404, detail="conversation not found") + return {"conversation_id": conversation_id, "messages": messages, "total": len(messages)} diff --git a/services/agents/app/routers/memory.py b/services/agents/app/routers/memory.py new file mode 100644 index 0000000..c8f18f5 --- /dev/null +++ b/services/agents/app/routers/memory.py @@ -0,0 +1,69 @@ +from typing import Any + +from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel + +from app.core.security import TenantContext, get_tenant_context +from app.memory.store import AgentMemoryStore + +router = APIRouter(prefix="/api/v1/agents/memory", tags=["Agent Memory"]) + +tenant_dependency = Depends(get_tenant_context) + + +class StoreMemoryRequest(BaseModel): + agent_id: str + key: str + value: Any + provenance: dict[str, Any] = {} + persist_long_term: bool = False + + +def get_memory_store() -> AgentMemoryStore: + from app.main import agent_memory_store + + return agent_memory_store + + +memory_store_dependency = Depends(get_memory_store) + + +@router.post("", status_code=201) +async def store_memory( + req: StoreMemoryRequest, + tenant: TenantContext = tenant_dependency, + store: AgentMemoryStore = memory_store_dependency, +) -> dict[str, Any]: + return await store.store( + tenant=tenant, + agent_id=req.agent_id, + key=req.key, + value=req.value, + provenance=req.provenance, + persist_long_term=req.persist_long_term, + ) + + +@router.get("/{agent_id}/{key}") +async def retrieve_memory( + agent_id: str, + key: str, + tenant: TenantContext = tenant_dependency, + store: AgentMemoryStore = memory_store_dependency, +) -> dict[str, Any]: + record = await store.retrieve(tenant=tenant, agent_id=agent_id, key=key) + if record is None: + raise HTTPException(status_code=404, detail="memory entry not found") + return record + + +@router.get("/{agent_id}") +async def search_memory( + agent_id: str, + query: str | None = None, + limit: int = 10, + tenant: TenantContext = tenant_dependency, + store: AgentMemoryStore = memory_store_dependency, +) -> dict[str, Any]: + records = await store.search(tenant=tenant, agent_id=agent_id, query=query, limit=limit) + return {"items": records, "total": len(records)} diff --git a/services/agents/requirements.txt b/services/agents/requirements.txt index db2b2dd..9848252 100644 --- a/services/agents/requirements.txt +++ b/services/agents/requirements.txt @@ -5,5 +5,6 @@ pydantic-settings==2.7.1 redis==5.2.1 httpx==0.28.1 python-json-logger==3.2.1 +pyjwt==2.10.1 pytest==8.3.4 pytest-asyncio==0.25.1 diff --git a/services/agents/tests/test_prompt8_memory.py b/services/agents/tests/test_prompt8_memory.py new file mode 100644 index 0000000..1fd9e27 --- /dev/null +++ b/services/agents/tests/test_prompt8_memory.py @@ -0,0 +1,245 @@ +"""Tests for Prompt 8 agent memory + conversation memory: unit tests +against a fake in-memory Redis (no real Redis required), plus API-level +tests through FastAPI's TestClient proving cross-organization and +cross-workspace isolation and basic request validation. +""" + +from __future__ import annotations + +import jwt +import pytest + +from app.core.config import get_settings +from app.core.security import TenantContext +from app.memory.store import AgentMemoryStore, ConversationMemoryStore + + +class FakeRedis: + """Minimal in-memory stand-in for redis.asyncio.Redis, covering only + the operations AgentMemoryStore/ConversationMemoryStore use.""" + + def __init__(self) -> None: + self._values: dict[str, str] = {} + self._sets: dict[str, set[str]] = {} + + async def set(self, key: str, value: str, ex: int | None = None) -> None: + self._values[key] = value + + async def get(self, key: str) -> str | None: + return self._values.get(key) + + async def sadd(self, key: str, value: str) -> None: + self._sets.setdefault(key, set()).add(value) + + async def expire(self, key: str, ttl: int) -> None: + pass + + async def smembers(self, key: str) -> set[str]: + return self._sets.get(key, set()) + + +ORG_A = TenantContext(organization_id="org-a", workspace_id="ws-1", user_id="user-1") +ORG_A_WS2 = TenantContext(organization_id="org-a", workspace_id="ws-2", user_id="user-2") +ORG_B = TenantContext(organization_id="org-b", workspace_id="ws-1", user_id="user-3") + + +# --------------------------------------------------------------------------- +# AgentMemoryStore +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_agent_memory_store_and_retrieve_roundtrip(): + store = AgentMemoryStore(FakeRedis()) + await store.store(tenant=ORG_A, agent_id="literature-agent", key="last_query", value={"q": "HER2"}) + + record = await store.retrieve(tenant=ORG_A, agent_id="literature-agent", key="last_query") + assert record is not None + assert record["value"] == {"q": "HER2"} + assert record["organization_id"] == "org-a" + assert record["workspace_id"] == "ws-1" + + +@pytest.mark.asyncio +async def test_agent_memory_search_filters_by_query_substring(): + store = AgentMemoryStore(FakeRedis()) + await store.store(tenant=ORG_A, agent_id="a1", key="k1", value="HER2 targeted therapy") + await store.store(tenant=ORG_A, agent_id="a1", key="k2", value="unrelated content") + + results = await store.search(tenant=ORG_A, agent_id="a1", query="her2") + assert len(results) == 1 + assert results[0]["key"] == "k1" + + +@pytest.mark.asyncio +async def test_agent_memory_is_isolated_across_organizations(): + store = AgentMemoryStore(FakeRedis()) + await store.store(tenant=ORG_A, agent_id="a1", key="secret", value="org-a-data") + + # Same agent_id/key, different organization -> nothing visible. + record = await store.retrieve(tenant=ORG_B, agent_id="a1", key="secret") + assert record is None + + results = await store.search(tenant=ORG_B, agent_id="a1") + assert results == [] + + +@pytest.mark.asyncio +async def test_agent_memory_is_isolated_across_workspaces_in_same_org(): + store = AgentMemoryStore(FakeRedis()) + await store.store(tenant=ORG_A, agent_id="a1", key="secret", value="ws-1-data") + + record = await store.retrieve(tenant=ORG_A_WS2, agent_id="a1", key="secret") + assert record is None + + +@pytest.mark.asyncio +async def test_agent_memory_long_term_persist_is_skipped_without_llm_wiki_url(): + store = AgentMemoryStore(FakeRedis()) + record = await store.store( + tenant=ORG_A, agent_id="a1", key="k1", value="v1", persist_long_term=True + ) + assert record["long_term"]["status"] == "skipped" + assert record["long_term"]["success"] is True + + +# --------------------------------------------------------------------------- +# ConversationMemoryStore +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_conversation_memory_add_and_get_messages(): + store = ConversationMemoryStore(FakeRedis()) + await store.add_message(tenant=ORG_A, conversation_id="conv-1", role="user", content="hello") + await store.add_message(tenant=ORG_A, conversation_id="conv-1", role="assistant", content="hi there") + + messages = await store.get_messages(tenant=ORG_A, conversation_id="conv-1") + assert messages is not None + assert [m["role"] for m in messages] == ["user", "assistant"] + assert messages[0]["content"] == "hello" + + +@pytest.mark.asyncio +async def test_conversation_memory_trims_to_max_messages(): + store = ConversationMemoryStore(FakeRedis(), max_messages=3) + for i in range(5): + await store.add_message(tenant=ORG_A, conversation_id="conv-1", role="user", content=f"msg-{i}") + + messages = await store.get_messages(tenant=ORG_A, conversation_id="conv-1") + assert len(messages) == 3 + assert [m["content"] for m in messages] == ["msg-2", "msg-3", "msg-4"] + + +@pytest.mark.asyncio +async def test_conversation_memory_is_isolated_across_organizations(): + store = ConversationMemoryStore(FakeRedis()) + await store.add_message(tenant=ORG_A, conversation_id="conv-1", role="user", content="org-a-secret") + + # A different org writing to the same conversation id is rejected... + result = await store.add_message(tenant=ORG_B, conversation_id="conv-1", role="user", content="hijack") + assert result is None + + # ...and cannot read it either. + messages = await store.get_messages(tenant=ORG_B, conversation_id="conv-1") + assert messages is None + + # The original organization's data is untouched. + owner_messages = await store.get_messages(tenant=ORG_A, conversation_id="conv-1") + assert len(owner_messages) == 1 + assert owner_messages[0]["content"] == "org-a-secret" + + +# --------------------------------------------------------------------------- +# API-level: auth, validation, and cross-tenant isolation through the routes +# --------------------------------------------------------------------------- + + +def _make_token(organization_id: str, workspace_id: str, user_id: str) -> str: + settings = get_settings() + return jwt.encode( + {"sub": user_id, "organization_id": organization_id, "workspace_id": workspace_id}, + settings.jwt_secret, + algorithm="HS256", + ) + + +@pytest.fixture +def api_client(monkeypatch): + from fastapi.testclient import TestClient + + import app.main as main_module + + fake_redis = FakeRedis() + monkeypatch.setattr(main_module, "agent_memory_store", AgentMemoryStore(fake_redis)) + monkeypatch.setattr(main_module, "conversation_memory_store", ConversationMemoryStore(fake_redis)) + return TestClient(main_module.app) + + +def _auth_headers(organization_id: str, workspace_id: str = "ws-1", user_id: str = "user-1") -> dict[str, str]: + return {"Authorization": f"Bearer {_make_token(organization_id, workspace_id, user_id)}"} + + +def test_memory_api_requires_authentication(api_client): + res = api_client.post("/api/v1/agents/memory", json={"agent_id": "a1", "key": "k1", "value": "v1"}) + assert res.status_code == 401 + + +def test_memory_api_store_and_retrieve(api_client): + headers = _auth_headers("org-a") + res = api_client.post( + "/api/v1/agents/memory", + json={"agent_id": "a1", "key": "k1", "value": {"score": 0.9}}, + headers=headers, + ) + assert res.status_code == 201 + + res = api_client.get("/api/v1/agents/memory/a1/k1", headers=headers) + assert res.status_code == 200 + assert res.json()["value"] == {"score": 0.9} + + +def test_memory_api_cross_organization_read_returns_404(api_client): + api_client.post( + "/api/v1/agents/memory", + json={"agent_id": "a1", "key": "k1", "value": "org-a-only"}, + headers=_auth_headers("org-a"), + ) + + res = api_client.get("/api/v1/agents/memory/a1/k1", headers=_auth_headers("org-b")) + assert res.status_code == 404 + + +def test_conversation_api_cross_organization_hijack_returns_404(api_client): + headers_a = _auth_headers("org-a") + headers_b = _auth_headers("org-b") + + res = api_client.post( + "/api/v1/agents/conversations/conv-1/messages", + json={"role": "user", "content": "org-a-secret"}, + headers=headers_a, + ) + assert res.status_code == 201 + + res = api_client.post( + "/api/v1/agents/conversations/conv-1/messages", + json={"role": "user", "content": "hijack-attempt"}, + headers=headers_b, + ) + assert res.status_code == 404 + + res = api_client.get("/api/v1/agents/conversations/conv-1/messages", headers=headers_b) + assert res.status_code == 404 + + res = api_client.get("/api/v1/agents/conversations/conv-1/messages", headers=headers_a) + assert res.status_code == 200 + assert res.json()["total"] == 1 + + +def test_memory_api_rejects_malformed_body(api_client): + res = api_client.post( + "/api/v1/agents/memory", + json={"agent_id": "a1"}, # missing required "key"/"value" + headers=_auth_headers("org-a"), + ) + assert res.status_code == 422 diff --git a/services/literature/app/core/config.py b/services/literature/app/core/config.py index b0cc54c..78893ce 100644 --- a/services/literature/app/core/config.py +++ b/services/literature/app/core/config.py @@ -1,5 +1,6 @@ from functools import lru_cache +from pydantic import AliasChoices, Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -17,8 +18,14 @@ class Settings(BaseSettings): opensearch_url: str = "http://opensearch:9200" jwt_secret: str = "change_this_dev_secret_before_deploying" - kg_service_url: str = "http://kg:8000" - okf_wiki_url: str | None = None + kg_service_url: str = "http://kg:8083" + # LLM_WIKI_URL is the repo-wide canonical env var (see root .env.example + # and services/search); OKF_WIKI_URL is kept as a fallback alias so + # existing literature-only deployments don't break. + okf_wiki_url: str | None = Field( + default=None, + validation_alias=AliasChoices("LLM_WIKI_URL", "OKF_WIKI_URL"), + ) okf_wiki_dir: str = "wiki-root" llm_provider: str | None = None @@ -37,7 +44,10 @@ class Settings(BaseSettings): kg_max_retries: int = 1 kg_backoff_seconds: float = 0.25 - wiki_api_key: str | None = None + wiki_api_key: str | None = Field( + default=None, + validation_alias=AliasChoices("LLM_WIKI_API_KEY", "WIKI_API_KEY"), + ) wiki_timeout: float = 5.0 wiki_max_retries: int = 1 wiki_backoff_seconds: float = 0.25 @@ -59,7 +69,7 @@ class Settings(BaseSettings): medrxiv_base_url: str = "https://api.biorxiv.org/details/medrxiv" # Services URLs and parameters - search_service_url: str = "http://search:8000" + search_service_url: str = "http://search:8084" kg_service_timeout_seconds: float = 5.0 kg_service_max_retries: int = 3 llmwiki_service_url: str = "http://wiki:8000" diff --git a/services/literature/app/core/security.py b/services/literature/app/core/security.py index 79b4578..c62aeed 100644 --- a/services/literature/app/core/security.py +++ b/services/literature/app/core/security.py @@ -1,8 +1,9 @@ import jwt -from fastapi import HTTPException, Security, status +from fastapi import Depends, HTTPException, Security, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from app.core.config import get_settings +from app.knowledge.models import TenantContext settings = get_settings() security = HTTPBearer(auto_error=False) @@ -37,3 +38,20 @@ def get_current_user( ) return payload + + +current_user_dependency = Depends(get_current_user) + + +def get_tenant_context( + auth_payload: dict[str, str] = current_user_dependency, +) -> TenantContext: + """Derive the caller's organization/workspace/project scope from the JWT. + + Mirrors services/auth's legacy claim shape (`organizationId`, `roles` — + see services/auth/src/tenantContext.ts::AuthPayload). workspaceId/ + projectId are not yet issued by the auth service's JWTs (see the auth + audit), so those fields are simply absent/None until that's added — + this function does not invent claims that aren't there. + """ + return TenantContext.from_claims(auth_payload) diff --git a/services/literature/app/integrations/kg_client.py b/services/literature/app/integrations/kg_client.py index 87c51d9..bd42fcf 100644 --- a/services/literature/app/integrations/kg_client.py +++ b/services/literature/app/integrations/kg_client.py @@ -6,6 +6,11 @@ import httpx +from app.knowledge.models import ( + deterministic_entity_id, + normalize_entity_label, + normalize_relationship_type, +) from app.observability.metrics import metrics logger = logging.getLogger(__name__) @@ -16,47 +21,85 @@ class KGClient: def __init__(self, config: dict[str, Any] | None = None): self.config = config or {} - self.base_url = self.config.get("kg_service_url", "http://localhost:8001").rstrip("/") + self.base_url = self.config.get("kg_service_url", "http://kg:8083").rstrip("/") self.timeout = float(self.config.get("kg_timeout", 5.0)) self.max_retries = int(self.config.get("kg_max_retries", 1)) self.backoff_seconds = float(self.config.get("kg_backoff_seconds", 0.25)) def update_knowledge_graph(self, entities: list[dict[str, Any]], relationships: list[dict[str, Any]]) -> dict[str, Any]: - """Convert extracted entities and relationships into KG update operations.""" - nodes_payload = [ - { - "label": e.get("label", "Entity").capitalize(), - "properties": { - "name": e.get("text"), - "category": e.get("category"), + """Convert extracted entities and relationships into a KG import request. + + Payload shape matches services/kg/app/schemas/imports.py::ImportJSONRequest + exactly (flat NodeCreate/RelationshipCreate, not a nested "properties" + envelope). Entity/relationship ids are derived deterministically (see + app.knowledge.models.deterministic_entity_id) so the same node id can + be referenced later from LLM Wiki chunk metadata without a round trip + to this service. Entities/relationships whose type isn't recognized by + the KG's label/relationship-type vocabulary are skipped rather than + sent and rejected as a batch (the KG's schema validation is atomic + across the whole request body). + """ + entity_id_by_text: dict[str, str] = {} + nodes_payload: list[dict[str, Any]] = [] + for e in entities: + text = e.get("text") + category = e.get("type") or e.get("category") + if not text or not category: + continue + label = normalize_entity_label(category) + if not label: + continue + entity_id = deterministic_entity_id(category, text) + entity_id_by_text[text.strip().lower()] = entity_id + nodes_payload.append( + { + "id": entity_id, + "label": label, + "name": text, "source": "literature_service", - }, - } - for e in entities - if e.get("text") - ] + "metadata": {"category": e.get("category")}, + } + ) - edges_payload = [ - { - "subject": r.get("subject"), - "predicate": r.get("predicate", "RELATED_TO").upper(), - "object": r.get("object"), - "properties": { - "confidence": r.get("confidence", 0.8), - "evidence": r.get("evidence", ""), - }, - } - for r in relationships - if r.get("subject") and r.get("object") - ] + edges_payload: list[dict[str, Any]] = [] + for r in relationships: + source_text = r.get("source_entity") or r.get("subject") + target_text = r.get("target_entity") or r.get("object") + predicate = r.get("predicate") + if not source_text or not target_text or not predicate: + continue + rel_type = normalize_relationship_type(predicate) + from_id = entity_id_by_text.get(str(source_text).strip().lower()) + to_id = entity_id_by_text.get(str(target_text).strip().lower()) + if not rel_type or not from_id or not to_id: + continue + edges_payload.append( + { + "from_node_id": from_id, + "to_node_id": to_id, + "type": rel_type, + "confidence": r.get("confidence"), + "source": "literature_service", + "evidence": (r.get("provenance") or {}).get("source_sentence"), + } + ) if not nodes_payload and not edges_payload: - return {"success": True, "updated_nodes": 0, "updated_edges": 0, "status": "no_op"} + return { + "success": True, + "updated_nodes": 0, + "updated_edges": 0, + "status": "no_op", + "entity_id_map": entity_id_by_text, + } for attempt in range(self.max_retries + 1): try: with httpx.Client(timeout=self.timeout) as client: - res = client.post(f"{self.base_url}/api/v1/graph/import", json={"nodes": nodes_payload, "relationships": edges_payload}) + res = client.post( + f"{self.base_url}/api/v1/graph/import/json", + json={"nodes": nodes_payload, "relationships": edges_payload}, + ) if res.status_code in (200, 201, 202): metrics.increment("literature.kg_update.success") return { @@ -64,6 +107,7 @@ def update_knowledge_graph(self, entities: list[dict[str, Any]], relationships: "updated_nodes": len(nodes_payload), "updated_edges": len(edges_payload), "status": "completed", + "entity_id_map": entity_id_by_text, } if res.status_code in (429, 500, 502, 503, 504) and attempt < self.max_retries: delay = self.backoff_seconds * (2**attempt) diff --git a/services/literature/app/integrations/wiki_client.py b/services/literature/app/integrations/wiki_client.py index f6890b1..2e9f0a5 100644 --- a/services/literature/app/integrations/wiki_client.py +++ b/services/literature/app/integrations/wiki_client.py @@ -9,6 +9,7 @@ import httpx +from app.knowledge.models import deterministic_entity_id from app.observability.metrics import metrics logger = logging.getLogger(__name__) @@ -20,14 +21,38 @@ class LLMWikiClient: def __init__(self, config: dict[str, Any] | None = None): self.config = config or {} self.wiki_dir = Path(self.config.get("wiki_dir", os.environ.get("OKF_WIKI_DIR", "wiki-root"))) - self.service_url = self.config.get("wiki_service_url", os.environ.get("OKF_WIKI_URL")) + self.service_url = self.config.get( + "wiki_service_url", + os.environ.get("LLM_WIKI_URL") or os.environ.get("OKF_WIKI_URL"), + ) self.wiki_api_key = self.config.get("wiki_api_key") self.timeout = float(self.config.get("wiki_timeout", 5.0)) self.max_retries = int(self.config.get("wiki_max_retries", 1)) self.backoff_seconds = float(self.config.get("wiki_backoff_seconds", 0.25)) - def update_wiki(self, document: dict[str, Any], entities: list[dict[str, Any]], summary: dict[str, Any]) -> dict[str, Any]: - """Compile literature intelligence into OKF markdown wiki concepts.""" + def update_wiki( + self, + document: dict[str, Any], + entities: list[dict[str, Any]], + summary: dict[str, Any], + *, + relationships: list[dict[str, Any]] | None = None, + evidence: list[dict[str, Any]] | None = None, + tenant: dict[str, Any] | None = None, + chunks: list[dict[str, Any]] | None = None, + ) -> dict[str, Any]: + """Compile literature intelligence into OKF markdown wiki concepts. + + relationships/evidence/tenant/chunks are optional and additive: all + default to None so existing callers (and the existing OKF-volume + test, which calls update_wiki with only document/entities/summary) + keep their prior behavior — with no tenant, pages are written to the + same untenanted paths as before. + """ + relationships = relationships or [] + evidence = evidence or [] + tenant = tenant or {} + if self.service_url: headers = {"Content-Type": "application/json"} if self.wiki_api_key: @@ -36,7 +61,19 @@ def update_wiki(self, document: dict[str, Any], entities: list[dict[str, Any]], for attempt in range(self.max_retries + 1): try: with httpx.Client(timeout=self.timeout) as client: - res = client.post(f"{self.service_url.rstrip('/')}/api/v1/wiki/compile", json={"document": document, "entities": entities, "summary": summary}, headers=headers) + res = client.post( + f"{self.service_url.rstrip('/')}/api/v1/wiki/compile", + json={ + "document": document, + "entities": entities, + "summary": summary, + "relationships": relationships, + "evidence": evidence, + "tenant": tenant, + "chunks": chunks or [], + }, + headers=headers, + ) if res.status_code in (200, 201): metrics.increment("literature.wiki_update.success") return {"success": True, "method": "http", "status": "completed"} @@ -64,21 +101,73 @@ def update_wiki(self, document: dict[str, Any], entities: list[dict[str, Any]], wiki_path.mkdir(parents=True, exist_ok=True) timestamp = datetime.now(timezone.utc).isoformat() + organization_id = tenant.get("organization_id") + workspace_id = tenant.get("workspace_id") + updated_files: list[str] = [] for entity in entities: text = entity.get("text") category = entity.get("category", "concepts") if not text: continue - category_dir = wiki_path / category + + # No tenant -> identical path to the pre-Prompt-8 behavior. + # With a tenant, pages are namespaced under org/workspace so + # one organization's wiki content is never written into (or + # read from) another's directory tree. + if organization_id: + category_dir = ( + wiki_path + / str(organization_id) + / str(workspace_id or "_shared") + / category + ) + else: + category_dir = wiki_path / category category_dir.mkdir(parents=True, exist_ok=True) - file_path = category_dir / f"{text.replace(' ', '_')}.md" + slug = text.replace(" ", "_") + file_path = category_dir / f"{slug}.md" + + version = self._next_version(file_path) + if file_path.exists(): + self._archive_previous_version(category_dir, slug, file_path) + + entity_id = deterministic_entity_id( + entity.get("type") or category, text + ) + + entity_relationships = [ + r + for r in relationships + if r.get("source_entity") == text or r.get("target_entity") == text + ] + entity_evidence = [e for e in evidence if e.get("entity") == text] + + relationships_section = "\n".join( + f"- {r.get('source_entity')} --{r.get('predicate')}--> " + f"{r.get('target_entity')} (confidence: {r.get('confidence', 'n/a')})" + for r in entity_relationships + ) or "No relationships recorded." + evidence_section = "\n".join( + f"- {e.get('entity')} (category: {e.get('category', 'unknown')}, " + f"score: {e.get('score', 'n/a')})" + for e in entity_evidence + ) or "No additional evidence recorded." + tenant_line = ( + f"- **Organization**: {organization_id}" + + (f" / **Workspace**: {workspace_id}" if workspace_id else "") + + "\n" + if organization_id + else "" + ) content = f"""# Concept: {text} - **Category**: {category} +- **Entity ID**: {entity_id} +- **Version**: {version} - **Last Updated**: {timestamp} - **Source**: {document.get('source')} ({document.get('source_id')}) - +{tenant_line} ## Primary Summary {summary.get('concise_summary', 'No summary provided.')} @@ -86,6 +175,12 @@ def update_wiki(self, document: dict[str, Any], entities: list[dict[str, Any]], - **Title**: {document.get('title')} - **DOI**: {document.get('doi')} - **URL**: {document.get('url')} + +## Relationships +{relationships_section} + +## Supporting Evidence +{evidence_section} """ file_path.write_text(content, encoding="utf-8") updated_files.append(str(file_path)) @@ -107,3 +202,42 @@ def update_wiki(self, document: dict[str, Any], entities: list[dict[str, Any]], logger.warning("OKF Wiki volume write failed: %s", exc) metrics.increment("literature.wiki_update.failure") return {"success": False, "error": str(exc), "retry_eligible": True, "status": "failed"} + + @staticmethod + def _next_version(file_path: Path) -> int: + """Read the current page's "- **Version**: N" line and return N + 1. + + Missing file, missing line, or an unparseable value all fall back + to version 1 rather than raising — versioning is best-effort + metadata on top of the markdown page, not a hard invariant. + """ + if not file_path.exists(): + return 1 + try: + content = file_path.read_text(encoding="utf-8") + except OSError: + return 1 + for line in content.splitlines(): + if line.startswith("- **Version**:"): + try: + return int(line.split(":", 1)[1].strip()) + 1 + except ValueError: + return 1 + return 1 + + @staticmethod + def _archive_previous_version(category_dir: Path, slug: str, file_path: Path) -> None: + """Copy the current page into a _versions/ subdirectory before overwrite. + + This is the "history where appropriate" versioning the concept + pages otherwise lack entirely (prior behavior was a full in-place + overwrite with no way to recover an earlier version of a concept). + """ + try: + versions_dir = category_dir / "_versions" + versions_dir.mkdir(parents=True, exist_ok=True) + previous_version = LLMWikiClient._next_version(file_path) - 1 + archive_path = versions_dir / f"{slug}_v{previous_version}.md" + archive_path.write_text(file_path.read_text(encoding="utf-8"), encoding="utf-8") + except OSError as exc: + logger.warning("Failed to archive previous wiki page version for %s: %s", slug, exc) diff --git a/services/literature/app/knowledge/__init__.py b/services/literature/app/knowledge/__init__.py new file mode 100644 index 0000000..1a93937 --- /dev/null +++ b/services/literature/app/knowledge/__init__.py @@ -0,0 +1,15 @@ +from app.knowledge.models import ( + KnowledgeMetadata, + TenantContext, + deterministic_entity_id, + normalize_entity_label, + normalize_relationship_type, +) + +__all__ = [ + "KnowledgeMetadata", + "TenantContext", + "deterministic_entity_id", + "normalize_entity_label", + "normalize_relationship_type", +] diff --git a/services/literature/app/knowledge/models.py b/services/literature/app/knowledge/models.py new file mode 100644 index 0000000..fba70f2 --- /dev/null +++ b/services/literature/app/knowledge/models.py @@ -0,0 +1,190 @@ +"""Canonical knowledge representation shared by the KG client, the LLM Wiki +client, and the chunking pipeline. + +This module intentionally does not introduce a new persistence layer or a +new entity taxonomy — it normalizes the literature pipeline's existing +NER/relationship output (see app/nlp/ner.py, app/nlp/relationship_extractor.py) +onto the entity labels and relationship types already defined by the +Knowledge Graph service (services/kg/app/schemas/{nodes,relationships}.py), +and defines the metadata/tenant shape that the LLM Wiki write path +(app/integrations/wiki_client.py) and the search hand-off +(app/services/search_integration.py) both need. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any + +# Fixed namespace so the same (category, text) pair always resolves to the +# same UUID across pipeline runs — this is what lets a wiki chunk's +# entity_ids be traced back to the matching Neo4j node id without a round +# trip to the KG service. +_ENTITY_ID_NAMESPACE = uuid.UUID("6f6d0f2e-6e0a-4f0b-9a6b-2e6b7f9c9e10") + +# NER `type` (see app/nlp/ner.py) -> Knowledge Graph node label +# (see services/kg/app/schemas/nodes.py::VALID_LABELS). Both singular +# ("drug", the NER `type` field) and plural ("drugs", the NER `category` +# field) spellings are mapped, since callers may only have one or the +# other (see app/nlp/ner.py, which emits both on the same entity dict). +ENTITY_TYPE_TO_KG_LABEL: dict[str, str] = { + "gene": "Gene", + "genes": "Gene", + "protein": "Protein", + "proteins": "Protein", + "disease": "Disease", + "diseases": "Disease", + "drug": "Drug", + "drugs": "Drug", + "target": "Target", + "targets": "Target", + "mutation": "Mutation", + "mutations": "Mutation", + "variant": "Mutation", + "variants": "Mutation", + "publication": "Publication", + "publications": "Publication", + "patent": "Patent", + "patents": "Patent", + "clinical_trial": "ClinicalTrial", + "clinical_trials": "ClinicalTrial", + "company": "Company", + "companies": "Company", + "organization": "Company", + "organizations": "Company", + "conference": "Conference", + "conferences": "Conference", + "biomarker": "Biomarker", + "biomarkers": "Biomarker", +} + +# Relationship `predicate` (see app/nlp/relationship_extractor.py) -> KG +# relationship type (see services/kg/app/schemas/relationships.py:: +# VALID_RELATIONSHIP_TYPES). The KG's relationship vocabulary is generic, +# so predicates without an exact semantic match fall back to INTERACTS +# rather than being silently dropped. +RELATIONSHIP_PREDICATE_TO_KG_TYPE: dict[str, str] = { + "treats": "TREATS", + "targets": "TARGETS", + "interacts_with": "INTERACTS", + "associated_with": "INTERACTS", + "presented_at": "PRESENTED_AT", + "published_in": "PUBLISHED_IN", + "owned_by": "OWNED_BY", + "competes_with": "COMPETES_WITH", + "validated_by": "VALIDATED_BY", +} + + +def deterministic_entity_id(category: str, text: str) -> str: + """Stable UUID for an entity, derived from its normalized text + category. + + Using a deterministic id (rather than a random one per pipeline run) + means the same concept always lands on the same KG node and the same + wiki page across repeated ingestions, which is what makes wiki + `entity_ids` traceable back to graph entities. + """ + normalized = f"{category.strip().lower()}:{text.strip().lower()}" + return str(uuid.uuid5(_ENTITY_ID_NAMESPACE, normalized)) + + +def normalize_entity_label(entity_type: str) -> str | None: + """Map a literature NER entity type to a valid KG node label, or None.""" + return ENTITY_TYPE_TO_KG_LABEL.get((entity_type or "").strip().lower()) + + +def normalize_relationship_type(predicate: str) -> str | None: + """Map a literature relationship predicate to a valid KG relationship type, or None.""" + return RELATIONSHIP_PREDICATE_TO_KG_TYPE.get((predicate or "").strip().lower()) + + +@dataclass(frozen=True) +class TenantContext: + """Organization/workspace/project/user scope for a request or pipeline run. + + Mirrors services/auth's tenant claim shape (organizationId/roles today, + workspaceId/projectId anticipated — see services/auth/src/abac.ts's + ResourceAttributes) so any isolation added here lines up with the + isolation model already used by the auth service's Postgres RLS. + """ + + organization_id: str | None = None + workspace_id: str | None = None + project_id: str | None = None + user_id: str | None = None + + def is_scoped(self) -> bool: + return bool(self.organization_id) + + def as_dict(self) -> dict[str, str]: + return { + k: v + for k, v in { + "organization_id": self.organization_id, + "workspace_id": self.workspace_id, + "project_id": self.project_id, + "user_id": self.user_id, + }.items() + if v + } + + @classmethod + def from_claims(cls, claims: dict[str, Any] | None) -> TenantContext: + claims = claims or {} + return cls( + organization_id=claims.get("organization_id") + or claims.get("organizationId"), + workspace_id=claims.get("workspace_id") or claims.get("workspaceId"), + project_id=claims.get("project_id") or claims.get("projectId"), + user_id=claims.get("user_id") or claims.get("sub"), + ) + + +@dataclass +class KnowledgeMetadata: + """Canonical metadata attached to every indexed document/chunk. + + Field names follow the Prompt 8 spec directly; source_type/source_id/ + provenance reuse the vocabulary already produced by + app/connectors (pubmed/biorxiv/clinicaltrials/patent/conference/company) + and app/services/evidence_ranking.py's provenance blocks. + """ + + document_id: str + source_type: str + source_id: str + title: str + entity_ids: list[str] = field(default_factory=list) + entity_types: list[str] = field(default_factory=list) + organization_id: str | None = None + workspace_id: str | None = None + project_id: str | None = None + version: int = 1 + created_at: str = field( + default_factory=lambda: datetime.now(timezone.utc).isoformat() + ) + updated_at: str = field( + default_factory=lambda: datetime.now(timezone.utc).isoformat() + ) + provenance: dict[str, Any] = field(default_factory=dict) + citation: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return { + "document_id": self.document_id, + "source_type": self.source_type, + "source_id": self.source_id, + "title": self.title, + "entity_ids": self.entity_ids, + "entity_types": self.entity_types, + "organization_id": self.organization_id, + "workspace_id": self.workspace_id, + "project_id": self.project_id, + "version": self.version, + "created_at": self.created_at, + "updated_at": self.updated_at, + "provenance": self.provenance, + "citation": self.citation, + } diff --git a/services/literature/app/nlp/embedding_service.py b/services/literature/app/nlp/embedding_service.py index f25e55b..616a803 100644 --- a/services/literature/app/nlp/embedding_service.py +++ b/services/literature/app/nlp/embedding_service.py @@ -89,6 +89,10 @@ def generate_embeddings(self, nlp_result: dict[str, Any]) -> dict[str, Any]: EMBEDDING_GENERATION_ERRORS_TOTAL.inc() raise + def embed_text(self, text: str) -> list[float]: + """Public entry point for embedding a single arbitrary string (e.g. a chunk).""" + return self._vectorize_text(text) + def _create_batches(self, items: list[str]) -> list[list[str]]: batches: list[list[str]] = [] for index in range(0, len(items), self.batch_size): diff --git a/services/literature/app/orchestrator/stages.py b/services/literature/app/orchestrator/stages.py index c559f2c..53ded9e 100644 --- a/services/literature/app/orchestrator/stages.py +++ b/services/literature/app/orchestrator/stages.py @@ -5,12 +5,24 @@ from app.integrations.kg_client import KGClient from app.integrations.wiki_client import LLMWikiClient +from app.knowledge.models import ( + TenantContext, + deterministic_entity_id, + normalize_entity_label, +) +from app.nlp.embedding_service import EmbeddingService from app.nlp.pipeline import LiteratureNLP from app.observability.metrics import metrics +from app.services.chunking import build_chunks +from app.services.search_integration import SearchIntegrationService logger = logging.getLogger(__name__) +def _tenant_from_payload(payload: dict[str, Any]) -> TenantContext: + return TenantContext.from_claims(payload.get("tenant")) + + class PipelineStage: """Base pipeline stage for reusable, callable stage implementations.""" @@ -156,6 +168,42 @@ def run(self, payload: dict[str, Any]) -> dict[str, Any]: return {**payload, "items": deduped, "duplicates": duplicates_found} +class ChunkingStage(PipelineStage): + """Splits each item's normalized text into provenance-carrying chunks. + + Runs after entity/relationship extraction (so entity_ids can be + attached) and before KGUpdateStage/WikiUpdateStage/SearchHandoffStage, + which all consume item["chunks"]. + """ + + name = "ChunkingStage" + + def run(self, payload: dict[str, Any]) -> dict[str, Any]: + tenant = _tenant_from_payload(payload) + for item in payload.get("items", []): + document = item.get("document", {}) + entities = item.get("structured_entities", []) + entity_ids: list[str] = [] + entity_types: list[str] = [] + for entity in entities: + category = entity.get("type") or entity.get("category") or "" + label = normalize_entity_label(category) + if not label or not entity.get("text"): + continue + entity_ids.append(deterministic_entity_id(category, entity["text"])) + entity_types.append(label) + + item["chunks"] = build_chunks( + document=document, + text=item.get("normalized_text") or document.get("content", ""), + source_type=document.get("source") or payload.get("source", "unknown"), + entity_ids=entity_ids, + entity_types=entity_types, + tenant=tenant, + ) + return payload + + class KGUpdateStage(PipelineStage): name = "KGUpdateStage" @@ -165,11 +213,13 @@ def __init__(self, kg_client: KGClient): def run(self, payload: dict[str, Any]) -> dict[str, Any]: kg_results: list[dict[str, Any]] = [] for item in payload.get("items", []): - kg_results.append( - self.kg_client.update_knowledge_graph( - item.get("structured_entities", []), item.get("relationships", []) - ) + result = self.kg_client.update_knowledge_graph( + item.get("structured_entities", []), item.get("relationships", []) ) + # Shares deterministic entity ids with the wiki chunks built in + # ChunkingStage, so WikiUpdateStage can link back to KG nodes. + item["kg_entity_id_map"] = result.get("entity_id_map", {}) + kg_results.append(result) return {**payload, "kg_updates": kg_results} @@ -180,6 +230,7 @@ def __init__(self, wiki_client: LLMWikiClient): self.wiki_client = wiki_client def run(self, payload: dict[str, Any]) -> dict[str, Any]: + tenant = _tenant_from_payload(payload).as_dict() wiki_results: list[dict[str, Any]] = [] for item in payload.get("items", []): wiki_results.append( @@ -187,6 +238,71 @@ def run(self, payload: dict[str, Any]) -> dict[str, Any]: item.get("document", {}), item.get("structured_entities", []), item.get("structured_summary", {}), + relationships=item.get("relationships", []), + evidence=item.get("evidence", []), + tenant=tenant, + chunks=item.get("chunks", []), ) ) return {**payload, "wiki_updates": wiki_results, "status": "completed"} + + +class SearchHandoffStage(PipelineStage): + """Feeds indexed chunks + embeddings to the search service. + + This is the previously-missing link between the literature pipeline + (Prompt 6) and search's semantic retrieval (Prompt 7): without it, + LLMWikiProvider's local QMD index never receives any content. Failures + are recorded per item rather than raised, matching the degrade-gracefully + behavior of KGUpdateStage/WikiUpdateStage's underlying clients — a + search outage should not fail literature ingestion. + """ + + name = "SearchHandoffStage" + + def __init__( + self, + search_client: SearchIntegrationService, + embedding_service: EmbeddingService, + ): + self.search_client = search_client + self.embedding_service = embedding_service + + def run(self, payload: dict[str, Any]) -> dict[str, Any]: + tenant = _tenant_from_payload(payload).as_dict() + handoff_results: list[dict[str, Any]] = [] + for item in payload.get("items", []): + chunks = item.get("chunks", []) + if not chunks: + continue + document_id = str( + chunks[0]["metadata"].get("document_id") + or item.get("document", {}).get("source_id") + or "" + ) + if not document_id: + continue + documents = [ + { + "id": chunk["chunk_id"], + "document_id": document_id, + "title": chunk["metadata"].get("title", ""), + "content": chunk["text"], + "source": chunk["metadata"].get("source_type", "literature_service"), + "embedding": self.embedding_service.embed_text(chunk["text"]), + } + for chunk in chunks + ] + try: + result = self.search_client.submit_embeddings( + document_id, documents, tenant=tenant + ) + except RuntimeError as exc: + metrics.increment("literature.search_handoff.failure") + result = { + "document_id": document_id, + "status": "failed", + "error": str(exc), + } + handoff_results.append(result) + return {**payload, "search_handoffs": handoff_results} diff --git a/services/literature/app/routers/ingestion.py b/services/literature/app/routers/ingestion.py index 70e4101..c585dbc 100644 --- a/services/literature/app/routers/ingestion.py +++ b/services/literature/app/routers/ingestion.py @@ -5,8 +5,9 @@ from fastapi import APIRouter, Depends, HTTPException -from app.core.security import get_current_user +from app.core.security import get_current_user, get_tenant_context from app.database.postgres import postgres_manager +from app.knowledge.models import TenantContext from app.orchestrator.manager import IngestionJob as OrchestrationJob from app.orchestrator.manager import JobStatus, orchestrator from app.schemas import IngestionJob as IngestionJobSchema @@ -15,6 +16,7 @@ router = APIRouter(prefix="/ingestion", tags=["Ingestion"]) auth_dependency = Depends(get_current_user) +tenant_dependency = Depends(get_tenant_context) def _use_postgres() -> bool: @@ -41,6 +43,7 @@ def _use_postgres() -> bool: async def start_ingestion( req: IngestionRequest, auth_payload: dict[str, str] = auth_dependency, + tenant: TenantContext = tenant_dependency, ) -> IngestionJobSchema: job_id = uuid4() if not _use_postgres(): @@ -54,7 +57,7 @@ async def start_ingestion( status="pending", ) job_store.save(job_state) - literature_service.ingest(req.source, req.query, job_id=job_id_str) + literature_service.ingest(req.source, req.query, job_id=job_id_str, tenant=tenant) updated_job = job_store.get(job_id_str) or job_state status_output = "completed" if updated_job.status == "completed" else updated_job.status created_at_dt = datetime.fromisoformat(updated_job.created_at) diff --git a/services/literature/app/services/chunking.py b/services/literature/app/services/chunking.py new file mode 100644 index 0000000..8efbf60 --- /dev/null +++ b/services/literature/app/services/chunking.py @@ -0,0 +1,129 @@ +"""Small, reusable scientific-document chunker. + +No chunking implementation existed anywhere in the repo prior to this +module (only unrelated I/O-buffer reads in app/parsing/duplicates.py and +app/parsing/parser.py). This is intentionally simple: paragraph-aware, +character-budget-based splitting with overlap, not a general NLP +framework. Every chunk carries enough metadata to be traced back to its +source document, entities, and tenant without a second lookup. +""" + +from __future__ import annotations + +import re +from typing import Any + +from app.knowledge.models import KnowledgeMetadata, TenantContext + +_PARAGRAPH_SPLIT = re.compile(r"\n\s*\n+") + + +def _split_paragraphs(text: str) -> list[str]: + paragraphs = [p.strip() for p in _PARAGRAPH_SPLIT.split(text or "") if p.strip()] + return paragraphs or ([text.strip()] if text and text.strip() else []) + + +def chunk_text( + text: str, + *, + max_chars: int = 1000, + overlap_chars: int = 150, +) -> list[str]: + """Split text into paragraph-respecting chunks of at most max_chars. + + Paragraphs are packed greedily; a paragraph longer than max_chars is + hard-split with overlap so no chunk ever exceeds the budget. + """ + if max_chars <= 0: + raise ValueError("max_chars must be positive") + + chunks: list[str] = [] + current = "" + + for paragraph in _split_paragraphs(text): + candidate = f"{current}\n\n{paragraph}".strip() if current else paragraph + if len(candidate) <= max_chars: + current = candidate + continue + + if current: + chunks.append(current) + current = "" + + if len(paragraph) <= max_chars: + current = paragraph + continue + + start = 0 + while start < len(paragraph): + end = min(start + max_chars, len(paragraph)) + chunks.append(paragraph[start:end]) + if end >= len(paragraph): + break + start = end - overlap_chars if end - overlap_chars > start else end + + if current: + chunks.append(current) + + return chunks + + +def build_chunks( + *, + document: dict[str, Any], + text: str, + source_type: str, + entity_ids: list[str] | None = None, + entity_types: list[str] | None = None, + tenant: TenantContext | None = None, + version: int = 1, + max_chars: int = 1000, + overlap_chars: int = 150, +) -> list[dict[str, Any]]: + """Chunk a document's text and attach canonical provenance/metadata. + + Every chunk is traceable to: document id, source type, title, entity + ids (where available), organization/workspace/project, and version — + the minimum set required by the Prompt 8 chunking requirement. + """ + tenant = tenant or TenantContext() + document_id = str( + document.get("id") or document.get("source_id") or document.get("doi") or "" + ) + source_id = str(document.get("source_id") or document_id) + title = str(document.get("title") or "") + + pieces = chunk_text(text, max_chars=max_chars, overlap_chars=overlap_chars) + chunks: list[dict[str, Any]] = [] + for index, piece in enumerate(pieces): + metadata = KnowledgeMetadata( + document_id=document_id, + source_type=source_type, + source_id=source_id, + title=title, + entity_ids=list(entity_ids or []), + entity_types=list(entity_types or []), + organization_id=tenant.organization_id, + workspace_id=tenant.workspace_id, + project_id=tenant.project_id, + version=version, + provenance={ + "source": document.get("source"), + "url": document.get("url"), + "doi": document.get("doi"), + }, + citation={ + "title": title, + "doi": document.get("doi"), + "url": document.get("url"), + }, + ) + chunks.append( + { + "chunk_id": f"{document_id or source_id}:{index}", + "chunk_index": index, + "text": piece, + "metadata": metadata.to_dict(), + } + ) + return chunks diff --git a/services/literature/app/services/literature_service.py b/services/literature/app/services/literature_service.py index 54bb92a..b6f4e4b 100644 --- a/services/literature/app/services/literature_service.py +++ b/services/literature/app/services/literature_service.py @@ -7,19 +7,24 @@ from app.database.models import IngestionJobState, job_store from app.integrations.kg_client import KGClient from app.integrations.wiki_client import LLMWikiClient +from app.knowledge.models import TenantContext +from app.nlp.embedding_service import EmbeddingService from app.nlp.pipeline import LiteratureNLP from app.observability.metrics import metrics from app.orchestrator.pipeline import PipelineRunner from app.orchestrator.stages import ( + ChunkingStage, DeduplicationStage, EvidenceRankingStage, KGUpdateStage, NERStage, ParsingStage, RelationshipExtractionStage, + SearchHandoffStage, SummarizationStage, WikiUpdateStage, ) +from app.services.search_integration import SearchIntegrationService logger = logging.getLogger(__name__) @@ -32,6 +37,8 @@ def __init__(self, config: dict[str, Any] | None = None): self.nlp = LiteratureNLP(self.config) self.kg_client = KGClient(self.config) self.wiki_client = LLMWikiClient(self.config) + self.search_client = SearchIntegrationService() + self.embedding_service = EmbeddingService() self.pipeline = PipelineRunner( stages=[ @@ -41,15 +48,24 @@ def __init__(self, config: dict[str, Any] | None = None): SummarizationStage(self.nlp), EvidenceRankingStage(self.nlp), DeduplicationStage(self.nlp), + ChunkingStage(), KGUpdateStage(self.kg_client), WikiUpdateStage(self.wiki_client), + SearchHandoffStage(self.search_client, self.embedding_service), ], retries=2, delay_seconds=0.05, ) - def ingest(self, source: str, query: str, job_id: str | None = None, **kwargs: Any) -> dict[str, Any]: + def ingest( + self, + source: str, + query: str, + job_id: str | None = None, + tenant: TenantContext | None = None, + **kwargs: Any, + ) -> dict[str, Any]: connector_config = {**self.config, **kwargs} connector = ConnectorFactory.create(source, connector_config) source_status = connector.connect() @@ -64,6 +80,7 @@ def ingest(self, source: str, query: str, job_id: str | None = None, **kwargs: A "items": results, "limitation": connector.get_limitation(), "source_status": source_status, + "tenant": (tenant or TenantContext()).as_dict(), } if job_id: diff --git a/services/literature/tests/test_prompt8_knowledge_layer.py b/services/literature/tests/test_prompt8_knowledge_layer.py new file mode 100644 index 0000000..0e9c7ad --- /dev/null +++ b/services/literature/tests/test_prompt8_knowledge_layer.py @@ -0,0 +1,257 @@ +"""Tests for the Prompt 8 LLM Wiki knowledge-layer additions: +chunking, canonical metadata, knowledge linking (KG <-> wiki entity ids), +tenant/workspace isolation in the wiki write path, and wiki versioning. + +Uses only local fakes/mocks for the external KG/search/LLM Wiki services — +no real credentials or running services required, matching the existing +test_prompt6_comprehensive.py / test_llmwiki_integration.py conventions. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from app.integrations.kg_client import KGClient +from app.integrations.wiki_client import LLMWikiClient +from app.knowledge.models import ( + TenantContext, + deterministic_entity_id, + normalize_entity_label, + normalize_relationship_type, +) +from app.services.chunking import build_chunks, chunk_text + + +# --------------------------------------------------------------------------- +# Chunking +# --------------------------------------------------------------------------- + + +def test_chunk_text_respects_max_chars_and_preserves_content(): + text = "\n\n".join([f"Paragraph {i} " + ("word " * 40) for i in range(5)]) + chunks = chunk_text(text, max_chars=250, overlap_chars=20) + + assert len(chunks) > 1 + assert all(len(c) <= 250 for c in chunks) + # No content is silently dropped. + assert "Paragraph 0" in chunks[0] + assert "Paragraph 4" in chunks[-1] + + +def test_chunk_text_hard_splits_a_single_oversized_paragraph(): + text = "x" * 1000 + chunks = chunk_text(text, max_chars=300, overlap_chars=50) + assert len(chunks) >= 3 + assert all(len(c) <= 300 for c in chunks) + + +def test_build_chunks_attaches_canonical_metadata_and_provenance(): + document = { + "id": "doc-42", + "source_id": "PMID42", + "title": "HER2 Targeted Therapy", + "source": "pubmed", + "doi": "10.1000/xyz", + "url": "https://example.com/42", + } + tenant = TenantContext(organization_id="org-1", workspace_id="ws-1") + + chunks = build_chunks( + document=document, + text="Trastuzumab is a HER2-targeted therapy. " * 10, + source_type="paper", + entity_ids=["entity-abc"], + entity_types=["Drug"], + tenant=tenant, + version=3, + max_chars=200, + ) + + assert len(chunks) >= 1 + for i, chunk in enumerate(chunks): + assert chunk["chunk_id"] == f"doc-42:{i}" + assert chunk["chunk_index"] == i + meta = chunk["metadata"] + assert meta["document_id"] == "doc-42" + assert meta["source_type"] == "paper" + assert meta["source_id"] == "PMID42" + assert meta["title"] == "HER2 Targeted Therapy" + assert meta["entity_ids"] == ["entity-abc"] + assert meta["entity_types"] == ["Drug"] + assert meta["organization_id"] == "org-1" + assert meta["workspace_id"] == "ws-1" + assert meta["version"] == 3 + assert meta["provenance"]["source"] == "pubmed" + assert meta["citation"]["doi"] == "10.1000/xyz" + + +def test_build_chunks_without_tenant_leaves_org_workspace_none(): + chunks = build_chunks( + document={"id": "doc-1", "title": "T"}, + text="short text", + source_type="paper", + ) + assert chunks[0]["metadata"]["organization_id"] is None + assert chunks[0]["metadata"]["workspace_id"] is None + + +# --------------------------------------------------------------------------- +# Knowledge linking (deterministic entity ids, KG label/relationship mapping) +# --------------------------------------------------------------------------- + + +def test_deterministic_entity_id_is_stable_and_case_insensitive(): + id1 = deterministic_entity_id("gene", "HER2") + id2 = deterministic_entity_id("gene", "her2") + id3 = deterministic_entity_id("gene", " HER2 ") + assert id1 == id2 == id3 + + id_different_category = deterministic_entity_id("disease", "HER2") + assert id_different_category != id1 + + +def test_normalize_entity_label_handles_singular_and_plural(): + assert normalize_entity_label("gene") == "Gene" + assert normalize_entity_label("genes") == "Gene" + assert normalize_entity_label("drugs") == "Drug" + assert normalize_entity_label("clinical_trials") == "ClinicalTrial" + assert normalize_entity_label("not_a_real_category") is None + + +def test_normalize_relationship_type_maps_known_predicates(): + assert normalize_relationship_type("treats") == "TREATS" + assert normalize_relationship_type("targets") == "TARGETS" + assert normalize_relationship_type("associated_with") == "INTERACTS" + assert normalize_relationship_type("nonsense") is None + + +def test_kg_client_sends_correct_url_and_flat_node_relationship_shape(): + entities = [ + {"text": "trastuzumab", "type": "drug"}, + {"text": "her2", "type": "gene"}, + ] + relationships = [ + {"source_entity": "trastuzumab", "target_entity": "her2", "predicate": "targets", "confidence": 0.9} + ] + + with patch("httpx.Client.post") as mock_post: + mock_response = MagicMock() + mock_response.status_code = 200 + mock_post.return_value = mock_response + + client = KGClient({"kg_service_url": "http://kg:8083"}) + result = client.update_knowledge_graph(entities, relationships) + + assert result["success"] is True + called_url = mock_post.call_args.args[0] + assert called_url == "http://kg:8083/api/v1/graph/import/json" + + payload = mock_post.call_args.kwargs["json"] + assert len(payload["nodes"]) == 2 + for node in payload["nodes"]: + assert set(node.keys()) >= {"id", "label", "name", "source", "metadata"} + assert node["label"] in {"Drug", "Gene"} + + assert len(payload["relationships"]) == 1 + rel = payload["relationships"][0] + assert rel["type"] == "TARGETS" + assert rel["from_node_id"] == deterministic_entity_id("drug", "trastuzumab") + assert rel["to_node_id"] == deterministic_entity_id("gene", "her2") + + +def test_kg_client_skips_unrecognized_entities_and_relationships_without_erroring(): + entities = [{"text": "some author fragment", "type": "unknown_type"}] + relationships = [{"source_entity": "a", "target_entity": "b", "predicate": "unmapped_predicate"}] + + client = KGClient({"kg_service_url": "http://kg:8083"}) + result = client.update_knowledge_graph(entities, relationships) + + assert result == { + "success": True, + "updated_nodes": 0, + "updated_edges": 0, + "status": "no_op", + "entity_id_map": {}, + } + + +def test_wiki_entity_id_matches_kg_entity_id_for_the_same_entity(tmp_path): + """The core knowledge-linking guarantee: a wiki concept page's Entity ID + is the same id KGClient would assign the matching KG node, so retrieved + wiki knowledge is traceable back to the graph without a lookup.""" + wiki_client = LLMWikiClient({"wiki_dir": str(tmp_path)}) + doc = {"source": "pubmed", "source_id": "PMID1", "title": "T"} + entities = [{"text": "her2", "category": "genes", "type": "gene"}] + summary = {"concise_summary": "summary"} + + wiki_client.update_wiki(doc, entities, summary) + + page = (tmp_path / "wiki" / "genes" / "her2.md").read_text(encoding="utf-8") + expected_id = deterministic_entity_id("gene", "her2") + assert f"**Entity ID**: {expected_id}" in page + + +# --------------------------------------------------------------------------- +# Tenant / workspace isolation in the wiki write path +# --------------------------------------------------------------------------- + + +def test_wiki_pages_are_namespaced_by_organization_and_workspace(tmp_path): + wiki_client = LLMWikiClient({"wiki_dir": str(tmp_path)}) + doc = {"source": "pubmed", "source_id": "PMID2", "title": "T"} + entities = [{"text": "aspirin", "category": "drugs"}] + summary = {"concise_summary": "s"} + + wiki_client.update_wiki( + doc, entities, summary, tenant={"organization_id": "org-a", "workspace_id": "ws-1"} + ) + wiki_client.update_wiki( + doc, entities, summary, tenant={"organization_id": "org-b", "workspace_id": "ws-1"} + ) + + org_a_page = tmp_path / "wiki" / "org-a" / "ws-1" / "drugs" / "aspirin.md" + org_b_page = tmp_path / "wiki" / "org-b" / "ws-1" / "drugs" / "aspirin.md" + assert org_a_page.exists() + assert org_b_page.exists() + # Writing org B's page must never touch org A's directory tree. + assert "org-a" not in org_b_page.read_text(encoding="utf-8").replace("org-a/ws-1", "") + + +def test_wiki_pages_without_tenant_use_pre_prompt8_untenanted_path(tmp_path): + """Backward compatibility: no tenant means the exact same path as before + tenant scoping was added (see test_wiki_integration_okf_volume_write in + test_prompt6_comprehensive.py, which asserts this same untenanted path).""" + wiki_client = LLMWikiClient({"wiki_dir": str(tmp_path)}) + doc = {"source": "pubmed", "source_id": "PMID3", "title": "T"} + entities = [{"text": "her2", "category": "genes"}] + summary = {"concise_summary": "s"} + + wiki_client.update_wiki(doc, entities, summary) + + assert (tmp_path / "wiki" / "genes" / "her2.md").exists() + assert not (tmp_path / "wiki" / "_none").exists() + + +# --------------------------------------------------------------------------- +# Versioning +# --------------------------------------------------------------------------- + + +def test_wiki_page_version_increments_and_archives_previous_content(tmp_path): + wiki_client = LLMWikiClient({"wiki_dir": str(tmp_path)}) + doc = {"source": "pubmed", "source_id": "PMID4", "title": "First Title"} + entities = [{"text": "her2", "category": "genes"}] + + wiki_client.update_wiki(doc, entities, {"concise_summary": "v1 summary"}) + page = tmp_path / "wiki" / "genes" / "her2.md" + assert "- **Version**: 1" in page.read_text(encoding="utf-8") + + doc2 = {"source": "pubmed", "source_id": "PMID5", "title": "Second Title"} + wiki_client.update_wiki(doc2, entities, {"concise_summary": "v2 summary"}) + content_v2 = page.read_text(encoding="utf-8") + assert "- **Version**: 2" in content_v2 + assert "v2 summary" in content_v2 + + archived = tmp_path / "wiki" / "genes" / "_versions" / "her2_v1.md" + assert archived.exists() + assert "v1 summary" in archived.read_text(encoding="utf-8") diff --git a/services/llm-wiki/Dockerfile b/services/llm-wiki/Dockerfile new file mode 100644 index 0000000..2ba4a65 --- /dev/null +++ b/services/llm-wiki/Dockerfile @@ -0,0 +1,33 @@ +FROM python:3.12-slim AS base + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 + +WORKDIR /app + +FROM base AS deps + +COPY services/llm-wiki/requirements.txt ./ + +RUN pip install \ + --upgrade pip \ + --default-timeout=1000 \ + --retries=20 \ + --no-cache-dir \ + -r requirements.txt + +FROM deps AS runtime + +RUN useradd --create-home --uid 1000 rxos + +COPY services/llm-wiki/app ./app +COPY services/llm-wiki/migrations ./migrations + +USER rxos + +EXPOSE 8092 + +ENV PORT=8092 + +CMD ["sh", "-c", "uvicorn app.main:app --host 0.0.0.0 --port ${PORT}"] diff --git a/services/llm-wiki/app/__init__.py b/services/llm-wiki/app/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/services/llm-wiki/app/core/__init__.py b/services/llm-wiki/app/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/services/llm-wiki/app/core/auth.py b/services/llm-wiki/app/core/auth.py new file mode 100644 index 0000000..0f40a40 --- /dev/null +++ b/services/llm-wiki/app/core/auth.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from fastapi import Header, HTTPException, status + +from app.core.config import get_settings + + +async def require_api_key(authorization: str | None = Header(default=None)) -> None: + """Service-to-service auth for literature/search -> llm-wiki calls. + + Matches the bearer scheme services/literature's LLMWikiClient and + services/search's LLMWikiProvider already send + (`Authorization: Bearer {LLM_WIKI_API_KEY}`). When LLM_WIKI_API_KEY is + not configured, requests are accepted unauthenticated -- this mirrors + the client's own behavior (it only attaches the header when it has a + key) and keeps local dev frictionless; app/main.py logs a warning at + startup when this dev posture is active. + """ + settings = get_settings() + if not settings.llm_wiki_api_key: + return + + if not authorization or not authorization.lower().startswith("bearer "): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="missing bearer token", + ) + + token = authorization.split(" ", 1)[1].strip() + if token != settings.llm_wiki_api_key: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="invalid API key", + ) diff --git a/services/llm-wiki/app/core/config.py b/services/llm-wiki/app/core/config.py new file mode 100644 index 0000000..36f7ee6 --- /dev/null +++ b/services/llm-wiki/app/core/config.py @@ -0,0 +1,30 @@ +from functools import lru_cache + +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class Settings(BaseSettings): + model_config = SettingsConfigDict(env_file=".env", extra="ignore") + + environment: str = "development" + log_level: str = "info" + + # Same shared Postgres instance every other AI-RxOS service uses. + # Wiki data lives in its own `llm_wiki` schema (see migrations/), not a + # separate database, so no new datastore is introduced. + database_url: str = "postgresql://ai_rxos:changeme@postgres:5432/ai_rxos" + + # Service-to-service auth. When unset, the service accepts + # unauthenticated requests (matches services/literature's + # LLMWikiClient, which only sends a bearer token when it has one) -- + # this is a dev-only posture and a startup warning is logged if unset. + llm_wiki_api_key: str | None = None + + # /llmwiki/query result size (see services/search's LLMWikiProvider). + query_default_limit: int = 10 + query_max_limit: int = 100 + + +@lru_cache +def get_settings() -> Settings: + return Settings() diff --git a/services/llm-wiki/app/db/__init__.py b/services/llm-wiki/app/db/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/services/llm-wiki/app/db/pool.py b/services/llm-wiki/app/db/pool.py new file mode 100644 index 0000000..3f552ad --- /dev/null +++ b/services/llm-wiki/app/db/pool.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import json +import logging +from contextlib import asynccontextmanager +from pathlib import Path +from typing import AsyncIterator + +import asyncpg + +from app.core.config import get_settings + +logger = logging.getLogger(__name__) + +MIGRATIONS_DIR = Path(__file__).resolve().parent.parent.parent / "migrations" + +_pool: asyncpg.Pool | None = None + + +async def _init_connection(conn: asyncpg.Connection) -> None: + # Let callers pass/receive Python dict/list values directly for jsonb + # columns instead of hand-rolling json.dumps + ::jsonb casts on every + # query. + await conn.set_type_codec( + "jsonb", + encoder=json.dumps, + decoder=json.loads, + schema="pg_catalog", + format="text", + ) + + +async def init_pool() -> asyncpg.Pool: + global _pool + settings = get_settings() + _pool = await asyncpg.create_pool( + settings.database_url, + min_size=1, + max_size=10, + init=_init_connection, + ) + await _run_migrations(_pool) + return _pool + + +async def close_pool() -> None: + global _pool + if _pool is not None: + await _pool.close() + _pool = None + + +def get_pool_dependency() -> asyncpg.Pool: + """FastAPI dependency returning the live pool. + + Kept as a plain callable (not a bare module import inside handlers) so + tests can swap it via `app.dependency_overrides` for a fake repository + without needing a real Postgres instance. + """ + if _pool is None: + raise RuntimeError("database pool is not initialized") + return _pool + + +async def _run_migrations(pool: asyncpg.Pool) -> None: + for path in sorted(MIGRATIONS_DIR.glob("*.sql")): + sql = path.read_text(encoding="utf-8") + async with pool.acquire() as conn: + await conn.execute(sql) + logger.info("applied migration %s", path.name) + + +@asynccontextmanager +async def tenant_connection( + pool: asyncpg.Pool, organization_id: str | None +) -> AsyncIterator[asyncpg.Connection]: + """Acquire a connection with the caller's tenant scope set for the + duration of one transaction. + + `set_config(..., true)` makes the GUC transaction-local, so it never + leaks onto the next request when the pooled connection is reused for a + different tenant. This is the second (database-layer) half of tenant + isolation -- see llm_wiki.wiki_pages_isolation in migrations/001; the + first half is the explicit WHERE clause every repository query also + applies. + """ + async with pool.acquire() as conn: + async with conn.transaction(): + await conn.execute( + "SELECT set_config('app.wiki_organization_id', $1, true)", + organization_id or "", + ) + yield conn diff --git a/services/llm-wiki/app/deps.py b/services/llm-wiki/app/deps.py new file mode 100644 index 0000000..4741f5c --- /dev/null +++ b/services/llm-wiki/app/deps.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +from fastapi import Depends + +from app.db.pool import get_pool_dependency +from app.repository import PostgresWikiRepository, WikiRepository + + +def get_repository(pool=Depends(get_pool_dependency)) -> WikiRepository: + """Overridden in tests via `app.dependency_overrides[get_repository]` + to swap in `InMemoryWikiRepository` -- since that override replaces + this callable outright, `get_pool_dependency` (and therefore the real + database) is never touched by the unit test suite.""" + return PostgresWikiRepository(pool) diff --git a/services/llm-wiki/app/main.py b/services/llm-wiki/app/main.py new file mode 100644 index 0000000..ffeaaee --- /dev/null +++ b/services/llm-wiki/app/main.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import logging +from contextlib import asynccontextmanager + +from fastapi import FastAPI + +from app.core.config import get_settings +from app.db.pool import close_pool, init_pool +from app.routers import health, wiki + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +settings = get_settings() + + +@asynccontextmanager +async def lifespan(_: FastAPI): + if not settings.llm_wiki_api_key: + logger.warning( + "LLM_WIKI_API_KEY is not set -- llm-wiki is accepting unauthenticated " + "requests. This is fine for local dev but must be set before deploying." + ) + await init_pool() + logger.info("llm-wiki ready, database pool initialized") + yield + await close_pool() + + +app = FastAPI( + title="AI-RxOS LLM Wiki Service", + description=( + "Persistent backend for the Open Knowledge Format (OKF) LLM Wiki. " + "Implements the write/read contract already used by " + "services/literature's LLMWikiClient (POST /api/v1/wiki/compile) " + "and services/search's LLMWikiProvider (POST /llmwiki/query)." + ), + version="0.1.0", + lifespan=lifespan, +) + +app.include_router(health.router) +app.include_router(wiki.router) diff --git a/services/llm-wiki/app/models.py b/services/llm-wiki/app/models.py new file mode 100644 index 0000000..f2eb2fb --- /dev/null +++ b/services/llm-wiki/app/models.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + + +class WikiCompileRequest(BaseModel): + """Exactly the payload services/literature's LLMWikiClient.update_wiki() + posts to POST /api/v1/wiki/compile -- see + services/literature/app/integrations/wiki_client.py:66-74.""" + + document: dict[str, Any] = Field(default_factory=dict) + entities: list[dict[str, Any]] = Field(default_factory=list) + summary: dict[str, Any] = Field(default_factory=dict) + relationships: list[dict[str, Any]] = Field(default_factory=list) + evidence: list[dict[str, Any]] = Field(default_factory=list) + tenant: dict[str, Any] = Field(default_factory=dict) + chunks: list[dict[str, Any]] = Field(default_factory=list) + + +class PageResult(BaseModel): + id: str + category: str + slug: str + entity_id: str + version: int + + +class WikiCompileResponse(BaseModel): + # The literature client never parses this body (any 2xx is treated as + # success -- wiki_client.py:77-79); shaped for the new read endpoints + # and for tests/tooling that do want the created page ids back. + success: bool = True + method: str = "http" + status: str = "completed" + pages: list[PageResult] = Field(default_factory=list) + + +class WikiQueryRequest(BaseModel): + """Exactly the payload services/search's LLMWikiProvider posts to + POST /llmwiki/query -- see + services/search/internal/search/providers_placeholder.go:64-69.""" + + embedding: list[float] = Field(default_factory=list) + limit: int = 10 + organization_id: str | None = None + workspace_id: str | None = None + + +class WikiHit(BaseModel): + """Field names/casing match services/search's Hit struct exactly + (services/search/internal/search/opensearch.go:161-169).""" + + id: str + score: float + title: str + snippet: str + source: str | None = None + citationCount: int | None = None + graphScore: float | None = None + rrfScore: float | None = None + + +class WikiQueryResponse(BaseModel): + items: list[WikiHit] = Field(default_factory=list) + + +class VersionSummary(BaseModel): + version: int + created_at: str + + +class PageVersionDetail(BaseModel): + version: int + document: dict[str, Any] + entity: dict[str, Any] + summary: dict[str, Any] + relationships: list[dict[str, Any]] + evidence: list[dict[str, Any]] + chunks: list[dict[str, Any]] + provenance: dict[str, Any] + created_at: str + + +class PageDetail(BaseModel): + id: str + organization_id: str | None + workspace_id: str | None + project_id: str | None + category: str + slug: str + entity_id: str | None + title: str | None + current_version: int + created_at: str + updated_at: str + latest_version: PageVersionDetail | None = None diff --git a/services/llm-wiki/app/repository.py b/services/llm-wiki/app/repository.py new file mode 100644 index 0000000..d4e4878 --- /dev/null +++ b/services/llm-wiki/app/repository.py @@ -0,0 +1,456 @@ +"""Persistence for LLM Wiki pages/versions. + +Two implementations sharing the same method signatures and the same +tenant-isolation semantics (`WikiRepository` is a structural protocol, not +a formal ABC -- there are exactly two implementations and adding +boilerplate for a third that doesn't exist isn't worth it): + +- `PostgresWikiRepository`: the production path, backed by the `llm_wiki` + Postgres schema (migrations/001_wiki_schema.sql). +- `InMemoryWikiRepository`: a test double with no external dependencies, + used by the fast unit test suite (tests/test_*.py) via + `app.dependency_overrides`. tests/test_postgres_integration.py and + tests/test_e2e_literature_roundtrip.py exercise the real Postgres path + instead, skipped automatically when no database is reachable. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any, Protocol + +import asyncpg + +from app.db.pool import tenant_connection + +# Must match services/literature/app/knowledge/models.py::_ENTITY_ID_NAMESPACE +# exactly -- this is what lets a wiki page's entity_id trace back to the +# same deterministic id services/kg would assign the matching graph node. +# Duplicated rather than imported: the two services deploy and run as +# separate processes/images, so there is no shared Python package boundary +# to import across. +_ENTITY_ID_NAMESPACE = uuid.UUID("6f6d0f2e-6e0a-4f0b-9a6b-2e6b7f9c9e10") + +MAX_FIELD_LEN = 512 + + +class ValidationError(ValueError): + pass + + +def deterministic_entity_id(category: str, text: str) -> str: + normalized = f"{category.strip().lower()}:{text.strip().lower()}" + return str(uuid.uuid5(_ENTITY_ID_NAMESPACE, normalized)) + + +def _clean(value: Any) -> str | None: + if value is None: + return None + value = str(value).strip() + return value or None + + +def _check_len(name: str, value: str | None) -> None: + if value and len(value) > MAX_FIELD_LEN: + raise ValidationError(f"{name} exceeds {MAX_FIELD_LEN} characters") + + +@dataclass(frozen=True) +class TenantScope: + organization_id: str | None = None + workspace_id: str | None = None + project_id: str | None = None + + @classmethod + def from_dict(cls, data: dict[str, Any] | None) -> "TenantScope": + data = data or {} + scope = cls( + organization_id=_clean(data.get("organization_id")), + workspace_id=_clean(data.get("workspace_id")), + project_id=_clean(data.get("project_id")), + ) + _check_len("organization_id", scope.organization_id) + _check_len("workspace_id", scope.workspace_id) + _check_len("project_id", scope.project_id) + return scope + + +def _entity_page_fields( + entity: dict[str, Any], document: dict[str, Any], relationships: list[dict[str, Any]], evidence: list[dict[str, Any]] +) -> tuple[str, str, str, list[dict[str, Any]], list[dict[str, Any]], dict[str, Any], str] | None: + text = entity.get("text") + if not text or not isinstance(text, str): + return None + category = str(entity.get("category") or "concepts") + slug = text.replace(" ", "_") + _check_len("category", category) + _check_len("slug", slug) + entity_id = deterministic_entity_id(str(entity.get("type") or category), text) + entity_relationships = [ + r for r in relationships if r.get("source_entity") == text or r.get("target_entity") == text + ] + entity_evidence = [e for e in evidence if e.get("entity") == text] + provenance = { + "source": document.get("source"), + "url": document.get("url"), + "doi": document.get("doi"), + } + title = str(document.get("title") or text) + return category, slug, entity_id, entity_relationships, entity_evidence, provenance, title + + +class WikiRepository(Protocol): + async def compile_pages( + self, + *, + tenant: TenantScope, + document: dict[str, Any], + entities: list[dict[str, Any]], + summary: dict[str, Any], + relationships: list[dict[str, Any]], + evidence: list[dict[str, Any]], + chunks: list[dict[str, Any]], + ) -> list[dict[str, Any]]: ... + + async def get_page(self, tenant: TenantScope, page_id: str) -> dict[str, Any] | None: ... + + async def list_versions(self, tenant: TenantScope, page_id: str) -> list[dict[str, Any]] | None: ... + + async def get_version(self, tenant: TenantScope, page_id: str, version: int) -> dict[str, Any] | None: ... + + async def find_page(self, tenant: TenantScope, *, category: str, slug: str) -> dict[str, Any] | None: ... + + async def query_pages(self, tenant: TenantScope, limit: int) -> list[dict[str, Any]]: ... + + +class PostgresWikiRepository: + def __init__(self, pool: asyncpg.Pool): + self._pool = pool + + async def compile_pages( + self, + *, + tenant: TenantScope, + document: dict[str, Any], + entities: list[dict[str, Any]], + summary: dict[str, Any], + relationships: list[dict[str, Any]], + evidence: list[dict[str, Any]], + chunks: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + results: list[dict[str, Any]] = [] + async with tenant_connection(self._pool, tenant.organization_id) as conn: + for entity in entities: + fields = _entity_page_fields(entity, document, relationships, evidence) + if fields is None: + continue + category, slug, entity_id, entity_relationships, entity_evidence, provenance, title = fields + + existing = await conn.fetchrow( + """ + SELECT id, current_version FROM llm_wiki.wiki_pages + WHERE organization_id IS NOT DISTINCT FROM $1 + AND workspace_id IS NOT DISTINCT FROM $2 + AND category = $3 AND slug = $4 + """, + tenant.organization_id, + tenant.workspace_id, + category, + slug, + ) + now = datetime.now(timezone.utc) + if existing is None: + page_id = await conn.fetchval( + """ + INSERT INTO llm_wiki.wiki_pages + (organization_id, workspace_id, project_id, category, slug, + entity_id, title, current_version, created_at, updated_at) + VALUES ($1,$2,$3,$4,$5,$6,$7,1,$8,$8) + RETURNING id + """, + tenant.organization_id, + tenant.workspace_id, + tenant.project_id, + category, + slug, + entity_id, + title, + now, + ) + version = 1 + else: + page_id = existing["id"] + version = existing["current_version"] + 1 + await conn.execute( + """ + UPDATE llm_wiki.wiki_pages + SET current_version = $2, title = $3, entity_id = $4, updated_at = $5 + WHERE id = $1 + """, + page_id, + version, + title, + entity_id, + now, + ) + + await conn.execute( + """ + INSERT INTO llm_wiki.wiki_page_versions + (page_id, version, document, entity, summary, relationships, evidence, chunks, provenance, created_at) + VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) + """, + page_id, + version, + document, + entity, + summary, + entity_relationships, + entity_evidence, + chunks, + provenance, + now, + ) + results.append( + { + "id": str(page_id), + "category": category, + "slug": slug, + "entity_id": entity_id, + "version": version, + } + ) + return results + + async def get_page(self, tenant: TenantScope, page_id: str) -> dict[str, Any] | None: + async with tenant_connection(self._pool, tenant.organization_id) as conn: + page = await conn.fetchrow( + """ + SELECT * FROM llm_wiki.wiki_pages + WHERE id = $1 AND (organization_id IS NULL OR organization_id = $2) + """, + page_id, + tenant.organization_id, + ) + if page is None: + return None + version = await conn.fetchrow( + "SELECT * FROM llm_wiki.wiki_page_versions WHERE page_id = $1 AND version = $2", + page_id, + page["current_version"], + ) + return {"page": dict(page), "version": dict(version) if version else None} + + async def list_versions(self, tenant: TenantScope, page_id: str) -> list[dict[str, Any]] | None: + async with tenant_connection(self._pool, tenant.organization_id) as conn: + page = await conn.fetchrow( + """ + SELECT id FROM llm_wiki.wiki_pages + WHERE id = $1 AND (organization_id IS NULL OR organization_id = $2) + """, + page_id, + tenant.organization_id, + ) + if page is None: + return None + rows = await conn.fetch( + "SELECT version, created_at FROM llm_wiki.wiki_page_versions WHERE page_id = $1 ORDER BY version DESC", + page_id, + ) + return [dict(r) for r in rows] + + async def get_version(self, tenant: TenantScope, page_id: str, version: int) -> dict[str, Any] | None: + async with tenant_connection(self._pool, tenant.organization_id) as conn: + page = await conn.fetchrow( + """ + SELECT id FROM llm_wiki.wiki_pages + WHERE id = $1 AND (organization_id IS NULL OR organization_id = $2) + """, + page_id, + tenant.organization_id, + ) + if page is None: + return None + row = await conn.fetchrow( + "SELECT * FROM llm_wiki.wiki_page_versions WHERE page_id = $1 AND version = $2", + page_id, + version, + ) + return dict(row) if row else None + + async def find_page(self, tenant: TenantScope, *, category: str, slug: str) -> dict[str, Any] | None: + async with tenant_connection(self._pool, tenant.organization_id) as conn: + row = await conn.fetchrow( + """ + SELECT * FROM llm_wiki.wiki_pages + WHERE category = $1 AND slug = $2 + AND organization_id IS NOT DISTINCT FROM $3 + AND workspace_id IS NOT DISTINCT FROM $4 + """, + category, + slug, + tenant.organization_id, + tenant.workspace_id, + ) + return dict(row) if row else None + + async def query_pages(self, tenant: TenantScope, limit: int) -> list[dict[str, Any]]: + async with tenant_connection(self._pool, tenant.organization_id) as conn: + rows = await conn.fetch( + """ + WITH scoped AS ( + SELECT p.id, p.title, p.category, p.slug, p.updated_at, v.summary, v.document + FROM llm_wiki.wiki_pages p + JOIN llm_wiki.wiki_page_versions v + ON v.page_id = p.id AND v.version = p.current_version + WHERE (p.organization_id IS NULL OR p.organization_id = $1) + AND ($2::text IS NULL OR p.workspace_id IS NULL OR p.workspace_id = $2) + ) + SELECT * FROM scoped ORDER BY updated_at DESC LIMIT $3 + """, + tenant.organization_id, + tenant.workspace_id, + limit, + ) + return [dict(r) for r in rows] + + +class InMemoryWikiRepository: + """Test double. Reimplements the same tenant-visibility rule as the + Postgres RLS policy in Python (untenanted OR own-org) so unit tests can + assert isolation behavior without a running database.""" + + def __init__(self) -> None: + self._pages: dict[str, dict[str, Any]] = {} + self._versions: dict[str, list[dict[str, Any]]] = {} + + @staticmethod + def _visible(tenant: TenantScope, page: dict[str, Any]) -> bool: + return page["organization_id"] is None or page["organization_id"] == tenant.organization_id + + async def compile_pages( + self, + *, + tenant: TenantScope, + document: dict[str, Any], + entities: list[dict[str, Any]], + summary: dict[str, Any], + relationships: list[dict[str, Any]], + evidence: list[dict[str, Any]], + chunks: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + results: list[dict[str, Any]] = [] + for entity in entities: + fields = _entity_page_fields(entity, document, relationships, evidence) + if fields is None: + continue + category, slug, entity_id, entity_relationships, entity_evidence, provenance, title = fields + + key = (tenant.organization_id or "", tenant.workspace_id or "", category, slug) + now = datetime.now(timezone.utc) + + existing_id = next( + ( + pid + for pid, p in self._pages.items() + if (p["organization_id"] or "", p["workspace_id"] or "", p["category"], p["slug"]) == key + ), + None, + ) + if existing_id is None: + page_id = str(uuid.uuid4()) + self._pages[page_id] = { + "id": page_id, + "organization_id": tenant.organization_id, + "workspace_id": tenant.workspace_id, + "project_id": tenant.project_id, + "category": category, + "slug": slug, + "entity_id": entity_id, + "title": title, + "current_version": 1, + "created_at": now, + "updated_at": now, + } + self._versions[page_id] = [] + version = 1 + else: + page_id = existing_id + version = self._pages[page_id]["current_version"] + 1 + self._pages[page_id].update( + current_version=version, title=title, entity_id=entity_id, updated_at=now + ) + + self._versions[page_id].append( + { + "page_id": page_id, + "version": version, + "document": document, + "entity": entity, + "summary": summary, + "relationships": entity_relationships, + "evidence": entity_evidence, + "chunks": chunks, + "provenance": provenance, + "created_at": now, + } + ) + results.append( + {"id": page_id, "category": category, "slug": slug, "entity_id": entity_id, "version": version} + ) + return results + + async def get_page(self, tenant: TenantScope, page_id: str) -> dict[str, Any] | None: + page = self._pages.get(page_id) + if page is None or not self._visible(tenant, page): + return None + versions = self._versions.get(page_id, []) + latest = next((v for v in versions if v["version"] == page["current_version"]), None) + return {"page": dict(page), "version": dict(latest) if latest else None} + + async def list_versions(self, tenant: TenantScope, page_id: str) -> list[dict[str, Any]] | None: + page = self._pages.get(page_id) + if page is None or not self._visible(tenant, page): + return None + return [ + {"version": v["version"], "created_at": v["created_at"]} + for v in sorted(self._versions.get(page_id, []), key=lambda v: -v["version"]) + ] + + async def get_version(self, tenant: TenantScope, page_id: str, version: int) -> dict[str, Any] | None: + page = self._pages.get(page_id) + if page is None or not self._visible(tenant, page): + return None + return next((dict(v) for v in self._versions.get(page_id, []) if v["version"] == version), None) + + async def find_page(self, tenant: TenantScope, *, category: str, slug: str) -> dict[str, Any] | None: + key = (tenant.organization_id or "", tenant.workspace_id or "", category, slug) + for page in self._pages.values(): + if (page["organization_id"] or "", page["workspace_id"] or "", page["category"], page["slug"]) == key: + return dict(page) + return None + + async def query_pages(self, tenant: TenantScope, limit: int) -> list[dict[str, Any]]: + visible = [p for p in self._pages.values() if self._visible(tenant, p)] + if tenant.workspace_id: + visible = [p for p in visible if p["workspace_id"] is None or p["workspace_id"] == tenant.workspace_id] + visible.sort(key=lambda p: p["updated_at"], reverse=True) + out = [] + for p in visible[:limit]: + latest = next( + (v for v in self._versions.get(p["id"], []) if v["version"] == p["current_version"]), None + ) + out.append( + { + "id": p["id"], + "title": p["title"], + "category": p["category"], + "slug": p["slug"], + "updated_at": p["updated_at"], + "summary": latest["summary"] if latest else {}, + "document": latest["document"] if latest else {}, + } + ) + return out diff --git a/services/llm-wiki/app/routers/__init__.py b/services/llm-wiki/app/routers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/services/llm-wiki/app/routers/health.py b/services/llm-wiki/app/routers/health.py new file mode 100644 index 0000000..f4d363a --- /dev/null +++ b/services/llm-wiki/app/routers/health.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +from fastapi import APIRouter, Response, status + +from app.db.pool import get_pool_dependency + +router = APIRouter() + + +@router.get("/healthz") +async def healthz() -> dict[str, str]: + return {"status": "ok", "service": "llm-wiki"} + + +@router.get("/readyz") +async def readyz(response: Response) -> dict[str, str]: + try: + pool = get_pool_dependency() + async with pool.acquire() as conn: + await conn.fetchval("SELECT 1") + return {"status": "ready", "service": "llm-wiki"} + except Exception: + response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE + return {"status": "not_ready", "service": "llm-wiki"} diff --git a/services/llm-wiki/app/routers/wiki.py b/services/llm-wiki/app/routers/wiki.py new file mode 100644 index 0000000..af48a0e --- /dev/null +++ b/services/llm-wiki/app/routers/wiki.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +import uuid +from datetime import datetime +from typing import Any + +from fastapi import APIRouter, Depends, HTTPException, Query, status + +from app.core.auth import require_api_key +from app.core.config import get_settings +from app.deps import get_repository +from app.models import ( + PageDetail, + PageVersionDetail, + VersionSummary, + WikiCompileRequest, + WikiCompileResponse, + WikiHit, + WikiQueryRequest, + WikiQueryResponse, +) +from app.repository import TenantScope, ValidationError, WikiRepository + +router = APIRouter(dependencies=[Depends(require_api_key)]) + + +def _iso(value: Any) -> str: + if isinstance(value, datetime): + return value.isoformat() + return str(value) + + +def _parse_page_id(page_id: str) -> str: + try: + return str(uuid.UUID(page_id)) + except ValueError as exc: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="page not found") from exc + + +def _tenant_from_query(organization_id: str | None, workspace_id: str | None) -> TenantScope: + try: + return TenantScope.from_dict({"organization_id": organization_id, "workspace_id": workspace_id}) + except ValidationError as exc: + raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)) from exc + + +@router.post("/api/v1/wiki/compile", response_model=WikiCompileResponse, status_code=status.HTTP_201_CREATED) +async def compile_wiki( + body: WikiCompileRequest, repo: WikiRepository = Depends(get_repository) +) -> WikiCompileResponse: + try: + tenant = TenantScope.from_dict(body.tenant) + pages = await repo.compile_pages( + tenant=tenant, + document=body.document, + entities=body.entities, + summary=body.summary, + relationships=body.relationships, + evidence=body.evidence, + chunks=body.chunks, + ) + except ValidationError as exc: + raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)) from exc + return WikiCompileResponse(pages=pages) + + +@router.post("/llmwiki/query", response_model=WikiQueryResponse) +async def query_wiki( + body: WikiQueryRequest, repo: WikiRepository = Depends(get_repository) +) -> WikiQueryResponse: + settings = get_settings() + limit = min(max(body.limit or settings.query_default_limit, 1), settings.query_max_limit) + try: + tenant = TenantScope.from_dict( + {"organization_id": body.organization_id, "workspace_id": body.workspace_id} + ) + except ValidationError as exc: + raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)) from exc + + rows = await repo.query_pages(tenant, limit) + items = [] + for i, row in enumerate(rows): + summary = row.get("summary") or {} + snippet = str(summary.get("concise_summary") or "")[:280] + items.append( + WikiHit( + id=str(row["id"]), + score=round(1.0 - i * 0.01, 4), + title=row.get("title") or row.get("slug") or "", + snippet=snippet, + source="llm-wiki", + ) + ) + return WikiQueryResponse(items=items) + + +@router.get("/api/v1/wiki/pages/{page_id}", response_model=PageDetail) +async def get_page( + page_id: str, + organization_id: str | None = Query(default=None), + workspace_id: str | None = Query(default=None), + repo: WikiRepository = Depends(get_repository), +) -> PageDetail: + page_id = _parse_page_id(page_id) + tenant = _tenant_from_query(organization_id, workspace_id) + result = await repo.get_page(tenant, page_id) + if result is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="page not found") + + page = result["page"] + version = result["version"] + latest_version = None + if version is not None: + latest_version = PageVersionDetail( + version=version["version"], + document=version["document"], + entity=version["entity"], + summary=version["summary"], + relationships=version["relationships"], + evidence=version["evidence"], + chunks=version["chunks"], + provenance=version["provenance"], + created_at=_iso(version["created_at"]), + ) + return PageDetail( + id=str(page["id"]), + organization_id=page["organization_id"], + workspace_id=page["workspace_id"], + project_id=page["project_id"], + category=page["category"], + slug=page["slug"], + entity_id=page["entity_id"], + title=page["title"], + current_version=page["current_version"], + created_at=_iso(page["created_at"]), + updated_at=_iso(page["updated_at"]), + latest_version=latest_version, + ) + + +@router.get("/api/v1/wiki/pages/{page_id}/versions", response_model=list[VersionSummary]) +async def list_page_versions( + page_id: str, + organization_id: str | None = Query(default=None), + workspace_id: str | None = Query(default=None), + repo: WikiRepository = Depends(get_repository), +) -> list[VersionSummary]: + page_id = _parse_page_id(page_id) + tenant = _tenant_from_query(organization_id, workspace_id) + versions = await repo.list_versions(tenant, page_id) + if versions is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="page not found") + return [VersionSummary(version=v["version"], created_at=_iso(v["created_at"])) for v in versions] + + +@router.get("/api/v1/wiki/pages/{page_id}/versions/{version}", response_model=PageVersionDetail) +async def get_page_version( + page_id: str, + version: int, + organization_id: str | None = Query(default=None), + workspace_id: str | None = Query(default=None), + repo: WikiRepository = Depends(get_repository), +) -> PageVersionDetail: + page_id = _parse_page_id(page_id) + tenant = _tenant_from_query(organization_id, workspace_id) + row = await repo.get_version(tenant, page_id, version) + if row is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="version not found") + return PageVersionDetail( + version=row["version"], + document=row["document"], + entity=row["entity"], + summary=row["summary"], + relationships=row["relationships"], + evidence=row["evidence"], + chunks=row["chunks"], + provenance=row["provenance"], + created_at=_iso(row["created_at"]), + ) + + +@router.get("/api/v1/wiki/pages", response_model=PageDetail) +async def find_page( + category: str = Query(...), + slug: str = Query(...), + organization_id: str | None = Query(default=None), + workspace_id: str | None = Query(default=None), + repo: WikiRepository = Depends(get_repository), +) -> PageDetail: + tenant = _tenant_from_query(organization_id, workspace_id) + page = await repo.find_page(tenant, category=category, slug=slug) + if page is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="page not found") + return await get_page(str(page["id"]), organization_id, workspace_id, repo) diff --git a/services/llm-wiki/migrations/001_wiki_schema.sql b/services/llm-wiki/migrations/001_wiki_schema.sql new file mode 100644 index 0000000..895b4a0 --- /dev/null +++ b/services/llm-wiki/migrations/001_wiki_schema.sql @@ -0,0 +1,104 @@ +-- LLM Wiki service schema. +-- +-- Additive only: lives in its own `llm_wiki` schema inside the shared +-- ai_rxos Postgres database (the same instance every other AI-RxOS service +-- already uses via DATABASE_URL) and never touches a table owned by +-- another service (auth's `public` schema tables, etc.). +CREATE SCHEMA IF NOT EXISTS llm_wiki; + +CREATE EXTENSION IF NOT EXISTS pgcrypto; + +-- Session-local tenant scope, set by the application on every connection +-- before it runs a query (see app/db/pool.py::tenant_connection). This +-- compares against a text GUC rather than a uuid foreign key the way +-- services/auth's app_current_tenant() does, because organization_id here +-- is an opaque caller-supplied string (from the wiki write/query payload), +-- not a row this database owns. +CREATE OR REPLACE FUNCTION llm_wiki.current_tenant() RETURNS text AS $$ + SELECT NULLIF(current_setting('app.wiki_organization_id', true), '') +$$ LANGUAGE sql STABLE; + +CREATE TABLE IF NOT EXISTS llm_wiki.wiki_pages ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + organization_id TEXT, + workspace_id TEXT, + project_id TEXT, + category TEXT NOT NULL, + slug TEXT NOT NULL, + entity_id TEXT, + title TEXT, + current_version INTEGER NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- NULLs are never equal under a plain UNIQUE constraint, so an untenanted +-- page (organization_id/workspace_id both NULL, matching wiki_client.py's +-- pre-Prompt-8 untenanted path) would otherwise never collide with itself +-- on repeated writes. Coalescing to '' makes the natural key +-- (org, workspace, category, slug) unique in every case. +CREATE UNIQUE INDEX IF NOT EXISTS wiki_pages_natural_key + ON llm_wiki.wiki_pages ( + COALESCE(organization_id, ''), + COALESCE(workspace_id, ''), + category, + slug + ); + +CREATE INDEX IF NOT EXISTS wiki_pages_org_idx ON llm_wiki.wiki_pages (organization_id); +CREATE INDEX IF NOT EXISTS wiki_pages_updated_idx ON llm_wiki.wiki_pages (updated_at DESC); + +CREATE TABLE IF NOT EXISTS llm_wiki.wiki_page_versions ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + page_id UUID NOT NULL REFERENCES llm_wiki.wiki_pages(id) ON DELETE CASCADE, + version INTEGER NOT NULL, + document JSONB NOT NULL DEFAULT '{}'::jsonb, + entity JSONB NOT NULL DEFAULT '{}'::jsonb, + summary JSONB NOT NULL DEFAULT '{}'::jsonb, + relationships JSONB NOT NULL DEFAULT '[]'::jsonb, + evidence JSONB NOT NULL DEFAULT '[]'::jsonb, + chunks JSONB NOT NULL DEFAULT '[]'::jsonb, + provenance JSONB NOT NULL DEFAULT '{}'::jsonb, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + UNIQUE (page_id, version) +); + +CREATE INDEX IF NOT EXISTS wiki_page_versions_page_idx + ON llm_wiki.wiki_page_versions (page_id, version DESC); + +ALTER TABLE llm_wiki.wiki_pages ENABLE ROW LEVEL SECURITY; +ALTER TABLE llm_wiki.wiki_pages FORCE ROW LEVEL SECURITY; +ALTER TABLE llm_wiki.wiki_page_versions ENABLE ROW LEVEL SECURITY; +ALTER TABLE llm_wiki.wiki_page_versions FORCE ROW LEVEL SECURITY; + +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_policies WHERE schemaname = 'llm_wiki' AND tablename = 'wiki_pages' AND policyname = 'wiki_pages_isolation' + ) THEN + -- A page is visible/writable when it's untenanted (global, pre-Prompt-8 + -- content) or when it belongs to the caller's own organization. Unlike + -- services/auth's admin-bypass RLS pattern, a request with no + -- organization in scope (current_tenant() IS NULL) does NOT see every + -- tenant's data -- it only ever sees untenanted pages. This is + -- defense-in-depth: the application layer applies the same filter + -- explicitly on every query regardless of RLS being enabled. + CREATE POLICY wiki_pages_isolation ON llm_wiki.wiki_pages FOR ALL + USING (organization_id IS NULL OR organization_id = llm_wiki.current_tenant()) + WITH CHECK (organization_id IS NULL OR organization_id = llm_wiki.current_tenant()); + END IF; + + IF NOT EXISTS ( + SELECT 1 FROM pg_policies WHERE schemaname = 'llm_wiki' AND tablename = 'wiki_page_versions' AND policyname = 'wiki_page_versions_isolation' + ) THEN + CREATE POLICY wiki_page_versions_isolation ON llm_wiki.wiki_page_versions FOR ALL + USING (EXISTS ( + SELECT 1 FROM llm_wiki.wiki_pages p WHERE p.id = wiki_page_versions.page_id + AND (p.organization_id IS NULL OR p.organization_id = llm_wiki.current_tenant()) + )) + WITH CHECK (EXISTS ( + SELECT 1 FROM llm_wiki.wiki_pages p WHERE p.id = wiki_page_versions.page_id + AND (p.organization_id IS NULL OR p.organization_id = llm_wiki.current_tenant()) + )); + END IF; +END $$; diff --git a/services/llm-wiki/pytest.ini b/services/llm-wiki/pytest.ini new file mode 100644 index 0000000..2f4c80e --- /dev/null +++ b/services/llm-wiki/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +asyncio_mode = auto diff --git a/services/llm-wiki/requirements.txt b/services/llm-wiki/requirements.txt new file mode 100644 index 0000000..17dfa3f --- /dev/null +++ b/services/llm-wiki/requirements.txt @@ -0,0 +1,8 @@ +fastapi==0.115.6 +uvicorn[standard]==0.34.0 +pydantic==2.10.4 +pydantic-settings==2.7.1 +asyncpg==0.30.0 +httpx==0.28.1 +pytest==8.3.4 +pytest-asyncio==0.25.1 diff --git a/services/llm-wiki/tests/__init__.py b/services/llm-wiki/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/services/llm-wiki/tests/conftest.py b/services/llm-wiki/tests/conftest.py new file mode 100644 index 0000000..1f352d6 --- /dev/null +++ b/services/llm-wiki/tests/conftest.py @@ -0,0 +1,39 @@ +import os +import sys +from pathlib import Path + +ROOT_DIR = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT_DIR)) + +os.environ["ENVIRONMENT"] = "test" + +import pytest # noqa: E402 +from fastapi.testclient import TestClient # noqa: E402 + +from app.core.auth import require_api_key # noqa: E402 +from app.deps import get_repository # noqa: E402 +from app.main import app # noqa: E402 +from app.repository import InMemoryWikiRepository # noqa: E402 + + +@pytest.fixture() +def repo() -> InMemoryWikiRepository: + return InMemoryWikiRepository() + + +@pytest.fixture() +def client(repo): + """Fast, DB-independent client: in-memory repository, auth bypassed. + + Deliberately does NOT use `with TestClient(app) as client:` -- entering + the context manager would run app/main.py's lifespan, which calls + init_pool() and requires a real Postgres connection. Every route here + gets its data access through the overridden `get_repository` + dependency, so the (never-initialized) real pool is never touched. + """ + app.dependency_overrides[get_repository] = lambda: repo + app.dependency_overrides[require_api_key] = lambda: None + try: + yield TestClient(app) + finally: + app.dependency_overrides.clear() diff --git a/services/llm-wiki/tests/test_auth.py b/services/llm-wiki/tests/test_auth.py new file mode 100644 index 0000000..a02eea1 --- /dev/null +++ b/services/llm-wiki/tests/test_auth.py @@ -0,0 +1,62 @@ +import pytest +from fastapi.testclient import TestClient + +from app.core.config import get_settings +from app.deps import get_repository +from app.main import app +from app.repository import InMemoryWikiRepository + +_COMPILE_BODY = {"document": {"title": "T"}, "entities": [], "summary": {}} + + +@pytest.fixture(autouse=True) +def _reset_settings(monkeypatch): + monkeypatch.delenv("LLM_WIKI_API_KEY", raising=False) + get_settings.cache_clear() + yield + get_settings.cache_clear() + + +@pytest.fixture() +def unauthed_client(): + """A client with data access faked out but auth NOT bypassed -- unlike + the `client` fixture in conftest.py, this exercises the real + require_api_key dependency.""" + app.dependency_overrides[get_repository] = lambda: InMemoryWikiRepository() + try: + yield TestClient(app) + finally: + app.dependency_overrides.clear() + + +def test_no_api_key_configured_allows_unauthenticated_requests(unauthed_client): + # Dev-mode posture: LLM_WIKI_API_KEY unset means no auth is enforced, + # matching services/literature's LLMWikiClient (which only ever sends + # a bearer token when it has one configured). + res = unauthed_client.post("/api/v1/wiki/compile", json=_COMPILE_BODY) + assert res.status_code == 201 + + +def test_missing_bearer_token_rejected_when_api_key_configured(unauthed_client, monkeypatch): + monkeypatch.setenv("LLM_WIKI_API_KEY", "secret-key") + get_settings.cache_clear() + res = unauthed_client.post("/api/v1/wiki/compile", json=_COMPILE_BODY) + assert res.status_code == 401 + + +def test_wrong_api_key_rejected(unauthed_client, monkeypatch): + monkeypatch.setenv("LLM_WIKI_API_KEY", "secret-key") + get_settings.cache_clear() + res = unauthed_client.post( + "/api/v1/wiki/compile", json=_COMPILE_BODY, headers={"Authorization": "Bearer wrong"} + ) + assert res.status_code == 401 + + +def test_correct_api_key_accepted(unauthed_client, monkeypatch): + monkeypatch.setenv("LLM_WIKI_API_KEY", "secret-key") + get_settings.cache_clear() + res = unauthed_client.post( + "/api/v1/wiki/compile", json=_COMPILE_BODY, headers={"Authorization": "Bearer secret-key"} + ) + assert res.status_code == 201 diff --git a/services/llm-wiki/tests/test_compile_and_retrieve.py b/services/llm-wiki/tests/test_compile_and_retrieve.py new file mode 100644 index 0000000..8f7991b --- /dev/null +++ b/services/llm-wiki/tests/test_compile_and_retrieve.py @@ -0,0 +1,113 @@ +"""Create/get/metadata/evidence/provenance round trip -- the core wiki +contract mirrored from services/literature/app/integrations/wiki_client.py's +own request shape (see WikiCompileRequest in app/models.py).""" + + +def _payload(**overrides): + base = { + "document": { + "source": "pubmed", + "source_id": "PMID1", + "title": "HER2 Targeted Therapy", + "doi": "10.1000/xyz", + "url": "https://example.com/42", + }, + "entities": [{"text": "trastuzumab", "category": "drugs", "type": "drug"}], + "summary": {"concise_summary": "s1"}, + "relationships": [ + { + "source_entity": "trastuzumab", + "target_entity": "her2", + "predicate": "targets", + "confidence": 0.9, + } + ], + "evidence": [{"entity": "trastuzumab", "category": "efficacy", "score": 0.8}], + "tenant": {}, + "chunks": [{"chunk_id": "doc-1:0", "chunk_index": 0, "text": "chunk text"}], + } + base.update(overrides) + return base + + +def test_create_page_returns_id_category_slug_version(client): + res = client.post("/api/v1/wiki/compile", json=_payload()) + assert res.status_code == 201 + body = res.json() + assert body["success"] is True + assert len(body["pages"]) == 1 + page_ref = body["pages"][0] + assert page_ref["category"] == "drugs" + assert page_ref["slug"] == "trastuzumab" + assert page_ref["version"] == 1 + + +def test_get_page_returns_metadata_evidence_provenance_relationships(client): + create_res = client.post("/api/v1/wiki/compile", json=_payload()) + page_id = create_res.json()["pages"][0]["id"] + + res = client.get(f"/api/v1/wiki/pages/{page_id}") + assert res.status_code == 200 + page = res.json() + + assert page["category"] == "drugs" + assert page["slug"] == "trastuzumab" + assert page["current_version"] == 1 + + lv = page["latest_version"] + assert lv["summary"]["concise_summary"] == "s1" + assert lv["provenance"] == { + "source": "pubmed", + "url": "https://example.com/42", + "doi": "10.1000/xyz", + } + assert lv["evidence"] == [{"entity": "trastuzumab", "category": "efficacy", "score": 0.8}] + assert lv["relationships"][0]["predicate"] == "targets" + assert lv["relationships"][0]["confidence"] == 0.9 + assert lv["chunks"][0]["chunk_id"] == "doc-1:0" + assert lv["document"]["doi"] == "10.1000/xyz" + + +def test_update_page_overwrites_metadata_and_bumps_version(client): + create_res = client.post("/api/v1/wiki/compile", json=_payload()) + page_id = create_res.json()["pages"][0]["id"] + + update_res = client.post( + "/api/v1/wiki/compile", json=_payload(summary={"concise_summary": "s2 updated"}) + ) + assert update_res.json()["pages"][0]["id"] == page_id + assert update_res.json()["pages"][0]["version"] == 2 + + page = client.get(f"/api/v1/wiki/pages/{page_id}").json() + assert page["current_version"] == 2 + assert page["latest_version"]["summary"]["concise_summary"] == "s2 updated" + + +def test_relationships_and_evidence_are_filtered_per_entity(client): + payload = _payload( + entities=[ + {"text": "trastuzumab", "category": "drugs"}, + {"text": "her2", "category": "genes"}, + ], + relationships=[ + { + "source_entity": "trastuzumab", + "target_entity": "her2", + "predicate": "targets", + "confidence": 0.9, + } + ], + evidence=[{"entity": "trastuzumab", "category": "efficacy", "score": 0.8}], + ) + res = client.post("/api/v1/wiki/compile", json=payload) + pages_by_slug = {p["slug"]: p["id"] for p in res.json()["pages"]} + + trastuzumab_page = client.get(f"/api/v1/wiki/pages/{pages_by_slug['trastuzumab']}").json() + her2_page = client.get(f"/api/v1/wiki/pages/{pages_by_slug['her2']}").json() + + # Both entities are relationship endpoints, so both pages see it... + assert len(trastuzumab_page["latest_version"]["relationships"]) == 1 + assert len(her2_page["latest_version"]["relationships"]) == 1 + # ...but evidence was only recorded against trastuzumab. + assert len(trastuzumab_page["latest_version"]["evidence"]) == 1 + assert len(her2_page["latest_version"]["evidence"]) == 0 diff --git a/services/llm-wiki/tests/test_e2e_literature_roundtrip.py b/services/llm-wiki/tests/test_e2e_literature_roundtrip.py new file mode 100644 index 0000000..524cd7b --- /dev/null +++ b/services/llm-wiki/tests/test_e2e_literature_roundtrip.py @@ -0,0 +1,195 @@ +"""The real end-to-end proof requested for this service: + + services/literature's actual LLMWikiClient + -> real HTTP call to a real running llm-wiki service + -> real Postgres (llm_wiki schema) + -> retrieved back through llm-wiki's own read API + +`wiki_client.py` is never modified or mocked -- it's imported and run +unmodified, in a **separate Python subprocess**, because both services' +top-level package is named `app`; importing both `services/literature/app` +and `services/llm-wiki/app` in one interpreter would collide in +sys.modules. A subprocess sidesteps that without touching either service's +code. + +Skipped automatically when Postgres isn't reachable or services/literature +isn't checked out alongside this service. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import subprocess +import sys +import threading +import time +import uuid +from pathlib import Path + +import asyncpg +import httpx +import pytest +import uvicorn + +from app.core.config import get_settings + +TEST_DATABASE_URL = os.environ.get( + "TEST_DATABASE_URL", "postgresql://ai_rxos:changeme@localhost:15432/ai_rxos" +) +LITERATURE_ROOT = Path(__file__).resolve().parents[2] / "literature" + + +def _postgres_available() -> bool: + async def _check() -> bool: + try: + conn = await asyncpg.connect(TEST_DATABASE_URL, timeout=2) + await conn.close() + return True + except Exception: + return False + + return asyncio.run(_check()) + + +pytestmark = [ + pytest.mark.skipif( + not _postgres_available(), + reason=f"Postgres not reachable at {TEST_DATABASE_URL} -- run `docker compose up -d postgres`", + ), + pytest.mark.skipif( + not (LITERATURE_ROOT / "app" / "integrations" / "wiki_client.py").exists(), + reason="services/literature not found next to services/llm-wiki", + ), +] + + +@pytest.fixture() +def running_llm_wiki_server(monkeypatch): + port = 18092 + api_key = f"e2e-key-{uuid.uuid4()}" + monkeypatch.setenv("DATABASE_URL", TEST_DATABASE_URL) + monkeypatch.setenv("LLM_WIKI_API_KEY", api_key) + get_settings.cache_clear() + + import importlib + + import app.main as main_module + + importlib.reload(main_module) # re-reads settings with the env vars set above + + config = uvicorn.Config(main_module.app, host="127.0.0.1", port=port, log_level="warning") + server = uvicorn.Server(config) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + + deadline = time.time() + 10 + while not server.started and time.time() < deadline: + time.sleep(0.05) + assert server.started, "llm-wiki test server did not start in time" + + yield f"http://127.0.0.1:{port}", api_key + + server.should_exit = True + thread.join(timeout=5) + get_settings.cache_clear() + + +def _run_literature_wiki_update( + base_url: str, + api_key: str, + *, + tenant: dict, + document: dict, + entities: list, + summary: dict, + relationships: list, + evidence: list, +) -> dict: + """Runs the real, unmodified LLMWikiClient.update_wiki() in a fresh + subprocess against `base_url`, and returns its parsed JSON result.""" + script = f""" +import json, sys +sys.path.insert(0, {str(LITERATURE_ROOT)!r}) +from app.integrations.wiki_client import LLMWikiClient + +client = LLMWikiClient({{ + "wiki_service_url": {base_url!r}, + "wiki_api_key": {api_key!r}, +}}) +result = client.update_wiki( + {document!r}, + {entities!r}, + {summary!r}, + relationships={relationships!r}, + evidence={evidence!r}, + tenant={tenant!r}, +) +print(json.dumps(result)) +""" + proc = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=30 + ) + assert proc.returncode == 0, f"literature-side subprocess failed:\n{proc.stderr}" + return json.loads(proc.stdout.strip().splitlines()[-1]) + + +def test_literature_client_writes_through_llm_wiki_to_postgres_and_back(running_llm_wiki_server): + base_url, api_key = running_llm_wiki_server + marker = f"e2e-{uuid.uuid4()}" + document = { + "source": "pubmed", + "source_id": "PMID-E2E", + "title": "E2E Title", + "doi": "10.9/e2e", + "url": "http://e2e.example", + } + entities = [{"text": "e2e-drug", "category": "drugs", "type": "drug"}] + summary = {"concise_summary": "e2e summary"} + relationships = [ + { + "source_entity": "e2e-drug", + "target_entity": "e2e-target", + "predicate": "targets", + "confidence": 0.5, + } + ] + evidence = [{"entity": "e2e-drug", "category": "efficacy", "score": 0.6}] + + try: + result = _run_literature_wiki_update( + base_url, + api_key, + tenant={"organization_id": marker}, + document=document, + entities=entities, + summary=summary, + relationships=relationships, + evidence=evidence, + ) + # method == "http" proves LLMWikiClient took the real remote path, + # not the local markdown wiki-root fallback. + assert result == {"success": True, "method": "http", "status": "completed"} + + res = httpx.get( + f"{base_url}/api/v1/wiki/pages", + params={"category": "drugs", "slug": "e2e-drug", "organization_id": marker}, + headers={"Authorization": f"Bearer {api_key}"}, + timeout=5, + ) + assert res.status_code == 200 + page = res.json() + assert page["organization_id"] == marker + assert page["latest_version"]["summary"]["concise_summary"] == "e2e summary" + assert page["latest_version"]["provenance"]["doi"] == "10.9/e2e" + assert page["latest_version"]["evidence"][0]["score"] == 0.6 + assert page["latest_version"]["relationships"][0]["predicate"] == "targets" + finally: + + async def _cleanup() -> None: + conn = await asyncpg.connect(TEST_DATABASE_URL) + await conn.execute("DELETE FROM llm_wiki.wiki_pages WHERE organization_id = $1", marker) + await conn.close() + + asyncio.run(_cleanup()) diff --git a/services/llm-wiki/tests/test_health.py b/services/llm-wiki/tests/test_health.py new file mode 100644 index 0000000..2149595 --- /dev/null +++ b/services/llm-wiki/tests/test_health.py @@ -0,0 +1,12 @@ +def test_healthz(client): + res = client.get("/healthz") + assert res.status_code == 200 + assert res.json() == {"status": "ok", "service": "llm-wiki"} + + +def test_readyz_reports_not_ready_without_a_pool(client): + # `client` never runs lifespan (see conftest.py), so the real pool is + # never initialized -- readyz must fail closed, not crash. + res = client.get("/readyz") + assert res.status_code == 503 + assert res.json()["status"] == "not_ready" diff --git a/services/llm-wiki/tests/test_postgres_integration.py b/services/llm-wiki/tests/test_postgres_integration.py new file mode 100644 index 0000000..186a542 --- /dev/null +++ b/services/llm-wiki/tests/test_postgres_integration.py @@ -0,0 +1,174 @@ +"""Real-Postgres integration tests -- create/get/update/version/tenant +isolation/persistence exercised against the actual llm_wiki schema +(migrations/001_wiki_schema.sql), not the InMemoryWikiRepository double the +fast unit suite uses. + +Skipped automatically when no Postgres is reachable. Point TEST_DATABASE_URL +elsewhere, or run `docker compose up -d postgres` for the default +localhost:15432 target (matches docker-compose.yml's host port mapping). +""" + +from __future__ import annotations + +import asyncio +import json +import os +import uuid + +import asyncpg +import pytest + +from app.db.pool import MIGRATIONS_DIR +from app.repository import PostgresWikiRepository, TenantScope + +TEST_DATABASE_URL = os.environ.get( + "TEST_DATABASE_URL", "postgresql://ai_rxos:changeme@localhost:15432/ai_rxos" +) + + +async def _init_connection(conn: asyncpg.Connection) -> None: + await conn.set_type_codec( + "jsonb", encoder=json.dumps, decoder=json.loads, schema="pg_catalog", format="text" + ) + + +def _postgres_available() -> bool: + async def _check() -> bool: + try: + conn = await asyncpg.connect(TEST_DATABASE_URL, timeout=2) + await conn.close() + return True + except Exception: + return False + + return asyncio.run(_check()) + + +pytestmark = pytest.mark.skipif( + not _postgres_available(), + reason=f"Postgres not reachable at {TEST_DATABASE_URL} -- run `docker compose up -d postgres`", +) + + +@pytest.fixture() +async def pg_repo(): + pool = await asyncpg.create_pool(TEST_DATABASE_URL, min_size=1, max_size=5, init=_init_connection) + for path in sorted(MIGRATIONS_DIR.glob("*.sql")): + async with pool.acquire() as conn: + await conn.execute(path.read_text(encoding="utf-8")) + repo = PostgresWikiRepository(pool) + marker = f"pytest-{uuid.uuid4()}" + yield repo, marker + async with pool.acquire() as conn: + await conn.execute("DELETE FROM llm_wiki.wiki_pages WHERE organization_id LIKE 'pytest-%'") + await pool.close() + + +async def test_real_create_and_retrieve_round_trip(pg_repo): + repo, marker = pg_repo + tenant = TenantScope(organization_id=marker) + + pages = await repo.compile_pages( + tenant=tenant, + document={"source": "pubmed", "source_id": "PMID1", "title": "T", "doi": "10.1/x"}, + entities=[{"text": "trastuzumab", "category": "drugs"}], + summary={"concise_summary": "s1"}, + relationships=[], + evidence=[{"entity": "trastuzumab", "category": "efficacy", "score": 0.7}], + chunks=[], + ) + assert len(pages) == 1 + page_id = pages[0]["id"] + + fetched = await repo.get_page(tenant, page_id) + assert fetched is not None + assert fetched["page"]["category"] == "drugs" + assert fetched["version"]["summary"]["concise_summary"] == "s1" + assert fetched["version"]["evidence"][0]["score"] == 0.7 + assert fetched["version"]["provenance"]["doi"] == "10.1/x" + + +async def test_real_version_increment_and_history(pg_repo): + repo, marker = pg_repo + tenant = TenantScope(organization_id=marker) + doc = {"source": "s", "source_id": "d1", "title": "T"} + + p1 = await repo.compile_pages( + tenant=tenant, + document=doc, + entities=[{"text": "her2", "category": "genes"}], + summary={"concise_summary": "v1"}, + relationships=[], + evidence=[], + chunks=[], + ) + page_id = p1[0]["id"] + assert p1[0]["version"] == 1 + + p2 = await repo.compile_pages( + tenant=tenant, + document=doc, + entities=[{"text": "her2", "category": "genes"}], + summary={"concise_summary": "v2"}, + relationships=[], + evidence=[], + chunks=[], + ) + assert p2[0]["id"] == page_id + assert p2[0]["version"] == 2 + + versions = await repo.list_versions(tenant, page_id) + assert [v["version"] for v in versions] == [2, 1] + + v1 = await repo.get_version(tenant, page_id, 1) + assert v1["summary"]["concise_summary"] == "v1" + + +async def test_real_tenant_a_cannot_read_tenant_b(pg_repo): + repo, marker = pg_repo + org_a = TenantScope(organization_id=f"{marker}-a") + org_b = TenantScope(organization_id=f"{marker}-b") + doc = {"source": "s", "source_id": "d1", "title": "T"} + + pages = await repo.compile_pages( + tenant=org_a, + document=doc, + entities=[{"text": "shared-slug", "category": "drugs"}], + summary={"concise_summary": "org-a data"}, + relationships=[], + evidence=[], + chunks=[], + ) + page_id = pages[0]["id"] + + assert await repo.get_page(org_b, page_id) is None + assert await repo.get_page(org_a, page_id) is not None + + +async def test_persistence_survives_a_fresh_connection_pool(pg_repo): + """Simulates a service restart: a brand-new pool (not the fixture's) + must still see data written earlier, proving persistence isn't an + artifact of connection-local state or in-process caching.""" + repo, marker = pg_repo + tenant = TenantScope(organization_id=marker) + pages = await repo.compile_pages( + tenant=tenant, + document={"source": "s", "source_id": "d1", "title": "T"}, + entities=[{"text": "restart-check", "category": "drugs"}], + summary={"concise_summary": "still here"}, + relationships=[], + evidence=[], + chunks=[], + ) + page_id = pages[0]["id"] + + fresh_pool = await asyncpg.create_pool( + TEST_DATABASE_URL, min_size=1, max_size=2, init=_init_connection + ) + try: + fresh_repo = PostgresWikiRepository(fresh_pool) + fetched = await fresh_repo.get_page(tenant, page_id) + assert fetched is not None + assert fetched["version"]["summary"]["concise_summary"] == "still here" + finally: + await fresh_pool.close() diff --git a/services/llm-wiki/tests/test_query_endpoint.py b/services/llm-wiki/tests/test_query_endpoint.py new file mode 100644 index 0000000..a877184 --- /dev/null +++ b/services/llm-wiki/tests/test_query_endpoint.py @@ -0,0 +1,38 @@ +"""POST /llmwiki/query response shape must match services/search's Hit +struct exactly (services/search/internal/search/opensearch.go:161-169) -- +`id`, `score`, `title`, `snippet` at minimum, since that's what +services/search's LLMWikiProvider.SimilaritySearchForTenant decodes.""" + + +def test_query_response_matches_hit_struct_field_names(client): + client.post( + "/api/v1/wiki/compile", + json={ + "document": {"source": "pubmed", "source_id": "P1", "title": "T"}, + "entities": [{"text": "aspirin", "category": "drugs"}], + "summary": {"concise_summary": "aspirin summary"}, + }, + ) + res = client.post("/llmwiki/query", json={"embedding": [0.1, 0.2], "limit": 5}) + assert res.status_code == 200 + items = res.json()["items"] + assert len(items) == 1 + hit = items[0] + assert set(hit.keys()) >= {"id", "score", "title", "snippet"} + assert hit["snippet"] == "aspirin summary" + + +def test_query_with_no_matching_tenant_returns_empty_items(client): + client.post( + "/api/v1/wiki/compile", + json={ + "document": {"source": "s", "source_id": "P2", "title": "T"}, + "entities": [{"text": "ibuprofen", "category": "drugs"}], + "summary": {"concise_summary": "s"}, + "tenant": {"organization_id": "org-a"}, + }, + ) + res = client.post( + "/llmwiki/query", json={"embedding": [], "limit": 10, "organization_id": "org-does-not-exist"} + ) + assert res.json() == {"items": []} diff --git a/services/llm-wiki/tests/test_tenant_isolation.py b/services/llm-wiki/tests/test_tenant_isolation.py new file mode 100644 index 0000000..420ebf7 --- /dev/null +++ b/services/llm-wiki/tests/test_tenant_isolation.py @@ -0,0 +1,64 @@ +def _compile(client, org, title): + return client.post( + "/api/v1/wiki/compile", + json={ + "document": {"source": "s", "source_id": org, "title": title}, + "entities": [{"text": "shared-drug-name", "category": "drugs"}], + "summary": {"concise_summary": title}, + "tenant": {"organization_id": org}, + }, + ) + + +def test_tenant_a_cannot_retrieve_tenant_b_page_by_id(client): + res_a = _compile(client, "org-a", "org-a page") + page_id_a = res_a.json()["pages"][0]["id"] + + # org-b explicitly asking for org-a's page id must not see it. + res_as_b = client.get(f"/api/v1/wiki/pages/{page_id_a}", params={"organization_id": "org-b"}) + assert res_as_b.status_code == 404 + + # org-a can read its own page. + res_as_a = client.get(f"/api/v1/wiki/pages/{page_id_a}", params={"organization_id": "org-a"}) + assert res_as_a.status_code == 200 + + # a caller with no organization scope at all also cannot see it. + res_no_org = client.get(f"/api/v1/wiki/pages/{page_id_a}") + assert res_no_org.status_code == 404 + + +def test_same_category_slug_creates_separate_pages_per_tenant(client): + res_a = _compile(client, "org-a", "org-a page") + res_b = _compile(client, "org-b", "org-b page") + id_a = res_a.json()["pages"][0]["id"] + id_b = res_b.json()["pages"][0]["id"] + assert id_a != id_b + + +def test_tenant_isolation_in_query_results(client): + id_a = _compile(client, "org-a", "org-a page").json()["pages"][0]["id"] + id_b = _compile(client, "org-b", "org-b page").json()["pages"][0]["id"] + + res = client.post( + "/llmwiki/query", json={"embedding": [], "limit": 10, "organization_id": "org-a"} + ) + returned_ids = {item["id"] for item in res.json()["items"]} + assert id_a in returned_ids + assert id_b not in returned_ids + + +def test_untenanted_page_is_visible_to_every_org(client): + global_res = client.post( + "/api/v1/wiki/compile", + json={ + "document": {"source": "s", "source_id": "g1", "title": "global page"}, + "entities": [{"text": "global-concept", "category": "concepts"}], + "summary": {"concise_summary": "global"}, + }, + ) + page_id = global_res.json()["pages"][0]["id"] + + for org in (None, "org-a", "org-b"): + params = {"organization_id": org} if org else {} + res = client.get(f"/api/v1/wiki/pages/{page_id}", params=params) + assert res.status_code == 200 diff --git a/services/llm-wiki/tests/test_validation.py b/services/llm-wiki/tests/test_validation.py new file mode 100644 index 0000000..2a3d9a1 --- /dev/null +++ b/services/llm-wiki/tests/test_validation.py @@ -0,0 +1,51 @@ +import uuid + + +def test_invalid_field_types_rejected(client): + res = client.post( + "/api/v1/wiki/compile", json={"document": {}, "entities": "not-a-list", "summary": {}} + ) + assert res.status_code == 422 + + +def test_missing_body_rejected(client): + res = client.post("/api/v1/wiki/compile") + assert res.status_code == 422 + + +def test_entities_without_text_are_skipped_not_erroring(client): + res = client.post( + "/api/v1/wiki/compile", + json={"document": {"title": "T"}, "entities": [{"category": "drugs"}], "summary": {}}, + ) + assert res.status_code == 201 + assert res.json()["pages"] == [] + + +def test_organization_id_too_long_rejected(client): + res = client.post( + "/api/v1/wiki/compile", + json={ + "document": {"title": "T"}, + "entities": [{"text": "x"}], + "summary": {}, + "tenant": {"organization_id": "a" * 600}, + }, + ) + assert res.status_code == 422 + + +def test_get_nonexistent_page_returns_404(client): + res = client.get(f"/api/v1/wiki/pages/{uuid.uuid4()}") + assert res.status_code == 404 + + +def test_get_page_with_malformed_id_returns_404_not_500(client): + res = client.get("/api/v1/wiki/pages/not-a-uuid") + assert res.status_code == 404 + + +def test_query_limit_is_clamped_not_rejected(client): + res = client.post("/llmwiki/query", json={"embedding": [], "limit": 10000}) + assert res.status_code == 200 + assert res.json() == {"items": []} diff --git a/services/llm-wiki/tests/test_versions.py b/services/llm-wiki/tests/test_versions.py new file mode 100644 index 0000000..0908277 --- /dev/null +++ b/services/llm-wiki/tests/test_versions.py @@ -0,0 +1,39 @@ +def _payload(summary_text): + return { + "document": {"source": "pubmed", "source_id": "PMID1", "title": "T1"}, + "entities": [{"text": "her2", "category": "genes"}], + "summary": {"concise_summary": summary_text}, + } + + +def test_version_history_and_archived_versions_are_retrievable(client): + res1 = client.post("/api/v1/wiki/compile", json=_payload("v1 summary")) + page_id = res1.json()["pages"][0]["id"] + assert res1.json()["pages"][0]["version"] == 1 + + res2 = client.post("/api/v1/wiki/compile", json=_payload("v2 summary")) + assert res2.json()["pages"][0]["version"] == 2 + + res3 = client.post("/api/v1/wiki/compile", json=_payload("v3 summary")) + assert res3.json()["pages"][0]["version"] == 3 + + versions = client.get(f"/api/v1/wiki/pages/{page_id}/versions").json() + assert [v["version"] for v in versions] == [3, 2, 1] + + v1 = client.get(f"/api/v1/wiki/pages/{page_id}/versions/1").json() + assert v1["summary"]["concise_summary"] == "v1 summary" + + v2 = client.get(f"/api/v1/wiki/pages/{page_id}/versions/2").json() + assert v2["summary"]["concise_summary"] == "v2 summary" + + # current page state reflects only the latest version + page = client.get(f"/api/v1/wiki/pages/{page_id}").json() + assert page["latest_version"]["summary"]["concise_summary"] == "v3 summary" + assert page["current_version"] == 3 + + +def test_unknown_version_returns_404(client): + res = client.post("/api/v1/wiki/compile", json=_payload("v1")) + page_id = res.json()["pages"][0]["id"] + missing = client.get(f"/api/v1/wiki/pages/{page_id}/versions/99") + assert missing.status_code == 404 diff --git a/services/search/cmd/search/main.go b/services/search/cmd/search/main.go index e8b87c1..a9b22fd 100644 --- a/services/search/cmd/search/main.go +++ b/services/search/cmd/search/main.go @@ -90,6 +90,7 @@ func main() { r.Get("/", h.Query) r.Post("/", h.Hybrid) r.Post("/index", h.Index) + r.Post("/context", h.Context) r.Get("/stream", h.StreamQuery) r.Post("/stream", h.StreamHybrid) }) diff --git a/services/search/internal/handlers/context.go b/services/search/internal/handlers/context.go new file mode 100644 index 0000000..1c8606f --- /dev/null +++ b/services/search/internal/handlers/context.go @@ -0,0 +1,181 @@ +package handlers + +import ( + "encoding/json" + "net/http" + "strings" + + "github.com/openhealthagents/ai-rxos/services/search/internal/search" +) + +// contextRequest is the input to POST /api/v1/search/context — the Prompt 8 +// context-optimization endpoint. It reuses the same retrieval/ranking legs +// as Hybrid (OpenSearch BM25 + LLM Wiki/QMD vector + graph/citation +// enrichment + RRF), but trims the output to compact, cited snippets +// instead of full documents, and supports the filters/scoping a caller +// building agent or conversation context needs. +type contextRequest struct { + Query string `json:"query"` + Embedding []float32 `json:"embedding,omitempty"` + TopK int `json:"top_k,omitempty"` + SourceFilters []string `json:"source_filters,omitempty"` + EntityFilters []string `json:"entity_filters,omitempty"` + OrganizationID string `json:"organization_id,omitempty"` + WorkspaceID string `json:"workspace_id,omitempty"` + ProjectID string `json:"project_id,omitempty"` + ConversationID string `json:"conversation_id,omitempty"` + AgentID string `json:"agent_id,omitempty"` + // MaxSnippetChars caps each returned snippet's length (default 320). + // The point of this endpoint is compact context, not full documents. + MaxSnippetChars int `json:"max_snippet_chars,omitempty"` +} + +func (r contextRequest) tenant() search.TenantScope { + return search.TenantScope{OrgID: r.OrganizationID, WorkspaceID: r.WorkspaceID} +} + +// contextItem is a single piece of retrieved context: a compact snippet +// plus enough provenance/citation metadata to trace it back to its source, +// never the full source document. +type contextItem struct { + ID string `json:"id"` + Title string `json:"title"` + Snippet string `json:"snippet"` + Source string `json:"source"` + Score float64 `json:"score"` + Citation string `json:"citation"` +} + +// Context handles POST /api/v1/search/context — builds compact, cited +// context for agent/conversation consumption from the same hybrid +// retrieval pipeline Hybrid uses, rather than a second ranking engine. +func (h *SearchHandler) Context(w http.ResponseWriter, r *http.Request) { + var req contextRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"code": "invalid_body", "message": err.Error()}) + return + } + if req.Query == "" && len(req.Embedding) == 0 { + writeJSON(w, http.StatusBadRequest, map[string]string{"code": "missing_query", "message": "query or embedding is required"}) + return + } + topK := req.TopK + if topK <= 0 { + topK = 10 + } + maxSnippet := req.MaxSnippetChars + if maxSnippet <= 0 { + maxSnippet = 320 + } + + var bm25Hits []search.Hit + var vectorHits []search.Hit + + if req.Query != "" && h.OpenSearch != nil { + if hits, err := h.OpenSearch.Query(r.Context(), req.Query, topK*2); err == nil { + for i := range hits { + if hits[i].Source == "" { + hits[i].Source = "opensearch" + } + } + bm25Hits = hits + } + } + + if len(req.Embedding) > 0 && h.Vectors != nil { + source := h.VectorSource + if source == "" { + source = search.ProviderLLMWiki + } + if hits, err := h.Vectors.SimilaritySearchForTenant(r.Context(), req.Embedding, topK*2, req.tenant()); err == nil { + for i := range hits { + if hits[i].Source == "" { + hits[i].Source = source + } + } + vectorHits = hits + } + } + + if h.Citations != nil { + bm25Hits = h.Citations.EnrichHits(r.Context(), bm25Hits) + vectorHits = h.Citations.EnrichHits(r.Context(), vectorHits) + } + if h.Graph != nil { + bm25Hits = h.Graph.EnrichHits(r.Context(), req.Query, bm25Hits) + vectorHits = h.Graph.EnrichHits(r.Context(), req.Query, vectorHits) + } + + ranker := h.Ranker + if ranker == nil { + ranker = search.NewResultRanker(60, 0.35, 0.35, 0.15, 0.15) + } + ranked := ranker.RankRRF(topK*2, bm25Hits, vectorHits) + ranked = applyContextFilters(ranked, req.SourceFilters, req.EntityFilters) + if len(ranked) > topK { + ranked = ranked[:topK] + } + + items := make([]contextItem, 0, len(ranked)) + for _, hit := range ranked { + items = append(items, contextItem{ + ID: hit.ID, + Title: hit.Title, + Snippet: truncateSnippet(hit.Snippet, maxSnippet), + Source: hit.Source, + Score: hit.Score, + Citation: hit.Title, + }) + } + + writeJSON(w, http.StatusOK, map[string]any{ + "items": items, + "total": len(items), + "scope": map[string]string{ + "organization_id": req.OrganizationID, + "workspace_id": req.WorkspaceID, + "project_id": req.ProjectID, + "conversation_id": req.ConversationID, + "agent_id": req.AgentID, + }, + }) +} + +func applyContextFilters(hits []search.Hit, sourceFilters, entityFilters []string) []search.Hit { + if len(sourceFilters) == 0 && len(entityFilters) == 0 { + return hits + } + sourceSet := make(map[string]bool, len(sourceFilters)) + for _, s := range sourceFilters { + sourceSet[strings.ToLower(s)] = true + } + + filtered := make([]search.Hit, 0, len(hits)) + for _, hit := range hits { + if len(sourceFilters) > 0 && !sourceSet[strings.ToLower(hit.Source)] { + continue + } + if len(entityFilters) > 0 { + haystack := strings.ToLower(hit.Title + " " + hit.Snippet) + matched := false + for _, entity := range entityFilters { + if strings.Contains(haystack, strings.ToLower(entity)) { + matched = true + break + } + } + if !matched { + continue + } + } + filtered = append(filtered, hit) + } + return filtered +} + +func truncateSnippet(snippet string, maxChars int) string { + if len(snippet) <= maxChars { + return snippet + } + return strings.TrimSpace(snippet[:maxChars]) + "..." +} diff --git a/services/search/internal/handlers/index.go b/services/search/internal/handlers/index.go index 8d474f7..e77ccfb 100644 --- a/services/search/internal/handlers/index.go +++ b/services/search/internal/handlers/index.go @@ -20,13 +20,21 @@ type indexItem struct { } type batchIndexPayload struct { - DocumentID string `json:"document_id"` - Title string `json:"title"` - Content string `json:"content"` - Source string `json:"source"` - Citations int `json:"citations"` - Documents []indexItem `json:"documents"` - Items []indexItem `json:"items"` + DocumentID string `json:"document_id"` + Title string `json:"title"` + Content string `json:"content"` + Source string `json:"source"` + Citations int `json:"citations"` + Documents []indexItem `json:"documents"` + Items []indexItem `json:"items"` + Tenant map[string]string `json:"tenant"` +} + +func tenantFromMap(tenant map[string]string) search.TenantScope { + return search.TenantScope{ + OrgID: tenant["organization_id"], + WorkspaceID: tenant["workspace_id"], + } } // Index handles POST /api/v1/search/index — indexing document text and vector embeddings @@ -39,6 +47,7 @@ func (h *SearchHandler) Index(w http.ResponseWriter, r *http.Request) { } var items []indexItem + var tenant search.TenantScope // Try decoding as JSON array first if err := json.Unmarshal(bodyBytes, &items); err != nil { @@ -52,6 +61,7 @@ func (h *SearchHandler) Index(w http.ResponseWriter, r *http.Request) { } items = []indexItem{single} } else { + tenant = tenantFromMap(wrapper.Tenant) if len(wrapper.Documents) > 0 { items = wrapper.Documents } else if len(wrapper.Items) > 0 { @@ -103,7 +113,7 @@ func (h *SearchHandler) Index(w http.ResponseWriter, r *http.Request) { }) if h.Vectors != nil { - _ = h.Vectors.Upsert(r.Context(), id, item.Title, item.Content, vec) + _ = h.Vectors.UpsertForTenant(r.Context(), id, item.Title, item.Content, vec, tenant) } if h.Citations != nil && item.Citations > 0 { h.Citations.SetCitationCount(id, item.Citations) diff --git a/services/search/internal/handlers/search.go b/services/search/internal/handlers/search.go index 01e3c86..14e6acc 100644 --- a/services/search/internal/handlers/search.go +++ b/services/search/internal/handlers/search.go @@ -18,9 +18,15 @@ type SearchHandler struct { } type hybridRequest struct { - Query string `json:"query"` - Embedding []float32 `json:"embedding,omitempty"` - Limit int `json:"limit,omitempty"` + Query string `json:"query"` + Embedding []float32 `json:"embedding,omitempty"` + Limit int `json:"limit,omitempty"` + OrganizationID string `json:"organization_id,omitempty"` + WorkspaceID string `json:"workspace_id,omitempty"` +} + +func (r hybridRequest) tenant() search.TenantScope { + return search.TenantScope{OrgID: r.OrganizationID, WorkspaceID: r.WorkspaceID} } func writeJSON(w http.ResponseWriter, status int, v any) { @@ -92,7 +98,7 @@ func (h *SearchHandler) Hybrid(w http.ResponseWriter, r *http.Request) { if source == "" { source = search.ProviderLLMWiki } - if hits, err := h.Vectors.SimilaritySearch(r.Context(), req.Embedding, req.Limit*2); err == nil { + if hits, err := h.Vectors.SimilaritySearchForTenant(r.Context(), req.Embedding, req.Limit*2, req.tenant()); err == nil { for i := range hits { if hits[i].Source == "" { hits[i].Source = source diff --git a/services/search/internal/handlers/stream.go b/services/search/internal/handlers/stream.go index 7db80e9..8af00ec 100644 --- a/services/search/internal/handlers/stream.go +++ b/services/search/internal/handlers/stream.go @@ -104,7 +104,7 @@ func (h *SearchHandler) StreamHybrid(w http.ResponseWriter, r *http.Request) { if source == "" { source = search.ProviderLLMWiki } - if hits, err := h.Vectors.SimilaritySearch(r.Context(), req.Embedding, req.Limit); err == nil { + if hits, err := h.Vectors.SimilaritySearchForTenant(r.Context(), req.Embedding, req.Limit, req.tenant()); err == nil { for i := range hits { if hits[i].Source == "" { hits[i].Source = source diff --git a/services/search/internal/search/okf_qmd.go b/services/search/internal/search/okf_qmd.go index 51feb2d..b9d14ce 100644 --- a/services/search/internal/search/okf_qmd.go +++ b/services/search/internal/search/okf_qmd.go @@ -9,6 +9,44 @@ import ( "unicode" ) +// TenantScope identifies the organization/workspace a document or query is +// scoped to. The zero value means "no tenant" (shared/public content, or an +// unscoped/admin query) — this keeps every pre-existing caller (including +// the QMD benchmark suite, which has no tenant concept at all) working +// unchanged after tenant scoping was added. +type TenantScope struct { + OrgID string + WorkspaceID string +} + +// matches reports whether a document scoped to docTenant is visible to a +// query scoped to t. +// +// - An unscoped query (t.OrgID == "") sees everything — back-compat for +// admin tooling and the benchmark suite. +// - Untenanted/shared documents (docTenant.OrgID == "") are visible to +// every tenant. +// - Otherwise OrgID must match exactly (cross-organization access is +// always blocked), and if the document also specifies a workspace, +// WorkspaceID must match exactly too (cross-workspace access to +// workspace-private content is always blocked). A document with an +// OrgID but no WorkspaceID is treated as org-shared. +func (t TenantScope) matches(docTenant TenantScope) bool { + if t.OrgID == "" { + return true + } + if docTenant.OrgID == "" { + return true + } + if t.OrgID != docTenant.OrgID { + return false + } + if docTenant.WorkspaceID == "" { + return true + } + return t.WorkspaceID == docTenant.WorkspaceID +} + // QMDDocument represents an item indexed in the local QMD engine over OKF concepts. type QMDDocument struct { ID string @@ -19,6 +57,7 @@ type QMDDocument struct { Citations int Length int Tokens map[string]int + Tenant TenantScope } // QMDEngine provides high-speed in-memory QMD (Query-Metadata-Document) hybrid search @@ -61,8 +100,17 @@ func tokenize(text string) []string { return tokens } -// IndexDocument adds or updates a document in the QMD inverted index and vector store. +// IndexDocument adds or updates a document in the QMD inverted index and +// vector store. It is untenanted (documents indexed this way are visible to +// every query) — kept for backward compatibility with existing callers +// (e.g. the benchmark suite). Use IndexDocumentForTenant to scope a +// document to an organization/workspace. func (e *QMDEngine) IndexDocument(id, title, content, source string, embedding []float32, citations int) { + e.IndexDocumentForTenant(id, title, content, source, embedding, citations, TenantScope{}) +} + +// IndexDocumentForTenant is IndexDocument with an explicit tenant scope. +func (e *QMDEngine) IndexDocumentForTenant(id, title, content, source string, embedding []float32, citations int, tenant TenantScope) { e.mu.Lock() defer e.mu.Unlock() @@ -94,6 +142,7 @@ func (e *QMDEngine) IndexDocument(id, title, content, source string, embedding [ Citations: citations, Length: docLen, Tokens: tokenCounts, + Tenant: tenant, } e.docs[id] = doc e.totalLength += docLen @@ -103,8 +152,15 @@ func (e *QMDEngine) IndexDocument(id, title, content, source string, embedding [ } } -// SearchBM25 calculates lexical BM25 scores across matching indexed documents. +// SearchBM25 calculates lexical BM25 scores across matching indexed +// documents. Untenanted — sees every indexed document regardless of tenant. +// Use SearchBM25ForTenant to enforce organization/workspace isolation. func (e *QMDEngine) SearchBM25(ctx context.Context, query string, limit int) ([]Hit, error) { + return e.SearchBM25ForTenant(ctx, query, limit, TenantScope{}) +} + +// SearchBM25ForTenant is SearchBM25 filtered to documents visible to tenant. +func (e *QMDEngine) SearchBM25ForTenant(ctx context.Context, query string, limit int, tenant TenantScope) ([]Hit, error) { e.mu.RLock() defer e.mu.RUnlock() @@ -154,6 +210,9 @@ func (e *QMDEngine) SearchBM25(ctx context.Context, query string, limit int) ([] var hits []Hit for id, score := range scores { doc := e.docs[id] + if !tenant.matches(doc.Tenant) { + continue + } hits = append(hits, Hit{ ID: doc.ID, Title: doc.Title, @@ -188,8 +247,15 @@ func cosineSimilarity(a, b []float32) float64 { return dot / (math.Sqrt(normA) * math.Sqrt(normB)) } -// SearchVector performs dense vector similarity ranking across indexed QMD documents. +// SearchVector performs dense vector similarity ranking across indexed QMD +// documents. Untenanted — sees every indexed document regardless of tenant. +// Use SearchVectorForTenant to enforce organization/workspace isolation. func (e *QMDEngine) SearchVector(ctx context.Context, embedding []float32, limit int) ([]Hit, error) { + return e.SearchVectorForTenant(ctx, embedding, limit, TenantScope{}) +} + +// SearchVectorForTenant is SearchVector filtered to documents visible to tenant. +func (e *QMDEngine) SearchVectorForTenant(ctx context.Context, embedding []float32, limit int, tenant TenantScope) ([]Hit, error) { e.mu.RLock() defer e.mu.RUnlock() @@ -202,6 +268,9 @@ func (e *QMDEngine) SearchVector(ctx context.Context, embedding []float32, limit if len(doc.Embedding) == 0 { continue } + if !tenant.matches(doc.Tenant) { + continue + } sim := cosineSimilarity(embedding, doc.Embedding) if sim > 0 { hits = append(hits, Hit{ @@ -220,10 +289,18 @@ func (e *QMDEngine) SearchVector(ctx context.Context, embedding []float32, limit return hits, nil } -// SearchHybrid combines BM25 lexical scores and local vector cosine similarity. +// SearchHybrid combines BM25 lexical scores and local vector cosine +// similarity. Untenanted — sees every indexed document regardless of +// tenant. Use SearchHybridForTenant to enforce organization/workspace +// isolation. func (e *QMDEngine) SearchHybrid(ctx context.Context, query string, embedding []float32, limit int) ([]Hit, error) { - bm25Hits, _ := e.SearchBM25(ctx, query, limit*2) - vectorHits, _ := e.SearchVector(ctx, embedding, limit*2) + return e.SearchHybridForTenant(ctx, query, embedding, limit, TenantScope{}) +} + +// SearchHybridForTenant is SearchHybrid filtered to documents visible to tenant. +func (e *QMDEngine) SearchHybridForTenant(ctx context.Context, query string, embedding []float32, limit int, tenant TenantScope) ([]Hit, error) { + bm25Hits, _ := e.SearchBM25ForTenant(ctx, query, limit*2, tenant) + vectorHits, _ := e.SearchVectorForTenant(ctx, embedding, limit*2, tenant) merged := make(map[string]*Hit) for _, h := range bm25Hits { diff --git a/services/search/internal/search/provider.go b/services/search/internal/search/provider.go index d54963f..6c8dc24 100644 --- a/services/search/internal/search/provider.go +++ b/services/search/internal/search/provider.go @@ -9,9 +9,16 @@ import ( // the semantic/vector leg of a search (as opposed to OpenSearch's keyword // leg). LLMWikiProvider and GoogleOKFProvider implement high-speed // retrieval using our Local QMD Engine and OKF microservices. +// +// SimilaritySearch/Upsert are untenanted (kept for backward compatibility +// with existing callers and tests); SimilaritySearchForTenant/ +// UpsertForTenant enforce organization/workspace isolation — see +// TenantScope.matches in okf_qmd.go for the exact isolation semantics. type RetrievalProvider interface { SimilaritySearch(ctx context.Context, embedding []float32, limit int) ([]Hit, error) + SimilaritySearchForTenant(ctx context.Context, embedding []float32, limit int, tenant TenantScope) ([]Hit, error) Upsert(ctx context.Context, id, title, content string, embedding []float32) error + UpsertForTenant(ctx context.Context, id, title, content string, embedding []float32, tenant TenantScope) error Close() } diff --git a/services/search/internal/search/providers_placeholder.go b/services/search/internal/search/providers_placeholder.go index 32748f5..4114c86 100644 --- a/services/search/internal/search/providers_placeholder.go +++ b/services/search/internal/search/providers_placeholder.go @@ -36,13 +36,21 @@ func NewLLMWikiProvider(cfg ProviderConfig) (*LLMWikiProvider, error) { } func (p *LLMWikiProvider) Upsert(ctx context.Context, id, title, content string, embedding []float32) error { - p.engine.IndexDocument(id, title, content, ProviderLLMWiki, embedding, 0) + return p.UpsertForTenant(ctx, id, title, content, embedding, TenantScope{}) +} + +func (p *LLMWikiProvider) UpsertForTenant(ctx context.Context, id, title, content string, embedding []float32, tenant TenantScope) error { + p.engine.IndexDocumentForTenant(id, title, content, ProviderLLMWiki, embedding, 0, tenant) return nil } func (p *LLMWikiProvider) SimilaritySearch(ctx context.Context, embedding []float32, limit int) ([]Hit, error) { + return p.SimilaritySearchForTenant(ctx, embedding, limit, TenantScope{}) +} + +func (p *LLMWikiProvider) SimilaritySearchForTenant(ctx context.Context, embedding []float32, limit int, tenant TenantScope) ([]Hit, error) { // Query local QMD index first - localHits, _ := p.engine.SearchVector(ctx, embedding, limit) + localHits, _ := p.engine.SearchVectorForTenant(ctx, embedding, limit, tenant) for i := range localHits { localHits[i].Source = ProviderLLMWiki } @@ -50,10 +58,14 @@ func (p *LLMWikiProvider) SimilaritySearch(ctx context.Context, embedding []floa return localHits, nil } - // Fall back to remote HTTP call if local index is empty + // Fall back to remote HTTP call if local index is empty. The remote + // LLM Wiki service is expected to enforce its own tenant isolation on + // org_id/workspace_id server-side — this client only forwards the scope. payload := map[string]any{ - "embedding": embedding, - "limit": limit, + "embedding": embedding, + "limit": limit, + "organization_id": tenant.OrgID, + "workspace_id": tenant.WorkspaceID, } body, err := json.Marshal(payload) if err != nil { @@ -116,12 +128,20 @@ func NewGoogleOKFProvider(cfg ProviderConfig) (*GoogleOKFProvider, error) { } func (p *GoogleOKFProvider) Upsert(ctx context.Context, id, title, content string, embedding []float32) error { - p.engine.IndexDocument(id, title, content, ProviderGoogleOKF, embedding, 0) + return p.UpsertForTenant(ctx, id, title, content, embedding, TenantScope{}) +} + +func (p *GoogleOKFProvider) UpsertForTenant(ctx context.Context, id, title, content string, embedding []float32, tenant TenantScope) error { + p.engine.IndexDocumentForTenant(id, title, content, ProviderGoogleOKF, embedding, 0, tenant) return nil } func (p *GoogleOKFProvider) SimilaritySearch(ctx context.Context, embedding []float32, limit int) ([]Hit, error) { - localHits, _ := p.engine.SearchVector(ctx, embedding, limit) + return p.SimilaritySearchForTenant(ctx, embedding, limit, TenantScope{}) +} + +func (p *GoogleOKFProvider) SimilaritySearchForTenant(ctx context.Context, embedding []float32, limit int, tenant TenantScope) ([]Hit, error) { + localHits, _ := p.engine.SearchVectorForTenant(ctx, embedding, limit, tenant) for i := range localHits { localHits[i].Source = ProviderGoogleOKF } @@ -129,7 +149,12 @@ func (p *GoogleOKFProvider) SimilaritySearch(ctx context.Context, embedding []fl return localHits, nil } - payload := map[string]any{"embedding": embedding, "limit": limit} + payload := map[string]any{ + "embedding": embedding, + "limit": limit, + "organization_id": tenant.OrgID, + "workspace_id": tenant.WorkspaceID, + } body, err := json.Marshal(payload) if err != nil { return nil, err diff --git a/services/search/internal/search/tenant_isolation_test.go b/services/search/internal/search/tenant_isolation_test.go new file mode 100644 index 0000000..6e1a5c4 --- /dev/null +++ b/services/search/internal/search/tenant_isolation_test.go @@ -0,0 +1,146 @@ +package search + +import ( + "context" + "testing" +) + +// These tests prove the Prompt 8 organization/workspace isolation +// guarantee for the LLM Wiki / QMD retrieval layer: a query scoped to one +// organization must never see another organization's documents, and a +// query scoped to one workspace must never see another workspace's +// private (workspace-scoped) documents within the same organization. + +func TestTenantScope_CrossOrganizationIsolation(t *testing.T) { + engine := NewQMDEngine() + ctx := context.Background() + embed := []float32{1.0, 0.0, 0.0} + + engine.IndexDocumentForTenant("doc-a", "Org A Secret", "confidential org A content", "okf_concept", embed, 0, TenantScope{OrgID: "org-a"}) + engine.IndexDocumentForTenant("doc-b", "Org B Secret", "confidential org B content", "okf_concept", embed, 0, TenantScope{OrgID: "org-b"}) + + hitsA, err := engine.SearchVectorForTenant(ctx, embed, 10, TenantScope{OrgID: "org-a"}) + if err != nil { + t.Fatalf("SearchVectorForTenant failed: %v", err) + } + for _, h := range hitsA { + if h.ID == "doc-b" { + t.Fatalf("organization A query must never see organization B's document, got %+v", hitsA) + } + } + found := false + for _, h := range hitsA { + if h.ID == "doc-a" { + found = true + } + } + if !found { + t.Fatalf("organization A query should see its own document, got %+v", hitsA) + } + + bm25HitsB, err := engine.SearchBM25ForTenant(ctx, "confidential", 10, TenantScope{OrgID: "org-b"}) + if err != nil { + t.Fatalf("SearchBM25ForTenant failed: %v", err) + } + for _, h := range bm25HitsB { + if h.ID == "doc-a" { + t.Fatalf("organization B query must never see organization A's document, got %+v", bm25HitsB) + } + } +} + +func TestTenantScope_CrossWorkspaceIsolationWithinSameOrganization(t *testing.T) { + engine := NewQMDEngine() + ctx := context.Background() + embed := []float32{0.0, 1.0, 0.0} + + engine.IndexDocumentForTenant("doc-ws1", "Workspace 1 private", "private workspace 1 content", "okf_concept", embed, 0, TenantScope{OrgID: "org-a", WorkspaceID: "ws-1"}) + engine.IndexDocumentForTenant("doc-ws2", "Workspace 2 private", "private workspace 2 content", "okf_concept", embed, 0, TenantScope{OrgID: "org-a", WorkspaceID: "ws-2"}) + // Org-shared (no workspace) content should be visible to every workspace in the org. + engine.IndexDocumentForTenant("doc-shared", "Org shared", "shared org content", "okf_concept", embed, 0, TenantScope{OrgID: "org-a"}) + + hitsWS1, err := engine.SearchVectorForTenant(ctx, embed, 10, TenantScope{OrgID: "org-a", WorkspaceID: "ws-1"}) + if err != nil { + t.Fatalf("SearchVectorForTenant failed: %v", err) + } + seen := map[string]bool{} + for _, h := range hitsWS1 { + seen[h.ID] = true + } + if seen["doc-ws2"] { + t.Fatalf("workspace 1 query must never see workspace 2's private document, got %+v", hitsWS1) + } + if !seen["doc-ws1"] { + t.Fatalf("workspace 1 query should see its own document, got %+v", hitsWS1) + } + if !seen["doc-shared"] { + t.Fatalf("workspace 1 query should see org-shared documents, got %+v", hitsWS1) + } +} + +func TestTenantScope_UnscopedQuerySeesEverything_BackwardCompat(t *testing.T) { + engine := NewQMDEngine() + ctx := context.Background() + embed := []float32{1.0, 1.0, 0.0} + + // Untenanted IndexDocument (pre-Prompt-8 call signature) must keep working. + engine.IndexDocument("legacy-doc", "Legacy", "legacy untenanted content", "benchmark", embed, 0) + engine.IndexDocumentForTenant("tenant-doc", "Tenant", "tenant scoped content", "okf_concept", embed, 0, TenantScope{OrgID: "org-a"}) + + hits, err := engine.SearchVector(ctx, embed, 10) + if err != nil { + t.Fatalf("SearchVector failed: %v", err) + } + seen := map[string]bool{} + for _, h := range hits { + seen[h.ID] = true + } + if !seen["legacy-doc"] || !seen["tenant-doc"] { + t.Fatalf("unscoped query should see both legacy and tenant-scoped documents, got %+v", hits) + } +} + +func TestTenantScope_UntenantedDocumentsAreSharedAcrossOrganizations(t *testing.T) { + engine := NewQMDEngine() + ctx := context.Background() + embed := []float32{0.5, 0.5, 0.5} + + // Content indexed before tenant scoping existed (no TenantScope) should + // remain visible to every tenant, rather than becoming orphaned. + engine.IndexDocument("pre-existing", "Pre-existing", "shared legacy content", "okf_concept", embed, 0) + + hits, err := engine.SearchVectorForTenant(ctx, embed, 10, TenantScope{OrgID: "org-a"}) + if err != nil { + t.Fatalf("SearchVectorForTenant failed: %v", err) + } + if len(hits) != 1 || hits[0].ID != "pre-existing" { + t.Fatalf("expected untenanted legacy document to remain visible, got %+v", hits) + } +} + +func TestLLMWikiProvider_UpsertForTenant_IsolatesSimilaritySearch(t *testing.T) { + ctx := context.Background() + provider, err := NewLLMWikiProvider(ProviderConfig{}) + if err != nil { + t.Fatalf("NewLLMWikiProvider failed: %v", err) + } + defer provider.Close() + + embed := []float32{0.2, 0.4, 0.6} + if err := provider.UpsertForTenant(ctx, "org-a-doc", "Org A", "org a content", embed, TenantScope{OrgID: "org-a"}); err != nil { + t.Fatalf("UpsertForTenant failed: %v", err) + } + if err := provider.UpsertForTenant(ctx, "org-b-doc", "Org B", "org b content", embed, TenantScope{OrgID: "org-b"}); err != nil { + t.Fatalf("UpsertForTenant failed: %v", err) + } + + hits, err := provider.SimilaritySearchForTenant(ctx, embed, 10, TenantScope{OrgID: "org-a"}) + if err != nil { + t.Fatalf("SimilaritySearchForTenant failed: %v", err) + } + for _, h := range hits { + if h.ID == "org-b-doc" { + t.Fatalf("LLMWikiProvider leaked org-b's document into org-a's search results: %+v", hits) + } + } +}