211 lines
8.7 KiB
Python
211 lines
8.7 KiB
Python
"""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 <working_dir>/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)
|