"""Knowledge-graph engines (optional graph stage). Two engines share one contract (``ingest_document`` + ``finalize``): * lightrag — per-document entity extraction into a merged graph (Ollama LLM + Ollama embeddings). The LightRAG package renamed its Ollama bindings across releases, so binding functions are resolved by stepping through the known candidate names. * graphify — stages extracted documents as Markdown, then one ``graphify extract --backend ollama`` pass builds graph.json, GRAPH_REPORT.md, and the interactive graph.html (Graphify-Labs). When an engine cannot be constructed — missing package, unknown API generation, unreachable service — it raises GraphUnavailable with the reason, and the stage reports unavailable with that reason. """ from __future__ import annotations import hashlib from collections.abc import Callable from pathlib import Path class GraphUnavailable(Exception): """Raised when a graph engine cannot be constructed.""" def _resolve_ollama_binding(): """Return (complete_fn, embed_fn) from whichever naming generation of lightrag is installed, or raise GraphUnavailable.""" try: import lightrag.llm.ollama as _ollama_mod except ImportError as e: raise GraphUnavailable( "lightrag not installed (or too old) — run: pip install 'thicket[graph]'" ) from e complete = next( (name for name in ("ollama_model_complete", "ollama_complete") if hasattr(_ollama_mod, name)), None ) embed = next( (name for name in ("ollama_embed", "ollama_embedding") if hasattr(_ollama_mod, name)), None ) if not complete or not embed: raise GraphUnavailable( "installed lightrag exposes no known Ollama binding " f"(found complete={complete!r}, embed={embed!r}) — pin " "lightrag-hku to a release matching this code" ) return getattr(_ollama_mod, complete), getattr(_ollama_mod, embed) class LightRAGGraphStore: """Extracts entities and relationships via LightRAG + Ollama.""" def __init__(self, working_dir: Path, llm_model: str = "llama3", embed_model: str = "nomic-embed-text", ollama_host: str = "localhost", log: Callable[[str], None] = lambda _msg: None): try: from lightrag import LightRAG from lightrag.utils import EmbeddingFunc except ImportError as e: raise GraphUnavailable( "lightrag not installed — run: pip install 'thicket[graph]'" ) from e self.working_dir = working_dir / ".lightrag" self.working_dir.mkdir(parents=True, exist_ok=True) complete_fn, embed_fn = _resolve_ollama_binding() # nomic-embed-text (the default Ollama embedding model) is 768-dim. # num_ctx 16384 keeps entity extraction from truncating mid-chunk. self._log = log self._embed_model = embed_model self.rag = LightRAG( working_dir=str(self.working_dir), llm_model_func=complete_fn, llm_model_name=llm_model, llm_model_kwargs={"options": {"num_ctx": 16384}}, embedding_func=EmbeddingFunc( embedding_dim=768, max_token_size=8192, func=lambda texts: embed_fn(texts, embed_model=embed_model), ), ) def ingest_document(self, title: str, text: str) -> None: """Insert one document through the full async lifecycle — pipeline status, storages initialize, ainsert, finalize — inside one event loop (LightRAG >= 1.5 binds its shared storage to the running loop).""" if not text.strip(): return import asyncio formatted = f"Document Title: {title}\n\n{text}" async def _run() -> None: try: from lightrag.kg.shared_storage import ( initialize_pipeline_status, ) await initialize_pipeline_status( workspace=str(self.working_dir)) except ImportError: pass # step-down: releases without the handshake need no init await self.rag.initialize_storages() try: await self.rag.ainsert(formatted) finally: await self.rag.finalize_storages() asyncio.run(_run()) # ────────────────────────────────────────────────────────────────── # Graphify engine # ────────────────────────────────────────────────────────────────── class GraphifyGraphStore: """Graphify-Labs graphify: extracted documents are staged as Markdown, then one ``graphify extract --backend ollama`` pass at finalize() builds the queryable graph (graph.json), report, and interactive HTML under /graphify-out/.""" def __init__(self, working_dir: Path, llm_model: str = "llama3", embed_model: str = "nomic-embed-text", log: Callable[[str], None] = lambda _msg: None): import importlib.util if importlib.util.find_spec("graphify") is None: raise GraphUnavailable( "graphify not installed — run: pip install 'thicket[graphify]'" ) self._log = log self._llm_model = llm_model self.working_dir = working_dir / ".graphify" self.corpus_dir = self.working_dir / "corpus" self.corpus_dir.mkdir(parents=True, exist_ok=True) self._staged = 0 def ingest_document(self, title: str, text: str) -> None: """Stage one document as Markdown for the batch build.""" if not text.strip(): return from slugify import slugify slug = slugify(title) or "untitled" target = self.corpus_dir / f"{slug}.md" digest = hashlib.md5(text.encode()).hexdigest()[:6] if target.exists(): # same title, different content — sibling file target = self.corpus_dir / f"{slug}-{digest}.md" target.write_text(f"# {title}\n\n{text}", encoding="utf-8") self._staged += 1 def finalize(self) -> None: """One graphify pass over the staged corpus (local Ollama LLM; concurrency 1 — local models do not parallelize well).""" if not self._staged: return import subprocess import sys cmd = [sys.executable, "-m", "graphify", "extract", str(self.corpus_dir), "--backend", "ollama", "--model", self._llm_model, "--max-concurrency", "1", "--out", str(self.working_dir)] self._log(f"Graphify: building graph from {self._staged} staged " f"document(s) (local Ollama pass)...") result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600) if result.returncode != 0: raise GraphUnavailable( f"graphify extract failed: " f"{(result.stderr or result.stdout).strip()[:300]}" ) graph_json = self.working_dir / "graphify-out" / "graph.json" self._log(f"Graphify: graph built at '{graph_json}' " f"(query it: graphify query \"...\" --graph {graph_json})") # ────────────────────────────────────────────────────────────────── # Engine registry — the single dispatch point # ────────────────────────────────────────────────────────────────── GRAPH_ENGINES: dict[str, dict] = { "lightrag": {"class": LightRAGGraphStore, "modules": ("lightrag",)}, "graphify": {"class": GraphifyGraphStore, "modules": ("graphify",)}, } DEFAULT_GRAPH_ENGINE = "lightrag" def create_graph(engine: str, working_dir: Path, llm_model: str, embed_model: str, log: Callable[[str], None] = lambda _msg: None): """Build the named engine; raises GraphUnavailable with the known choices for an unknown name.""" spec = GRAPH_ENGINES.get(engine) if spec is None: known = ", ".join(sorted(GRAPH_ENGINES)) raise GraphUnavailable( f"unknown graph engine '{engine}' — known: {known}" ) return spec["class"](working_dir=working_dir, llm_model=llm_model, embed_model=embed_model, log=log)