From f1ccb3926aaf16ecec5d6abb23393e6244703485 Mon Sep 17 00:00:00 2001 From: jjk9090 Date: Sun, 25 Oct 2026 10:25:58 +0000 Subject: [PATCH] feat(graph): add TerminusDB backend with incremental updates and versioning - Add TerminusDB as a graph storage backend - Implement incremental graph updates and version control - Update graph and ingest layers to support version-aware operations - Add integration tests for graph versioning and rollback in `test_graph_versioning.py` - Ensure all tests pass --- .../capabilities/graph/providers/__init__.py | 1 + .../graph/providers/terminusdb_store.py | 1171 +++++++++++++++++ datamind/capabilities/graph/service.py | 78 ++ datamind/capabilities/graph/tools.py | 192 ++- datamind/capabilities/ingest/service.py | 135 +- datamind/capabilities/ingest/tools.py | 31 +- .../capabilities/kb/providers/chroma_store.py | 1 - .../kb/providers/lancedb_store.py | 1 - datamind/config.py | 1 + datamind/core/registry.py | 1 + datamind/direct_tools.py | 4 +- datamind/tests/test_graph_versioning.py | 284 ++++ 12 files changed, 1884 insertions(+), 16 deletions(-) create mode 100644 datamind/capabilities/graph/providers/terminusdb_store.py create mode 100644 datamind/tests/test_graph_versioning.py diff --git a/datamind/capabilities/graph/providers/__init__.py b/datamind/capabilities/graph/providers/__init__.py index 4118048..94f82f0 100644 --- a/datamind/capabilities/graph/providers/__init__.py +++ b/datamind/capabilities/graph/providers/__init__.py @@ -2,5 +2,6 @@ from __future__ import annotations from . import networkx_store # noqa: F401 +from . import terminusdb_store # noqa: F401 __all__: list[str] = [] diff --git a/datamind/capabilities/graph/providers/terminusdb_store.py b/datamind/capabilities/graph/providers/terminusdb_store.py new file mode 100644 index 0000000..d1775f7 --- /dev/null +++ b/datamind/capabilities/graph/providers/terminusdb_store.py @@ -0,0 +1,1171 @@ +"""TerminusDB-backed graph storage with versioned reads and rollback support. + +The provider persists graph entities and edges in TerminusDB, while a local +SQLite history records the DataMind ingest operations associated with each +TerminusDB commit. +""" +from __future__ import annotations + +import asyncio +import difflib +import json +import os +import re +import threading +import sqlite3 +import shutil +import subprocess +import time +import httpx + +from pathlib import Path +from typing import Any, Sequence + +from datamind.core.logging import get_logger +from datamind.core.protocols import Edge, Entity, GraphPath, GraphTriple +from datamind.core.registry import graph_registry + +_log = get_logger("graph.terminusdb") + + +@graph_registry.register("terminusdb") +class TerminusDBGraphStore: + + def __init__( + self, + *, + dsn: str, + password: str, + persist_path: str | Path, + ) -> None: + self._dsn = dsn + self._password = password + self._path = Path(persist_path) + self._path.mkdir(parents=True, exist_ok=True) + self._terminusdb_home = None + self._db = "kg_default" + self._team = "admin" + self._init_terminusdb() + self._client = self._get_client() + self._schema() + self._dirty = False + self._mutation_revision = 0 + self._state_lock = threading.RLock() + self._persist_lock = asyncio.Lock() + self._persist_worker_lock = threading.Lock() + + self._history_path = self._path / "version_history.sqlite3" + with sqlite3.connect(self._history_path) as conn: + conn.execute("CREATE TABLE IF NOT EXISTS history (database TEXT, version INTEGER, payload TEXT, PRIMARY KEY(database, version))") + + _log.info( + "graph_loaded", + extra={ + "path": str(self._path), + }, + ) + + def _is_terminusdb_running(self) -> bool: + """Check whether TerminusDB is already available.""" + try: + response = httpx.get( + f"{self._dsn.rstrip('/')}/api/ok", + timeout=2.0, + ) + return response.is_success + except httpx.HTTPError: + return False + + def _init_terminusdb(self): + """Clone official Compose configuration and start TerminusDB.""" + if self._is_terminusdb_running(): + return + + repo_url = "https://github.com/terminusdb/terminusdb" + home = Path.home() / ".datamind" / "terminusdb" + home = home.expanduser().resolve() + home.parent.mkdir(parents=True, exist_ok=True) + self._terminusdb_home = home + + # 1. Check prerequisites. + if not shutil.which("docker"): + raise RuntimeError("Docker is not installed or not in PATH.") + + subprocess.run( + ["docker", "compose", "version"], + check=True, + capture_output=True, + text=True, + ) + + # 2. Clone once; preserve the existing checkout and data. + if not (home / ".git").exists(): + if home.exists() and any(home.iterdir()): + raise RuntimeError( + f"{home} exists but is not a TerminusDB checkout. " + "Choose another directory." + ) + subprocess.run( + ["git", "clone", "--depth", "1", repo_url, str(home)], + check=True, + ) + + # 3. Create credentials if this is a new installation. + env_file = home / ".env" + if not env_file.exists(): + port = self._dsn.split(":")[-1] + fd = os.open( + env_file, + os.O_WRONLY | os.O_CREAT | os.O_EXCL, + 0o600, + ) + with os.fdopen(fd, "w") as f: + f.write(f"TERMINUSDB_SERVER_DB_PATH={str(self._path)}\n") + f.write(f"TERMINUSDB_SERVER_PORT={port}\n") + + # 4. Start in the background. Compose's storage directories + # are relative to the checkout directory. + subprocess.run( + ["docker", "compose", "up", "-d"], + cwd=home, + check=True, + ) + + # Wait for TerminusDB to become available (up to 60 seconds). + timeout = 60 + interval = 2 + deadline = time.monotonic() + timeout + + while time.monotonic() < deadline: + if self._is_terminusdb_running(): + print(f"TerminusDB is ready at {self._dsn}") + break + time.sleep(interval) + else: + raise RuntimeError( + f"TerminusDB failed to become ready within {timeout} seconds " + f"at {self._dsn}. Check the Docker logs with " + f"'docker compose -f {home / 'docker-compose.yml'} logs'." + ) + + def _get_client(self): + from terminusdb_client import Client + if not self._dsn: + raise ValueError("Missing TerminusDB DSN. Please configure DATAMIND__GRAPH__DSN.") + + if not self._password: + raise ValueError("Missing TerminusDB password. Please configure DATAMIND__GRAPH__PASSWORD.") + + try: + client = Client(self._dsn) + client.connect( + team=self._team, + user=self._team, + key=self._password, + server_url=self._dsn, + ) + except Exception as e: + raise RuntimeError( + f"TerminusDB connection failed: " + f"dsn={self._dsn}, password={self._password}, " + f"error_type={type(e).__name__}, " + f"error={e}" + ) from e + + if not client.has_database(self._db, team=self._team): + client.create_database( + self._db, + label=self._db, + description="Graph of DataMind stores in TerminusDB", + include_schema=False, + ) + + client.connect( + team=self._team, + db=self._db, + user=self._team, + key=self._password, + server_url=self._dsn, + ) + return client + + def _normalize(self, schema: dict) -> str: + return json.dumps( + schema, + sort_keys=True, + ensure_ascii=False, + separators=(",", ":"), + ) + + def _schema(self) -> None: + schema_docs = [ + { + "@type": "Class", + "@id": "entity", + "label": "xsd:string", + "outgoing_edge": { + "@type": "Set", + "@class": "GraphEdge", + }, + }, + { + "@type": "Class", + "@id": "GraphEdge", + "subject": "entity", + "object": "entity", + "relation": "xsd:string", + "key": "xsd:string", + "weight": "xsd:double", + "source": "xsd:string", + "properties": "sys:JSON", + "incoming_edge": { + "@type": "Optional", + "@class": "entity", + }, + }, + ] + + to_sync = [] + + for schema in schema_docs: + doc_id = schema["@id"] + + if not self._has_doc(doc_id, graph_type="schema"): + to_sync.append(schema) + continue + + existing = self._client.get_document( + doc_id, + graph_type="schema", + ) + if self._normalize(existing) != self._normalize(schema): + to_sync.append(schema) + + if to_sync: + self._client.replace_document( + to_sync, + graph_type="schema", + create=True, + commit_msg="Synchronize graph schema", + ) + + for schema in schema_docs: + doc_id = schema["@id"] + if not self._has_doc(doc_id, graph_type="schema"): + raise RuntimeError( + f"TerminusDB schema registration failed: {doc_id}" + ) + + def _has_doc(self, doc_id: str, graph_type=None): + if graph_type is None: + graph_type = 'instance' + + client = self._client + client._check_connection() + + response = client._session.get( + client._documents_url(), + headers=client._default_headers, + params={ + "id": doc_id, + "graph_type": graph_type, + }, + auth=client._auth(), + ) + if response.status_code == 200: + return True + if response.status_code == 404: + try: + body = response.json() + except ValueError: + return False + error = body.get("api:error", {}) + if error.get("@type") == "api:DocumentNotFound": + return False + + response.raise_for_status() + response.raise_for_status() + return False + + async def aclose(self) -> None: + """Stop TerminusDB only if this instance started it.""" + if not self._is_terminusdb_running() or not self._terminusdb_home: + # manually close + return + + process = await asyncio.create_subprocess_exec( + "docker", + "compose", + "stop", + cwd=self._terminusdb_home, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + + try: + stdout, stderr = await asyncio.wait_for( + process.communicate(), + timeout=60, + ) + except asyncio.TimeoutError: + process.kill() + await process.wait() + raise TimeoutError("Timed out stopping TerminusDB") from None + + if process.returncode != 0: + raise RuntimeError( + f"Failed to stop TerminusDB: " + f"{stderr.decode(errors='replace').strip()}" + ) + + # ------------------------------------------------------------ mutation + + def _upsert_documents(self, documents, commit_msg): + from terminusdb_client.errors import DatabaseError + + to_insert = [] + to_replace = [] + for doc in documents: + doc_id = doc["@id"] + if not self._has_doc(doc): + to_insert.append(doc) + continue + existing = self._client.get_document(doc_id) + + if self._normalize(existing) != self._normalize(doc): + to_replace.append(doc) + + if to_insert: + self._client.insert_document(to_insert, commit_msg=commit_msg) + if to_replace: + self._client.replace_document(to_replace, commit_msg=commit_msg) + + return { + "inserted": len(to_insert), + "replaced": len(to_replace), + } + + def _edge_id(self, t): + profile_managed = bool((t.properties or {}).get("_profile_managed")) + origin = "profile" if profile_managed else (t.source or "runtime") + edge_key = f"{t.relation}\x1f{origin}" + subject = ( + t.subject[len(t.subject_type):].lstrip("/") + if t.subject.startswith(t.subject_type) + else t.subject + ) + object = ( + t.object[len(t.object_type):].lstrip("/") + if t.object.startswith(t.object_type) + else t.object + ) + return f"GraphEdge/{subject}_{object}_{edge_key}", edge_key + + async def upsert_triples(self, triples: Sequence[GraphTriple]) -> None: + from terminusdb_client.woqlquery import WOQLQuery + added_triples, updated_triples = [], [] + with self._state_lock: + for t in triples: + # Nodes + nodes = [] + node_replace = False + edge_replace = False + for side, node_id, node_type in ( + ("subject", t.subject, t.subject_type), + ("object", t.object, t.object_type), + ): + entity_id = f"entity/{node_id}" + node = {"@type": node_type, "@id": entity_id, "label": node_id} + result = self._upsert_documents([node], commit_msg=f"Add Node {entity_id}") + node_replace = bool(result["replaced"]) + nodes.append(node) + + # Profile snapshots and runtime writes have separate identities; + # exact duplicates within either origin overwrite deterministically. + edge_id, edge_key = self._edge_id(t) + subject_id = f"entity/{t.subject}" + object_id = f"entity/{t.object}" + edge_doc = { + "@type": "GraphEdge", + "@id": edge_id, + # TerminusDB requires ``@type: @id`` to distinguish a + # document reference from an embedded document. + "subject": {"@id": subject_id, "@type": "@id"}, + "object": {"@id": object_id, "@type": "@id"}, + "relation": t.relation, + "key": edge_key, + "weight": float(t.confidence), + "source": t.source or "", + "properties": {f"p_{k}": v for k, v in (t.properties or {}).items()}, + } + + result = self._upsert_documents([edge_doc], commit_msg=f"Add Edge {t.subject} --{t.relation}--> {t.object}") + node_replace = bool(result["replaced"]) + WOQLQuery().add_triple( + nodes[0]["@id"], + "outgoing_edge", + edge_doc["@id"], + ).execute(self._client) + + WOQLQuery().add_triple( + edge_doc["@id"], + "incoming_edge", + nodes[1]["@id"], + ).execute(self._client) + if node_replace or edge_replace: + updated_triples.append(t) + else: + added_triples.append(t) + self._mutation_revision += 1 + self._dirty = True + return self.get_triples_info(added=added_triples, updated=updated_triples) + + def _validate_key(self, doc, key, value=None): + props = doc.get("properties") or {} + if key == "profile": + return bool(props.get("p__profile_managed")) + elif key == "source": + if not value: + return False + source_path = str(props.get("p__source_path") or "") + return ( + source_path == value + or re.fullmatch( + re.escape(value) + r"#part=\d+", + source_path, + ) is not None + ) + elif key == "lineage": + return props.get("p__lineage_root") == value + return True + + def _get_triples(self, key="profile", value=None) -> list[GraphTriple]: + with self._state_lock: + docs = self._client.get_documents_by_type( + "GraphEdge", + as_list=True, + ) + + result = [] + + for doc in docs: + if not self._validate_key(doc, key, value): continue + props = doc.get("properties") or {} + subject = doc.get("subject") + obj = doc.get("object") + if isinstance(subject, dict): + subject = subject.get("@id") + if isinstance(obj, dict): + obj = obj.get("@id") + if not subject or not obj: + continue + result.append( + GraphTriple( + subject=subject, + relation=doc["relation"], + object=obj, + subject_type="entity", + object_type="entity", + confidence=float(doc.get("weight", 1.0)), + source=doc.get("source", ""), + properties=props, + ) + ) + return result + + def _diff_triples(self, previous_triples, triples): + def edge_identity(t: GraphTriple) -> tuple: + source_path = getattr(t, "source_path", "") or "" + + return ( + t.subject, + t.relation, + t.object, + source_path, + ) + + def edge_payload(t: GraphTriple) -> str: + payload = { + "confidence": float(t.confidence), + "source": getattr(t, "source", "") or "", + "properties": dict(t.properties or {}), + } + + return json.dumps( + payload, + sort_keys=True, + ensure_ascii=False, + separators=(",", ":"), + default=str, + ) + old = {edge_identity(t): t for t in previous_triples} + new = {edge_identity(t): t for t in triples} + + added = [new[k] for k in new.keys() - old.keys()] + removed = [old[k] for k in old.keys() - new.keys()] + + updated = [ + new[k] + for k in old.keys() & new.keys() + if edge_payload(old[k]) != edge_payload(new[k]) + ] + return added, updated, removed + + def _delete_edge_document_and_links(self, edge_doc: dict) -> None: + from terminusdb_client.woqlquery import WOQLQuery + edge_id = edge_doc["@id"] + subject = edge_doc.get("subject") + obj = edge_doc.get("object") + + if isinstance(subject, dict): + subject = subject.get("@id") + if isinstance(obj, dict): + obj = obj.get("@id") + + if not subject or not obj: + raise ValueError( + f"GraphEdge {edge_id} Lack of effective subject/object: " + f"subject={subject}, object={obj}" + ) + + WOQLQuery().delete_triple( + subject, + "outgoing_edge", + edge_id, + ).execute(self._client) + WOQLQuery().delete_triple( + edge_id, + "incoming_edge", + obj, + ).execute(self._client) + if self._has_doc(edge_id): + self._client.delete_document( + edge_id, + commit_msg=f"Delete GraphEdge {edge_id}", + ) + + def _delete_orphaned_nodes( + self, + candidate_node_ids: set[str], + *, + key: str = "profile", + value: str | None = None, + ) -> list[str]: + if not candidate_node_ids: + return [] + + connected_node_ids = set() + docs = self._client.get_all_documents(as_list=True) + for doc in docs: + if not isinstance(doc, dict): + continue + doc_id = doc.get("@id", "") + if not doc_id.startswith("GraphEdge/"): + continue + + subject = doc.get("subject") + obj = doc.get("object") + + if subject: + if self._validate_key(subject, key, value): + connected_node_ids.add(subject.get("@id")) + if obj: + if self._validate_key(obj, key, value): + connected_node_ids.add(obj.get("@id")) + + orphaned_node_ids = sorted( + node_id + for node_id in candidate_node_ids + if node_id not in connected_node_ids + ) + + deleted_node_ids = [] + for node_id in orphaned_node_ids: + if not self._has_doc(node_id): + continue + self._client.delete_document( + node_id, + commit_msg=f"Delete orphaned profile node {node_id}", + ) + deleted_node_ids.append(node_id) + return deleted_node_ids + + async def _delete_triples(self, triples, key, value=None): + candidate_nodes = set() + for t in triples: + candidate_nodes.add(t.subject) + candidate_nodes.add(t.object) + for t in triples: + edge_id, _ = self._edge_id(t) + edge_doc = self._client.get_document(edge_id) + if edge_doc: + self._delete_edge_document_and_links(edge_doc) + + if candidate_nodes: + self._delete_orphaned_nodes(candidate_nodes, key=key, value=value) + + def _docs_to_triples(self, edge_docs: list[dict]) -> list[GraphTriple]: + triples = [] + for doc in edge_docs: + subject = doc.get("subject") + obj = doc.get("object") + if isinstance(subject, dict): + subject = subject.get("@id") + if isinstance(obj, dict): + obj = obj.get("@id") + if not subject or not obj: + continue + triples.append( + GraphTriple( + subject=subject, + relation=doc["relation"], + object=obj, + subject_type="entity", + object_type="entity", + confidence=float(doc.get("weight", 1.0)), + source=doc.get("source", ""), + properties=doc.get("properties") or {}, + ) + ) + return triples + + def get_triples_info(self, added=[], updated=[], removed=[]): + def to_dict(triple: GraphTriple) -> dict: + props = triple.properties or {} + return { + "subject": triple.subject, + "relation": triple.relation, + "object": triple.object, + "source": props.get("_source_path") or triple.source or "", + } + + def serialize(triples) -> list[dict]: + return sorted( + (to_dict(triple) for triple in triples), + key=lambda t: (t["source"], t["subject"], t["relation"], t["object"]), + ) + + return { + "triples_added": len(added), + "triples_updated": len(updated), + "triples_deleted": len(removed), + "added_triples": serialize(added), + "updated_triples": serialize(updated), + "deleted_triples": serialize(removed), + } + + async def reconcile_profile_triples(self, triples): + previous_triples = self._get_triples(key="profile") + added, updated, removed = self._diff_triples( + previous_triples, + triples, + ) + await self._delete_triples(removed, key="profile") + changed = [*added, *updated] + if changed: + await self.upsert_triples(changed) + + return self.get_triples_info(added, updated, removed) + + async def reconcile_source_triples( + self, source: str, triples: Sequence[GraphTriple] + ) -> None: + """Replace edges produced from one source file while preserving others.""" + previous_triples = self._get_triples(key="source", value=source) + added, updated, removed = self._diff_triples( + previous_triples, + triples, + ) + await self._delete_triples(removed, key="source", value=source) + changed = [*added, *updated] + if changed: + await self.upsert_triples(changed) + + return self.get_triples_info(added, updated, removed) + + def source_paths(self) -> set[str]: + """List document-managed origins, excluding lineage and explicit writes.""" + with self._state_lock: + docs = self._client.get_documents_by_type( + "GraphEdge", + as_list=True, + ) + + paths = set() + for doc in docs: + props = doc.get("properties") or {} + source_path = props.get("p__source_path") + + if not props.get("p__source_managed") or not source_path: + continue + + path = re.sub(r"#part=\d+$", "", str(source_path)) + paths.add(path) + + return paths + + async def reconcile_lineage_triples( + self, root: str, triples: Sequence[GraphTriple] + ) -> None: + """Replace all deterministic lineage edges for one workspace root.""" + previous_triples = self._get_triples(key="lineage", value=root) + added, updated, removed = self._diff_triples( + previous_triples, + triples, + ) + await self._delete_triples(removed, key="lineage", value=root) + changed = [*added, *updated] + if changed: + await self.upsert_triples(changed) + + return self.get_triples_info(added, updated, removed) + + async def delete_from_source( + self, + source: str, + key: str, + *, + value: str | None = None, + ) -> dict: + if not source: + raise ValueError("Source cannot be empty") + + with self._state_lock: + edge_docs = self._client.get_documents_by_type( + "GraphEdge", + as_list=True, + ) + def matches_source(doc): + if doc.get("source") == source: + return True + props = doc.get("properties") or {} + if props.get("p__source_path") == source: + return True + return False + + target_edges = [ + doc for doc in edge_docs + if matches_source(doc) + ] + target_triples = self._docs_to_triples(target_edges) + await self._delete_triples(target_triples, key, value) + + return { + "triples_deleted": len(target_triples) + } + + async def reset(self) -> None: + await asyncio.to_thread(self._reset_sync) + + def _reset_sync(self) -> None: + from terminusdb_client import Client + + admin_client = Client(self._dsn) + admin_client.connect( + team=self._team, + user=self._team, + key=self._password, + server_url=self._dsn, + ) + + if admin_client.has_database(self._db, team=self._team): + admin_client.delete_database(self._db, team=self._team) + + admin_client.create_database( + self._db, + team=self._team, + label=self._db, + description="Graph of DataMind stores in TerminusDB", + include_schema=False, + ) + + self._client = Client(self._dsn) + self._client.connect( + team=self._team, + db=self._db, + user=self._team, + key=self._password, + server_url=self._dsn, + ) + + self._schema() + with sqlite3.connect(self._history_path) as conn: + conn.execute("DELETE FROM history WHERE database = ?", (self._db,)) + + # ------------------------------------------------------------- lookup + + def _documents_at_version(self, version: str | None) -> list[dict]: + """Read instance documents from HEAD or an immutable commit.""" + with self._state_lock: + previous_ref = self._client.ref + try: + self._client.ref = version + return self._client.get_all_documents(as_list=True) + finally: + self._client.ref = previous_ref + + @staticmethod + def _entity_id(document_id: str) -> str: + return document_id.removeprefix("entity/") + + async def search_entities(self, query: str, *, top_k: int = 5, version: str | None = None,) -> list[Entity]: + q = query.lower().strip() + if not q: + return [] + scored: list[tuple[float, str, dict]] = [] + docs = self._documents_at_version(version) + for data in docs: + if data.get("@type") == "GraphEdge": + continue + nid = data["@id"] + public_id = self._entity_id(nid) + label = str(data.get("label", public_id)) + # Exact and substring hits beat fuzzy. + if public_id.lower() == q or label.lower() == q: + score = 1.0 + elif q in public_id.lower() or q in label.lower(): + score = 0.7 + else: + ratio = difflib.SequenceMatcher(None, q, label.lower()).ratio() + if ratio < 0.4: + continue + score = 0.4 + 0.4 * ratio # cap below substring score + scored.append((score, public_id, data)) + scored.sort(key=lambda x: -x[0]) + return [ + Entity( + id=nid, + label=str(data.get("label", nid)), + type=str(data.get("@type", "entity")), + score=score, + properties={ + k: v for k, v in data.items() if k not in {"label", "@type"} + }, + ) + for score, nid, data in scored[:top_k] + ] + + async def neighbors( + self, + entity: str, + *, + direction: str = "both", + relation_filter: list[str] | None = None, + limit: int = 100, + version: str | None = None, + ) -> list[Edge]: + if direction not in {"in", "out", "both"}: + raise ValueError(f"Invalid direction: {direction}") + if limit <= 0: + return [] + + docs = self._documents_at_version(version) + + allowed = set(relation_filter) if relation_filter else None + edges: list[Edge] = [] + + for doc in docs: + if doc.get("@type") != "GraphEdge": + continue + subject = doc.get("subject") + obj = doc.get("object") + if isinstance(subject, dict): + subject = subject.get("@id") + if isinstance(obj, dict): + obj = obj.get("@id") + if not isinstance(subject, str) or not isinstance(obj, str): + continue + subject_id = self._entity_id(subject) + object_id = self._entity_id(obj) + relation = doc.get("relation", "related") + + if allowed is not None and relation not in allowed: + continue + matches_out = ( + direction in {"out", "both"} and subject_id == entity + ) + matches_in = ( + direction in {"in", "both"} and object_id == entity + ) + + if matches_out or matches_in: + edges.append(self._edge(subject_id, object_id, doc)) + if len(edges) >= limit: + break + return edges + + async def traverse( + self, + start: str, + *, + max_hops: int = 3, + relation_filter: list[str] | None = None, + max_results: int = 100, + version: str | None = None + ) -> list[GraphPath]: + from terminusdb_client import WOQLQuery + + if max_hops < 1 or max_results <= 0: + return [] + + internal_start = ( + start if start.startswith("entity/") else f"entity/{start}" + ) + allowed = set(relation_filter) if relation_filter else None + # Each business edge of the jump operation = Entity -> GraphEdge -> Entity + pattern = ( + f"(outgoing_edge,incoming_edge){{1,{max_hops}}}" + ) + if version is None: + query = WOQLQuery().path( + internal_start, + pattern, + "v:End", + "v:Path", + ) + else: + query = WOQLQuery().using( + f"{self._team}/{self._db}/local/commit/{version}" + ).path( + internal_start, + pattern, + "v:End", + "v:Path", + ) + + docs = self._documents_at_version(version) + + result = self._client.query(query) + bindings = result.get("bindings", []) + + edge_doc_cache: dict[str, dict] = {} + def get_edge_doc(edge_id: str) -> dict: + if edge_id not in edge_doc_cache: + doc = next( + (d for d in docs if d.get("@id") == edge_id), + None, + ) + if not doc: + raise ValueError(f"GraphEdge document not found: {edge_id}") + edge_doc_cache[edge_id] = doc + return edge_doc_cache[edge_id] + + candidates = [] + for binding in bindings: + raw_path = binding.get("Path", []) + if len(raw_path) % 2 != 0: + continue + internal_nodes = [internal_start] + nodes = [self._entity_id(internal_start)] + edges = [] + seen = {internal_start} + bfs_keys = [] + valid = True + for i in range(0, len(raw_path), 2): + out_edge = raw_path[i] + in_edge = raw_path[i + 1] + + if not ( + out_edge.get("woql:predicate", "").endswith( + "outgoing_edge" + ) + and in_edge.get("woql:predicate", "").endswith( + "incoming_edge" + ) + ): + valid = False + break + + source_id = out_edge.get("woql:subject") + edge_id = out_edge.get("woql:object") + intermediate_id = in_edge.get("woql:subject") + target_id = in_edge.get("woql:object") + + if not edge_id or intermediate_id != edge_id: + valid = False + break + if source_id != internal_nodes[-1] or not target_id: + valid = False + break + if target_id in seen: + valid = False + break + + edge_doc = get_edge_doc(edge_id) + relation = edge_doc.get("relation", "related") + if allowed is not None and relation not in allowed: + valid = False + break + edge_obj = self._edge( + self._entity_id(source_id), + self._entity_id(target_id), + edge_doc, + ) + edges.append(edge_obj) + internal_nodes.append(target_id) + nodes.append(self._entity_id(target_id)) + seen.add(target_id) + + edge_key = str( + edge_doc.get("key") + or edge_doc.get("@id") + or edge_id + ) + bfs_keys.append( + (target_id, str(relation), edge_key) + ) + + if not valid or not edges: + continue + if len(edges) > max_hops: + continue + score = sum(edge.weight for edge in edges) / len(edges) + path_obj = GraphPath( + nodes=nodes, + edges=edges, + score=score, + ) + candidates.append( + (len(edges), tuple(bfs_keys), path_obj) + ) + + candidates.sort(key=lambda item: (item[0], item[1])) + selected = candidates[:max_results] + selected.sort(key=lambda item: -item[2].score) + return [item[2] for item in selected] + + async def rollback(self, version: str): + before_version = self._get_latest_version() + diff = self._client.diff_version(before_version, version) + if not diff: + return { + "success": True, + "before_version": before_version, + "target_version": version, + "message": "Current status matches the target version; no rollback is required.", + "info": { + "triples_added": 0, + "triples_updated": 0, + "triples_deleted": 0, + }, + "diff": [], + } + + message = f"Rollback to version {version}" + self._client.apply( + before_version=before_version, + after_version=version, + message=message, + ) + current = self._get_latest_version() + info = { + "success": True, + "message": message, + "version": current, + "diff": diff, + } + await self.add_version(info) + return info + + async def add_version(self, info: dict): + version = self._get_latest_version() + with sqlite3.connect(self._history_path) as conn: + conn.execute("INSERT OR REPLACE INTO history VALUES (?, ?, ?)", + (self._db, version, json.dumps(info, default=str))) + + async def list_versions(self): + with sqlite3.connect(self._history_path) as conn: + return [ + {"version": row[0], "payload": json.loads(row[1])} + for row in conn.execute( + "SELECT version, payload FROM history " + "WHERE database = ? ORDER BY rowid", + (self._db,), + ) + ] + # -------------------------------------------------------------- helpers + + @staticmethod + def _edge(u: str, v: str, data: dict) -> Edge: + return Edge( + source=u, + target=v, + relation=str(data.get("relation", "related")), + weight=float(data.get("weight", 1.0)), + properties={ + k[2:] if k.startswith("p_") else k: val + for k, val in data.items() + if k not in {"relation", "weight"} + }, + ) + + def stats(self) -> dict[str, int]: + """Return the number of graph nodes and edges in TerminusDB.""" + with self._state_lock: + edge_docs = self._client.get_documents_by_type( + "GraphEdge", + as_list=True, + ) + + node_ids: set[str] = set() + edge_count = 0 + + for doc in edge_docs: + if not isinstance(doc, dict): + continue + + subject = doc.get("subject") + obj = doc.get("object") + + # TerminusDB may return references as {"@id": "..."}. + if isinstance(subject, dict): + subject = subject.get("@id") + if isinstance(obj, dict): + obj = obj.get("@id") + + # Count only valid graph edges. + if not isinstance(subject, str) or not subject: + continue + if not isinstance(obj, str) or not obj: + continue + + node_ids.add(subject) + node_ids.add(obj) + edge_count += 1 + + return { + "nodes": len(node_ids), + "edges": edge_count, + } + + def _get_latest_version(self): + history = self._client.get_commit_history(max_history=1) + + if not history: + raise RuntimeError("No commit history found") + + if isinstance(history, dict): + first_key = next(iter(history)) + first_value = history[first_key] + + if isinstance(first_value, dict): + return ( + first_value.get("identifier") + or first_value.get("commit_id") + or first_key + ) + + return first_key + + if isinstance(history, list): + item = history[0] + return ( + item.get("identifier") + or item.get("commit_id") + or item.get("commit") + ) + + raise TypeError(f"Unexpected commit history: {type(history)}") + + + diff --git a/datamind/capabilities/graph/service.py b/datamind/capabilities/graph/service.py index b7d6ab2..8d728d8 100644 --- a/datamind/capabilities/graph/service.py +++ b/datamind/capabilities/graph/service.py @@ -11,6 +11,7 @@ from datamind.core.protocols import GraphEvidence, GraphStore, GraphTriple from .evidence import persist_input from datamind.core.registry import graph_registry +from datamind.core.errors import CapabilityError # Importing providers populates graph_registry. from . import providers # noqa: F401 @@ -82,6 +83,13 @@ async def load_from_profile(self) -> dict[str, int]: async def search_entities(self, query: str, *, top_k: int = 5) -> list[dict[str, Any]]: return [e.model_dump() for e in await self.store.search_entities(query, top_k=top_k)] + + async def search_entities_at_version(self, query: str, *, top_k: int = 5, version: str = None) -> list[dict[str, Any]]: + if self.store.name != "terminusdb": + raise CapabilityError("kb", "rollback is only supported by the LanceDB vector store") + if not isinstance(version, str): + raise CapabilityError("kb", f"version must be an string, got {type(version).__name__}") + return [e.model_dump() for e in await self.store.search_entities(query, top_k=top_k, version=version)] async def traverse( self, @@ -96,6 +104,25 @@ async def traverse( max_results=max_results, ) return [p.model_dump() for p in paths] + + async def traverse_at_version( + self, + start: str, + *, + max_hops: int = 2, + relation_filter: list[str] | None = None, + max_results: int = 100, + version: str = None, + ) -> list[dict[str, Any]]: + if self.store.name != "terminusdb": + raise CapabilityError("kb", "rollback is only supported by the LanceDB vector store") + if not isinstance(version, str): + raise CapabilityError("kb", f"version must be an string, got {type(version).__name__}") + paths = await self.store.traverse( + start, max_hops=max_hops, relation_filter=relation_filter, + max_results=max_results, version=version, + ) + return [p.model_dump() for p in paths] async def neighbors( self, @@ -109,6 +136,20 @@ async def neighbors( entity, direction=direction, relation_filter=relation_filter, limit=limit, ) return [e.model_dump() for e in edges] + + async def neighbors_at_version( + self, + entity: str, + *, + direction: str = "both", + relation_filter: list[str] | None = None, + limit: int = 100, + version: str = None, + ) -> list[dict[str, Any]]: + edges = await self.store.neighbors( + entity, direction=direction, relation_filter=relation_filter, limit=limit, version=version + ) + return [e.model_dump() for e in edges] async def upsert(self, triples: list[dict[str, Any]]) -> dict[str, int]: parsed = [GraphTriple(**t) for t in triples] @@ -127,8 +168,38 @@ async def upsert(self, triples: list[dict[str, Any]]) -> dict[str, int]: persist = getattr(self.store, "persist", None) if callable(persist): await persist() + + add_version = getattr(self.store, "add_version", None) + if add_version and callable(add_version): + info = self.store.get_triples_info(added=parsed) + await add_version(info) + return {"upserted": len(parsed)} + async def rollback(self, version: str = None) -> dict[str, Any]: + if self.store.name != "terminusdb": + raise CapabilityError("kb", "rollback is only supported by the LanceDB vector store") + if not isinstance(version, str): + raise CapabilityError("kb", f"version must be an string, got {type(version).__name__}") + + versions = await self.list_versions() + if not any(v["version"] == version for v in versions): + raise CapabilityError("kb", f"version {version} does not exist in this LanceDB table") + + result = await self.store.rollback(version=version) + return { + "status": "ok", + "message": f"Successfully rolled back to version {version}", + "rollback_version": version, # Target + "new_version": result["version"], # New Commit + } + + async def list_versions(self) -> dict[str, Any]: + if self.store.name != "terminusdb": + raise CapabilityError("kb", "list_versions is only supported by the LanceDB vector store") + + return await self.store.list_versions() + def stats(self) -> dict[str, int]: stats = getattr(self.store, "stats", None) return stats() if callable(stats) else {} @@ -148,6 +219,13 @@ def build_graph_service(settings: Settings) -> GraphService: "networkx", persist_path=str(storage_dir / "graph.json"), ) + elif backend == "terminusdb": + store = graph_registry.create( + "terminusdb", + dsn=settings.graph.dsn, + password=settings.graph.password, + persist_path=str(storage_dir / "graph"), + ) else: # DSN-based remote backends (neo4j, ...) — future work. store = graph_registry.create(backend, dsn=settings.graph.dsn) diff --git a/datamind/capabilities/graph/tools.py b/datamind/capabilities/graph/tools.py index ddaf41e..ad56497 100644 --- a/datamind/capabilities/graph/tools.py +++ b/datamind/capabilities/graph/tools.py @@ -12,6 +12,10 @@ def build_graph_tools(graph: GraphService) -> list[ToolSpec]: async def _search(query: str, top_k: int = 5) -> dict: entities = await graph.search_entities(query, top_k=top_k) return {"query": query, "count": len(entities), "entities": entities} + + async def _search_at_version(query: str, top_k: int = 5, version: str = None) -> dict: + entities = await graph.search_entities(query, top_k=top_k, version=version) + return {"query": query, "count": len(entities), "entities": entities} async def _traverse( start: str, @@ -30,6 +34,25 @@ async def _traverse( "total_count": len(paths), "truncated": truncated, "next_cursor": max_results if truncated else None, "paths": visible, } + + async def _traverse_at_version( + start: str, + max_hops: int = 2, + relation_filter: list[str] | None = None, + max_results: int = 100, + version: str = None, + ) -> dict: + paths = await graph.traverse( + start, max_hops=max_hops, relation_filter=relation_filter, + max_results=max_results + 1, version=version, + ) + truncated = len(paths) > max_results + visible = paths[:max_results] + return { + "start": start, "max_hops": max_hops, "count": len(visible), + "total_count": len(paths), "truncated": truncated, + "next_cursor": max_results if truncated else None, "paths": visible, + } async def _neighbors( entity: str, @@ -47,6 +70,24 @@ async def _neighbors( "total_count": len(edges), "truncated": truncated, "next_cursor": limit if truncated else None, "edges": visible, } + + async def _neighbors_at_version( + entity: str, + direction: str = "both", + relation_filter: list[str] | None = None, + limit: int = 100, + version: str = None + ) -> dict: + edges = await graph.neighbors( + entity, direction=direction, relation_filter=relation_filter, limit=limit + 1, version=version, + ) + truncated = len(edges) > limit + visible = edges[:limit] + return { + "entity": entity, "direction": direction, "count": len(visible), + "total_count": len(edges), "truncated": truncated, + "next_cursor": limit if truncated else None, "edges": visible, + } async def _upsert(triples: list[dict]) -> dict: return await graph.upsert(triples) @@ -54,7 +95,13 @@ async def _upsert(triples: list[dict]) -> dict: async def _schema(): return graph.store.schema() - return [ + async def _list_versions(): + return await graph.list_versions() + + async def _rollback(version: str = None): + return await graph.rollback(version=version) + + tools = [ ToolSpec(name='graph_schema', description='Discover observed node types and relation names before graph search.', input_schema={'type': 'object', 'properties': {}}, handler=_schema, metadata={'surface': 'graph', 'access': 'read'}), @@ -178,7 +225,148 @@ async def _schema(): }, ), ] - + + if ( + graph is not None + and getattr(getattr(graph, "store", None), "name", None) == "terminusdb" + ): + tools.extend([ + ToolSpec( + name="graph_search_entities_at_version", + description=( + "Search for knowledge graph entities by name or alias in a specific " + "graph version. Use this tool first when the query mentions an entity " + "by name and its entity ID is unknown. Returns matching entities from " + "the specified historical state without changing the current graph state." + ), + input_schema={ + "type": "object", + "properties": { + "query": {"type": "string", "description": "Entity name or alias."}, + "top_k": {"type": "integer", "minimum": 1, "maximum": 20, "default": 5}, + "version": { + "type": "string", + "description": ( + "The exact graph version or commit identifier to query. " + "Use an identifier returned by list_versions. " + ), + }, + }, + "required": ["query", "version"], + }, + handler=_search_at_version, + metadata={"group": "graph", "surface": "graph", "access": "read"}, + ), + ToolSpec( + name="graph_traverse_at_version", + description=( + "Traverse the knowledge graph from a starting entity in a specific " + "historical version. Explores paths up to max_hops and returns a bounded " + "set of paths ranked by average edge weight. Use relation_filter to " + "restrict traversal to exact relation types. The query does not change " + "the current graph state." + ), + input_schema={ + "type": "object", + "properties": { + "start": {"type": "string", "description": "Entity id to start from."}, + "max_hops": {"type": "integer", "minimum": 1, "maximum": 5, "default": 2}, + "relation_filter": { + "type": "array", + "items": {"type": "string"}, + "description": "Optional allow-list of relation names.", + }, + "max_results": { + "type": "integer", "minimum": 1, "maximum": 500, "default": 100, + "description": "Hard cap on returned paths.", + }, + "version": { + "type": "string", + "description": ( + "The exact graph version or commit identifier to query. " + "Use an identifier returned by list_versions. " + ), + }, + }, + "required": ["start", "version"], + }, + handler=_traverse_at_version, + metadata={"group": "graph", "surface": "graph", "access": "read"}, + ), + ToolSpec( + name="graph_neighbors_at_version", + description=( + "Retrieve edges connected to an entity in a specific historical graph " + "version. Use direction='out' for outgoing edges, 'in' for incoming " + "edges, or 'both' for both directions. Optionally restrict results using " + "relation_filter and cap the number of returned edges with limit. " + "Use this tool to inspect an entity's historical relationships without " + "changing the current graph state." + ), + input_schema={ + "type": "object", + "properties": { + "entity": {"type": "string"}, + "direction": { + "type": "string", + "enum": ["out", "in", "both"], + "default": "both", + }, + "relation_filter": { + "type": "array", "items": {"type": "string"}, + "description": "Only return these exact relation types.", + }, + "limit": { + "type": "integer", "minimum": 1, "maximum": 500, "default": 100, + }, + "version": { + "type": "string", + "description": ( + "The exact graph version or commit identifier to query. " + "Use an identifier returned by list_versions. " + ), + }, + }, + "required": ["entity", "version"], + }, + handler=_neighbors_at_version, + metadata={"group": "graph", "surface": "graph", "access": "read"}, + ), + ToolSpec( + name="list_versions", + description=( + "List available knowledge graph versions or commits. Use this tool to " + "discover valid version identifiers before calling historical entity " + "search, traversal, or neighbor tools. Use the returned identifiers " + "as the version argument for historical queries." + ), + input_schema={}, + handler=_list_versions, + metadata={"group": "graph", "surface": "graph", "access": "read"}, + ), + ToolSpec( + name="rollback", + description=( + "Restore the knowledge graph to a specified historical version. " + "This is a state-changing operation, unlike historical read tools. " + "Only use it when explicitly requested to roll back the graph; do not " + "use it merely to query historical data." + ), + input_schema={ + "version": { + "type": "string", + "description": ( + "The exact graph version or commit identifier to query. " + "Use an identifier returned by list_versions. " + ), + }, + }, + handler=_rollback, + metadata={"group": "graph", "surface": "graph", "access": "write", "destructive": True,}, + ), + ]) + + return tools @tool_provider_registry.register("graph") class _GraphToolProvider: diff --git a/datamind/capabilities/ingest/service.py b/datamind/capabilities/ingest/service.py index 4f9d709..f7ad0d6 100644 --- a/datamind/capabilities/ingest/service.py +++ b/datamind/capabilities/ingest/service.py @@ -1856,17 +1856,29 @@ async def _extract_graph_triples( triples.append(triple) return triples - async def _commit_graph_source(self, source: str | None, triples: list[GraphTriple]) -> None: + async def _commit_graph_source(self, source: str | None, triples: list[GraphTriple], version_info: bool = False) -> None: + """ + Args: + version_info: Whether to record this update as a new graph version + when the store supports versioning. + """ + reconcile = getattr(self._graph.store, "reconcile_source_triples", None) if source and callable(reconcile): - await reconcile(source, triples) + info = await reconcile(source, triples) else: - await self._graph.store.upsert_triples(triples) + info = await self._graph.store.upsert_triples(triples) persist = getattr(self._graph.store, "persist", None) if callable(persist): result = persist() if hasattr(result, "__await__"): await result + + if version_info: + add_version = getattr(self._graph.store, "add_version", None) + if add_version and callable(add_version): + await add_version(info) + return info or {} async def graph_add_triples_from_text( self, *, text: str, max_triples: int = 30, source: str | None = None, @@ -1884,7 +1896,8 @@ async def graph_add_triples_from_text( source_version=source_version, locator={"uri": str(evidence_path), "char_start": 0, "char_end": len(text)}, ) - await self._commit_graph_source(source, triples) + await self._commit_graph_source(source, triples, version_info=True) + return {"triples_added": len(triples), "source": source, "samples": [t.model_dump(mode="json") for t in triples[:8]]} @@ -1916,6 +1929,7 @@ async def graph_add_path( raise CapabilityError("ingest", f"path does not exist: {resolved}") processed, skipped, removed = [], [], [] + infos = {} for candidate in candidates: try: candidate = _resolve_safe_path(str(candidate), self._allowed_roots) @@ -1945,11 +1959,19 @@ async def graph_add_path( results.append({"source": part_source, "triples_added": len(extracted)}) if self._file_hash(candidate) != raw_version: raise CapabilityError("ingest", "source changed during extraction; retry") - await self._commit_graph_source(str(candidate), triples) + info = await self._commit_graph_source(str(candidate), triples) + processed.append({"path": str(candidate), "source": str(candidate), "source_version": raw_version, "triples_added": len(triples), "parts_processed": len(results), "parts": results, - "truncated": len(results) < len(parts) or len(triples) >= max_triples_per_file}) + "truncated": len(results) < len(parts) or len(triples) >= max_triples_per_file, **info}) + + for key, value in info.items(): + if isinstance(value, (int, float)): + infos[key] = infos.get(key, 0) + value + elif isinstance(value, list): + infos.setdefault(key, []).extend(value) + except Exception as exc: skipped.append(f"{candidate}: {type(exc).__name__}: {exc}") @@ -1958,12 +1980,107 @@ async def graph_add_path( for old_source in sorted(sources()): old = Path(old_source) if old.is_relative_to(resolved) and (recursive or old.parent == resolved) and not old.exists(): - await self._commit_graph_source(old_source, []) + info = await self._commit_graph_source(old_source, []) + for key, value in info.items(): + if isinstance(value, (int, float)): + infos[key] = infos.get(key, 0) + value + elif isinstance(value, list): + infos.setdefault(key, []).extend(value) removed.append(old_source) + + add_version = getattr(self._graph.store, "add_version", None) + if add_version and callable(add_version): + await add_version(infos) return {"path": str(resolved), "files_processed": len(processed), - "triples_added": sum(p["triples_added"] for p in processed), "files": processed, - "removed_sources": removed, "skipped": skipped[:20], "skipped_count": len(skipped)} + "triples_added": sum(p["triples_added"] for p in processed), + "triples_deleted": sum(p.get("triples_deleted") or 0 for p in processed), + "files": processed, "removed_sources": removed, + "skipped": skipped[:20], "skipped_count": len(skipped)} + + async def graph_delete_path( + self, + *, + path: str, + recursive: bool = True, + ) -> dict[str, Any]: + """ + Delete graph data associated with a file or directory and remove source files from disk. + The source files must not be deleted by the caller before invoking this method. + This method is responsible for coordinating graph cleanup and filesystem deletion. + """ + + if self._graph is None: + raise CapabilityError("ingest", "Graph surface is disabled") + if not path or not path.strip(): + raise CapabilityError("ingest", "path is required") + if getattr(self._graph.store, "name", None) != "terminusdb": + raise CapabilityError("ingest", "graph_delete_path is only supported by the TerminusDB graph backend") + + resolved = _resolve_safe_path(path, self._allowed_roots) + if resolved.is_file(): + if resolved.suffix.lower() not in _GRAPH_DOCUMENT_EXTS: + raise CapabilityError( + "ingest", + f"unsupported extension '{resolved.suffix}'. Supported: {sorted(_GRAPH_DOCUMENT_EXTS)}", + ) + candidates = [resolved] + elif resolved.is_dir(): + iterator = resolved.rglob("*") if recursive else resolved.glob("*") + candidates = sorted( + item for item in iterator + if item.is_file() and item.suffix.lower() in _GRAPH_DOCUMENT_EXTS + ) + else: + raise CapabilityError("ingest", f"path does not exist: {resolved}") + + processed: list[dict[str, Any]] = [] + skipped: list[str] = [] + total = 0 + for candidate in candidates: + try: + if candidate.suffix.lower() in _TEXT_EXTS: + content = candidate.read_text(encoding="utf-8", errors="replace") + else: + content = extract_document(candidate).text + if not content.strip(): + skipped.append(f"{candidate}: empty") + candidate.unlink() + continue + result = await self._graph.store.delete_from_source( + source=str(candidate), key="source", value=str(candidate) + ) + deleted = int(result.get("triples_deleted", 0)) + total += deleted + + # Actually delete the source files from disk. + candidate.unlink() + processed.append({ + "path": str(candidate), + "triples_deleted": deleted, + }) + except Exception as exc: + _log.warning( + "graph_delete_file_skipped", + extra={"path": str(candidate), "error": str(exc)}, + ) + skipped.append(f"{candidate}: {type(exc).__name__}: {exc}") + + + result = { + "path": str(resolved), + "files_deleted": len(processed), + "triples_deleted": total, + "files": processed, + "skipped": skipped[:20], + "skipped_count": len(skipped), + } + add_version = getattr(self._graph.store, "add_version", None) + if add_version and callable(add_version): + await add_version(result) + + return result + @staticmethod def _parse_triples_json(raw: str) -> list[dict[str, Any]]: """Valid empty output retracts edges; malformed output must never do so.""" diff --git a/datamind/capabilities/ingest/tools.py b/datamind/capabilities/ingest/tools.py index 23b683e..d229c7e 100644 --- a/datamind/capabilities/ingest/tools.py +++ b/datamind/capabilities/ingest/tools.py @@ -666,7 +666,7 @@ async def _graph_add_path( ), ] - if svc.ingest_mode == "cocoindex": + if svc and svc.ingest_mode == "cocoindex": tools.extend([ ToolSpec( name="kb_update_workspace", @@ -681,6 +681,35 @@ async def _graph_add_path( ), ]) + if ( + svc is not None + and getattr(svc, "_graph", None) is not None + and getattr(svc._graph, "store", None) is not None + and svc._graph.store.name == "terminusdb" + ): + tools.extend([ + ToolSpec( + name="graph_delete_path", + description=( + "Delete a file or directory from the dist workspace and remove " + "the graph triples ingested from that path. " + "When recursive is true, recursively delete directory contents. " + "Only delete paths inside the configured dist workspace. " + "Preserve graph data unrelated to the deleted path." + ), + input_schema={ + "type": "object", + "properties": { + "path": {"type": "string", "description": "Path to the source file or directory to delete."}, + "recursive": {"type": "boolean", "default": True}, + }, + "required": ["path"], + }, + handler=_kb_update_worksapce, + metadata={"group": "ingest", "surface": "graph", "access": "write"}, + ), + ]) + return tools __all__ = ["build_ingest_tools"] diff --git a/datamind/capabilities/kb/providers/chroma_store.py b/datamind/capabilities/kb/providers/chroma_store.py index a21c97a..5204e16 100644 --- a/datamind/capabilities/kb/providers/chroma_store.py +++ b/datamind/capabilities/kb/providers/chroma_store.py @@ -25,7 +25,6 @@ def __init__( dimension: int, ) -> None: import chromadb # type: ignore - self.name = "chroma" self.dimension = dimension self._persist_dir = Path(persist_dir) self._persist_dir.mkdir(parents=True, exist_ok=True) diff --git a/datamind/capabilities/kb/providers/lancedb_store.py b/datamind/capabilities/kb/providers/lancedb_store.py index 48d2284..0247117 100644 --- a/datamind/capabilities/kb/providers/lancedb_store.py +++ b/datamind/capabilities/kb/providers/lancedb_store.py @@ -33,7 +33,6 @@ def __init__( except ImportError as exc: from datamind.core.errors import ConfigError raise ConfigError("LanceDB requires pip install datamind[versioning]") from exc - self.name = "lancedb" self._write_lock = threading.RLock() self.dimension = dimension self._persist_dir = Path(persist_dir) diff --git a/datamind/config.py b/datamind/config.py index f119ca2..5e21c12 100644 --- a/datamind/config.py +++ b/datamind/config.py @@ -123,6 +123,7 @@ class GraphConfig(BaseModel): backend: str = "networkx" # networkx | neo4j | ... dsn: str | None = None # e.g. bolt://user:pw@host:7687 + password: str | None = None # for terminusdb | ... embed_entities: bool = False diff --git a/datamind/core/registry.py b/datamind/core/registry.py index a28ec23..60008eb 100644 --- a/datamind/core/registry.py +++ b/datamind/core/registry.py @@ -46,6 +46,7 @@ def deco(cls: type[T]) -> type[T]: f"new={cls.__module__}.{cls.__name__})" ) self._entries[name] = cls + cls.name = name # Set the backend identifier for runtime backend detection. return cls return deco diff --git a/datamind/direct_tools.py b/datamind/direct_tools.py index 989186e..0107297 100644 --- a/datamind/direct_tools.py +++ b/datamind/direct_tools.py @@ -204,12 +204,12 @@ async def replace_source(path, expected_version, triples): triple.evidence = [GraphEvidence(source_id='file:' + hashlib.sha256(str(source).encode()).hexdigest(), source_version=version, locator={'uri': str(evidence_path), 'original_uri': str(source)}, method='extracted', pipeline_version='codex-explicit-v1')] - await ingest._commit_graph_source(str(source), parsed) + await ingest._commit_graph_source(str(source), parsed, version_info=True) return {'source': str(source), 'source_version': version, 'triples_added': len(parsed)} async def forget_source(path): source = _resolve_safe_path(path, ingest._allowed_roots) - await ingest._commit_graph_source(str(source), []) + await ingest._commit_graph_source(str(source), [], version_info=True) return {'source': str(source), 'retracted': True, 'source_file_deleted': False} return [ diff --git a/datamind/tests/test_graph_versioning.py b/datamind/tests/test_graph_versioning.py new file mode 100644 index 0000000..4e84961 --- /dev/null +++ b/datamind/tests/test_graph_versioning.py @@ -0,0 +1,284 @@ +"""Integration tests for graph versioning backed by TerminusDB. + +Data is deliberately written through ``IngestService``. That is the public +path which both reconciles a document and registers the resulting TerminusDB +commit in DataMind's graph-version history. + +Run these tests against a disposable TerminusDB instance with:: + + DATAMIND_TEST_TERMINUSDB_DSN=http://127.0.0.1:6363 \ + DATAMIND_TEST_TERMINUSDB_PASSWORD=xxxxxxx pytest -q \ + datamind/tests/test_graph_versioning.py + +The configured database is reset between tests, so the suite is opt-in. +""" +from __future__ import annotations + +import inspect +import importlib.util +import os +import sqlite3 +from pathlib import Path + +import pytest +import pytest_asyncio + +from datamind.capabilities.graph.providers.terminusdb_store import ( + TerminusDBGraphStore, +) +from datamind.capabilities.graph.service import GraphService +from datamind.capabilities.ingest.service import IngestService + + +class _Model: + """Deterministic extractor; the entity is selected from the input text.""" + + async def generate_text(self, prompt: str, **kwargs) -> str: + if "Bob" in prompt: + return '[{"subject":"Bob","relation":"owns","object":"Project"}]' + return '[{"subject":"Alice","relation":"owns","object":"Project"}]' + + +async def _maybe_await(value): + """TerminusDB store methods are synchronous while services are async.""" + return await value if inspect.isawaitable(value) else value + + +@pytest.fixture(scope="session") +def terminusdb_config() -> tuple[str, str]: + if importlib.util.find_spec("terminusdb_client") is None: + pytest.skip("terminusdb_client is not installed") + dsn = os.getenv("DATAMIND_TEST_TERMINUSDB_DSN") + if not dsn: + pytest.skip( + "set DATAMIND_TEST_TERMINUSDB_DSN to run destructive TerminusDB tests" + ) + return dsn, os.getenv("DATAMIND_TEST_TERMINUSDB_PASSWORD", None) + + +@pytest_asyncio.fixture +async def graph(terminusdb_config, tmp_path: Path): + dsn, password = terminusdb_config + try: + store = TerminusDBGraphStore( + dsn=dsn, + password=password, + persist_path=tmp_path / "terminusdb", + ) + except Exception as exc: # pragma: no cover - depends on local service + pytest.skip(f"TerminusDB is unavailable: {type(exc).__name__}: {exc}") + + # The provider currently owns one fixed database (kg_default). Requiring + # an explicit test DSN above prevents accidental resets of normal setups. + await _maybe_await(store.reset()) + service = GraphService( + store=store, + data_dir=tmp_path / "profile", + storage_dir=tmp_path / "storage", + ) + yield service + await _maybe_await(service.aclose()) + + +@pytest.fixture +def profile(tmp_path: Path) -> Path: + path = tmp_path / "profile" + (path / "uploads").mkdir(parents=True) + return path + + +@pytest.fixture +def ingest(graph: GraphService, profile: Path) -> IngestService: + return IngestService( + kb=None, + db=None, + graph=graph, + llm_client=_Model(), + llm_model="test", + profile_data_dir=profile, + chunk_size=512, + chunk_overlap=64, + ingest_mode="direct", + ) + + +def _write(profile: Path, text: str, name: str = "ownership.md") -> Path: + path = profile / "uploads" / name + path.write_text(text, encoding="utf-8") + return path + + +async def _versions(graph: GraphService) -> list[dict]: + return await _maybe_await(graph.store.list_versions()) + + +async def _latest_version(graph: GraphService) -> str: + versions = await _versions(graph) + assert versions, "ingest did not register a graph version" + version = versions[-1].get("version") + assert isinstance(version, str) and version + return version + + +async def test_list_versions_uses_ingest_order_not_commit_hash_order(tmp_path: Path): + store = TerminusDBGraphStore.__new__(TerminusDBGraphStore) + store._db = "kg_default" + store._history_path = tmp_path / "version_history.sqlite3" + with sqlite3.connect(store._history_path) as conn: + conn.execute( + "CREATE TABLE history (" + "database TEXT, version TEXT, payload TEXT, " + "PRIMARY KEY(database, version))" + ) + conn.execute( + "INSERT INTO history VALUES (?, ?, ?)", + (store._db, "ffff", '{"step": 1}'), + ) + conn.execute( + "INSERT INTO history VALUES (?, ?, ?)", + (store._db, "0000", '{"step": 2}'), + ) + + versions = await store.list_versions() + + assert [item["version"] for item in versions] == ["ffff", "0000"] + assert versions[-1]["payload"] == {"step": 2} + + +@pytest.mark.asyncio +async def test_text_ingest_registers_a_graph_version(ingest, graph): + before = len(await _versions(graph)) + + result = await ingest.graph_add_triples_from_text(text="Alice owns Project") + + assert result["triples_added"] == 1 + versions = await _versions(graph) + assert len(versions) == before + 1 + assert versions[-1]["payload"]["triples_added"] == 1 + assert versions[-1]["version"] + + +@pytest.mark.asyncio +async def test_path_ingest_records_added_and_deleted_triples( + ingest, graph, profile, +): + source = _write(profile, "Alice owns Project") + first = await ingest.graph_add_path(path=str(source)) + assert first["triples_added"] == 1 + + source.write_text("Bob owns Project", encoding="utf-8") + second = await ingest.graph_add_path(path=str(source)) + + assert second["triples_added"] == 1 + version = (await _versions(graph))[-1] + assert version["payload"]["triples_added"] == 1 + assert version["payload"]["triples_deleted"] == 1 + + +@pytest.mark.asyncio +async def test_graph_delete_path_removes_file_edges_and_registers_version( + ingest, graph, profile, +): + source = _write(profile, "Alice owns Project") + await ingest.graph_add_path(path=str(source)) + count_before = len(await _versions(graph)) + + result = await ingest.graph_delete_path(path=str(source)) + + assert result["files_deleted"] == 1 + assert result["triples_deleted"] == 1 + assert result["skipped_count"] == 0 + assert not source.exists() + assert await graph.neighbors("Alice") == [] + versions = await _versions(graph) + assert len(versions) == count_before + 1 + assert versions[-1]["payload"]["triples_deleted"] == 1 + + +@pytest.mark.asyncio +async def test_versioned_search_reads_the_ingest_snapshot(ingest, graph, profile): + source = _write(profile, "Alice owns Project") + await ingest.graph_add_path(path=str(source)) + alice_version = await _latest_version(graph) + + source.write_text("Bob owns Project", encoding="utf-8") + await ingest.graph_add_path(path=str(source)) + + current = await graph.search_entities("Bob", top_k=10) + historical = await graph.search_entities_at_version( + "Alice", top_k=10, version=alice_version + ) + missing_historically = await graph.search_entities_at_version( + "Bob", top_k=10, version=alice_version + ) + + assert {item["id"] for item in current} == {"Bob"} + assert {item["id"] for item in historical} == {"Alice"} + assert missing_historically == [] + + +@pytest.mark.asyncio +async def test_neighbors_at_version_reads_the_ingest_snapshot( + ingest, graph, profile, +): + source = _write(profile, "Alice owns Project") + await ingest.graph_add_path(path=str(source)) + alice_version = await _latest_version(graph) + + source.write_text("Bob owns Project", encoding="utf-8") + await ingest.graph_add_path(path=str(source)) + + historical = await graph.neighbors_at_version( + "Alice", direction="out", version=alice_version + ) + assert [(edge["source"], edge["relation"], edge["target"]) + for edge in historical] == [("Alice", "owns", "Project")] + assert await graph.neighbors_at_version( + "Bob", direction="out", version=alice_version + ) == [] + + +@pytest.mark.asyncio +async def test_traverse_at_version_reads_the_ingest_snapshot( + ingest, graph, profile, +): + source = _write(profile, "Alice owns Project") + await ingest.graph_add_path(path=str(source)) + alice_version = await _latest_version(graph) + + source.write_text("Bob owns Project", encoding="utf-8") + await ingest.graph_add_path(path=str(source)) + + paths = await graph.traverse_at_version( + "Alice", max_hops=1, version=alice_version + ) + assert [path["nodes"] for path in paths] == [["Alice", "Project"]] + assert await graph.traverse_at_version( + "Bob", max_hops=1, version=alice_version + ) == [] + + +@pytest.mark.asyncio +async def test_rollback_restores_snapshot_and_appends_history( + ingest, graph, profile, +): + source = _write(profile, "Alice owns Project") + await ingest.graph_add_path(path=str(source)) + alice_version = await _latest_version(graph) + + source.write_text("Bob owns Project", encoding="utf-8") + await ingest.graph_add_path(path=str(source)) + count_before = len(await _versions(graph)) + + result = await graph.rollback(alice_version) + + assert result["status"] == "ok" + assert result["rollback_version"] == alice_version + assert result["new_version"] != alice_version + + assert len(await _versions(graph)) == count_before + 1 + assert {item["id"] for item in await graph.search_entities("Alice")} == { + "Alice" + } + assert await graph.search_entities("Bob") == [] +