From acd07eabaed79357b5fd12889080ca29d2d89076 Mon Sep 17 00:00:00 2001 From: Rene Zander <96420681+renezander030@users.noreply.github.com> Date: Wed, 23 Sep 2026 09:49:37 +0000 Subject: [PATCH] Store each group in its own Ladybug file, and expose the library helpers database.ladybug.layout: per-group keeps every configured group in a separate file named by its SHA-256, so no id can address a path outside the directory. Ingest, drain and restore write into the group's file; the MCP tools and kg reads open only the file of the group they address, and a token scoped to one group never opens another. single stays the default. Library helpers: a path-keyed embedder record, public BM25-only search configs, an extension check for the running Ladybug engine, and the contained per-group path. kg export on Ladybug now exports every record when a named group holds no edges. --- CHANGELOG.md | 28 +++ README.md | 26 +++ config/ladybug.example.yaml | 3 + docs/reference.md | 36 +++ src/kg_mcp/cli.py | 97 ++++---- src/kg_mcp/config.py | 57 ++++- src/kg_mcp/demo.py | 6 +- src/kg_mcp/doctor.py | 18 +- src/kg_mcp/export.py | 5 + src/kg_mcp/fingerprint.py | 106 ++++++++- src/kg_mcp/ingest.py | 39 +++- src/kg_mcp/ladybug.py | 44 ++++ src/kg_mcp/restore.py | 58 ++++- src/kg_mcp/retrieval.py | 41 ++++ src/kg_mcp/server.py | 116 ++++++++-- src/kg_mcp/verify.py | 18 +- tests/test_v05_multi_user.py | 413 +++++++++++++++++++++++++++++++++++ 17 files changed, 1028 insertions(+), 83 deletions(-) create mode 100644 tests/test_v05_multi_user.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 259b27d..8a96fbb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,33 @@ # Changelog +## 0.5.0 (unreleased) + +### Added + +- `database.ladybug.layout: per-group` stores each configured group in its own Ladybug + file inside `database.ladybug.directory`, named by the SHA-256 of the group so no id + can address a path outside it. Ingest, drain and restore write into the group's + file; the MCP tools, `kg ask` and the other reads open only the file of the group + they address. Bearer tokens may be scoped to a subset of groups in this layout, and + a token scoped to one group never opens another group's file. `single` remains the + default. +- `kg_mcp.fingerprint.record_embedder`, `recorded_embedder` and `embedder_drift` + record and check the embedder that wrote a Ladybug file from its path alone, in + `.embedder.json`. Configured writes also keep this record, and the configured + drift check reads it. +- `kg_mcp.retrieval.keyword_edge_search_config` and `keyword_node_search_config` + return the BM25-only search configuration behind `kg ask --keyword` and `kg-demo`. +- `kg_mcp.ladybug.extension_status` reports whether the FTS and VECTOR extensions are + installed for the running Ladybug engine version, with the command that installs + them. +- `kg_mcp.ladybug.group_database_path` returns the contained per-group file path. +- `docs/reference.md` lists the library helpers. + +### Fixed + +- `kg export` on a Ladybug graph exports every record when a named group holds no + edges. + ## 0.4.1 - Use absolute HTTPS URLs for the README logo and demo so both images render in the diff --git a/README.md b/README.md index b5edbd9..76468ff 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,32 @@ To review extraction before it touches the configured graph, run printed restore command to promote those exact records. Searches return current facts by default; `kg ask --history` is the explicit historical view. +## Multiple users on Ladybug + +Give each user a group, and each group its own Ladybug file: + +```yaml +graph: + groups: [alice, bob] +database: + provider: ladybug + ladybug: + layout: per-group # default: single, one file for every group + directory: ./workspace/groups +server: + transport: streamable-http + auth: + tokens: + - {name: alice, token: "${ALICE_TOKEN}", groups: [alice]} + - {name: bob, token: "${BOB_TOKEN}", groups: [bob]} +``` + +Each file is named by the SHA-256 of its group, so no group id can point outside the +directory. Reads, `kg ask` and the drain open only the file of the group they address; +a token granted `alice` never opens Bob's file. A read names one group, or the token's +single group is used. For a tenant boundary inside one database server, use FalkorDB +or Neo4j. + ## Is it a fit? Use it for local agent memory with explicit human review. Skip it if you need diff --git a/config/ladybug.example.yaml b/config/ladybug.example.yaml index c77b6f5..7c5d151 100644 --- a/config/ladybug.example.yaml +++ b/config/ladybug.example.yaml @@ -23,6 +23,9 @@ database: provider: ladybug ladybug: path: ./workspace/example.ladybug + # One file per group instead, named by the group's SHA-256 inside `directory`: + # layout: per-group + # directory: ./workspace/groups llm: model: ${KG_LLM_MODEL:gpt-4.1-mini} diff --git a/docs/reference.md b/docs/reference.md index 4aa5ec0..ceb4dfd 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -92,6 +92,42 @@ indexes before the read-only server opens it. `kg-ladybug-setup --apply` prepare them explicitly. Opening a reader never installs extensions automatically. For a database created by 0.2.x, run setup again to add the indexes. +With `database.ladybug.layout: per-group`, each configured group is a separate file in +`database.ladybug.directory`, named `.ladybug`. Ingest, drain and +restore write each record into its group's file, which the first write creates, and +record the embedder beside that file. `kg ask`, `kg nodes`, `kg episodes`, `kg export` +and the MCP search tools read one group per call; `kg status`, `kg edge`, +`get_entity_edge` and `get_status` visit every group file the caller may read. Tokens +may be scoped to a subset of groups in this layout. + +### Library use + +These helpers work from a database path and need no configuration file: + +| Name | Purpose | +| --- | --- | +| `kg_mcp.ladybug.build_ladybug_driver(path, read_only=...)` | A graphiti driver on one Ladybug file | +| `kg_mcp.ladybug.group_database_path(directory, group)` | The contained SHA-256 file path for a group or user id | +| `kg_mcp.ladybug.extension_status(path=None)` | `engine_version`, `installed`, `missing`, `ok` and the `fix` command for the FTS and VECTOR extensions of the running engine | +| `kg_mcp.fingerprint.record_embedder(path, model=..., dimensions=...)` | Record the embedder that wrote a file, in `.embedder.json` | +| `kg_mcp.fingerprint.recorded_embedder(path)` | The recorded embedder, or `None` before the first write | +| `kg_mcp.fingerprint.embedder_drift(path, model=..., dimensions=...)` | A refusal message when a different embedder would write, else `None` | +| `kg_mcp.retrieval.keyword_edge_search_config(limit)` | A BM25-only `SearchConfig` for facts; no model call | +| `kg_mcp.retrieval.keyword_node_search_config(limit)` | The same for entities | +| `kg_mcp.write_lock.ladybug_path_lock(path, timeout)` | The cross-process writer lock the CLI uses | + +```python +from kg_mcp.ladybug import build_ladybug_driver, extension_status, group_database_path +from kg_mcp.retrieval import keyword_edge_search_config + +status = extension_status() +if not status["ok"]: + raise SystemExit(f"missing {status['missing']}: {status['fix']}") +path = group_database_path("./memory", user_id) +driver = build_ladybug_driver(str(path), read_only=True) +# Graphiti(graph_driver=driver, ...).search_(query, config=keyword_edge_search_config(10)) +``` + ## Explicit ingestion Ingestion is a separate command and is dry-run by default: diff --git a/src/kg_mcp/cli.py b/src/kg_mcp/cli.py index 325a314..b91b90c 100644 --- a/src/kg_mcp/cli.py +++ b/src/kg_mcp/cli.py @@ -82,47 +82,6 @@ def _doctor_command(args: argparse.Namespace) -> dict[str, Any]: return {"status": worst_status(results), "checks": results} -def _keyword_edge_config(limit: int): - """Edge search on BM25 alone, reranked by reciprocal rank fusion. - - The vector half of the default hybrid search embeds the query, which means an - embedding backend has to be reachable before any question can be answered. Keyword - search needs no such call, so a shipped graph stays readable on a machine with - nothing running. Reciprocal rank fusion is kept because it orders a single result - list unchanged and costs no extra call. - """ - from graphiti_core.search.search_config import ( - EdgeReranker, - EdgeSearchConfig, - EdgeSearchMethod, - SearchConfig, - ) - - return SearchConfig( - edge_config=EdgeSearchConfig( - search_methods=[EdgeSearchMethod.bm25], reranker=EdgeReranker.rrf - ), - limit=limit, - ) - - -def _keyword_node_config(limit: int): - """Node equivalent of :func:`_keyword_edge_config`.""" - from graphiti_core.search.search_config import ( - NodeReranker, - NodeSearchConfig, - NodeSearchMethod, - SearchConfig, - ) - - return SearchConfig( - node_config=NodeSearchConfig( - search_methods=[NodeSearchMethod.bm25], reranker=NodeReranker.rrf - ), - limit=limit, - ) - - async def _edge_search( graph: Any, args: argparse.Namespace, @@ -134,6 +93,7 @@ async def _edge_search( candidate_limit, compatible_query, current_edges, + keyword_edge_search_config, rerank_edges, ) @@ -144,7 +104,7 @@ async def _edge_search( else: results = await graph.search_( query, - config=_keyword_edge_config(candidates), + config=keyword_edge_search_config(candidates), group_ids=groups, ) edges = results.edges[:candidates] @@ -155,7 +115,50 @@ async def _edge_search( async def _graph_command(args: argparse.Namespace) -> dict[str, Any]: + from kg_mcp.config import per_group_ladybug, settings_for_group, single_group + settings = load_config() + if not per_group_ladybug(settings): + return await _graph_command_on(args, settings) + if args.command in ("status", "edge"): + return await _each_group_file(args, settings) + # One file per group: a read opens only the file of the one group it names. + requested = list(getattr(args, "groups", None) or []) + allowed_groups(requested or None, settings) + group = single_group(requested or settings.graph.groups, settings) + args.groups = [group] + return await _graph_command_on(args, settings_for_group(settings, group)) + + +async def _each_group_file(args: argparse.Namespace, settings: Any) -> dict[str, Any]: + """``status`` and ``edge`` visit every configured group's file that exists.""" + from pathlib import Path + + from kg_mcp.config import settings_for_group + + existing = [ + group + for group in settings.graph.groups + if Path(settings.database.ladybug.path_for(group)).exists() + ] + if args.command == "status": + counts = {group: 0 for group in settings.graph.groups} + for group in existing: + payload = await _graph_command_on(args, settings_for_group(settings, group)) + counts[group] = payload["groups"][0]["nodes"] + return { + "provider": "ladybug", + "groups": [{"group": group, "nodes": count} for group, count in counts.items()], + } + for group in existing: + try: + return await _graph_command_on(args, settings_for_group(settings, group)) + except CommandError: + continue + raise CommandError(f"edge not found: {args.uuid}") + + +async def _graph_command_on(args: argparse.Namespace, settings: Any) -> dict[str, Any]: from kg_mcp.runtime import build_graphiti timeout = settings.graph.query_timeout_seconds @@ -185,18 +188,22 @@ async def _graph_command(args: argparse.Namespace) -> dict[str, Any]: ], "suppressed_invalidated": suppressed, "ranker": settings.reranker.provider, - "pending": _pending_items(groups), + "pending": _pending_items(groups or settings.graph.groups), } if args.command == "nodes": groups = allowed_groups(args.groups, settings) - from kg_mcp.retrieval import candidate_limit, compatible_query + from kg_mcp.retrieval import ( + candidate_limit, + compatible_query, + keyword_node_search_config, + ) candidate_limit(args.limit, settings) query = compatible_query(args.query, settings) from graphiti_core.search.search_config_recipes import NODE_HYBRID_SEARCH_RRF node_config = ( - _keyword_node_config(args.limit) if args.keyword else NODE_HYBRID_SEARCH_RRF + keyword_node_search_config(args.limit) if args.keyword else NODE_HYBRID_SEARCH_RRF ) result = await bounded( graph.search_(query, config=node_config, group_ids=groups), diff --git a/src/kg_mcp/config.py b/src/kg_mcp/config.py index 45b6f90..f2bec6c 100644 --- a/src/kg_mcp/config.py +++ b/src/kg_mcp/config.py @@ -185,7 +185,26 @@ class Neo4jConfig(BaseModel): class LadybugConfig(BaseModel): + """Where the embedded graph lives. + + ``single`` keeps every group in one file at ``path``. ``per-group`` gives each + configured group its own file inside ``directory``, named by the SHA-256 of the + group, so a reader bound to one group never opens another group's file. + """ + path: str = "./workspace/graph.ladybug" + layout: Literal["single", "per-group"] = "single" + directory: str = "./workspace/groups" + + def path_for(self, group: str | None = None) -> str: + """The database file that holds ``group``.""" + if self.layout == "single": + return self.path + if not group: + raise ValueError("the per-group Ladybug layout needs a group to choose a file") + from kg_mcp.ladybug import group_database_path + + return str(group_database_path(self.directory, group)) class DatabaseConfig(BaseModel): @@ -214,12 +233,14 @@ def validate_auth_scopes(self) -> Settings: raise ValueError(f"server auth token grants an unconfigured group: {forbidden[0]}") if ( self.database.provider == "ladybug" + and self.database.ladybug.layout == "single" and grant.groups and set(grant.groups) != configured ): raise ValueError( "Ladybug is a single embedded graph and cannot enforce a partial token " - "group scope; grant all configured groups or use FalkorDB/Neo4j" + "group scope; grant all configured groups, set " + "database.ladybug.layout: per-group, or use FalkorDB/Neo4j" ) return self @@ -256,3 +277,37 @@ def allowed_groups(requested: str | list[str] | None, settings: Settings) -> lis if settings.database.provider == "ladybug": return None return groups + + +def per_group_ladybug(settings: Settings) -> bool: + """Whether each group lives in its own Ladybug file.""" + database = settings.database + return database.provider == "ladybug" and database.ladybug.layout == "per-group" + + +def settings_for_group(settings: Settings, group: str) -> Settings: + """A view of ``settings`` bound to the one file that holds ``group``. + + With the per-group Ladybug layout every reader and writer goes through this, so the + rest of the code sees an ordinary single-file graph that contains only that group. + Any other configuration is returned unchanged. + """ + if not per_group_ladybug(settings): + return settings + allowed_groups(group, settings) + ladybug = settings.database.ladybug.model_copy( + update={"path": settings.database.ladybug.path_for(group), "layout": "single"} + ) + database = settings.database.model_copy(update={"ladybug": ladybug}) + graph = settings.graph.model_copy(update={"groups": [group]}) + return settings.model_copy(update={"database": database, "graph": graph}) + + +def single_group(groups: list[str], settings: Settings) -> str: + """The one group a per-group Ladybug read addresses, or a refusal naming the choice.""" + if len(groups) != 1: + raise ValueError( + "the per-group Ladybug layout reads one group per call; name one of: " + + ", ".join(groups or settings.graph.groups) + ) + return groups[0] diff --git a/src/kg_mcp/demo.py b/src/kg_mcp/demo.py index b6d450b..be3fcbc 100644 --- a/src/kg_mcp/demo.py +++ b/src/kg_mcp/demo.py @@ -141,18 +141,16 @@ def _ensure_graph(home: Path) -> Path: async def _ask(question: str, limit: int) -> list[dict[str, Any]]: - from kg_mcp.cli import _keyword_edge_config from kg_mcp.config import load_config + from kg_mcp.retrieval import current_edges, keyword_edge_search_config from kg_mcp.runtime import build_graphiti settings = load_config() graph = build_graphiti(settings, read_only=True) try: results = await graph.search_( - question, config=_keyword_edge_config(limit), group_ids=[GROUP] + question, config=keyword_edge_search_config(limit), group_ids=[GROUP] ) - from kg_mcp.retrieval import current_edges - edges, _ = current_edges(list(results.edges), include_invalidated=False) return [ {"fact": edge.fact, "valid_at": edge.valid_at.isoformat() if edge.valid_at else None} diff --git a/src/kg_mcp/doctor.py b/src/kg_mcp/doctor.py index 1f8b40d..2799599 100644 --- a/src/kg_mcp/doctor.py +++ b/src/kg_mcp/doctor.py @@ -275,6 +275,20 @@ def check_embedding_model_known(settings: Settings) -> dict[str, str]: ) +def _per_group(check: Any, settings: Settings) -> list[dict[str, str]]: + """Run a per-graph check once per group file when each group has its own file.""" + from kg_mcp.config import per_group_ladybug, settings_for_group + + if not per_group_ladybug(settings): + return [check(settings)] + results = [] + for group in settings.graph.groups: + result = check(settings_for_group(settings, group)) + result["check"] = f"{result['check']}[{group}]" + results.append(result) + return results + + def run_checks(explicit: str | None = None, *, offline: bool = False) -> list[dict[str, str]]: results = [check_versions()] config_results, settings = check_config(explicit) @@ -286,13 +300,13 @@ def run_checks(explicit: str | None = None, *, offline: bool = False) -> list[di results.append(check_reranker(settings)) results.append(check_embedding_model_known(settings)) results.append(check_embedding_dim_binding(settings)) - results.append(check_embedder_fingerprint(settings)) + results.extend(_per_group(check_embedder_fingerprint, settings)) if offline: results.append(_check("database", WARN, "skipped (--offline)")) results.append(_check("llm", WARN, "skipped (--offline)")) results.append(_check("embedder", WARN, "skipped (--offline)")) return results - results.append(check_database(settings)) + results.extend(_per_group(check_database, settings)) results.append(check_llm(settings)) results.append(check_embedder(settings)) return results diff --git a/src/kg_mcp/export.py b/src/kg_mcp/export.py index f9bcbfd..cc1c65b 100644 --- a/src/kg_mcp/export.py +++ b/src/kg_mcp/export.py @@ -35,6 +35,7 @@ def default_output() -> Path: async def _collect(driver: Any, groups: list[str]) -> list[dict[str, Any]]: from graphiti_core.edges import EntityEdge, EpisodicEdge + from graphiti_core.errors import GroupsEdgesNotFoundError from graphiti_core.nodes import EntityNode, EpisodicNode kinds = ( @@ -47,6 +48,10 @@ async def _collect(driver: Any, groups: list[str]) -> list[dict[str, Any]]: for model, kind in kinds: try: records = await model.get_by_group_ids(driver, groups) + except GroupsEdgesNotFoundError: + # graphiti reports "no edges in these groups" as an error; it is an empty + # result, not a partial collection. + records = [] except Exception as exc: raise RuntimeError(f"snapshot export could not collect {kind}: {exc}") from exc collected.extend(_serialize(record, kind) for record in records or []) diff --git a/src/kg_mcp/fingerprint.py b/src/kg_mcp/fingerprint.py index b407f78..d5146fa 100644 --- a/src/kg_mcp/fingerprint.py +++ b/src/kg_mcp/fingerprint.py @@ -5,6 +5,11 @@ search still returns results, and the ranking is quietly wrong. So the first ingest records the embedder for the database it wrote to, later writes must match it, and ``kg doctor`` compares the configuration against the record. + +For a Ladybug file the record is also kept beside the file, as ``.embedder.json``. +:func:`recorded_embedder`, :func:`record_embedder` and :func:`embedder_drift` read and +write that record from a database path alone, for callers that use the graph as a +library without a configuration file or workspace. """ from __future__ import annotations @@ -18,6 +23,72 @@ from kg_mcp.config import Settings, embedding_model_base FILE = "embedders.json" +SIDECAR_SUFFIX = ".embedder.json" + + +def sidecar_path(database_path: str | Path) -> Path: + target = Path(database_path).expanduser().resolve() + return target.with_name(target.name + SIDECAR_SUFFIX) + + +def recorded_embedder(database_path: str | Path) -> dict[str, Any] | None: + """The embedder recorded for a Ladybug file, or ``None`` before its first write.""" + path = sidecar_path(database_path) + if not path.exists(): + return None + try: + data = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + return None + return data if isinstance(data, dict) else None + + +def record_embedder( + database_path: str | Path, + *, + model: str, + dimensions: int, + api_url: str = "", +) -> dict[str, Any]: + """Record the embedder that wrote a Ladybug file. Call it after a write lands.""" + from filelock import FileLock + + path = sidecar_path(database_path) + path.parent.mkdir(parents=True, exist_ok=True) + entry = { + "model": embedding_model_base(model), + "dimensions": dimensions, + "api_url": api_url, + "recorded_at": datetime.now(timezone.utc).isoformat(), + } + with FileLock(str(path) + ".lock"): + temporary = path.with_name(path.name + ".tmp") + temporary.write_text(json.dumps(entry, indent=2, ensure_ascii=False), encoding="utf-8") + os.replace(temporary, path) + return entry + + +def _drift_message(previous: dict[str, Any], model: str, dimensions: int) -> str | None: + if ( + previous.get("model") == embedding_model_base(model) + and previous.get("dimensions") == dimensions + ): + return None + return ( + f"the stored vectors were produced by {previous.get('model')} " + f"({previous.get('dimensions')} dimensions) but the embedder is now " + f"{model} ({dimensions}); mixing embedders corrupts search silently. Restore the " + "previous embedder, or rebuild the graph under the new one: kg export, an empty " + "database, then kg-ingest SNAPSHOT --restore --apply" + ) + + +def embedder_drift(database_path: str | Path, *, model: str, dimensions: int) -> str | None: + """A refusal message when ``model`` is not the embedder that wrote the file.""" + previous = recorded_embedder(database_path) + if previous is None: + return None + return _drift_message(previous, model, dimensions) def database_key(settings: Settings) -> str: @@ -59,7 +130,21 @@ def _load() -> dict[str, Any]: def recorded(settings: Settings) -> dict[str, Any] | None: entry = _load().get(database_key(settings)) - return entry if isinstance(entry, dict) else None + if isinstance(entry, dict): + return entry + path = _ladybug_file(settings) + return recorded_embedder(path) if path else None + + +def _ladybug_file(settings: Settings) -> str | None: + """The one Ladybug file ``settings`` addresses, if it addresses exactly one.""" + if settings.database.provider != "ladybug": + return None + groups = settings.graph.groups + try: + return settings.database.ladybug.path_for(groups[0] if len(groups) == 1 else None) + except ValueError: + return None def record(settings: Settings) -> dict[str, Any]: @@ -78,6 +163,14 @@ def record(settings: Settings) -> dict[str, Any]: temporary = target.with_suffix(".json.tmp") temporary.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8") os.replace(temporary, target) + path = _ladybug_file(settings) + if path: + record_embedder( + path, + model=settings.embedder.model, + dimensions=settings.embedder.dimensions, + api_url=settings.embedder.api_url, + ) return entry @@ -86,16 +179,7 @@ def drift(settings: Settings) -> str | None: previous = recorded(settings) if previous is None: return None - now = current(settings) - if previous.get("model") == now["model"] and previous.get("dimensions") == now["dimensions"]: - return None - return ( - f"the stored vectors were produced by {previous.get('model')} " - f"({previous.get('dimensions')} dimensions) but embedder.model is " - f"{settings.embedder.model} ({now['dimensions']}); mixing embedders corrupts search " - "silently. Restore the previous embedder, or rebuild the graph under the new one: " - "kg export, an empty database, then kg-ingest SNAPSHOT --restore --apply" - ) + return _drift_message(previous, settings.embedder.model, settings.embedder.dimensions) def check(settings: Settings) -> tuple[str, str]: diff --git a/src/kg_mcp/ingest.py b/src/kg_mcp/ingest.py index 071985b..25fec33 100644 --- a/src/kg_mcp/ingest.py +++ b/src/kg_mcp/ingest.py @@ -159,6 +159,43 @@ async def ingest_records( domain = str(record.get("domain") or settings.graph.groups[0]) allowed_groups(domain, settings) + from kg_mcp.config import per_group_ladybug, settings_for_group + + if per_group_ladybug(settings): + # One file per group: each domain's records go to its own file, under its own + # writer lock and embedder record, and never open another group's file. + by_domain: dict[str, list[dict[str, Any]]] = {} + for record in records: + by_domain.setdefault( + str(record.get("domain") or settings.graph.groups[0]), [] + ).append(record) + merged: dict[str, Any] = { + "applied": apply, + "ingested": 0, + "skipped": 0, + "failed": [], + "interrupted": False, + "planned": [], + } + for domain, subset in by_domain.items(): + part = await ingest_records( + subset, + apply=apply, + resume=resume, + fail_fast=fail_fast, + ledger=ledger, + settings=settings_for_group(settings, domain), + track_fingerprint=track_fingerprint, + ) + for key in ("ingested", "skipped"): + merged[key] += part[key] + merged["failed"].extend(part["failed"]) + merged["planned"].extend(part["planned"]) + merged["interrupted"] = merged["interrupted"] or part["interrupted"] + if part["interrupted"] or (fail_fast and part["failed"]): + break + return merged + seen = completed_keys() if (resume and ledger) else set() planned, skipped = [], 0 for record in records: @@ -278,7 +315,7 @@ async def stage_extraction( update={ "provider": "ladybug", "ladybug": settings.database.ladybug.model_copy( - update={"path": str(database_path)} + update={"path": str(database_path), "layout": "single"} ), } ) diff --git a/src/kg_mcp/ladybug.py b/src/kg_mcp/ladybug.py index dabfbcb..3897b63 100644 --- a/src/kg_mcp/ladybug.py +++ b/src/kg_mcp/ladybug.py @@ -14,6 +14,7 @@ from __future__ import annotations import argparse +import hashlib import os import re import sys @@ -33,6 +34,49 @@ def _alias_kuzu() -> Any: return ladybug +def group_database_path(directory: str | Path, group: str) -> Path: + """The file that holds one group in the per-group layout. + + The name is the SHA-256 of the group, so no group or user id, however hostile, can + name a path outside ``directory``: ``../x``, an absolute path, a NUL byte and a + 5000-character id all map to 64 hex characters inside it. + """ + if not isinstance(group, str) or not group: + raise ValueError("group must be a non-empty string") + root = Path(directory).expanduser().resolve() + digest = hashlib.sha256(group.encode("utf-8", "surrogatepass")).hexdigest() + target = root / f"{digest}.ladybug" + if target.parent != root: # cannot happen with a hex name; kept as a second guard + raise ValueError(f"group file escapes {root}") + return target + + +def extension_status(database: str | Path | None = None) -> dict[str, Any]: + """Whether the FTS and VECTOR extensions are installed for the running engine. + + Ladybug installs extensions per engine version, so an upgrade of the ``ladybug`` + package leaves them missing until they are installed again. The check loads them + into a throwaway in-memory database and never touches a graph file. ``fix`` is the + command that installs them, or ``None`` when nothing is missing. + """ + ladybug = _alias_kuzu() + database_handle = ladybug.Database(":memory:") + connection = ladybug.Connection(database_handle) + try: + missing = load_extensions(connection) + finally: + connection.close() + database_handle.close() + target = str(Path(database).expanduser().resolve()) if database else "PATH" + return { + "engine_version": getattr(ladybug, "__version__", "unknown"), + "installed": [name for name in EXTENSIONS if name not in missing], + "missing": missing, + "ok": not missing, + "fix": f"kg-ladybug-setup --database {target} --apply" if missing else None, + } + + def fulltext_indexes() -> list[tuple[str, str, str]]: """``(table, index, statement)`` for every full-text index upstream search queries. diff --git a/src/kg_mcp/restore.py b/src/kg_mcp/restore.py index 29fdaea..85075de 100644 --- a/src/kg_mcp/restore.py +++ b/src/kg_mcp/restore.py @@ -80,6 +80,18 @@ def target_group(record: dict[str, Any], override: str | None, settings: Setting return group +def file_group(record: dict[str, Any], override: str | None, settings: Settings) -> str: + """The group whose file a record is restored into, with the per-group Ladybug layout.""" + group = override or str(record.get("group_id") or "") + if not group: + raise ValueError( + f"record {record.get('uuid')} has no group_id; pass --group to choose the " + "group file it is restored into" + ) + allowed_groups(group, settings) + return group + + def build_model(record: dict[str, Any], group: str) -> Any: from graphiti_core.edges import EntityEdge, EpisodicEdge from graphiti_core.nodes import EntityNode, EpisodicNode @@ -139,15 +151,25 @@ async def restore_snapshot( group: str | None = None, fail_fast: bool = False, ) -> dict[str, Any]: + from kg_mcp.config import per_group_ladybug, settings_for_group + header, records = read_snapshot(path) settings = load_config() if group: allowed_groups(group, settings) - planned = [(record, target_group(record, group, settings)) for record in records] + per_group = per_group_ladybug(settings) + if per_group: + # One file per group: the record's group chooses the file; inside it the + # record is stored like any other Ladybug record. + files = [(record, file_group(record, group, settings)) for record in records] + planned = [(record, "") for record, _ in files] + groups = sorted({target for _, target in files}) + else: + planned = [(record, target_group(record, group, settings)) for record in records] + groups = sorted({target for _, target in planned}) counts = {kind: 0 for kind in KINDS} for record, _ in planned: counts[record["kind"]] += 1 - groups = sorted({target for _, target in planned}) summary = { "source": { "provider": header.get("provider"), @@ -160,6 +182,36 @@ async def restore_snapshot( if not apply: return {"applied": False, "restored": {}, "failed": [], **summary} + if not per_group: + restored, failures = await _restore_into(settings, planned, fail_fast=fail_fast) + return {"applied": True, "restored": restored, "failed": failures, **summary} + + bound = {target: settings_for_group(settings, target) for target in groups} + from kg_mcp import fingerprint + + for target in bound.values(): + message = fingerprint.drift(target) + if message: + raise ValueError(message) + restored = {kind: 0 for kind in KINDS} + failures: list[dict[str, Any]] = [] + for target in groups: + subset = [(record, "") for record, chosen in files if chosen == target] + part, failed = await _restore_into(bound[target], subset, fail_fast=fail_fast) + for kind, count in part.items(): + restored[kind] += count + failures.extend(failed) + if fail_fast and failures: + break + return {"applied": True, "restored": restored, "failed": failures, **summary} + + +async def _restore_into( + settings: Settings, + planned: list[tuple[dict[str, Any], str]], + *, + fail_fast: bool, +) -> tuple[dict[str, int], list[dict[str, Any]]]: from kg_mcp import fingerprint from kg_mcp.runtime import build_graphiti from kg_mcp.write_lock import graph_writer_lock @@ -214,7 +266,7 @@ async def restore_snapshot( fingerprint.record(settings) finally: await graph.close() - return {"applied": True, "restored": restored, "failed": failures, **summary} + return restored, failures def human_lines(result: dict[str, Any]) -> list[str]: diff --git a/src/kg_mcp/retrieval.py b/src/kg_mcp/retrieval.py index 9606a1a..a45e44e 100644 --- a/src/kg_mcp/retrieval.py +++ b/src/kg_mcp/retrieval.py @@ -71,3 +71,44 @@ def candidate_limit(limit: int, settings: Settings) -> int: if limit < 1 or limit > 100: raise ValueError("result limit must be between 1 and 100") return min(1000, limit * settings.reranker.candidate_multiplier) + + +def keyword_edge_search_config(limit: int = 10): + """A graphiti ``SearchConfig`` for facts that matches on keywords alone. + + BM25 over the full-text index, ordered by reciprocal rank fusion. The vector half of + the default hybrid search embeds the query, so an embedding backend has to answer + before any question can; this config makes no model call at all, which keeps a graph + readable on a machine with nothing running. Pass it as ``config=`` to + ``Graphiti.search_``; the facts are in ``result.edges``. + """ + from graphiti_core.search.search_config import ( + EdgeReranker, + EdgeSearchConfig, + EdgeSearchMethod, + SearchConfig, + ) + + return SearchConfig( + edge_config=EdgeSearchConfig( + search_methods=[EdgeSearchMethod.bm25], reranker=EdgeReranker.rrf + ), + limit=limit, + ) + + +def keyword_node_search_config(limit: int = 10): + """Entity equivalent of :func:`keyword_edge_search_config`; results in ``result.nodes``.""" + from graphiti_core.search.search_config import ( + NodeReranker, + NodeSearchConfig, + NodeSearchMethod, + SearchConfig, + ) + + return SearchConfig( + node_config=NodeSearchConfig( + search_methods=[NodeSearchMethod.bm25], reranker=NodeReranker.rrf + ), + limit=limit, + ) diff --git a/src/kg_mcp/server.py b/src/kg_mcp/server.py index 4926ad5..7d5aa57 100644 --- a/src/kg_mcp/server.py +++ b/src/kg_mcp/server.py @@ -15,7 +15,15 @@ from mcp.server.fastmcp import FastMCP -from kg_mcp.config import Settings, TokenGrant, allowed_groups, load_config +from kg_mcp.config import ( + Settings, + TokenGrant, + allowed_groups, + load_config, + per_group_ladybug, + settings_for_group, + single_group, +) from kg_mcp.runtime import bounded logger = logging.getLogger(__name__) @@ -65,6 +73,14 @@ def _edge(edge: Any, score: float | None = None) -> dict[str, Any]: def authorized_groups(requested: str | list[str] | None, settings: Settings) -> list[str] | None: """Apply both the configured graph allow-list and the authenticated token scope.""" + groups = granted_groups(requested, settings) + if settings.database.provider == "ladybug": + return None + return groups + + +def granted_groups(requested: str | list[str] | None, settings: Settings) -> list[str]: + """The groups a request may read: the allow-list narrowed by the token scope.""" token_scope = _request_group_scope.get() requested_groups = [requested] if isinstance(requested, str) else requested if not requested_groups: @@ -78,9 +94,7 @@ def authorized_groups(requested: str | list[str] | None, settings: Settings) -> forbidden = sorted(set(requested_groups) - token_scope) if forbidden: raise ValueError(f"group is not granted by the bearer token: {forbidden[0]}") - if settings.database.provider == "ladybug": - return None - return requested_groups + return list(requested_groups) def _parse_time(value: str | None) -> datetime | None: @@ -143,13 +157,17 @@ async def lifespan(_: FastMCP): from kg_mcp.runtime import build_graphiti active = settings or load_config() - graph = build_graphiti(active, read_only=True) - state.update(settings=active, graph=graph) + # With one Ladybug file per group, a group's file is opened on its first read. + graphs: dict[str, Any] = {} + if not per_group_ladybug(active): + graphs[""] = build_graphiti(active, read_only=True) + state.update(settings=active, graphs=graphs) try: yield state finally: state.clear() - await graph.close() + for graph in graphs.values(): + await graph.close() selected = settings mcp = FastMCP( @@ -161,10 +179,35 @@ async def lifespan(_: FastMCP): transport_security=transport_security(selected), ) - def runtime() -> tuple[Any, Settings, float]: - if "graph" not in state: + def configured() -> Settings: + if "graphs" not in state: raise RuntimeError("Graphiti Local runtime is not initialized") - graph, active = state["graph"], state["settings"] + return state["settings"] + + def runtime(requested: str | list[str] | None = None) -> tuple[Any, Settings, float]: + """The graph a request reads, and the settings it reads under. + + With one Ladybug file per group the request must resolve to exactly one granted + group, and only that group's file is opened: a token scoped to one group cannot + reach another group's file. The returned settings are bound to that group. + """ + active = configured() + key = "" + if per_group_ladybug(active): + key = single_group(granted_groups(requested, active), active) + active = settings_for_group(active, key) + graphs = state["graphs"] + if key not in graphs: + from kg_mcp.runtime import build_graphiti + + path = Path(active.database.ladybug.path) + if not path.exists(): + raise LookupError( + f"no graph for group '{key}' yet; the first ingest or drain for it " + "creates its file" + ) + graphs[key] = build_graphiti(active, read_only=True) + graph = graphs[key] # An embedded read-only handle sees the file as it was at open time; a drain that # landed since must become visible without restarting the server. reopen = getattr(graph.driver, "reopen_if_changed", None) @@ -172,6 +215,18 @@ def runtime() -> tuple[Any, Settings, float]: reopen() return graph, active, active.graph.query_timeout_seconds + def each_granted(requested: str | list[str] | None = None): + """Every granted group's graph that exists, for calls without a group argument.""" + active = configured() + if not per_group_ladybug(active): + yield runtime(requested) + return + for group in granted_groups(requested, active): + try: + yield runtime(group) + except LookupError: + continue + @mcp.tool() async def search_nodes( query: str, @@ -184,7 +239,7 @@ async def search_nodes( if max_nodes < 1: return {"error": "max_nodes must be a positive integer"} try: - graph, active, timeout = runtime() + graph, active, timeout = runtime(group_ids) groups = authorized_groups(group_ids, active) from kg_mcp.retrieval import compatible_query @@ -260,7 +315,7 @@ async def search_memory_facts( if max_facts < 1: return {"error": "max_facts must be a positive integer"} try: - graph, active, timeout = runtime() + graph, active, timeout = runtime(group_ids) groups = authorized_groups(group_ids, active) from kg_mcp.retrieval import ( candidate_limit, @@ -309,9 +364,22 @@ async def find_facts() -> tuple[list[tuple[Any, float]], int]: async def get_entity_edge(uuid: str) -> dict[str, Any]: """Get one fact edge by UUID from an allow-listed graph.""" try: + from graphiti_core.edges import EntityEdge + + if per_group_ladybug(configured()): + for graph, _, timeout in each_granted(): + try: + edge = await bounded( + EntityEdge.get_by_uuid(graph.driver, uuid), timeout, "get_entity_edge" + ) + except TimeoutError: + raise + except Exception: + continue + return _edge(edge) + raise LookupError(f"edge not found: {uuid}") graph, active, timeout = runtime() groups = authorized_groups(None, active) - from graphiti_core.edges import EntityEdge if active.database.provider == "falkordb": for group in groups or active.graph.groups: @@ -346,7 +414,7 @@ async def get_episodes( if max_episodes < 1: return {"error": "max_episodes must be a positive integer"} try: - graph, active, timeout = runtime() + graph, active, timeout = runtime(group_ids) groups = authorized_groups(group_ids, active) episodes = await bounded( graph.retrieve_episodes( @@ -387,7 +455,7 @@ async def get_episode_entities( if not episode_uuids: return {"error": "episode_uuids must contain at least one UUID"} try: - graph, active, timeout = runtime() + graph, active, timeout = runtime(group_id) groups = authorized_groups(group_id, active) from graphiti_core.edges import EntityEdge from graphiti_core.nodes import EpisodicNode @@ -436,6 +504,24 @@ async def trace(driver: Any) -> tuple[list[Any], list[Any]]: async def get_status() -> dict[str, Any]: """Check the server and its configured database connection.""" try: + if per_group_ladybug(configured()): + granted = granted_groups(None, configured()) + opened = [] + for graph, active, timeout in each_granted(): + await bounded( + graph.driver.execute_query("MATCH (n) RETURN count(n) AS count"), + timeout, + "get_status", + ) + opened.extend(active.graph.groups) + return { + "status": "ok", + "message": ( + f"Graphiti Local is connected to ladybug; {len(opened)} of " + f"{len(granted)} group file(s) exist" + ), + "groups": granted, + } graph, active, timeout = runtime() groups = authorized_groups(None, active) driver = graph.driver diff --git a/src/kg_mcp/verify.py b/src/kg_mcp/verify.py index 7b4af4c..bdfc73f 100644 --- a/src/kg_mcp/verify.py +++ b/src/kg_mcp/verify.py @@ -7,6 +7,7 @@ from __future__ import annotations from datetime import datetime, timezone +from pathlib import Path from typing import Any from kg_mcp.doctor import FAIL, OK, WARN, _check, run_checks, worst_status @@ -50,11 +51,26 @@ async def check_tool_surface() -> list[dict[str, str]]: async def check_retrieval(query: str) -> list[dict[str, Any]]: - from kg_mcp.config import allowed_groups, load_config + from kg_mcp.config import ( + allowed_groups, + load_config, + per_group_ladybug, + settings_for_group, + ) from kg_mcp.retrieval import current_edges from kg_mcp.runtime import bounded, build_graphiti settings = load_config() + if per_group_ladybug(settings): + # Probe the first group that has a file; every group file has the same schema. + existing = [ + group + for group in settings.graph.groups + if Path(settings.database.ladybug.path_for(group)).exists() + ] + if not existing: + return [_check("retrieval", WARN, "no group file exists yet; nothing to search")] + settings = settings_for_group(settings, existing[0]) timeout = settings.graph.query_timeout_seconds try: graph = build_graphiti(settings, read_only=True) diff --git a/tests/test_v05_multi_user.py b/tests/test_v05_multi_user.py new file mode 100644 index 0000000..2198b08 --- /dev/null +++ b/tests/test_v05_multi_user.py @@ -0,0 +1,413 @@ +"""0.5: one Ladybug file per group, and the library helpers around an embedded graph. + +The isolation tests run against real Ladybug files and need the FTS and VECTOR +extensions on the host; they skip otherwise, like the rest of the Ladybug suite. No +model or network is contacted: embeddings come from a deterministic fake. +""" + +from __future__ import annotations + +import asyncio +import os +import uuid +from datetime import datetime, timezone +from pathlib import Path + +import pytest + +from kg_mcp.config import Settings + +HOSTILE_IDS = [ + "../escape", + "../../etc/passwd", + "/etc/passwd", + "C:\\Windows\\System32", + "a/../../b", + "nul\x00byte", + "x" * 5000, + ".", + "..", + "~root", + "spaces and\nnewlines", + "\udcff lone surrogate", +] + + +def _extensions_available() -> bool: + from kg_mcp.ladybug import extension_status + + try: + return extension_status()["ok"] + except Exception: + return False + + +needs_extensions = pytest.mark.skipif( + not _extensions_available(), reason="Ladybug search extensions are not installed" +) + + +def _use(settings: Settings, tmp_path: Path, monkeypatch) -> Settings: + """Make ``settings`` the loaded configuration, as the CLI and the workspace read it.""" + import yaml + + path = tmp_path / "config.yaml" + path.write_text(yaml.safe_dump(settings.model_dump()), encoding="utf-8") + monkeypatch.setenv("GRAPHITI_LOCAL_CONFIG", str(path)) + monkeypatch.delenv("KG_WORKSPACE_DIR", raising=False) + return settings + + +def _settings(tmp_path: Path, layout: str = "per-group", **server) -> Settings: + return Settings.model_validate( + { + "server": server, + "graph": {"groups": ["alice", "bob"], "workspace_dir": str(tmp_path / "ws")}, + "database": { + "provider": "ladybug", + "ladybug": { + "path": str(tmp_path / "graph.ladybug"), + "layout": layout, + "directory": str(tmp_path / "groups"), + }, + }, + "embedder": {"model": "nomic-embed-text", "dimensions": 768}, + } + ) + + +# --- file layout ------------------------------------------------------------------- + + +def test_single_layout_is_the_default_and_unchanged(tmp_path: Path): + from kg_mcp.config import per_group_ladybug, settings_for_group + + settings = Settings.model_validate( + {"database": {"provider": "ladybug", "ladybug": {"path": "/data/graph.ladybug"}}} + ) + assert settings.database.ladybug.layout == "single" + assert not per_group_ladybug(settings) + assert settings.database.ladybug.path_for("anyone") == "/data/graph.ladybug" + assert settings_for_group(settings, "main") is settings + + +def test_each_group_gets_its_own_file_named_by_sha256(tmp_path: Path): + import hashlib + + settings = _settings(tmp_path) + alice = Path(settings.database.ladybug.path_for("alice")) + bob = Path(settings.database.ladybug.path_for("bob")) + assert alice != bob + assert alice.parent == bob.parent == (tmp_path / "groups").resolve() + assert alice.name == hashlib.sha256(b"alice").hexdigest() + ".ladybug" + with pytest.raises(ValueError, match="needs a group"): + settings.database.ladybug.path_for(None) + + +@pytest.mark.parametrize("hostile", HOSTILE_IDS) +def test_hostile_ids_stay_inside_the_directory(tmp_path: Path, hostile: str): + from kg_mcp.ladybug import group_database_path + + root = (tmp_path / "groups").resolve() + target = group_database_path(root, hostile) + assert target.parent == root + assert len(target.stem) == 64 and all(c in "0123456789abcdef" for c in target.stem) + + +def test_an_empty_group_is_refused(tmp_path: Path): + from kg_mcp.ladybug import group_database_path + + with pytest.raises(ValueError, match="non-empty"): + group_database_path(tmp_path, "") + + +def test_binding_narrows_settings_to_one_group_file(tmp_path: Path): + from kg_mcp.config import per_group_ladybug, settings_for_group + + settings = _settings(tmp_path) + bound = settings_for_group(settings, "alice") + assert bound.graph.groups == ["alice"] + assert bound.database.ladybug.path == settings.database.ladybug.path_for("alice") + assert not per_group_ladybug(bound) + with pytest.raises(ValueError, match="not allowed"): + settings_for_group(settings, "mallory") + + +def test_per_group_layout_accepts_a_partial_token_scope(tmp_path: Path): + grants = {"auth": {"tokens": [{"token": "alice-token-0123456", "groups": ["alice"]}]}} + settings = _settings(tmp_path, transport="streamable-http", **grants) + assert settings.server.auth.grants()[0].groups == ["alice"] + with pytest.raises(ValueError, match="per-group"): + _settings(tmp_path, layout="single", transport="streamable-http", **grants) + + +# --- isolation on real files ------------------------------------------------------- + + +def _fake_embedder_class(): + from graphiti_core.embedder.client import EmbedderClient + + class FakeEmbedder(EmbedderClient): + def __init__(self, config=None): + del config + + async def create(self, input_data): + return [0.5] * 8 + + async def create_batch(self, input_data_list): + return [[0.5] * 8 for _ in input_data_list] + + return FakeEmbedder + + +def _write_fact(path: str, fact: str) -> None: + """Write one fact straight into a file, as a drain would leave it.""" + import ladybug + + from kg_mcp.ladybug import build_ladybug_driver, load_extensions + + Path(path).parent.mkdir(parents=True, exist_ok=True) + asyncio.run(build_ladybug_driver(path, read_only=False).close()) + database = ladybug.Database(path) + connection = ladybug.Connection(database) + assert load_extensions(connection) == [] + now = datetime.now(timezone.utc) + a, b = str(uuid.uuid4()), str(uuid.uuid4()) + for node in (a, b): + connection.execute( + "CREATE (:Entity {uuid: $u, name: $n, group_id: '', created_at: $t, " + "name_embedding: $v, summary: '', labels: ['Entity']})", + {"u": node, "n": f"node {node[:4]}", "t": now, "v": [0.5] * 8}, + ) + connection.execute( + "MATCH (x:Entity {uuid: $a}), (y:Entity {uuid: $b}) " + "CREATE (x)-[:RELATES_TO]->(:RelatesToNode_ {uuid: $e, name: 'NOTE', fact: $f, " + "group_id: '', created_at: $t, fact_embedding: $v, episodes: []})-[:RELATES_TO]->(y)", + {"a": a, "b": b, "e": str(uuid.uuid4()), "f": fact, "t": now, "v": [0.5] * 8}, + ) + connection.execute("CHECKPOINT") + connection.close() + database.close() + + +def _raw_facts(path: str) -> list[str]: + import ladybug + + database = ladybug.Database(path, read_only=True) + connection = ladybug.Connection(database) + try: + rows = connection.execute("MATCH (e:RelatesToNode_) RETURN e.fact AS fact") + return sorted(row["fact"] for row in rows.rows_as_dict()) + finally: + connection.close() + database.close() + + +@pytest.fixture +def two_users(tmp_path: Path, monkeypatch): + os.environ.setdefault("OPENAI_API_KEY", "test-only") + import graphiti_core.embedder.openai as openai_embedder + + monkeypatch.setattr(openai_embedder, "OpenAIEmbedder", _fake_embedder_class()) + settings = _use(_settings(tmp_path), tmp_path, monkeypatch) + _write_fact(settings.database.ladybug.path_for("alice"), "Alice keeps the invoice ledger") + _write_fact(settings.database.ladybug.path_for("bob"), "Bob keeps the invoice archive") + return settings + + +async def _call(server, name: str, arguments: dict): + return await server._tool_manager.call_tool(name, arguments) + + +@needs_extensions +def test_two_users_two_files_each_holding_only_its_own_facts(two_users: Settings): + ladybug = two_users.database.ladybug + assert _raw_facts(ladybug.path_for("alice")) == ["Alice keeps the invoice ledger"] + assert _raw_facts(ladybug.path_for("bob")) == ["Bob keeps the invoice archive"] + assert not Path(ladybug.path).exists() # the single-layout file is never created + + +@needs_extensions +def test_a_token_for_one_user_reads_only_that_users_file(two_users: Settings): + from kg_mcp.server import _request_group_scope, create_server + + server = create_server(two_users) + + async def scenario(): + marker = _request_group_scope.set(frozenset({"alice"})) + try: + async with server.settings.lifespan(server) as state: + facts = await _call(server, "search_memory_facts", {"query": "invoice"}) + denied = await _call( + server, "search_memory_facts", {"query": "invoice", "group_ids": ["bob"]} + ) + status = await _call(server, "get_status", {}) + opened = sorted(state["graphs"]) + return facts, denied, status, opened + finally: + _request_group_scope.reset(marker) + + facts, denied, status, opened = asyncio.run(scenario()) + assert [item["fact"] for item in facts["facts"]] == ["Alice keeps the invoice ledger"] + assert "not granted" in denied["error"] + assert status["status"] == "ok" and status["groups"] == ["alice"] + assert opened == ["alice"] # bob's file was never opened + + +@needs_extensions +def test_an_unscoped_reader_must_name_one_group(two_users: Settings): + from kg_mcp.server import create_server + + server = create_server(two_users) + + async def scenario(): + async with server.settings.lifespan(server): + ambiguous = await _call(server, "search_memory_facts", {"query": "invoice"}) + bob = await _call( + server, "search_memory_facts", {"query": "invoice", "group_ids": "bob"} + ) + return ambiguous, bob + + ambiguous, bob = asyncio.run(scenario()) + assert "one group per call" in ambiguous["error"] + assert [item["fact"] for item in bob["facts"]] == ["Bob keeps the invoice archive"] + + +@needs_extensions +def test_kg_ask_reads_the_named_users_file(two_users: Settings): + from kg_mcp import cli + + args = cli.parser().parse_args(["ask", "invoice", "bob", "--keyword"]) + payload = asyncio.run(cli._dispatch(args)) + assert payload["groups"] == ["bob"] + assert [item["fact"] for item in payload["facts"]] == ["Bob keeps the invoice archive"] + + ambiguous = cli.parser().parse_args(["ask", "invoice", "--keyword"]) + with pytest.raises(ValueError, match="one group per call"): + asyncio.run(cli._dispatch(ambiguous)) + + +def test_ingest_routes_each_domain_to_its_own_file(tmp_path: Path, monkeypatch): + from kg_mcp import ingest, runtime + + settings = _use(_settings(tmp_path), tmp_path, monkeypatch) + written: dict[str, list[str]] = {} + + class FakeGraph: + def __init__(self, bound): + self.path = bound.database.ladybug.path + + async def build_indices_and_constraints(self): + return None + + async def add_episode(self, *, name, **kwargs): + written.setdefault(self.path, []).append(name) + + async def close(self): + return None + + monkeypatch.setattr(runtime, "build_graphiti", lambda bound, read_only: FakeGraph(bound)) + records = [ + {"name": "a1", "body": "Alice fact", "domain": "alice"}, + {"name": "b1", "body": "Bob fact", "domain": "bob"}, + {"name": "a2", "body": "Alice again", "domain": "alice"}, + ] + result = asyncio.run( + ingest.ingest_records(records, apply=True, ledger=False, settings=settings) + ) + assert result["ingested"] == 3 and result["failed"] == [] + alice = settings.database.ladybug.path_for("alice") + bob = settings.database.ladybug.path_for("bob") + assert written == {alice: ["a1", "a2"], bob: ["b1"]} + + from kg_mcp.fingerprint import recorded_embedder + + assert recorded_embedder(alice)["model"] == "nomic-embed-text" + assert recorded_embedder(bob)["dimensions"] == 768 + + +# --- library helpers --------------------------------------------------------------- + + +def test_embedder_record_is_keyed_by_database_path(tmp_path: Path): + from kg_mcp.fingerprint import embedder_drift, record_embedder, recorded_embedder + + path = tmp_path / "user.ladybug" + assert recorded_embedder(path) is None + assert embedder_drift(path, model="nomic-embed-text", dimensions=768) is None + entry = record_embedder(path, model="nomic-embed-text:latest", dimensions=768) + assert entry["model"] == "nomic-embed-text" + assert (tmp_path / "user.ladybug.embedder.json").is_file() + assert recorded_embedder(path)["dimensions"] == 768 + assert embedder_drift(path, model="nomic-embed-text:latest", dimensions=768) is None + message = embedder_drift(path, model="mxbai-embed-large", dimensions=1024) + assert message and "mixing embedders" in message + other = tmp_path / "other.ladybug" + assert recorded_embedder(other) is None + + +def test_configured_drift_check_sees_a_record_written_as_a_library(tmp_path: Path, monkeypatch): + from kg_mcp import fingerprint + from kg_mcp.config import settings_for_group + + _use(_settings(tmp_path), tmp_path, monkeypatch) + bound = settings_for_group(_settings(tmp_path), "alice") + fingerprint.record_embedder( + bound.database.ladybug.path, model="mxbai-embed-large", dimensions=1024 + ) + assert "mixing embedders" in fingerprint.drift(bound) + + +def test_keyword_search_config_makes_no_model_call(): + from graphiti_core.search.search_config import EdgeSearchMethod, NodeSearchMethod + + from kg_mcp.retrieval import keyword_edge_search_config, keyword_node_search_config + + edges = keyword_edge_search_config(7) + assert edges.limit == 7 + assert edges.edge_config.search_methods == [EdgeSearchMethod.bm25] + assert edges.node_config is None + nodes = keyword_node_search_config() + assert nodes.limit == 10 + assert nodes.node_config.search_methods == [NodeSearchMethod.bm25] + assert nodes.edge_config is None + + +def test_extension_status_names_the_fix_when_something_is_missing(tmp_path: Path, monkeypatch): + import ladybug + + from kg_mcp import ladybug as adapter + + status = adapter.extension_status() + assert status["engine_version"] == ladybug.__version__ + assert set(status["installed"]) | set(status["missing"]) == {"FTS", "VECTOR"} + assert status["ok"] is (not status["missing"]) + + monkeypatch.setattr(adapter, "load_extensions", lambda connection: ["VECTOR"]) + target = tmp_path / "user.ladybug" + missing = adapter.extension_status(target) + assert missing["ok"] is False + assert missing["missing"] == ["VECTOR"] and missing["installed"] == ["FTS"] + assert missing["fix"] == f"kg-ladybug-setup --database {target.resolve()} --apply" + + +@needs_extensions +def test_export_and_restore_move_facts_between_group_files(two_users: Settings, tmp_path: Path): + from kg_mcp import cli + from kg_mcp.restore import restore_snapshot + + snapshot = tmp_path / "alice.jsonl" + args = cli.parser().parse_args(["export", "alice", "--output", str(snapshot)]) + exported = asyncio.run(cli._dispatch(args)) + assert exported["edges"] == 1 + + with pytest.raises(ValueError, match="pass --group"): + asyncio.run(restore_snapshot(snapshot, apply=False)) + result = asyncio.run(restore_snapshot(snapshot, apply=True, group="bob")) + assert result["groups"] == ["bob"] and result["failed"] == [] + ladybug = two_users.database.ladybug + assert _raw_facts(ladybug.path_for("bob")) == [ + "Alice keeps the invoice ledger", + "Bob keeps the invoice archive", + ] + assert _raw_facts(ladybug.path_for("alice")) == ["Alice keeps the invoice ledger"]