From c9b0edc00fd3c66f7b1b96f3bc00430f5a598302 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 12 Jul 2026 13:25:49 +0300 Subject: [PATCH 01/10] fix: sync __version__ with pyproject (1.3.0) __version__ was left at 1.2.0 when 1.3.0 shipped; the value is stamped into every graph as __GraphRAGConfig__.sdk_version, so 1.3.0 installs were mislabeling graphs. Adds a drift test against pyproject.toml. Co-Authored-By: Claude Fable 5 --- graphrag_sdk/src/graphrag_sdk/__init__.py | 2 +- graphrag_sdk/tests/test_version_sync.py | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) create mode 100644 graphrag_sdk/tests/test_version_sync.py diff --git a/graphrag_sdk/src/graphrag_sdk/__init__.py b/graphrag_sdk/src/graphrag_sdk/__init__.py index 2f8d8ee4..b703b1da 100644 --- a/graphrag_sdk/src/graphrag_sdk/__init__.py +++ b/graphrag_sdk/src/graphrag_sdk/__init__.py @@ -11,7 +11,7 @@ # Adaptability — Optimization-ready core, strategies are swappable. # Velocity — Production-grade throughput. -__version__ = "1.2.0" +__version__ = "1.3.0" # ── API Surface (Facade) ──────────────────────────────────────── from graphrag_sdk.api.main import GraphRAG diff --git a/graphrag_sdk/tests/test_version_sync.py b/graphrag_sdk/tests/test_version_sync.py new file mode 100644 index 00000000..a084000e --- /dev/null +++ b/graphrag_sdk/tests/test_version_sync.py @@ -0,0 +1,14 @@ +"""Guard against src/graphrag_sdk/__init__.py __version__ drifting from pyproject.toml.""" +from __future__ import annotations + +import re +from pathlib import Path + +import graphrag_sdk + + +def test_version_matches_pyproject() -> None: + pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" + m = re.search(r'^version\s*=\s*"([^"]+)"', pyproject.read_text(encoding="utf-8"), re.M) + assert m, "version not found in pyproject.toml" + assert graphrag_sdk.__version__ == m.group(1) From c9210922e21eece51963f5125cdd6ca0057d31ae Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 12 Jul 2026 13:28:36 +0300 Subject: [PATCH 02/10] feat(tools): result models with budgeted to_llm_text rendering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Typed pydantic-v2 results for the agent toolkit (SearchResult, AnswerResult, SchemaResult, CypherResult, EntityResult, RememberResult + shared parts). to_llm_text() renders deterministic plain text bounded by max_chars, truncating only at item boundaries with an explicit …(N more) marker; every string passes the ingestion control-char sanitizer. Adds ReadOnlyViolation to core exceptions. Golden files pin the rendering contract. Co-Authored-By: Claude Fable 5 --- .../src/graphrag_sdk/core/exceptions.py | 20 + .../src/graphrag_sdk/tools/__init__.py | 42 ++ graphrag_sdk/src/graphrag_sdk/tools/models.py | 362 ++++++++++++++++++ .../tests/golden/tools/answer_result.txt | 7 + .../tests/golden/tools/cypher_result.txt | 5 + .../tests/golden/tools/entity_result.txt | 9 + .../tests/golden/tools/remember_result.txt | 2 + .../tests/golden/tools/schema_result.txt | 6 + .../tests/golden/tools/search_result.txt | 12 + .../golden/tools/search_result_truncated.txt | 6 + graphrag_sdk/tests/test_tools_models.py | 234 +++++++++++ 11 files changed, 705 insertions(+) create mode 100644 graphrag_sdk/src/graphrag_sdk/tools/__init__.py create mode 100644 graphrag_sdk/src/graphrag_sdk/tools/models.py create mode 100644 graphrag_sdk/tests/golden/tools/answer_result.txt create mode 100644 graphrag_sdk/tests/golden/tools/cypher_result.txt create mode 100644 graphrag_sdk/tests/golden/tools/entity_result.txt create mode 100644 graphrag_sdk/tests/golden/tools/remember_result.txt create mode 100644 graphrag_sdk/tests/golden/tools/schema_result.txt create mode 100644 graphrag_sdk/tests/golden/tools/search_result.txt create mode 100644 graphrag_sdk/tests/golden/tools/search_result_truncated.txt create mode 100644 graphrag_sdk/tests/test_tools_models.py diff --git a/graphrag_sdk/src/graphrag_sdk/core/exceptions.py b/graphrag_sdk/src/graphrag_sdk/core/exceptions.py index 99c3c379..efc17012 100644 --- a/graphrag_sdk/src/graphrag_sdk/core/exceptions.py +++ b/graphrag_sdk/src/graphrag_sdk/core/exceptions.py @@ -124,6 +124,26 @@ class SchemaValidationError(GraphRAGError): pass +# ── Read-Only Surface Errors ───────────────────────────────────── + + +class ReadOnlyViolation(GraphRAGError): + """Raised when a write is attempted through a read-only surface. + + Used by :mod:`graphrag_sdk.tools` both for guarded ``cypher_read`` + queries containing write clauses and for ``remember``/``flush`` + calls on a toolkit constructed with ``read_only=True``. + + Attributes: + offending_token: The specific rejected token (e.g. ``"MERGE"``, + ``"CALL apoc.load.json"``), or None for non-query violations. + """ + + def __init__(self, message: str, *, offending_token: str | None = None) -> None: + super().__init__(message) + self.offending_token = offending_token + + # ── Configuration Errors ───────────────────────────────────────── diff --git a/graphrag_sdk/src/graphrag_sdk/tools/__init__.py b/graphrag_sdk/src/graphrag_sdk/tools/__init__.py new file mode 100644 index 00000000..d0d40018 --- /dev/null +++ b/graphrag_sdk/src/graphrag_sdk/tools/__init__.py @@ -0,0 +1,42 @@ +"""Agentic GraphRAG toolkit — framework-neutral agent surface over GraphRAG. + +Canonical import: ``from graphrag_sdk.tools import GraphRAGToolkit``. +""" + +from __future__ import annotations + +from graphrag_sdk.core.exceptions import ReadOnlyViolation +from graphrag_sdk.tools.models import ( + AnswerResult, + ChunkRef, + Citation, + CypherResult, + DocumentRef, + EntityCard, + EntityResult, + EntityTypeInfo, + RelationTriple, + RelationTypeInfo, + RememberResult, + SchemaResult, + SearchResult, + ToolResult, +) + +__all__ = [ + "AnswerResult", + "ChunkRef", + "Citation", + "CypherResult", + "DocumentRef", + "EntityCard", + "EntityResult", + "EntityTypeInfo", + "ReadOnlyViolation", + "RelationTriple", + "RelationTypeInfo", + "RememberResult", + "SchemaResult", + "SearchResult", + "ToolResult", +] diff --git a/graphrag_sdk/src/graphrag_sdk/tools/models.py b/graphrag_sdk/src/graphrag_sdk/tools/models.py new file mode 100644 index 00000000..674e99fc --- /dev/null +++ b/graphrag_sdk/src/graphrag_sdk/tools/models.py @@ -0,0 +1,362 @@ +# GraphRAG SDK — Tools: result models + LLM-text rendering +# Typed results for the agent toolkit. Every model is pydantic v2 and +# renders a deterministic, budget-bounded plain-text form via to_llm_text(). + +from __future__ import annotations + +import json +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + +from graphrag_sdk.storage.graph_store import GraphStore + +_ELLIPSIS = "…" +_SNIPPET_CHARS = 240 + + +def _clean(value: Any) -> str: + """Sanitize a value for LLM output: strip control chars (reuses the + ingestion sanitizer) and collapse all whitespace runs to single spaces.""" + return " ".join(GraphStore._sanitize_string(str(value)).split()) + + +def _snippet(value: Any, limit: int = _SNIPPET_CHARS) -> str: + """A cleaned, length-capped one-line excerpt ending with an ellipsis.""" + text = _clean(value) + return text if len(text) <= limit else text[: limit - 1] + _ELLIPSIS + + +def _render( + preamble: list[str], + sections: list[tuple[str, list[str]]], + *, + max_chars: int, +) -> str: + """Assemble preamble lines + (header, items) sections into text ≤ max_chars. + + Truncation happens only at item boundaries; a dropped tail is marked + with the exact marker ``…(N more)``. Deterministic for equal inputs. + """ + if max_chars < 1: + return "" + lines: list[str] = [] + used = 0 + + def try_add(line: str) -> bool: + nonlocal used + cost = len(line) + (1 if lines else 0) + if used + cost > max_chars: + return False + lines.append(line) + used += cost + return True + + for line in preamble: + if not try_add(line): + if not lines: # always say something + lines.append(line[: max_chars - 1] + _ELLIPSIS) + return "\n".join(lines) + + for header, items in sections: + if not items: + continue + header_line = f"{header} ({len(items)}):" + full_marker = f" {_ELLIPSIS}({len(items)} more)" + # Only start a section if the header plus at least a drop-marker fit — + # otherwise a bare header would dangle with nothing under it. Marker + # strings never grow as the remaining count shrinks, so this reserve + # guarantees every emitted header is followed by an item or a marker. + needed = used + (1 if lines else 0) + len(header_line) + 1 + len(full_marker) + if needed > max_chars: + try_add(f"{header}: {_ELLIPSIS}({len(items)} items)") + continue + try_add(header_line) + for idx, item in enumerate(items): + marker = f" {_ELLIPSIS}({len(items) - idx} more)" + reserve = (len(marker) + 1) if idx < len(items) - 1 else 0 + if used + len(item) + 1 + reserve <= max_chars: + try_add(item) + else: + try_add(marker) + break + return "\n".join(lines) + + +class ToolResult(BaseModel): + """Base class for toolkit results: strict fields + LLM-text rendering.""" + + model_config = ConfigDict(extra="forbid") + + def to_llm_text(self, *, max_chars: int = 4000) -> str: + """Render a compact, deterministic plain-text form bounded by max_chars.""" + return _render(self._preamble(), self._sections(), max_chars=max_chars) + + def _preamble(self) -> list[str]: # pragma: no cover - overridden + return [] + + def _sections(self) -> list[tuple[str, list[str]]]: # pragma: no cover + return [] + + +class DocumentRef(BaseModel): + """A source document reference.""" + + model_config = ConfigDict(extra="forbid") + document_id: str + document_path: str = "" + + +class EntityCard(BaseModel): + """A knowledge-graph entity with its user-facing properties.""" + + model_config = ConfigDict(extra="forbid") + name: str + label: str = "" + description: str | None = None + properties: dict[str, Any] = Field(default_factory=dict) + + +class RelationTriple(BaseModel): + """A directed relationship ``source -[type]-> target`` with optional evidence.""" + + model_config = ConfigDict(extra="forbid") + source: str + type: str + target: str + fact: str | None = None + + +class ChunkRef(BaseModel): + """A retrieved source chunk with provenance ids.""" + + model_config = ConfigDict(extra="forbid") + chunk_id: str + document_id: str = "" + document_path: str = "" + text: str + + +class Citation(BaseModel): + """A provenance citation attached to a generated answer.""" + + model_config = ConfigDict(extra="forbid") + document_id: str + document_path: str = "" + chunk_id: str + snippet: str + + +class EntityTypeInfo(BaseModel): + """A declared or observed entity label with a live node count.""" + + model_config = ConfigDict(extra="forbid") + label: str + description: str | None = None + count: int = 0 + properties: list[str] = Field(default_factory=list) + + +class RelationTypeInfo(BaseModel): + """A declared or observed relation type with patterns and a live count.""" + + model_config = ConfigDict(extra="forbid") + label: str + description: str | None = None + patterns: list[tuple[str, str]] = Field(default_factory=list) + count: int = 0 + + +def _entity_line(e: EntityCard) -> str: + line = f"- {_clean(e.name)}" + if e.label: + line += f" [{_clean(e.label)}]" + if e.description: + line += f": {_snippet(e.description, 160)}" + return line + + +def _relation_line(r: RelationTriple) -> str: + line = f"- {_clean(r.source)} -[{_clean(r.type)}]-> {_clean(r.target)}" + if r.fact: + line += f": {_snippet(r.fact, 160)}" + return line + + +class SearchResult(ToolResult): + """Ranked, typed retrieval context for a query (no LLM generation).""" + + query: str + entities: list[EntityCard] = Field(default_factory=list) + relations: list[RelationTriple] = Field(default_factory=list) + facts: list[str] = Field(default_factory=list) + chunks: list[ChunkRef] = Field(default_factory=list) + documents: list[DocumentRef] = Field(default_factory=list) + + def _preamble(self) -> list[str]: + return [f"Query: {_clean(self.query)}"] + + def _sections(self) -> list[tuple[str, list[str]]]: + return [ + ("Entities", [_entity_line(e) for e in self.entities]), + ("Relations", [_relation_line(r) for r in self.relations]), + ("Facts", [f"- {_snippet(f)}" for f in self.facts]), + ( + "Chunks", + [ + f"- [{_clean(c.document_path or c.document_id)}#{_clean(c.chunk_id)}] " + f"{_snippet(c.text)}" + for c in self.chunks + ], + ), + ( + "Documents", + [ + f"- {_clean(d.document_id)} ({_clean(d.document_path)})" + if d.document_path + else f"- {_clean(d.document_id)}" + for d in self.documents + ], + ), + ] + + +class AnswerResult(ToolResult): + """A generated answer plus provenance citations.""" + + answer: str + citations: list[Citation] = Field(default_factory=list) + entities_touched: list[str] = Field(default_factory=list) + cypher_used: str | None = None + + def _preamble(self) -> list[str]: + return [_clean(line) for line in self.answer.splitlines() if line.strip()] + + def _sections(self) -> list[tuple[str, list[str]]]: + sections = [ + ( + "Citations", + [ + f"- [{_clean(c.document_path or c.document_id)}#{_clean(c.chunk_id)}] " + f"{_snippet(c.snippet)}" + for c in self.citations + ], + ), + ("Entities", [f"- {_clean(n)}" for n in self.entities_touched]), + ] + if self.cypher_used: + sections.append(("Cypher used", [f"- {_clean(self.cypher_used)}"])) + return sections + + +class SchemaResult(ToolResult): + """The graph's entity labels and relation types with live counts.""" + + entities: list[EntityTypeInfo] = Field(default_factory=list) + relations: list[RelationTypeInfo] = Field(default_factory=list) + node_count: int = 0 + edge_count: int = 0 + + def _preamble(self) -> list[str]: + return [f"Nodes: {self.node_count} | Edges: {self.edge_count}"] + + def _sections(self) -> list[tuple[str, list[str]]]: + ent_lines = [] + for e in self.entities: + line = f"- {_clean(e.label)}: {e.count}" + if e.description: + line += f" — {_snippet(e.description, 100)}" + if e.properties: + line += f" (props: {', '.join(_clean(p) for p in e.properties)})" + ent_lines.append(line) + rel_lines = [] + for r in self.relations: + line = f"- {_clean(r.label)}: {r.count}" + if r.patterns: + pats = ", ".join(f"{_clean(a)}->{_clean(b)}" for a, b in r.patterns) + line += f" [{pats}]" + rel_lines.append(line) + return [("Entity labels", ent_lines), ("Relation types", rel_lines)] + + +class CypherResult(ToolResult): + """Rows returned by a guarded read-only Cypher query.""" + + columns: list[str] = Field(default_factory=list) + rows: list[list[Any]] = Field(default_factory=list) + row_count: int = 0 + truncated: bool = False + + def _preamble(self) -> list[str]: + suffix = " (truncated)" if self.truncated else "" + return [ + f"Columns: {', '.join(_clean(c) for c in self.columns)}", + f"Rows: {self.row_count}{suffix}", + ] + + def _sections(self) -> list[tuple[str, list[str]]]: + return [ + ( + "Rows", + [ + f"- {_clean(json.dumps(row, ensure_ascii=False, default=str))}" + for row in self.rows + ], + ) + ] + + +class EntityResult(ToolResult): + """Entity card: best-match node, neighbors, and source documents.""" + + query: str + found: bool = False + entity: EntityCard | None = None + neighbors: list[RelationTriple] = Field(default_factory=list) + nearby: list[str] = Field(default_factory=list) + documents: list[DocumentRef] = Field(default_factory=list) + + def _preamble(self) -> list[str]: + if not self.found or self.entity is None: + return [ + f"No entity found matching '{_clean(self.query)}'. " + f"Try graph_search for fuzzy discovery." + ] + e = self.entity + lines = [f"Entity: {_clean(e.name)}" + (f" [{_clean(e.label)}]" if e.label else "")] + if e.description: + lines.append(f"Description: {_snippet(e.description)}") + if e.properties: + props = "; ".join(f"{_clean(k)}={_clean(v)}" for k, v in sorted(e.properties.items())) + lines.append(f"Properties: {props}") + return lines + + def _sections(self) -> list[tuple[str, list[str]]]: + return [ + ("Neighbors", [_relation_line(r) for r in self.neighbors]), + ("Documents", [f"- {_clean(d.document_id)}" for d in self.documents]), + ("Nearby", [f"- {_clean(n)}" for n in self.nearby]), + ] + + +class RememberResult(ToolResult): + """Outcome of storing text into the graph via graph_remember.""" + + document_id: str + chunks_indexed: int = 0 + nodes_created: int = 0 + relationships_created: int = 0 + finalized: bool = False + + def _preamble(self) -> list[str]: + lines = [ + f"Stored document '{_clean(self.document_id)}' " + f"({self.chunks_indexed} chunks, {self.nodes_created} nodes, " + f"{self.relationships_created} relations)." + ] + lines.append( + "Finalized." + if self.finalized + else "Pending finalize — call graph_flush (or GraphRAG.finalize())." + ) + return lines diff --git a/graphrag_sdk/tests/golden/tools/answer_result.txt b/graphrag_sdk/tests/golden/tools/answer_result.txt new file mode 100644 index 00000000..f3daf961 --- /dev/null +++ b/graphrag_sdk/tests/golden/tools/answer_result.txt @@ -0,0 +1,7 @@ +Alice works at Acme Corp. +She is an engineer. +Citations (1): +- [docs/a.md#c1] Alice works at Acme Corp. +Entities (2): +- Alice +- Acme Corp diff --git a/graphrag_sdk/tests/golden/tools/cypher_result.txt b/graphrag_sdk/tests/golden/tools/cypher_result.txt new file mode 100644 index 00000000..e0e4e335 --- /dev/null +++ b/graphrag_sdk/tests/golden/tools/cypher_result.txt @@ -0,0 +1,5 @@ +Columns: name, n +Rows: 2 (truncated) +Rows (2): +- ["Alice", 3] +- ["Bob", 1] diff --git a/graphrag_sdk/tests/golden/tools/entity_result.txt b/graphrag_sdk/tests/golden/tools/entity_result.txt new file mode 100644 index 00000000..f9e98c3e --- /dev/null +++ b/graphrag_sdk/tests/golden/tools/entity_result.txt @@ -0,0 +1,9 @@ +Entity: Alice [Person] +Description: Engineer +Properties: seniority=senior +Neighbors (1): +- Alice -[WORKS_AT]-> Acme Corp +Documents (1): +- doc-a +Nearby (1): +- Alice Smith diff --git a/graphrag_sdk/tests/golden/tools/remember_result.txt b/graphrag_sdk/tests/golden/tools/remember_result.txt new file mode 100644 index 00000000..17cae038 --- /dev/null +++ b/graphrag_sdk/tests/golden/tools/remember_result.txt @@ -0,0 +1,2 @@ +Stored document 'text-abc123' (1 chunks, 2 nodes, 1 relations). +Pending finalize — call graph_flush (or GraphRAG.finalize()). diff --git a/graphrag_sdk/tests/golden/tools/schema_result.txt b/graphrag_sdk/tests/golden/tools/schema_result.txt new file mode 100644 index 00000000..67374b9a --- /dev/null +++ b/graphrag_sdk/tests/golden/tools/schema_result.txt @@ -0,0 +1,6 @@ +Nodes: 7 | Edges: 9 +Entity labels (2): +- Person: 2 — A human (props: seniority) +- Organization: 1 +Relation types (1): +- WORKS_AT: 1 [Person->Organization] diff --git a/graphrag_sdk/tests/golden/tools/search_result.txt b/graphrag_sdk/tests/golden/tools/search_result.txt new file mode 100644 index 00000000..2fa97ab7 --- /dev/null +++ b/graphrag_sdk/tests/golden/tools/search_result.txt @@ -0,0 +1,12 @@ +Query: Who works at Acme? +Entities (2): +- Alice [Person]: Engineer at Acme +- Acme Corp [Organization]: A tech company +Relations (1): +- Alice -[WORKS_AT]-> Acme Corp: Alice is employed at Acme Corp +Facts (1): +- Alice —[WORKS_AT]→ Acme Corp: employment +Chunks (1): +- [docs/a.md#c1] Alice works at Acme Corp. +Documents (1): +- doc-a (docs/a.md) diff --git a/graphrag_sdk/tests/golden/tools/search_result_truncated.txt b/graphrag_sdk/tests/golden/tools/search_result_truncated.txt new file mode 100644 index 00000000..aec958b6 --- /dev/null +++ b/graphrag_sdk/tests/golden/tools/search_result_truncated.txt @@ -0,0 +1,6 @@ +Query: Who works at Acme? +Entities (2): +- Alice [Person]: Engineer at Acme +- Acme Corp [Organization]: A tech company +Relations (1): + …(1 more) diff --git a/graphrag_sdk/tests/test_tools_models.py b/graphrag_sdk/tests/test_tools_models.py new file mode 100644 index 00000000..fae3fdd3 --- /dev/null +++ b/graphrag_sdk/tests/test_tools_models.py @@ -0,0 +1,234 @@ +"""Rendering and shape tests for graphrag_sdk.tools result models.""" + +from __future__ import annotations + +import json +import os +import re +from pathlib import Path + +from graphrag_sdk.core.exceptions import GraphRAGError, ReadOnlyViolation +from graphrag_sdk.tools import ( + AnswerResult, + ChunkRef, + Citation, + CypherResult, + DocumentRef, + EntityCard, + EntityResult, + EntityTypeInfo, + RelationTriple, + RelationTypeInfo, + RememberResult, + SchemaResult, + SearchResult, +) + +GOLDEN = Path(__file__).parent / "golden" / "tools" + + +def check_golden(name: str, actual: str) -> None: + path = GOLDEN / name + if os.getenv("UPDATE_GOLDEN") == "1": + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(actual + "\n", encoding="utf-8") + return + assert actual + "\n" == path.read_text(encoding="utf-8"), f"golden mismatch: {name}" + + +def _search_result() -> SearchResult: + return SearchResult( + query="Who works at Acme?", + entities=[ + EntityCard( + name="Alice", + label="Person", + description="Engineer at Acme", + properties={"seniority": "senior"}, + ), + EntityCard( + name="Acme Corp", + label="Organization", + description="A tech company", + properties={}, + ), + ], + relations=[ + RelationTriple( + source="Alice", + type="WORKS_AT", + target="Acme Corp", + fact="Alice is employed at Acme Corp", + ) + ], + facts=["Alice —[WORKS_AT]→ Acme Corp: employment"], + chunks=[ + ChunkRef( + chunk_id="c1", + document_id="doc-a", + document_path="docs/a.md", + text="Alice works at Acme Corp.", + ) + ], + documents=[DocumentRef(document_id="doc-a", document_path="docs/a.md")], + ) + + +def test_readonly_violation_hierarchy(): + err = ReadOnlyViolation("no", offending_token="MERGE") + assert isinstance(err, GraphRAGError) + assert err.offending_token == "MERGE" + assert ReadOnlyViolation("no").offending_token is None + + +def test_model_dump_json_round_trip(): + sr = _search_result() + assert SearchResult.model_validate(json.loads(sr.model_dump_json())) == sr + + +def test_to_llm_text_deterministic_and_bounded(): + sr = _search_result() + assert sr.to_llm_text() == sr.to_llm_text() + for max_chars in (10, 50, 120, 4000): + out = sr.to_llm_text(max_chars=max_chars) + assert len(out) <= max_chars + + +def test_truncation_marker_at_item_boundary(): + sr = SearchResult( + query="q", + entities=[ + EntityCard(name=f"E{i:02d}", label="Person", description="d" * 30, properties={}) + for i in range(20) + ], + relations=[], + facts=[], + chunks=[], + documents=[], + ) + out = sr.to_llm_text(max_chars=300) + assert "…(" in out and "more)" in out + assert len(out) <= 300 + # never truncates mid-entity: every emitted entity line is complete + for line in out.splitlines(): + if line.startswith("- E"): + assert line.endswith("d" * 30) + + +def test_no_dangling_section_header(): + """Every emitted section header is followed by an item or a drop-marker.""" + sr = _search_result() + for max_chars in range(40, 400, 7): + out = sr.to_llm_text(max_chars=max_chars) + lines = out.splitlines() + for i, line in enumerate(lines): + if re.match(r"^[A-Z][A-Za-z ]+ \(\d+\):$", line): + assert i + 1 < len(lines), f"dangling header at max_chars={max_chars}: {out!r}" + nxt = lines[i + 1] + assert nxt.startswith("- ") or nxt.startswith(f" {chr(0x2026)}"), ( + f"header not followed by content at max_chars={max_chars}: {out!r}" + ) + + +def test_control_characters_stripped(): + sr = SearchResult( + query="bad\x00query\x07", entities=[], relations=[], facts=[], chunks=[], documents=[] + ) + out = sr.to_llm_text() + assert "\x00" not in out and "\x07" not in out and "badquery" in out + + +def test_golden_search(): + check_golden("search_result.txt", _search_result().to_llm_text()) + + +def test_golden_search_truncated(): + check_golden("search_result_truncated.txt", _search_result().to_llm_text(max_chars=160)) + + +def test_golden_answer(): + ar = AnswerResult( + answer="Alice works at Acme Corp.\nShe is an engineer.", + citations=[ + Citation( + document_id="doc-a", + document_path="docs/a.md", + chunk_id="c1", + snippet="Alice works at Acme Corp.", + ) + ], + entities_touched=["Alice", "Acme Corp"], + cypher_used=None, + ) + check_golden("answer_result.txt", ar.to_llm_text()) + + +def test_golden_schema(): + sr = SchemaResult( + entities=[ + EntityTypeInfo( + label="Person", description="A human", count=2, properties=["seniority"] + ), + EntityTypeInfo(label="Organization", description=None, count=1, properties=[]), + ], + relations=[ + RelationTypeInfo( + label="WORKS_AT", + description=None, + patterns=[("Person", "Organization")], + count=1, + ) + ], + node_count=7, + edge_count=9, + ) + check_golden("schema_result.txt", sr.to_llm_text()) + + +def test_golden_cypher(): + cr = CypherResult( + columns=["name", "n"], rows=[["Alice", 3], ["Bob", 1]], row_count=2, truncated=True + ) + check_golden("cypher_result.txt", cr.to_llm_text()) + + +def test_golden_entity_found_and_not_found(): + er = EntityResult( + query="Alice", + found=True, + entity=EntityCard( + name="Alice", + label="Person", + description="Engineer", + properties={"seniority": "senior"}, + ), + neighbors=[RelationTriple(source="Alice", type="WORKS_AT", target="Acme Corp", fact=None)], + nearby=["Alice Smith"], + documents=[DocumentRef(document_id="doc-a", document_path="docs/a.md")], + ) + check_golden("entity_result.txt", er.to_llm_text()) + missing = EntityResult( + query="Zorp", found=False, entity=None, neighbors=[], nearby=[], documents=[] + ) + assert "No entity" in missing.to_llm_text() and "Zorp" in missing.to_llm_text() + + +def test_golden_remember(): + rr = RememberResult( + document_id="text-abc123", + chunks_indexed=1, + nodes_created=2, + relationships_created=1, + finalized=False, + ) + check_golden("remember_result.txt", rr.to_llm_text()) + assert ( + "Finalized." + in RememberResult( + document_id="d", + chunks_indexed=0, + nodes_created=0, + relationships_created=0, + finalized=True, + ).to_llm_text() + ) From 5a60b0bb9a8cfef8be7a3884933a14f274c2ef1b Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 12 Jul 2026 13:29:02 +0300 Subject: [PATCH 03/10] feat(tools): fail-closed read-only Cypher guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-pass lexer strips comments and masks string/backtick literals (regex-order stripping is exploitable via quote/comment interleaving), then scans both the raw and NFKC-normalized copies for write keywords, multi-statement separators, disallowed start keywords, and CALLs to procedures outside a full-name read-safe allowlist. Rejections raise ReadOnlyViolation naming the offending token. apply_limit() injects LIMIT only when the (noise-stripped) query lacks one. The original query text — never the normalized copy — is what executes. Co-Authored-By: Claude Fable 5 --- .../src/graphrag_sdk/tools/cypher_guard.py | 141 ++++++++++++++++++ graphrag_sdk/tests/test_tools_cypher_guard.py | 91 +++++++++++ 2 files changed, 232 insertions(+) create mode 100644 graphrag_sdk/src/graphrag_sdk/tools/cypher_guard.py create mode 100644 graphrag_sdk/tests/test_tools_cypher_guard.py diff --git a/graphrag_sdk/src/graphrag_sdk/tools/cypher_guard.py b/graphrag_sdk/src/graphrag_sdk/tools/cypher_guard.py new file mode 100644 index 00000000..df1ded7c --- /dev/null +++ b/graphrag_sdk/src/graphrag_sdk/tools/cypher_guard.py @@ -0,0 +1,141 @@ +# GraphRAG SDK — Tools: read-only Cypher guard +# Fail-closed validation for agent-supplied Cypher. Strings/comments are +# removed with a single-pass lexer state machine (ordering regex passes is +# exploitable), then write keywords are scanned on both the raw and the +# NFKC-normalized copies. The ORIGINAL query text is what gets executed. + +from __future__ import annotations + +import re +import unicodedata + +from graphrag_sdk.core.exceptions import ReadOnlyViolation + +_WRITE_TOKENS = ("CREATE", "MERGE", "DELETE", "DETACH", "SET", "REMOVE", "DROP", "FOREACH") +_START_KEYWORDS = ("MATCH", "OPTIONAL", "UNWIND", "WITH", "RETURN", "CALL") +# Full procedure names only — prefixes are unsafe (db.idx.fulltext.createNodeIndex +# is a WRITE). Compared lowercase. +READ_SAFE_PROCEDURES = frozenset( + { + "db.labels", + "db.relationshiptypes", + "db.propertykeys", + "db.indexes", + "db.idx.fulltext.querynodes", + "db.idx.fulltext.queryrelationships", + "db.idx.vector.querynodes", + "db.idx.vector.queryrelationships", + } +) + + +def _strip_noise(text: str) -> str: + """Remove comments and mask string/backtick literals in one lexer pass. + + Comments are removed as EMPTY string (not a space) on purpose: it is + fail-closed against mid-token splitting (``Cr/**/eate`` reassembles to + ``Create`` and is caught), at the cost of over-rejecting queries that + rely on a comment as the only token separator — acceptable for a guard. + String/backtick literals become single spaces so their content can never + trip (or hide) a keyword. + """ + out: list[str] = [] + i, n = 0, len(text) + state = "code" + while i < n: + ch = text[i] + nxt = text[i + 1] if i + 1 < n else "" + if state == "code": + if ch == "/" and nxt == "/": + state = "line" + i += 2 + continue + if ch == "/" and nxt == "*": + state = "block" + i += 2 + continue + if ch in ("'", '"', "`"): + state = ch + out.append(" ") + i += 1 + continue + out.append(ch) + i += 1 + continue + if state == "line": + if ch == "\n": + state = "code" + out.append("\n") + i += 1 + continue + if state == "block": + if ch == "*" and nxt == "/": + state = "code" + i += 2 + continue + i += 1 + continue + # inside a quoted literal (state is the quote char) + if ch == "\\" and state in ("'", '"'): + i += 2 + continue + if ch == state: + state = "code" + out.append(" ") + i += 1 + return "".join(out) + + +def _scan(stripped: str) -> None: + """Reject write constructs in noise-stripped Cypher text.""" + first = re.match(r"\s*([A-Za-z]+)", stripped) + if first and first.group(1).upper() not in _START_KEYWORDS: + raise ReadOnlyViolation( + f"Query must start with one of {', '.join(_START_KEYWORDS)}; got '{first.group(1)}'.", + offending_token=first.group(1), + ) + body = stripped.rstrip().rstrip(";") + if ";" in body: + raise ReadOnlyViolation("Multiple Cypher statements are not allowed.", offending_token=";") + for token in _WRITE_TOKENS: + if re.search(rf"\b{token}\b", stripped, re.IGNORECASE): + raise ReadOnlyViolation( + f"Write operation '{token}' is not allowed — cypher_read is read-only.", + offending_token=token, + ) + if re.search(r"\bLOAD\s+CSV\b", stripped, re.IGNORECASE): + raise ReadOnlyViolation("LOAD CSV is not allowed.", offending_token="LOAD CSV") + for match in re.finditer(r"\bCALL\b(\s*)([A-Za-z0-9_.]*)", stripped, re.IGNORECASE): + proc = match.group(2) + if not proc: + rest = stripped[match.end() :].lstrip() + if rest.startswith("{"): + continue # CALL { subquery }: inner writes caught by the scan above + raise ReadOnlyViolation("Bare CALL is not allowed.", offending_token="CALL") + if proc.lower() not in READ_SAFE_PROCEDURES: + raise ReadOnlyViolation( + f"Procedure '{proc}' is not on the read-safe allowlist " + f"({', '.join(sorted(READ_SAFE_PROCEDURES))}).", + offending_token=f"CALL {proc}", + ) + + +def ensure_read_only(query: str) -> None: + """Raise :class:`ReadOnlyViolation` unless *query* is a read-only statement.""" + if not query or not query.strip(): + raise ReadOnlyViolation("Empty Cypher query.", offending_token=None) + _scan(_strip_noise(query)) + _scan(_strip_noise(unicodedata.normalize("NFKC", query))) + + +def apply_limit(query: str, limit: int) -> tuple[str, bool]: + """Append ``LIMIT {limit}`` when the query has no LIMIT clause. + + LIMIT detection runs on the noise-stripped copy so a literal string + containing the word LIMIT does not suppress injection. Returns the + (possibly rewritten) query and whether injection happened. + """ + trimmed = query.rstrip().rstrip(";") + if re.search(r"\bLIMIT\b", _strip_noise(trimmed), re.IGNORECASE): + return trimmed, False + return f"{trimmed}\nLIMIT {int(limit)}", True diff --git a/graphrag_sdk/tests/test_tools_cypher_guard.py b/graphrag_sdk/tests/test_tools_cypher_guard.py new file mode 100644 index 00000000..9d0aea01 --- /dev/null +++ b/graphrag_sdk/tests/test_tools_cypher_guard.py @@ -0,0 +1,91 @@ +"""Table-driven tests for the tools read-only Cypher guard.""" + +from __future__ import annotations + +import pytest + +from graphrag_sdk.core.exceptions import ReadOnlyViolation +from graphrag_sdk.tools.cypher_guard import apply_limit, ensure_read_only + +BENIGN = [ + "MATCH (n:Person) RETURN n.name LIMIT 5", + "MATCH (n) WHERE n.name = 'DELETE ME' RETURN n", # write kw inside string + 'MATCH (n) WHERE n.note = "please MERGE later" RETURN n', + "MATCH (n) WHERE n.created > 2020 RETURN n", # substring identifier + "MATCH (a)-[r:RELATES]->(b) RETURN a.name, r.rel_type, b.name", + "UNWIND $ids AS i MATCH (n {id: i}) RETURN n.name", + "WITH 1 AS x RETURN x", + "RETURN 1", + "CALL db.labels()", + "CALL db.idx.fulltext.queryNodes('Entity', 'alice') YIELD node RETURN node.name", + "MATCH (n) // trailing comment\nRETURN n.name", + "MATCH (n) /* block comment */ RETURN n LIMIT 3", + "MATCH (n) WHERE n.url = 'http://x.com//path' RETURN n", # // inside string + "MATCH (n) WHERE n.x = 'semi;colon' RETURN n", # ; inside string + "MATCH (n) RETURN n;", # single trailing semicolon + "UNWIND $kws AS kw CALL { WITH kw MATCH (e:__Entity__) WHERE e.name = kw " + "RETURN e LIMIT 1 } RETURN e", # CALL subquery + "OPTIONAL MATCH (n:Chunk) RETURN count(n)", + "MATCH (n) RETURN n.`create`", # backtick identifier +] + +REJECTED = [ # (query, offending_token) — token = FIRST check that fires: + # start-keyword check runs before the write scan, and the write scan + # checks tokens in tuple order (DELETE before DETACH, SET before FOREACH). + ("CREATE (n:Person {name:'X'})", "CREATE"), + ("MATCH (n) SET n.x = 1 RETURN n", "SET"), + ("MATCH (n) DETACH DELETE n", "DELETE"), + ("MATCH (n) DELETE n", "DELETE"), + ("MERGE (n:Person {name:'X'}) RETURN n", "MERGE"), + ("merge(n) return n", "MERGE"), + ("MATCH (n) REMOVE n.x RETURN n", "REMOVE"), + ("DROP INDEX ON :Person(name)", "DROP"), + ("LOAD CSV FROM 'file:///x' AS row RETURN row", "LOAD"), # start-keyword check + ("load csv from 'x' as r return r", "LOAD"), # start-keyword check + ("MATCH (n) WITH n LOAD CSV FROM 'x' AS r RETURN r", "LOAD CSV"), # embedded LOAD CSV + ("MATCH (n) FOREACH (x IN [1] | SET n.y = x)", "SET"), + ("Cr/**/eate (n)", "CREATE"), # comment-split keyword reassembles + ("ＣＲＥＡＴＥ (n)", "CREATE"), # fullwidth unicode + ("CALL db.idx.fulltext.createNodeIndex('L','p')", "CALL db.idx.fulltext.createnodeindex"), + ("CALL apoc.load.json('u')", "CALL apoc.load.json"), + ("MATCH (n) RETURN n; MATCH (m) RETURN m", ";"), # multi-statement + ("PROFILE MATCH (n) CREATE (m) RETURN n", "PROFILE"), # start-keyword check fires first + ("EXPLAIN MATCH (n) RETURN n", "EXPLAIN"), # disallowed start keyword + ("", "empty"), + (" ", "empty"), +] + + +@pytest.mark.parametrize("query", BENIGN) +def test_benign_queries_pass(query): + ensure_read_only(query) # must not raise + + +@pytest.mark.parametrize("query,token", REJECTED) +def test_write_queries_rejected_with_token(query, token): + with pytest.raises(ReadOnlyViolation) as exc_info: + ensure_read_only(query) + if token != "empty": + assert exc_info.value.offending_token is not None + assert token.lower() in exc_info.value.offending_token.lower() + assert str(exc_info.value) # actionable message + + +def test_apply_limit_injects_when_absent(): + q, injected = apply_limit("MATCH (n) RETURN n", 100) + assert injected and q.endswith("LIMIT 100") + + +def test_apply_limit_respects_existing(): + q, injected = apply_limit("MATCH (n) RETURN n LIMIT 3", 100) + assert not injected and q == "MATCH (n) RETURN n LIMIT 3" + + +def test_apply_limit_ignores_limit_inside_string(): + q, injected = apply_limit("MATCH (n) WHERE n.x = 'no LIMIT here' RETURN n", 50) + assert injected and q.endswith("LIMIT 50") + + +def test_apply_limit_strips_trailing_semicolon(): + q, injected = apply_limit("MATCH (n) RETURN n;", 10) + assert injected and ";" not in q From e99c3c6406352b9edbea2ef264db38de673d5d69 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 12 Jul 2026 13:33:28 +0300 Subject: [PATCH 04/10] feat(retrieval): expose structured provenance in MultiPath result metadata RetrieverResult.metadata now carries a 'provenance' dict with entity ids/names/descriptions, kept chunk ids/texts/document paths (in rerank order, captured before [Source:] tagging), fact strings, and relationship strings. Additive only: assemble_raw_result inputs are untouched and a byte-stability test pins the exact section rendering for a fixed input both before and after this change. Consumed by the graphrag_sdk.tools agent toolkit for citations. Approved deviation from the Task-01 out-of-scope fence (user ruling 2026-07-12). Co-Authored-By: Claude Fable 5 --- .../retrieval/strategies/multi_path.py | 34 +++++++- .../tests/test_multi_path_retrieval.py | 82 +++++++++++++++++++ 2 files changed, 114 insertions(+), 2 deletions(-) diff --git a/graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py b/graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py index fe8e70da..53cb8704 100644 --- a/graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py +++ b/graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py @@ -319,6 +319,34 @@ async def _execute( ctx=ctx, ) + # Structured provenance for graphrag_sdk.tools — additive only; does + # not alter section content. Captured before [Source:] tagging so + # chunk texts still match candidate_chunks values exactly. + text_to_cid: dict[str, str] = {} + for cid, text in candidate_chunks.items(): + text_to_cid.setdefault(text, cid) + kept_cids = [text_to_cid[p] for p in source_passages if p in text_to_cid] + provenance: dict[str, Any] = { + "entities": [ + { + "id": eid, + "name": einfo.get("name", ""), + "description": einfo.get("description", ""), + } + for eid, einfo in entity_list + ], + "chunks": [ + { + "id": cid, + "text": candidate_chunks[cid], + "document_path": chunk_doc_map.get(cid, ""), + } + for cid in kept_cids + ], + "facts": list(fact_strings), + "relationships": list(relationship_strings), + } + # Tag with source docs text_to_doc: dict[str, str] = { candidate_chunks[cid]: doc_name @@ -333,7 +361,7 @@ async def _execute( # 9. Detect question type + assemble ctx.ensure_budget("MultiPath result assembly") q_type_hint = detect_question_type(query) - return assemble_raw_result( + raw = assemble_raw_result( entity_list, relationship_strings, fact_strings, @@ -341,6 +369,8 @@ async def _execute( q_type_hint, cypher_results=cypher_facts if cypher_facts else None, ) + raw.metadata["provenance"] = provenance + return raw def _format(self, raw: RawSearchResult) -> RetrieverResult: """Produce RetrieverResultItems as markdown sections.""" @@ -394,7 +424,7 @@ async def _extract_keywords( async def _search_relates_edges( self, query_vector: list[float] - ) -> tuple[list[tuple[str, float]], dict[str, dict]]: + ) -> tuple[list[tuple[str, float]], dict[str, dict[str, str]]]: """Backward-compat wrapper — delegates to module function.""" return await search_relates_edges( self._vector, diff --git a/graphrag_sdk/tests/test_multi_path_retrieval.py b/graphrag_sdk/tests/test_multi_path_retrieval.py index 5eeabcfd..cf04eb73 100644 --- a/graphrag_sdk/tests/test_multi_path_retrieval.py +++ b/graphrag_sdk/tests/test_multi_path_retrieval.py @@ -839,3 +839,85 @@ async def capture_query(cypher, params=None): sibling_queries = [q for q in queries_seen if "hub" in q.lower() and "sibling" in q.lower()] assert len(sibling_queries) == 0 + + +# -- Provenance metadata (graphrag_sdk.tools contract) -- + + +def _rows(rows): + result = MagicMock() + result.result_set = rows + return result + + +class TestProvenanceMetadata: + """metadata['provenance'] carries typed ids; section items stay byte-stable.""" + + @pytest.fixture + def prov_strategy(self, mp_vector_store, mp_embedder): + graph_store = MagicMock() + + async def route_query(cypher, params=None): + if "toLower(e.name) = toLower(kw)" in cypher: # exact entity match + return _rows([["e1", "Alice", "Engineer"]]) + if "CONTAINS toLower(kw)" in cypher: # CONTAINS discovery path + return _rows([]) + if "MENTIONED_IN" in cypher and "vec.cosineDistance" in cypher: + return _rows([["e1", "c1", "Alice works at Acme."]]) + if "r.rel_type" in cypher and "RELATES" in cypher: # 1-hop relationships + return _rows([["Alice", "WORKS_AT", "Acme Corp", "employment"]]) + if "PART_OF" in cypher: # chunk -> document path + return _rows([["c1", "docs/a.md"]]) + return _rows([]) + + graph_store.query_raw = AsyncMock(side_effect=route_query) + mp_vector_store.search_chunks = AsyncMock( + return_value=[{"id": "c1", "text": "Alice works at Acme."}] + ) + return MultiPathRetrieval( + graph_store=graph_store, + vector_store=mp_vector_store, + embedder=mp_embedder, + llm=MockLLM(responses=["Alice"]), + ) + + async def test_provenance_present_and_typed(self, prov_strategy): + result = await prov_strategy.search("Who is Alice?") + prov = result.metadata["provenance"] + assert {"id": "e1", "name": "Alice", "description": "Engineer"} in prov["entities"] + assert { + "id": "c1", + "text": "Alice works at Acme.", + "document_path": "docs/a.md", + } in prov["chunks"] + assert isinstance(prov["facts"], list) + assert prov["relationships"] == ["Alice —[WORKS_AT]→ Acme Corp: employment"] + + async def test_sections_byte_stable(self, prov_strategy): + """Pins the exact pre-patch section rendering for this fixed input. + + These assertions must pass both BEFORE and AFTER the provenance + patch — they are the guarantee that the patch is non-behavioral. + """ + result = await prov_strategy.search("Who is Alice?") + by_section = {i.metadata["section"]: i.content for i in result.items} + assert by_section["entities"] == "## Key Entities\n- Alice: Engineer" + assert by_section["relationships"] == ( + "## Entity Relationships\n- Alice —[WORKS_AT]→ Acme Corp: employment" + ) + assert by_section["passages"] == ( + "## Source Document Passages\n[Source: docs/a.md]\nAlice works at Acme." + ) + assert result.metadata["strategy"] == "multi_path" + + async def test_empty_graph_provenance(self, mp_graph_store, mp_vector_store, mp_embedder): + s = MultiPathRetrieval( + graph_store=mp_graph_store, + vector_store=mp_vector_store, + embedder=mp_embedder, + llm=MockLLM(responses=["Nobody"]), + ) + result = await s.search("Anything at all") + prov = result.metadata["provenance"] + assert prov["entities"] == [] and prov["chunks"] == [] + assert prov["facts"] == [] and prov["relationships"] == [] From 714253e609a7ca2781f42b65cd5c7de52bc66435 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 12 Jul 2026 13:35:31 +0300 Subject: [PATCH 05/10] feat(tools): tool_specs registry with JSON-Schema input models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Seven pydantic input models (extra='forbid' -> additionalProperties: false) and a static registry binding tool names to toolkit methods with LLM-ready descriptions and output hints. build_tool_specs() filters by read_only / finalize_policy / include. Adapters and MCP servers generate their tool definitions from this — the single source of truth. Adds jsonschema (dev) for Draft 2020-12 meta-validation in tests. Co-Authored-By: Claude Fable 5 --- graphrag_sdk/pyproject.toml | 1 + .../src/graphrag_sdk/tools/__init__.py | 2 + graphrag_sdk/src/graphrag_sdk/tools/specs.py | 217 ++++++++++++++++++ graphrag_sdk/tests/test_tools_specs.py | 77 +++++++ 4 files changed, 297 insertions(+) create mode 100644 graphrag_sdk/src/graphrag_sdk/tools/specs.py create mode 100644 graphrag_sdk/tests/test_tools_specs.py diff --git a/graphrag_sdk/pyproject.toml b/graphrag_sdk/pyproject.toml index df918370..2389862f 100644 --- a/graphrag_sdk/pyproject.toml +++ b/graphrag_sdk/pyproject.toml @@ -70,6 +70,7 @@ dev = [ "pytest-cov>=5.0", "ruff>=0.4", "mypy>=1.10", + "jsonschema>=4.0", ] [project.urls] diff --git a/graphrag_sdk/src/graphrag_sdk/tools/__init__.py b/graphrag_sdk/src/graphrag_sdk/tools/__init__.py index d0d40018..b94b2b01 100644 --- a/graphrag_sdk/src/graphrag_sdk/tools/__init__.py +++ b/graphrag_sdk/src/graphrag_sdk/tools/__init__.py @@ -22,6 +22,7 @@ SearchResult, ToolResult, ) +from graphrag_sdk.tools.specs import ToolSpec __all__ = [ "AnswerResult", @@ -39,4 +40,5 @@ "SchemaResult", "SearchResult", "ToolResult", + "ToolSpec", ] diff --git a/graphrag_sdk/src/graphrag_sdk/tools/specs.py b/graphrag_sdk/src/graphrag_sdk/tools/specs.py new file mode 100644 index 00000000..50d9cc83 --- /dev/null +++ b/graphrag_sdk/src/graphrag_sdk/tools/specs.py @@ -0,0 +1,217 @@ +# GraphRAG SDK — Tools: tool specifications (single source of truth) +# Adapters (pydantic-ai, LangGraph, MCP) generate their tool definitions +# from tool_specs() — descriptions and schemas are never duplicated downstream. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + + +class SearchInput(BaseModel): + """Arguments for graph_search.""" + + model_config = ConfigDict(extra="forbid") + query: str = Field( + min_length=1, description="Natural-language query about entities/relationships." + ) + top_k: int = Field( + default=8, ge=1, le=25, description="How many entities and chunks to return." + ) + expand_hops: int = Field( + default=1, ge=1, le=3, description="Relationship expansion depth from found entities." + ) + include_chunks: bool = Field( + default=True, description="Include source text passages (heavier output)." + ) + + +class AnswerInput(BaseModel): + """Arguments for graph_answer.""" + + model_config = ConfigDict(extra="forbid") + question: str = Field(min_length=1, description="The question to answer from the graph.") + top_k: int = Field( + default=8, ge=1, le=25, description="Retrieval breadth used to build the answer context." + ) + + +class SchemaInput(BaseModel): + """Arguments for graph_schema (none).""" + + model_config = ConfigDict(extra="forbid") + + +class CypherReadInput(BaseModel): + """Arguments for cypher_read.""" + + model_config = ConfigDict(extra="forbid") + query: str = Field(min_length=1, description="Read-only Cypher. Write clauses are rejected.") + params: dict[str, Any] | None = Field( + default=None, description="Query parameters — always prefer over inlining values." + ) + limit: int = Field( + default=100, ge=1, le=1000, description="Row cap injected as LIMIT when the query has none." + ) + timeout_ms: int = Field( + default=5000, ge=100, le=60000, description="Server-side query timeout in milliseconds." + ) + + +class EntityInput(BaseModel): + """Arguments for graph_entity.""" + + model_config = ConfigDict(extra="forbid") + name: str = Field(min_length=1, description="Entity name (exact or partial).") + hops: int = Field( + default=1, ge=1, le=3, description="Neighborhood depth to include around the entity." + ) + + +class RememberInput(BaseModel): + """Arguments for graph_remember.""" + + model_config = ConfigDict(extra="forbid") + text: str = Field( + min_length=1, + max_length=200_000, + description="Text to extract into the knowledge graph.", + ) + document_id: str | None = Field( + default=None, description="Stable document id; auto-generated when omitted." + ) + + +class FlushInput(BaseModel): + """Arguments for graph_flush (none).""" + + model_config = ConfigDict(extra="forbid") + + +class ToolSpec(BaseModel): + """Machine-readable tool definition consumed by agent-framework adapters.""" + + model_config = ConfigDict(frozen=True) + name: str + description: str + input_schema: dict[str, Any] + output_hint: str + + +@dataclass(frozen=True) +class ToolDef: + """Internal registry row binding a tool name to a toolkit method.""" + + name: str + method: str + input_model: type[BaseModel] + description: str + output_hint: str + is_write: bool = False + manual_only: bool = False + + +_TOOL_REGISTRY: tuple[ToolDef, ...] = ( + ToolDef( + "graph_search", + "search", + SearchInput, + "Search the knowledge graph for entities, relationships, facts, and source " + "passages relevant to a query. Use this when you will compose the reply " + "yourself (prefer it over graph_answer for multi-step reasoning), and cite " + "sources with the returned document_id/chunk_id values.", + "SearchResult{query, entities[], relations[], facts[], chunks[], documents[]}; " + "call .to_llm_text() for prompt-ready text.", + ), + ToolDef( + "graph_answer", + "answer", + AnswerInput, + "Ask the knowledge graph a natural-language question and get a fully generated " + "answer with citations. Use for one-shot Q&A when you do not need to inspect " + "the raw context yourself.", + "AnswerResult{answer, citations[], entities_touched[], cypher_used}; " + "call .to_llm_text() for prompt-ready text.", + ), + ToolDef( + "graph_schema", + "schema", + SchemaInput, + "List the graph's entity labels, relationship types, directional patterns, and " + "live counts. Call once before other graph tools to learn what the graph " + "contains and plan queries.", + "SchemaResult{entities[], relations[], node_count, edge_count}; " + "call .to_llm_text() for prompt-ready text.", + ), + ToolDef( + "graph_entity", + "entity", + EntityInput, + "Look up one entity by name and get its properties, relationships up to `hops` " + "away, and source documents. Use when the user asks about a specific named " + "person, organization, or thing.", + "EntityResult{query, found, entity, neighbors[], nearby[], documents[]}; " + "call .to_llm_text() for prompt-ready text.", + ), + ToolDef( + "cypher_read", + "cypher_read", + CypherReadInput, + "Run a read-only Cypher query against the knowledge graph. Use ONLY for " + "aggregations or precise filters graph_search cannot express (counts, sorting, " + "property predicates). Write clauses are rejected; a LIMIT is added if missing.", + "CypherResult{columns[], rows[], row_count, truncated}; " + "call .to_llm_text() for prompt-ready text.", + ), + ToolDef( + "graph_remember", + "remember", + RememberInput, + "Store new text (a fact, note, or document) into the knowledge graph so future " + "searches can find it. Use when the user tells you something worth remembering.", + "RememberResult{document_id, chunks_indexed, nodes_created, " + "relationships_created, finalized}; call .to_llm_text() for prompt-ready text.", + is_write=True, + ), + ToolDef( + "graph_flush", + "flush", + FlushInput, + "Run graph finalization (entity dedup, embeddings, indexes) after one or more " + "graph_remember calls. Expensive — O(graph size); call once at the end of a " + "write session, never after every write.", + "Returns null.", + is_write=True, + manual_only=True, + ), +) + +TOOL_NAMES: tuple[str, ...] = tuple(td.name for td in _TOOL_REGISTRY) + + +def build_tool_specs( + *, + read_only: bool, + finalize_policy: str, + include: frozenset[str] | None, +) -> list[ToolSpec]: + """Build the advertised ToolSpec list for a toolkit configuration.""" + specs: list[ToolSpec] = [] + for td in _TOOL_REGISTRY: + if read_only and td.is_write: + continue + if td.manual_only and finalize_policy != "manual": + continue + if include is not None and td.name not in include: + continue + specs.append( + ToolSpec( + name=td.name, + description=td.description, + input_schema=td.input_model.model_json_schema(), + output_hint=td.output_hint, + ) + ) + return specs diff --git a/graphrag_sdk/tests/test_tools_specs.py b/graphrag_sdk/tests/test_tools_specs.py new file mode 100644 index 00000000..c7d2990e --- /dev/null +++ b/graphrag_sdk/tests/test_tools_specs.py @@ -0,0 +1,77 @@ +"""tool_specs contract: JSON-Schema validity, filtering, stable order.""" + +from __future__ import annotations + +import json + +import jsonschema +import pytest + +from graphrag_sdk.tools import ToolSpec +from graphrag_sdk.tools.specs import _TOOL_REGISTRY, TOOL_NAMES, build_tool_specs + +ALL_NAMES = [ + "graph_search", + "graph_answer", + "graph_schema", + "graph_entity", + "cypher_read", + "graph_remember", + "graph_flush", +] + + +def test_registry_names_and_order(): + assert list(TOOL_NAMES) == ALL_NAMES + + +def test_default_build_has_all_tools(): + specs = build_tool_specs(read_only=False, finalize_policy="manual", include=None) + assert [s.name for s in specs] == ALL_NAMES + + +def test_read_only_removes_write_tools(): + specs = build_tool_specs(read_only=True, finalize_policy="manual", include=None) + names = [s.name for s in specs] + assert "graph_remember" not in names and "graph_flush" not in names + assert "graph_search" in names + + +@pytest.mark.parametrize("policy", ["on_write", "never"]) +def test_non_manual_policy_hides_flush(policy): + specs = build_tool_specs(read_only=False, finalize_policy=policy, include=None) + names = [s.name for s in specs] + assert "graph_flush" not in names and "graph_remember" in names + + +def test_include_filters_and_preserves_order(): + specs = build_tool_specs( + read_only=False, + finalize_policy="manual", + include=frozenset({"graph_schema", "graph_search"}), + ) + assert [s.name for s in specs] == ["graph_search", "graph_schema"] + + +def test_schemas_are_valid_draft202012_and_strict(): + for spec in build_tool_specs(read_only=False, finalize_policy="manual", include=None): + jsonschema.Draft202012Validator.check_schema(spec.input_schema) + assert spec.input_schema.get("additionalProperties") is False, spec.name + + +def test_specs_round_trip_through_json(): + specs = build_tool_specs(read_only=False, finalize_policy="manual", include=None) + dumped = json.dumps([s.model_dump() for s in specs]) + assert [ToolSpec.model_validate(d) for d in json.loads(dumped)] == specs + + +def test_descriptions_are_llm_ready(): + for td in _TOOL_REGISTRY: + assert len(td.description) >= 40, td.name # says when to use it + assert td.output_hint and len(td.output_hint) <= 200, td.name + + +def test_field_descriptions_present(): + for td in _TOOL_REGISTRY: + for fname, field in td.input_model.model_fields.items(): + assert field.description, f"{td.name}.{fname} missing description" From e1562403b450b1e88e2dac042df209ae045b7f42 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 12 Jul 2026 13:38:52 +0300 Subject: [PATCH 06/10] feat(tools): GraphRAGToolkit facade with policies, dispatch, tenancy The seven async operations (search/answer/schema/entity/cypher_read/ remember/flush) over a wrapped GraphRAG. search()/answer() drive a per-call tuned MultiPathRetrieval through the real retrieve()/ completion() paths and map provenance metadata to typed results with document/chunk-id citations. cypher_read guards via ensure_read_only + LIMIT injection and enforces timeout_ms on the connection. call(name, arguments) validates through the registry input models for adapters/ MCP. for_tenant() derives an owned tenant-scoped graph (validated tenant_id); aclose()/async-with manage its lifecycle. Root package re-exports GraphRAGToolkit/ToolSpec/ReadOnlyViolation. graph_ops centralizes toolkit-owned parameterized Cypher; entity/chunk bulky props (embedding, source_chunk_ids) never reach agent output. Co-Authored-By: Claude Fable 5 --- graphrag_sdk/src/graphrag_sdk/__init__.py | 8 + .../src/graphrag_sdk/tools/__init__.py | 3 + .../src/graphrag_sdk/tools/graph_ops.py | 192 ++++++++ .../src/graphrag_sdk/tools/toolkit.py | 440 ++++++++++++++++++ graphrag_sdk/tests/test_tools_toolkit.py | 148 ++++++ 5 files changed, 791 insertions(+) create mode 100644 graphrag_sdk/src/graphrag_sdk/tools/graph_ops.py create mode 100644 graphrag_sdk/src/graphrag_sdk/tools/toolkit.py create mode 100644 graphrag_sdk/tests/test_tools_toolkit.py diff --git a/graphrag_sdk/src/graphrag_sdk/__init__.py b/graphrag_sdk/src/graphrag_sdk/__init__.py index b703b1da..dadfb5be 100644 --- a/graphrag_sdk/src/graphrag_sdk/__init__.py +++ b/graphrag_sdk/src/graphrag_sdk/__init__.py @@ -23,6 +23,7 @@ DocumentNotFoundError, GraphRAGError, LatencyBudgetExceededError, + ReadOnlyViolation, ) from graphrag_sdk.core.models import ( ApplyChangesResult, @@ -126,6 +127,9 @@ ) from graphrag_sdk.storage.vector_store import VectorStore +# ── Agent Toolkit ─────────────────────────────────────────────── +from graphrag_sdk.tools import GraphRAGToolkit, ToolSpec + __all__ = [ # Version "__version__", @@ -209,6 +213,10 @@ "OntologyModificationNotAllowedError", "OntologyStore", "VectorStore", + # Agent toolkit + "GraphRAGToolkit", + "ReadOnlyViolation", + "ToolSpec", ] diff --git a/graphrag_sdk/src/graphrag_sdk/tools/__init__.py b/graphrag_sdk/src/graphrag_sdk/tools/__init__.py index b94b2b01..f9d6b5d0 100644 --- a/graphrag_sdk/src/graphrag_sdk/tools/__init__.py +++ b/graphrag_sdk/src/graphrag_sdk/tools/__init__.py @@ -23,8 +23,11 @@ ToolResult, ) from graphrag_sdk.tools.specs import ToolSpec +from graphrag_sdk.tools.toolkit import FinalizePolicy, GraphRAGToolkit __all__ = [ + "FinalizePolicy", + "GraphRAGToolkit", "AnswerResult", "ChunkRef", "Citation", diff --git a/graphrag_sdk/src/graphrag_sdk/tools/graph_ops.py b/graphrag_sdk/src/graphrag_sdk/tools/graph_ops.py new file mode 100644 index 00000000..f55dc70d --- /dev/null +++ b/graphrag_sdk/src/graphrag_sdk/tools/graph_ops.py @@ -0,0 +1,192 @@ +# GraphRAG SDK — Tools: toolkit-owned graph queries and result conversion +# Parameterized Cypher only — user values NEVER interpolated into query text. + +from __future__ import annotations + +from typing import Any + +from graphrag_sdk.storage.graph_store import GraphStore +from graphrag_sdk.tools.models import CypherResult, DocumentRef, EntityCard, RelationTriple + +_BULKY_PROPS = frozenset({"embedding", "source_chunk_ids"}) + + +def _card_from_row(row: list[Any]) -> tuple[str, EntityCard]: + """(id, name, description, labels, properties) row -> (id, EntityCard).""" + eid, name, desc, labels, props = (list(row) + [None] * 5)[:5] + label = next((str(lbl) for lbl in (labels or []) if lbl != "__Entity__"), "") + clean_props = { + k: v + for k, v in (props or {}).items() + if k not in _BULKY_PROPS and k not in ("id", "name", "description") + } + return str(eid), EntityCard( + name=str(name or ""), label=label, description=desc or None, properties=clean_props + ) + + +async def enrich_entities(store: GraphStore, entity_ids: list[str]) -> dict[str, EntityCard]: + """Fetch label + user properties for entity ids (bulky props excluded).""" + if not entity_ids: + return {} + result = await store.query_raw( + "UNWIND $ids AS eid MATCH (e:__Entity__ {id: eid}) " + "RETURN e.id, e.name, e.description, labels(e), properties(e)", + {"ids": entity_ids}, + ) + return dict(_card_from_row(row) for row in result.result_set) + + +async def expand_triples( + store: GraphStore, + seed_ids: list[str], + *, + hops: int, + cap: int = 60, + per_hop_limit: int = 25, +) -> list[RelationTriple]: + """Frontier expansion over RELATES edges, hop by hop. + + Never builds a variable-length pattern from user input. Direction comes + from the edge's own src_name/tgt_name properties, so one undirected + query per hop suffices. + """ + triples: list[RelationTriple] = [] + seen: set[tuple[str, str, str]] = set() + visited = set(seed_ids) + frontier = list(seed_ids) + for _ in range(max(1, hops)): + if not frontier or len(triples) >= cap: + break + result = await store.query_raw( + "MATCH (a:__Entity__)-[r:RELATES]-(b:__Entity__) " + "WHERE a.id IN $ids " + "RETURN r.src_name, r.rel_type, r.tgt_name, " + "COALESCE(r.fact, r.description, ''), b.id " + "ORDER BY r.src_name, r.rel_type, r.tgt_name " + "LIMIT $limit", + {"ids": frontier, "limit": per_hop_limit}, + ) + next_frontier: list[str] = [] + for row in result.result_set: + src, rel_type, tgt = row[0] or "", row[1] or "", row[2] or "" + fact = row[3] if len(row) > 3 else "" + other = row[4] if len(row) > 4 else "" + key = (src.lower(), rel_type, tgt.lower()) + if src and rel_type and tgt and key not in seen: + seen.add(key) + triples.append( + RelationTriple(source=src, type=rel_type, target=tgt, fact=fact or None) + ) + if other and other not in visited: + visited.add(other) + next_frontier.append(other) + frontier = next_frontier + return triples[:cap] + + +async def chunk_documents(store: GraphStore, chunk_ids: list[str]) -> dict[str, tuple[str, str]]: + """chunk_id -> (document_id, document_path) via PART_OF.""" + if not chunk_ids: + return {} + result = await store.query_raw( + "UNWIND $ids AS cid MATCH (d:Document)-[:PART_OF]->(c:Chunk {id: cid}) " + "RETURN c.id, d.id, d.path", + {"ids": chunk_ids}, + ) + return {row[0]: (str(row[1] or ""), str(row[2] or "")) for row in result.result_set if row[0]} + + +async def find_entity_matches( + store: GraphStore, name: str, limit: int = 5 +) -> list[tuple[str, EntityCard]]: + """Ranked (exact > case-insensitive > substring; shorter first) name matches.""" + result = await store.query_raw( + "MATCH (e:__Entity__) WHERE toLower(e.name) CONTAINS toLower($name) " + "RETURN e.id, e.name, e.description, labels(e), properties(e), " + "CASE WHEN e.name = $name THEN 0 " + "WHEN toLower(e.name) = toLower($name) THEN 1 ELSE 2 END AS rank " + "ORDER BY rank, size(e.name), e.name LIMIT $limit", + {"name": name, "limit": limit}, + ) + return [_card_from_row(list(row)[:5]) for row in result.result_set] + + +async def entity_documents(store: GraphStore, entity_id: str) -> list[DocumentRef]: + """Distinct source documents mentioning the entity.""" + result = await store.query_raw( + "MATCH (e:__Entity__ {id: $eid})-[:MENTIONED_IN]->(:Chunk)" + "<-[:PART_OF]-(d:Document) " + "RETURN DISTINCT d.id, d.path ORDER BY d.id LIMIT 10", + {"eid": entity_id}, + ) + return [ + DocumentRef(document_id=str(row[0] or ""), document_path=str(row[1] or "")) + for row in result.result_set + if row[0] + ] + + +async def schema_counts(store: GraphStore) -> tuple[dict[str, int], dict[str, int]]: + """(entity-label -> count, RELATES rel_type -> count) from the live graph.""" + labels = await store.query_raw( + "MATCH (n:__Entity__) UNWIND labels(n) AS l " + "WITH l, count(*) AS c WHERE l <> '__Entity__' " + "RETURN l, c ORDER BY l" + ) + rels = await store.query_raw( + "MATCH (:__Entity__)-[r:RELATES]->(:__Entity__) " + "WITH r.rel_type AS t, count(*) AS c WHERE t IS NOT NULL " + "RETURN t, c ORDER BY t" + ) + return ( + {row[0]: int(row[1]) for row in labels.result_set}, + {row[0]: int(row[1]) for row in rels.result_set}, + ) + + +def _to_jsonable(value: Any) -> Any: + """Convert falkordb values to JSON-safe data; strips bulky/binary props. + + falkordb attribute contract (verified against falkordb-py 1.x): + ``Node{labels, properties}``, ``Edge{relation, src_node, dest_node, + properties}``. Duck-typed so plain scalars/containers pass through. + """ + if value is None or isinstance(value, (bool, int, float)): + return value + if isinstance(value, str): + return GraphStore._sanitize_string(value) + if isinstance(value, (list, tuple)): + return [_to_jsonable(v) for v in value] + if isinstance(value, dict): + return {str(k): _to_jsonable(v) for k, v in value.items()} + props = getattr(value, "properties", None) + if isinstance(props, dict): + clean = {k: _to_jsonable(v) for k, v in props.items() if k not in _BULKY_PROPS} + labels = getattr(value, "labels", None) + if labels is not None: + return {"labels": [str(lbl) for lbl in labels], "properties": clean} + relation = getattr(value, "relation", None) + if relation is not None: + return { + "type": str(relation), + "src": getattr(value, "src_node", None), + "dst": getattr(value, "dest_node", None), + "properties": clean, + } + return GraphStore._sanitize_string(str(value)) + + +def convert_query_result(result: Any, *, limit: int, limit_injected: bool) -> CypherResult: + """FalkorDB QueryResult -> CypherResult (columns from header pairs).""" + header = getattr(result, "header", None) or [] + columns = [ + str(h[1]) if isinstance(h, (list, tuple)) and len(h) >= 2 else str(h) for h in header + ] + rows = [[_to_jsonable(v) for v in row] for row in (result.result_set or [])] + return CypherResult( + columns=columns, + rows=rows, + row_count=len(rows), + truncated=bool(limit_injected and len(rows) == limit), + ) diff --git a/graphrag_sdk/src/graphrag_sdk/tools/toolkit.py b/graphrag_sdk/src/graphrag_sdk/tools/toolkit.py new file mode 100644 index 00000000..f06ff4a1 --- /dev/null +++ b/graphrag_sdk/src/graphrag_sdk/tools/toolkit.py @@ -0,0 +1,440 @@ +# GraphRAG SDK — Tools: GraphRAGToolkit facade +# Framework-neutral agent surface over GraphRAG. Retrieval rides the real +# MultiPathRetrieval pipeline (with provenance metadata); QA rides completion(). + +from __future__ import annotations + +import dataclasses +import re +from collections.abc import Sequence +from types import TracebackType +from typing import TYPE_CHECKING, Any, Literal + +from graphrag_sdk.core.connection import ConnectionConfig +from graphrag_sdk.core.context import Context +from graphrag_sdk.core.exceptions import ConfigError, ReadOnlyViolation +from graphrag_sdk.core.models import Ontology +from graphrag_sdk.core.providers import Embedder, LLMInterface +from graphrag_sdk.retrieval.strategies.multi_path import MultiPathRetrieval +from graphrag_sdk.tools import graph_ops +from graphrag_sdk.tools.cypher_guard import apply_limit, ensure_read_only +from graphrag_sdk.tools.models import ( + AnswerResult, + ChunkRef, + Citation, + CypherResult, + DocumentRef, + EntityCard, + EntityResult, + EntityTypeInfo, + RelationTypeInfo, + RememberResult, + SchemaResult, + SearchResult, +) +from graphrag_sdk.tools.specs import ( + _TOOL_REGISTRY, + AnswerInput, + CypherReadInput, + EntityInput, + RememberInput, + SearchInput, + ToolSpec, + build_tool_specs, +) + +if TYPE_CHECKING: + from graphrag_sdk.api.main import GraphRAG + +FinalizePolicy = Literal["manual", "on_write", "never"] +_POLICIES = ("manual", "on_write", "never") +_TENANT_RE = re.compile(r"^[A-Za-z0-9_-]{1,64}$") +_SNIPPET_CHARS = 200 + + +class GraphRAGToolkit: + """Framework-neutral agent toolkit over an async :class:`GraphRAG` instance. + + Exposes a small, stable set of async operations with LLM-friendly typed + results, plus :meth:`tool_specs` — the machine-readable contract adapters + (pydantic-ai, LangGraph, MCP) generate their tool definitions from. + + Args: + rag: The GraphRAG instance to wrap (binds to its graph/tenant). + finalize_policy: ``"manual"`` (default — call :meth:`flush` yourself), + ``"on_write"`` (finalize after every remember; demos only — + finalize is O(graph size)), or ``"never"`` (you call + ``rag.finalize()`` yourself). + read_only: Disable remember/flush and drop them from tool_specs(). + include: Optional subset of tool names to advertise via + tool_specs()/call(). Direct method calls are not affected. + tenant_id: Stamped into the Context of every operation. + owns_rag: When True, :meth:`aclose` closes the wrapped GraphRAG. + """ + + def __init__( + self, + rag: GraphRAG, + *, + finalize_policy: FinalizePolicy = "manual", + read_only: bool = False, + include: Sequence[str] | None = None, + tenant_id: str = "default", + owns_rag: bool = False, + ) -> None: + if finalize_policy not in _POLICIES: + raise ValueError(f"finalize_policy must be one of {_POLICIES}, got {finalize_policy!r}") + valid = {td.name for td in _TOOL_REGISTRY} + if include is not None: + unknown = sorted(set(include) - valid) + if unknown: + raise ValueError(f"Unknown tool names in include={unknown}; valid: {sorted(valid)}") + self._rag = rag + self._finalize_policy: FinalizePolicy = finalize_policy + self._read_only = read_only + self._include = frozenset(include) if include is not None else None + self._tenant_id = tenant_id + self._owns_rag = owns_rag + + # ── Lifecycle ──────────────────────────────────────────────── + + @property + def rag(self) -> GraphRAG: + """The wrapped GraphRAG instance.""" + return self._rag + + @classmethod + def for_tenant( + cls, + base_config: ConnectionConfig, + tenant_id: str, + *, + llm: LLMInterface, + embedder: Embedder, + ontology: Ontology | None = None, + embedding_dimension: int = 256, + **toolkit_kwargs: Any, + ) -> GraphRAGToolkit: + """Build a toolkit bound to a tenant-scoped graph. + + Derives ``graph_name = f"{base_config.graph_name}__{tenant_id}"`` and + constructs a dedicated GraphRAG the toolkit owns (closed by + :meth:`aclose` / ``async with``). + """ + if not _TENANT_RE.match(tenant_id): + raise ValueError( + "tenant_id must match ^[A-Za-z0-9_-]{1,64}$ (it becomes part " + f"of the graph name); got {tenant_id!r}" + ) + from graphrag_sdk.api.main import GraphRAG # local import: avoid cycle + + config = dataclasses.replace( + base_config, graph_name=f"{base_config.graph_name}__{tenant_id}" + ) + rag = GraphRAG( + connection=config, + llm=llm, + embedder=embedder, + ontology=ontology, + embedding_dimension=embedding_dimension, + ) + return cls(rag, tenant_id=tenant_id, owns_rag=True, **toolkit_kwargs) + + async def aclose(self) -> None: + """Close the wrapped GraphRAG if this toolkit owns it.""" + if self._owns_rag: + await self._rag.close() + + async def __aenter__(self) -> GraphRAGToolkit: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + await self.aclose() + + # ── Specs & dispatch ───────────────────────────────────────── + + def tool_specs(self) -> list[ToolSpec]: + """Machine-readable tool definitions for this toolkit configuration.""" + return build_tool_specs( + read_only=self._read_only, + finalize_policy=self._finalize_policy, + include=self._include, + ) + + async def call(self, name: str, arguments: dict[str, Any] | None = None) -> Any: + """Validate *arguments* against the tool's input model and invoke it. + + The generic entry point for adapters/MCP servers. Enforces the same + read_only/include gates as :meth:`tool_specs`. + """ + by_name = {td.name: td for td in _TOOL_REGISTRY} + td = by_name.get(name) + if td is None: + raise ValueError(f"Unknown tool {name!r}; valid: {sorted(by_name)}") + if self._include is not None and name not in self._include: + enabled = [s.name for s in self.tool_specs()] + raise ValueError(f"Tool {name!r} is not enabled; enabled: {enabled}") + if self._read_only and td.is_write: + raise ReadOnlyViolation(f"{name} is disabled: toolkit is read-only.") + if td.manual_only and self._finalize_policy != "manual": + raise ValueError( + f"Tool {name!r} is unavailable under finalize_policy={self._finalize_policy!r}." + ) + model = td.input_model(**(arguments or {})) + method = getattr(self, td.method) + return await method(**model.model_dump()) + + def _new_ctx(self) -> Context: + return Context(tenant_id=self._tenant_id) + + def _make_strategy(self, top_k: int) -> MultiPathRetrieval: + """A per-call MultiPathRetrieval tuned to top_k (ctor is pure assignment).""" + return MultiPathRetrieval( + graph_store=self._rag._graph_store, + vector_store=self._rag._vector_store, + embedder=self._rag.embedder, + llm=self._rag.llm, + chunk_top_k=top_k, + rel_top_k=top_k, + max_entities=max(2 * top_k, 10), + ontology=self._rag._global_ontology, + ) + + # ── Read path ──────────────────────────────────────────────── + + async def search( + self, + query: str, + *, + top_k: int = 8, + expand_hops: int = 1, + include_chunks: bool = True, + ctx: Context | None = None, + ) -> SearchResult: + """Typed retrieval only — no generation. The default agent mode: + returns ranked entities/relations/facts/chunks the host LLM composes + from, with document/chunk ids for citations.""" + inp = SearchInput( + query=query, top_k=top_k, expand_hops=expand_hops, include_chunks=include_chunks + ) + ctx = ctx or self._new_ctx() + store = self._rag._graph_store + rr = await self._rag.retrieve(inp.query, strategy=self._make_strategy(inp.top_k), ctx=ctx) + prov: dict[str, Any] = rr.metadata.get("provenance") or {} + seeds = [e for e in (prov.get("entities") or []) if e.get("id")][: inp.top_k] + seed_ids = [e["id"] for e in seeds] + cards_by_id = await graph_ops.enrich_entities(store, seed_ids) + entities = [ + cards_by_id.get( + e["id"], + EntityCard(name=e.get("name", ""), description=e.get("description") or None), + ) + for e in seeds + ] + relations = await graph_ops.expand_triples( + store, seed_ids, hops=inp.expand_hops, cap=min(4 * inp.top_k, 40) + ) + chunks: list[ChunkRef] = [] + documents: list[DocumentRef] = [] + if inp.include_chunks: + prov_chunks = (prov.get("chunks") or [])[: inp.top_k] + doc_map = await graph_ops.chunk_documents(store, [c["id"] for c in prov_chunks]) + seen_docs: set[str] = set() + for c in prov_chunks: + doc_id, doc_path = doc_map.get(c["id"], ("", c.get("document_path", ""))) + chunks.append( + ChunkRef( + chunk_id=c["id"], + document_id=doc_id, + document_path=doc_path, + text=c.get("text", ""), + ) + ) + if doc_id and doc_id not in seen_docs: + seen_docs.add(doc_id) + documents.append(DocumentRef(document_id=doc_id, document_path=doc_path)) + return SearchResult( + query=inp.query, + entities=entities, + relations=relations, + facts=(prov.get("facts") or [])[: 2 * inp.top_k], + chunks=chunks, + documents=documents, + ) + + async def answer( + self, question: str, *, top_k: int = 8, ctx: Context | None = None + ) -> AnswerResult: + """Full RAG: the real completion() pipeline plus provenance citations. + + ``cypher_used`` is currently always None (the experimental + text-to-Cypher path does not surface its query); the field exists + for forward compatibility. + """ + inp = AnswerInput(question=question, top_k=top_k) + ctx = ctx or self._new_ctx() + rag_result = await self._rag.completion( + inp.question, + strategy=self._make_strategy(inp.top_k), + return_context=True, + ctx=ctx, + ) + prov: dict[str, Any] = {} + if rag_result.retriever_result is not None: + prov = rag_result.retriever_result.metadata.get("provenance") or {} + prov_chunks = (prov.get("chunks") or [])[: inp.top_k] + doc_map = await graph_ops.chunk_documents( + self._rag._graph_store, [c["id"] for c in prov_chunks] + ) + citations = [] + for c in prov_chunks: + doc_id, doc_path = doc_map.get(c["id"], ("", c.get("document_path", ""))) + text = c.get("text", "") + snippet = text if len(text) <= _SNIPPET_CHARS else text[: _SNIPPET_CHARS - 1] + "…" + citations.append( + Citation( + document_id=doc_id, + document_path=doc_path, + chunk_id=c["id"], + snippet=snippet, + ) + ) + seen: set[str] = set() + touched: list[str] = [] + for e in prov.get("entities") or []: + n = e.get("name", "") + if n and n.lower() not in seen: + seen.add(n.lower()) + touched.append(n) + return AnswerResult( + answer=rag_result.answer, + citations=citations, + entities_touched=touched, + cypher_used=None, + ) + + async def schema(self, *, ctx: Context | None = None) -> SchemaResult: + """Entity labels + relation types with declared metadata and live counts.""" + ontology = await self._rag.get_ontology() + label_counts, rel_counts = await graph_ops.schema_counts(self._rag._graph_store) + stats = await self._rag.get_statistics() + declared_e = {e.label: e for e in ontology.entities} + declared_r = {r.label: r for r in ontology.relations} + entity_infos = [ + EntityTypeInfo( + label=label, + description=declared_e[label].description if label in declared_e else None, + count=count, + properties=[a.name for a in declared_e[label].properties] + if label in declared_e + else [], + ) + for label, count in sorted({**{e: 0 for e in declared_e}, **label_counts}.items()) + ] + relation_infos = [ + RelationTypeInfo( + label=label, + description=declared_r[label].description if label in declared_r else None, + patterns=list(declared_r[label].patterns) if label in declared_r else [], + count=count, + ) + for label, count in sorted({**{r: 0 for r in declared_r}, **rel_counts}.items()) + ] + return SchemaResult( + entities=entity_infos, + relations=relation_infos, + node_count=int(stats.get("node_count", 0)), + edge_count=int(stats.get("edge_count", 0)), + ) + + async def cypher_read( + self, + query: str, + params: dict[str, Any] | None = None, + *, + limit: int = 100, + timeout_ms: int = 5000, + ctx: Context | None = None, + ) -> CypherResult: + """Guarded read-only Cypher. Rejects writes (:class:`ReadOnlyViolation`), + injects LIMIT when absent, enforces a server-side timeout. + + The connection retries transient failures up to ``retry_count`` times, + so worst-case wall time is about ``retry_count * timeout_ms``. + """ + inp = CypherReadInput(query=query, params=params, limit=limit, timeout_ms=timeout_ms) + ensure_read_only(inp.query) + final_query, injected = apply_limit(inp.query, inp.limit) + result = await self._rag._conn.query(final_query, params=inp.params, timeout=inp.timeout_ms) + return graph_ops.convert_query_result(result, limit=inp.limit, limit_injected=injected) + + async def entity(self, name: str, *, hops: int = 1, ctx: Context | None = None) -> EntityResult: + """Entity card: best name match, neighbors up to *hops*, source documents.""" + inp = EntityInput(name=name, hops=hops) + store = self._rag._graph_store + matches = await graph_ops.find_entity_matches(store, inp.name) + if not matches: + return EntityResult(query=inp.name, found=False) + (eid, card), rest = matches[0], matches[1:] + neighbors = await graph_ops.expand_triples(store, [eid], hops=inp.hops) + documents = await graph_ops.entity_documents(store, eid) + return EntityResult( + query=inp.name, + found=True, + entity=card, + neighbors=neighbors, + nearby=[c.name for _, c in rest], + documents=documents, + ) + + # ── Write path ─────────────────────────────────────────────── + + def _ensure_writable(self, tool: str) -> None: + if self._read_only: + raise ReadOnlyViolation(f"{tool} is disabled: toolkit is read-only.") + + async def remember( + self, text: str, *, document_id: str | None = None, ctx: Context | None = None + ) -> RememberResult: + """Ingest raw text into the graph (agent memory / fact capture). + + Under ``finalize_policy="on_write"`` this also runs finalize — + which is O(graph size); use "manual" + :meth:`flush` in production. + """ + self._ensure_writable("graph_remember") + inp = RememberInput(text=text, document_id=document_id) + ctx = ctx or self._new_ctx() + result = await self._rag.ingest(text=inp.text, document_id=inp.document_id, ctx=ctx) + finalized = False + if self._finalize_policy == "on_write": + await self._rag.finalize() + finalized = True + return RememberResult( + document_id=result.document_info.uid, + chunks_indexed=result.chunks_indexed, + nodes_created=result.nodes_created, + relationships_created=result.relationships_created, + finalized=finalized, + ) + + async def flush(self, *, ctx: Context | None = None) -> None: + """Run ``GraphRAG.finalize()`` under the "manual" policy. + + No-op under "on_write" (each remember already finalized); raises + :class:`ConfigError` under "never" (call ``rag.finalize()`` yourself). + """ + self._ensure_writable("graph_flush") + if self._finalize_policy == "manual": + await self._rag.finalize() + return + if self._finalize_policy == "on_write": + return + raise ConfigError( + 'finalize_policy="never": the toolkit never finalizes — call ' + "GraphRAG.finalize() yourself." + ) diff --git a/graphrag_sdk/tests/test_tools_toolkit.py b/graphrag_sdk/tests/test_tools_toolkit.py new file mode 100644 index 00000000..345a944e --- /dev/null +++ b/graphrag_sdk/tests/test_tools_toolkit.py @@ -0,0 +1,148 @@ +"""GraphRAGToolkit unit tests (stubbed GraphRAG — no server, no LLM).""" + +from __future__ import annotations + +import inspect +import subprocess +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from graphrag_sdk.core.connection import ConnectionConfig +from graphrag_sdk.core.exceptions import ConfigError, ReadOnlyViolation +from graphrag_sdk.core.models import DocumentInfo, FinalizeResult, IngestionResult +from graphrag_sdk.tools import GraphRAGToolkit, RememberResult +from graphrag_sdk.tools.specs import _TOOL_REGISTRY + + +def make_stub_rag() -> MagicMock: + rag = MagicMock() + rag.ingest = AsyncMock( + return_value=IngestionResult( + document_info=DocumentInfo(uid="doc-1"), + nodes_created=2, + relationships_created=1, + chunks_indexed=1, + ) + ) + rag.finalize = AsyncMock(return_value=FinalizeResult()) + rag.close = AsyncMock() + return rag + + +def test_ctor_validates_policy_and_include(): + rag = make_stub_rag() + with pytest.raises(ValueError, match="finalize_policy"): + GraphRAGToolkit(rag, finalize_policy="sometimes") # type: ignore[arg-type] + with pytest.raises(ValueError, match="graph_serch"): + GraphRAGToolkit(rag, include=["graph_serch"]) + + +async def test_remember_manual_policy_defers_finalize(): + rag = make_stub_rag() + tk = GraphRAGToolkit(rag, finalize_policy="manual") + result = await tk.remember("Alice works at Acme.", document_id="doc-1") + assert isinstance(result, RememberResult) + assert result.document_id == "doc-1" and result.finalized is False + rag.finalize.assert_not_awaited() + rag.ingest.assert_awaited_once() + assert rag.ingest.await_args.kwargs["text"] == "Alice works at Acme." + await tk.flush() + rag.finalize.assert_awaited_once() + + +async def test_remember_on_write_policy_finalizes(): + rag = make_stub_rag() + tk = GraphRAGToolkit(rag, finalize_policy="on_write") + result = await tk.remember("x") + assert result.finalized is True + rag.finalize.assert_awaited_once() + await tk.flush() # documented no-op + rag.finalize.assert_awaited_once() + + +async def test_never_policy_flush_raises(): + tk = GraphRAGToolkit(make_stub_rag(), finalize_policy="never") + with pytest.raises(ConfigError): + await tk.flush() + + +async def test_read_only_blocks_writes(): + rag = make_stub_rag() + tk = GraphRAGToolkit(rag, read_only=True) + with pytest.raises(ReadOnlyViolation): + await tk.remember("x") + with pytest.raises(ReadOnlyViolation): + await tk.flush() + with pytest.raises(ReadOnlyViolation): + await tk.call("graph_remember", {"text": "x"}) + rag.ingest.assert_not_awaited() + assert "graph_remember" not in [s.name for s in tk.tool_specs()] + + +async def test_call_dispatch_and_validation(): + rag = make_stub_rag() + tk = GraphRAGToolkit(rag) + result = await tk.call("graph_remember", {"text": "hi", "document_id": "d"}) + assert isinstance(result, RememberResult) + with pytest.raises(ValueError, match="Unknown tool"): + await tk.call("nope", {}) + with pytest.raises(Exception): # pydantic ValidationError on extra arg + await tk.call("graph_remember", {"text": "hi", "bogus": 1}) + + +async def test_call_flush_unavailable_under_non_manual_policy(): + tk = GraphRAGToolkit(make_stub_rag(), finalize_policy="on_write") + with pytest.raises(ValueError, match="unavailable"): + await tk.call("graph_flush", {}) + + +async def test_include_gates_call_and_specs(): + tk = GraphRAGToolkit(make_stub_rag(), include=["graph_schema"]) + assert [s.name for s in tk.tool_specs()] == ["graph_schema"] + with pytest.raises(ValueError, match="not enabled"): + await tk.call("graph_remember", {"text": "x"}) + + +def test_signature_matches_input_model(): + for td in _TOOL_REGISTRY: + method = getattr(GraphRAGToolkit, td.method) + params = { + n: p + for n, p in inspect.signature(method).parameters.items() + if n not in {"self", "ctx"} + } + fields = td.input_model.model_fields + assert set(params) == set(fields), td.name + for pname, param in params.items(): + field = fields[pname] + expected = inspect.Parameter.empty if field.is_required() else field.default + assert param.default == expected, f"{td.name}.{pname} default drift" + + +def test_for_tenant_derives_graph_name_lazily(): + base = ConnectionConfig(graph_name="app") + tk = GraphRAGToolkit.for_tenant(base, "acme", llm=MagicMock(), embedder=MagicMock()) + assert tk.rag._conn.config.graph_name == "app__acme" + assert base.graph_name == "app" # base untouched + with pytest.raises(ValueError, match="tenant_id"): + GraphRAGToolkit.for_tenant(base, "bad tenant!", llm=MagicMock(), embedder=MagicMock()) + + +async def test_aclose_closes_owned_rag_only(): + rag = make_stub_rag() + tk = GraphRAGToolkit(rag) + await tk.aclose() + rag.close.assert_not_awaited() + owned = GraphRAGToolkit(rag, owns_rag=True) + async with owned: + pass + rag.close.assert_awaited_once() + + +def test_module_imports_cleanly(): + proc = subprocess.run( + [sys.executable, "-c", "import graphrag_sdk.tools"], capture_output=True, text=True + ) + assert proc.returncode == 0, proc.stderr From 3a5e68f9d89a798b1f33e33f53cf132d13bc3cb6 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 12 Jul 2026 13:40:52 +0300 Subject: [PATCH 07/10] test(tools): graph ops, schema merge, guarded cypher_read, entity card Covers falkordb Node/Edge conversion with real driver objects (embedding and source_chunk_ids stripped, truncated heuristic), frontier expansion dedupe/caps, guard-before-connection ordering, LIMIT injection with params/timeout forwarding, ontology+live-count merging, and ranked entity matching with nearby suggestions. Co-Authored-By: Claude Fable 5 --- graphrag_sdk/tests/test_tools_graph_ops.py | 276 +++++++++++++++++++++ 1 file changed, 276 insertions(+) create mode 100644 graphrag_sdk/tests/test_tools_graph_ops.py diff --git a/graphrag_sdk/tests/test_tools_graph_ops.py b/graphrag_sdk/tests/test_tools_graph_ops.py new file mode 100644 index 00000000..e7917ad5 --- /dev/null +++ b/graphrag_sdk/tests/test_tools_graph_ops.py @@ -0,0 +1,276 @@ +"""graph_ops behavior + toolkit schema/cypher_read/entity (stubbed stores).""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from falkordb import Edge, Node + +from graphrag_sdk.core.exceptions import ReadOnlyViolation +from graphrag_sdk.core.models import Attribute, Entity, Ontology, Relation +from graphrag_sdk.tools import GraphRAGToolkit +from graphrag_sdk.tools.graph_ops import ( + chunk_documents, + convert_query_result, + enrich_entities, + expand_triples, + find_entity_matches, +) + + +def _res(rows, header=None): + return SimpleNamespace(result_set=rows, header=header or []) + + +# ── convert_query_result ───────────────────────────────────────── + + +def test_convert_query_result_nodes_edges_and_columns(): + node = Node( + node_id=1, + alias="n", + labels=["Person", "__Entity__"], + properties={"name": "Alice", "embedding": [0.1] * 4, "source_chunk_ids": ["c1"]}, + ) + edge = Edge( + src_node=1, + relation="RELATES", + dest_node=2, + edge_id=9, + properties={"fact": "x", "embedding": [0.2] * 4}, + ) + res = SimpleNamespace( + header=[[1, "n"], [1, "r"], [1, "k"]], result_set=[[node, edge, "plain\x00text"]] + ) + cr = convert_query_result(res, limit=10, limit_injected=True) + assert cr.columns == ["n", "r", "k"] + n_dict, e_dict, s = cr.rows[0] + assert n_dict["labels"] == ["Person", "__Entity__"] + assert "embedding" not in n_dict["properties"] + assert "source_chunk_ids" not in n_dict["properties"] + assert n_dict["properties"]["name"] == "Alice" + assert e_dict["type"] == "RELATES" and e_dict["src"] == 1 and e_dict["dst"] == 2 + assert "embedding" not in e_dict["properties"] + assert s == "plaintext" # control char stripped + assert cr.truncated is False # 1 row != limit + assert cr.model_dump_json() # JSON-serializable end to end + + +def test_convert_truncated_heuristic(): + res = SimpleNamespace(header=[[1, "x"]], result_set=[[1], [2]]) + assert convert_query_result(res, limit=2, limit_injected=True).truncated is True + assert convert_query_result(res, limit=2, limit_injected=False).truncated is False + assert convert_query_result(res, limit=5, limit_injected=True).truncated is False + + +# ── graph_ops helpers ──────────────────────────────────────────── + + +async def test_enrich_entities_maps_rows_and_short_circuits(): + store = MagicMock() + store.query_raw = AsyncMock( + return_value=_res( + [ + [ + "e1", + "Alice", + "Engineer", + ["Person", "__Entity__"], + { + "id": "e1", + "name": "Alice", + "description": "Engineer", + "seniority": "senior", + "embedding": [0.1], + "source_chunk_ids": ["c1"], + }, + ] + ] + ) + ) + cards = await enrich_entities(store, ["e1"]) + card = cards["e1"] + assert card.name == "Alice" and card.label == "Person" + assert card.properties == {"seniority": "senior"} # bulky/system props dropped + + store.query_raw.reset_mock() + assert await enrich_entities(store, []) == {} + store.query_raw.assert_not_awaited() + + +async def test_expand_triples_dedupes_frontier_and_caps(): + store = MagicMock() + store.query_raw = AsyncMock( + return_value=_res( + [ + ["A", "REL", "B", "", "e2"], + ["A", "REL", "B", "", "e2"], # duplicate triple + ["B", "REL2", "C", "evidence", "e3"], + ] + ) + ) + triples = await expand_triples(store, ["e1"], hops=1, cap=10) + assert len(triples) == 2 + assert triples[0].fact is None and triples[1].fact == "evidence" + assert store.query_raw.await_count == 1 + + store.query_raw.reset_mock() + triples = await expand_triples(store, ["e1"], hops=2, cap=10) + assert store.query_raw.await_count == 2 # one query per hop + # second hop frontier excludes already-visited ids + second_ids = store.query_raw.await_args_list[1].args[1]["ids"] + assert "e1" not in second_ids and set(second_ids) == {"e2", "e3"} + + capped = await expand_triples(store, ["e1"], hops=1, cap=1) + assert len(capped) == 1 + + +async def test_chunk_documents_maps_ids(): + store = MagicMock() + store.query_raw = AsyncMock(return_value=_res([["c1", "doc-a", "docs/a.md"]])) + mapping = await chunk_documents(store, ["c1"]) + assert mapping == {"c1": ("doc-a", "docs/a.md")} + store.query_raw.reset_mock() + assert await chunk_documents(store, []) == {} + store.query_raw.assert_not_awaited() + + +async def test_find_entity_matches_passes_rank_ordering_through(): + store = MagicMock() + store.query_raw = AsyncMock( + return_value=_res( + [ + ["e1", "Alice", "Engineer", ["Person", "__Entity__"], {}, 0], + ["e9", "Alice Smith", None, ["Person", "__Entity__"], {}, 2], + ] + ) + ) + matches = await find_entity_matches(store, "Alice") + assert [eid for eid, _ in matches] == ["e1", "e9"] + assert matches[0][1].name == "Alice" and matches[1][1].name == "Alice Smith" + + +# ── toolkit read methods ───────────────────────────────────────── + + +def _stub_rag() -> MagicMock: + rag = MagicMock() + rag._conn.query = AsyncMock(return_value=_res([])) + rag._graph_store.query_raw = AsyncMock(return_value=_res([])) + return rag + + +async def test_cypher_read_guard_fires_before_connection(): + rag = _stub_rag() + tk = GraphRAGToolkit(rag) + with pytest.raises(ReadOnlyViolation): + await tk.cypher_read("CREATE (n) RETURN n") + rag._conn.query.assert_not_awaited() + + +async def test_cypher_read_injects_limit_and_forwards_kwargs(): + rag = _stub_rag() + rag._conn.query = AsyncMock( + return_value=SimpleNamespace(header=[[1, "name"]], result_set=[["Alice"], ["Bob"]]) + ) + tk = GraphRAGToolkit(rag) + result = await tk.cypher_read("MATCH (n) RETURN n.name", {"x": 1}, limit=2, timeout_ms=1234) + query_arg = rag._conn.query.await_args.args[0] + assert query_arg.endswith("LIMIT 2") + assert rag._conn.query.await_args.kwargs == {"params": {"x": 1}, "timeout": 1234} + assert result.columns == ["name"] and result.row_count == 2 + assert result.truncated is True + + +async def test_cypher_read_respects_existing_limit(): + rag = _stub_rag() + tk = GraphRAGToolkit(rag) + await tk.cypher_read("MATCH (n) RETURN n LIMIT 3") + assert rag._conn.query.await_args.args[0] == "MATCH (n) RETURN n LIMIT 3" + + +async def test_schema_merges_ontology_and_live_counts(): + rag = _stub_rag() + rag.get_ontology = AsyncMock( + return_value=Ontology( + entities=[ + Entity( + label="Person", + description="A human", + properties=[Attribute(name="seniority")], + ), + Entity(label="Location"), # declared but absent from the live graph + ], + relations=[Relation(label="WORKS_AT", patterns=[("Person", "Organization")])], + ) + ) + rag.get_statistics = AsyncMock(return_value={"node_count": 7, "edge_count": 9}) + + async def route(cypher, params=None): + if "UNWIND labels(n)" in cypher: + return _res([["Person", 2], ["Organization", 1]]) + if "r.rel_type AS t" in cypher: + return _res([["WORKS_AT", 1]]) + return _res([]) + + rag._graph_store.query_raw = AsyncMock(side_effect=route) + tk = GraphRAGToolkit(rag) + schema = await tk.schema() + by_label = {e.label: e for e in schema.entities} + assert by_label["Person"].count == 2 + assert by_label["Person"].description == "A human" + assert by_label["Person"].properties == ["seniority"] + assert by_label["Organization"].count == 1 and by_label["Organization"].description is None + assert by_label["Location"].count == 0 # declared labels always listed + assert schema.relations[0].label == "WORKS_AT" + assert schema.relations[0].patterns == [("Person", "Organization")] + assert schema.node_count == 7 and schema.edge_count == 9 + + +async def test_entity_found_with_nearby_hops_and_documents(): + rag = _stub_rag() + calls: list[str] = [] + + async def route(cypher, params=None): + calls.append(cypher) + if "CONTAINS toLower($name)" in cypher: + return _res( + [ + [ + "e1", + "Alice", + "Engineer", + ["Person", "__Entity__"], + {"seniority": "senior", "embedding": [0.1]}, + 0, + ], + ["e9", "Alice Smith", None, ["Person", "__Entity__"], {}, 2], + ] + ) + if "r.src_name" in cypher: + return _res([["Alice", "WORKS_AT", "Acme Corp", "employment", "e2"]]) + if "MENTIONED_IN" in cypher: + return _res([["doc-a", "docs/a.md"]]) + return _res([]) + + rag._graph_store.query_raw = AsyncMock(side_effect=route) + tk = GraphRAGToolkit(rag) + er = await tk.entity("Alice", hops=2) + assert er.found and er.entity is not None + assert er.entity.name == "Alice" and er.entity.label == "Person" + assert er.entity.properties == {"seniority": "senior"} + assert er.nearby == ["Alice Smith"] + assert any(t.type == "WORKS_AT" and t.target == "Acme Corp" for t in er.neighbors) + assert er.documents[0].document_id == "doc-a" + assert len([c for c in calls if "r.src_name" in c]) == 2 # hops=2 -> two frontier queries + + +async def test_entity_not_found(): + rag = _stub_rag() + tk = GraphRAGToolkit(rag) + er = await tk.entity("Zorp") + assert er.found is False and er.entity is None + assert er.to_llm_text().startswith("No entity") + assert rag._graph_store.query_raw.await_count == 1 # only the match query ran From 55b9de34ff114fe36d3c32ca74cda5a7d3ff3a49 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 12 Jul 2026 13:43:52 +0300 Subject: [PATCH 08/10] test(tools): search/answer provenance mapping and strategy tuning Asserts per-call MultiPathRetrieval tuning (chunk_top_k/rel_top_k/ max_entities from top_k), provenance-order entity mapping with enrichment fallback, top_k caps, include_chunks/expand_hops behavior, citation construction (chunk id, document id/path, 200-char snippet), case-insensitive entities_touched dedupe, and the graceful no-provenance degradation path. Co-Authored-By: Claude Fable 5 --- graphrag_sdk/tests/test_tools_toolkit.py | 154 ++++++++++++++++++++++- 1 file changed, 153 insertions(+), 1 deletion(-) diff --git a/graphrag_sdk/tests/test_tools_toolkit.py b/graphrag_sdk/tests/test_tools_toolkit.py index 345a944e..2e3292d6 100644 --- a/graphrag_sdk/tests/test_tools_toolkit.py +++ b/graphrag_sdk/tests/test_tools_toolkit.py @@ -11,7 +11,14 @@ from graphrag_sdk.core.connection import ConnectionConfig from graphrag_sdk.core.exceptions import ConfigError, ReadOnlyViolation -from graphrag_sdk.core.models import DocumentInfo, FinalizeResult, IngestionResult +from graphrag_sdk.core.models import ( + DocumentInfo, + FinalizeResult, + IngestionResult, + RagResult, + RetrieverResult, +) +from graphrag_sdk.retrieval.strategies.multi_path import MultiPathRetrieval from graphrag_sdk.tools import GraphRAGToolkit, RememberResult from graphrag_sdk.tools.specs import _TOOL_REGISTRY @@ -146,3 +153,148 @@ def test_module_imports_cleanly(): [sys.executable, "-c", "import graphrag_sdk.tools"], capture_output=True, text=True ) assert proc.returncode == 0, proc.stderr + + +# ── search() / answer() over provenance metadata ───────────────── + + +_PROVENANCE = { + "entities": [ + {"id": "e1", "name": "Alice", "description": "Engineer"}, + {"id": "e2", "name": "Ghost", "description": ""}, + {"id": "e3", "name": "Extra", "description": ""}, + ], + "chunks": [ + {"id": "c1", "text": "Alice works at Acme.", "document_path": "docs/a.md"}, + {"id": "c2", "text": "Bob leads data.", "document_path": "docs/b.md"}, + {"id": "c3", "text": "HQ is in Berlin.", "document_path": "docs/c.md"}, + ], + "facts": [f"fact {i}" for i in range(6)], + "relationships": ["Alice —[WORKS_AT]→ Acme Corp"], +} + + +def _rows(rows): + result = MagicMock() + result.result_set = rows + return result + + +def make_search_rag() -> MagicMock: + rag = make_stub_rag() + rag.retrieve = AsyncMock( + return_value=RetrieverResult(items=[], metadata={"provenance": _PROVENANCE}) + ) + rag.completion = AsyncMock( + return_value=RagResult( + answer="Alice and Bob work at Acme.", + retriever_result=RetrieverResult(items=[], metadata={"provenance": _PROVENANCE}), + metadata={}, + ) + ) + + async def route(cypher, params=None): + if "UNWIND $ids AS eid MATCH (e:__Entity__" in cypher: # enrichment + return _rows( + [["e1", "Alice", "Engineer", ["Person", "__Entity__"], {"seniority": "senior"}]] + ) + if "r.src_name" in cypher: # triples expansion + return _rows([["Alice", "WORKS_AT", "Acme Corp", "employment", "e9"]]) + if "PART_OF" in cypher: # chunk -> document ids + return _rows([["c1", "doc-a", "docs/a.md"]]) + return _rows([]) + + rag._graph_store.query_raw = AsyncMock(side_effect=route) + return rag + + +async def test_search_tunes_strategy_and_maps_provenance(): + rag = make_search_rag() + tk = GraphRAGToolkit(rag) + sr = await tk.search("Who works at Acme?", top_k=2) + + strategy = rag.retrieve.await_args.kwargs["strategy"] + assert isinstance(strategy, MultiPathRetrieval) + assert strategy._chunk_top_k == 2 and strategy._rel_top_k == 2 + assert strategy._max_entities == 10 # max(2*top_k, 10) + + # entities: provenance order, capped to top_k, enrichment fallback for e2 + assert [e.name for e in sr.entities] == ["Alice", "Ghost"] + assert sr.entities[0].label == "Person" + assert sr.entities[0].properties == {"seniority": "senior"} + assert sr.entities[1].label == "" and sr.entities[1].description is None + + assert [c.chunk_id for c in sr.chunks] == ["c1", "c2"] + assert sr.chunks[0].document_id == "doc-a" + assert sr.chunks[1].document_id == "" and sr.chunks[1].document_path == "docs/b.md" + assert [d.document_id for d in sr.documents] == ["doc-a"] + + assert len(sr.facts) == 4 # 2 * top_k + assert any(r.type == "WORKS_AT" for r in sr.relations) + + +async def test_search_include_chunks_false_skips_chunk_queries(): + rag = make_search_rag() + tk = GraphRAGToolkit(rag) + sr = await tk.search("q", top_k=2, include_chunks=False) + assert sr.chunks == [] and sr.documents == [] + part_of_calls = [ + c.args[0] for c in rag._graph_store.query_raw.await_args_list if "PART_OF" in c.args[0] + ] + assert part_of_calls == [] + + +async def test_search_expand_hops_controls_frontier_queries(): + rag = make_search_rag() + tk = GraphRAGToolkit(rag) + await tk.search("q", top_k=2, expand_hops=2) + rel_calls = [ + c.args[0] for c in rag._graph_store.query_raw.await_args_list if "r.src_name" in c.args[0] + ] + assert len(rel_calls) == 2 + + +async def test_answer_builds_citations_and_entities_touched(): + rag = make_search_rag() + long_text = "X" * 250 + provenance = { + "entities": [ + {"id": "e1", "name": "Alice"}, + {"id": "e2", "name": "alice"}, # case-insensitive duplicate + {"id": "e3", "name": "Bob"}, + ], + "chunks": [{"id": "c1", "text": long_text, "document_path": "docs/a.md"}], + "facts": [], + "relationships": [], + } + rag.completion = AsyncMock( + return_value=RagResult( + answer="Alice works at Acme.", + retriever_result=RetrieverResult(items=[], metadata={"provenance": provenance}), + metadata={}, + ) + ) + tk = GraphRAGToolkit(rag) + ar = await tk.answer("Who is Alice?", top_k=3) + + assert rag.completion.await_args.kwargs["return_context"] is True + strategy = rag.completion.await_args.kwargs["strategy"] + assert isinstance(strategy, MultiPathRetrieval) and strategy._chunk_top_k == 3 + + assert ar.answer == "Alice works at Acme." + assert len(ar.citations) == 1 + citation = ar.citations[0] + assert citation.chunk_id == "c1" and citation.document_id == "doc-a" + assert len(citation.snippet) == 200 and citation.snippet.endswith("…") + assert ar.entities_touched == ["Alice", "Bob"] + assert ar.cypher_used is None + + +async def test_answer_degrades_without_provenance(): + rag = make_search_rag() + rag.completion = AsyncMock( + return_value=RagResult(answer="ok", retriever_result=None, metadata={}) + ) + tk = GraphRAGToolkit(rag) + ar = await tk.answer("q") + assert ar.answer == "ok" and ar.citations == [] and ar.entities_touched == [] From ce7c2965c5ad786e3c6f1f0f1e0bbb4e0de3d6d5 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 12 Jul 2026 13:49:21 +0300 Subject: [PATCH 09/10] test(tools): real-FalkorDB integration suite + CI wiring Round-trip over 3 scripted docs + remember/flush: typed search with document ids, entity card with WORKS_AT neighbors and provenance docs, schema counts, parameterized cypher_read with write rejection and live LIMIT injection, answer() citations, and for_tenant graph isolation with owned-connection cleanup. Real-LLM smoke stays behind OPENAI_API_KEY. The CI integration job now runs this file alongside test_integration.py (its test path is hardcoded). Co-Authored-By: Claude Fable 5 --- .github/workflows/ci.yml | 1 + graphrag_sdk/tests/test_tools_integration.py | 224 +++++++++++++++++++ 2 files changed, 225 insertions(+) create mode 100644 graphrag_sdk/tests/test_tools_integration.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7ca8ebc0..33376d3a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -90,3 +90,4 @@ jobs: python -m pytest -v -m integration tests/test_integration.py + tests/test_tools_integration.py diff --git a/graphrag_sdk/tests/test_tools_integration.py b/graphrag_sdk/tests/test_tools_integration.py new file mode 100644 index 00000000..ea639a72 --- /dev/null +++ b/graphrag_sdk/tests/test_tools_integration.py @@ -0,0 +1,224 @@ +"""Real-FalkorDB integration tests for graphrag_sdk.tools (RUN_INTEGRATION-gated).""" + +from __future__ import annotations + +import json +import os +from uuid import uuid4 + +import pytest + +from graphrag_sdk.core.exceptions import ReadOnlyViolation +from graphrag_sdk.tools import GraphRAGToolkit + +# Both markers are load-bearing: `-m integration` selects the file in the CI +# integration job; the skipif keeps it out of the plain unit run. +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif( + os.getenv("RUN_INTEGRATION") != "1", + reason="Set RUN_INTEGRATION=1 to run real-FalkorDB tests", + ), +] + + +def _step2(entities, relationships=()): + """One scripted GraphExtraction step-2 response (entities + relationships).""" + return json.dumps( + { + "entities": [{"name": n, "type": t, "description": d} for n, t, d in entities], + "relationships": [ + { + "source": s, + "target": o, + "type": r, + "description": d, + "keywords": "", + "weight": 0.9, + } + for s, r, o, d in relationships + ], + } + ) + + +async def test_toolkit_round_trip(real_falkordb_rag_factory): + from graphrag_sdk.ingestion.resolution_strategies.exact_match import ExactMatchResolution + + from .conftest import MockLLM + + resolver = ExactMatchResolution() + llm = MockLLM( + responses=[ + _step2( + [ + ("Alice", "Person", "Engineer at Acme"), + ("Acme Corp", "Organization", "A tech company"), + ], + [("Alice", "WORKS_AT", "Acme Corp", "Alice is employed at Acme")], + ), + _step2( + [ + ("Bob", "Person", "Data lead at Acme"), + ("Acme Corp", "Organization", "A tech company"), + ], + [("Bob", "WORKS_AT", "Acme Corp", "Bob works at Acme")], + ), + _step2( + [ + ("Acme Corp", "Organization", "A tech company"), + ("Berlin", "Location", "Capital of Germany"), + ], + [("Acme Corp", "HEADQUARTERED_IN", "Berlin", "Acme HQ is in Berlin")], + ), + _step2([("Carol", "Person", "CFO of Acme")]), # consumed by remember() + "Alice, Acme Corp", # keyword-extraction / completion clamp + ] + ) + rag = real_falkordb_rag_factory(llm=llm, resolver=resolver) + for doc_id, text in [ + ("doc-alice", "Alice is a software engineer at Acme Corp."), + ("doc-bob", "Bob leads the data team at Acme Corp."), + ("doc-hq", "Acme Corp is headquartered in Berlin."), + ]: + await rag.ingest(text=text, document_id=doc_id, resolver=resolver) + + tk = GraphRAGToolkit(rag, finalize_policy="manual") + + remembered = await tk.remember("Carol is the CFO of Acme Corp.", document_id="doc-carol") + assert remembered.document_id == "doc-carol" and not remembered.finalized + await tk.flush() # dedup + embeddings + fulltext/vector indexes + + sr = await tk.search("Who works at Acme Corp?", top_k=8) + names = {e.name for e in sr.entities} + assert "Alice" in names or "Acme Corp" in names + assert any(r.type == "WORKS_AT" for r in sr.relations) + assert sr.chunks and all(c.document_id for c in sr.chunks) + assert sr.to_llm_text() + assert len(sr.to_llm_text(max_chars=500)) <= 500 + + er = await tk.entity("Alice", hops=2) + assert er.found and er.entity is not None and er.entity.label == "Person" + assert any(t.type == "WORKS_AT" and t.target == "Acme Corp" for t in er.neighbors) + assert any(d.document_id == "doc-alice" for d in er.documents) + + sch = await tk.schema() + labels = {e.label: e.count for e in sch.entities} + assert labels.get("Person", 0) >= 2 and labels.get("Organization", 0) >= 1 + assert any(r.label == "WORKS_AT" and r.count >= 1 for r in sch.relations) + assert sch.node_count > 0 + + cr = await tk.cypher_read( + "MATCH (e:__Entity__) WHERE e.name = $name RETURN e.name", {"name": "Alice"} + ) + assert cr.rows == [["Alice"]] and cr.columns + with pytest.raises(ReadOnlyViolation): + await tk.cypher_read("CREATE (n:Hack) RETURN n") + + ar = await tk.answer("Who works at Acme Corp?", top_k=5) + assert ar.answer.strip() and len(ar.citations) >= 1 + assert all(c.chunk_id and c.document_id for c in ar.citations) + + +async def test_cypher_read_limit_injection_live(real_falkordb_rag_factory): + from graphrag_sdk.ingestion.resolution_strategies.exact_match import ExactMatchResolution + + from .conftest import MockLLM + + resolver = ExactMatchResolution() + llm = MockLLM( + responses=[ + _step2( + [ + ("D1", "Concept", "one"), + ("D2", "Concept", "two"), + ("D3", "Concept", "three"), + ] + ) + ] + ) + rag = real_falkordb_rag_factory(llm=llm, resolver=resolver) + await rag.ingest(text="D1 and D2 and D3 are concepts.", document_id="doc-d", resolver=resolver) + + tk = GraphRAGToolkit(rag) + cr = await tk.cypher_read("MATCH (e:__Entity__) RETURN e.name", limit=2) + assert cr.row_count <= 2 + if cr.row_count == 2: + assert cr.truncated is True + + +async def test_for_tenant_isolated_graphs(embedder): + from graphrag_sdk.core.connection import ConnectionConfig + + from .conftest import MockLLM + + base = ConnectionConfig( + host=os.getenv("FALKOR_HOST", "localhost"), + port=int(os.getenv("FALKOR_PORT", "6379")), + username=os.getenv("FALKOR_USERNAME") or None, + password=os.getenv("FALKOR_PASSWORD") or None, + graph_name=f"test_{uuid4().hex[:8]}", + ) + llm_a = MockLLM(responses=[_step2([("Alice", "Person", "Engineer")])]) + tk_a = GraphRAGToolkit.for_tenant( + base, "tenant-a", llm=llm_a, embedder=embedder, embedding_dimension=embedder.dimension + ) + tk_b = GraphRAGToolkit.for_tenant( + base, + "tenant-b", + llm=MockLLM(responses=["Alice"]), + embedder=embedder, + embedding_dimension=embedder.dimension, + ) + try: + assert tk_a.rag._conn.config.graph_name == f"{base.graph_name}__tenant-a" + assert tk_b.rag._conn.config.graph_name == f"{base.graph_name}__tenant-b" + + await tk_a.remember("Alice is an engineer.", document_id="doc-a") + await tk_a.flush() + found_a = await tk_a.entity("Alice") + assert found_a.found + + # tenant-b's graph never saw Alice + found_b = await tk_b.entity("Alice") + assert found_b.found is False + finally: + for tk in (tk_a, tk_b): + try: + await tk.rag._graph_store.delete_all() + finally: + await tk.aclose() + + +@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="needs OPENAI_API_KEY") +async def test_real_llm_smoke(real_falkordb_rag_factory): + """One end-to-end pass with a real LLM (kept tiny: one 2-sentence doc).""" + pytest.importorskip("litellm") + from graphrag_sdk import LiteLLM, LiteLLMEmbedder + from graphrag_sdk.api.main import GraphRAG + from graphrag_sdk.core.connection import ConnectionConfig + + config = ConnectionConfig( + host=os.getenv("FALKOR_HOST", "localhost"), + port=int(os.getenv("FALKOR_PORT", "6379")), + graph_name=f"test_{uuid4().hex[:8]}", + ) + llm = LiteLLM(model="openai/gpt-4o-mini") + embedder = LiteLLMEmbedder(model="openai/text-embedding-3-small", dimensions=256) + rag = GraphRAG(connection=config, llm=llm, embedder=embedder, embedding_dimension=256) + try: + tk = GraphRAGToolkit(rag) + await tk.remember( + "Ada Lovelace wrote the first computer program. " + "She collaborated with Charles Babbage on the Analytical Engine.", + document_id="doc-ada", + ) + await tk.flush() + ar = await tk.answer("Who wrote the first computer program?") + assert ar.answer.strip() + assert len(ar.citations) >= 1 + finally: + try: + await rag._graph_store.delete_all() + finally: + await rag.close() From 2f0396792dcc69724948068d12d41ba92532eb96 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 12 Jul 2026 13:52:35 +0300 Subject: [PATCH 10/10] docs(tools): agentic guide, example 11, README section, changelog docs/agentic.md covers the tool surface, tool_specs contract for adapter authors, finalize policies (finalize is O(graph size)), the cypher_read guard semantics + allowlist and the treat-agent-Cypher-as- untrusted guidance, and for_tenant. Wired into mkdocs nav (strict build verified). README gains a 'Use with Agents' section + example table row and marks the 2026-Q3 roadmap item; fixes the MENTIONS -> MENTIONED_IN edge name. Also repairs a pre-existing broken relative link in incremental-updates.md that failed mkdocs --strict (the docs workflow triggers on any docs/** change). Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 41 ++++++ README.md | 22 ++- docs/agentic.md | 168 ++++++++++++++++++++++ docs/incremental-updates.md | 2 +- graphrag_sdk/examples/11_agent_toolkit.py | 91 ++++++++++++ mkdocs.yml | 1 + 6 files changed, 322 insertions(+), 3 deletions(-) create mode 100644 docs/agentic.md create mode 100644 graphrag_sdk/examples/11_agent_toolkit.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ff0a9a56..9280e603 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,47 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +Agentic GraphRAG: a framework-neutral agent toolkit +(`graphrag_sdk.tools`) so any agent framework — pydantic-ai, LangGraph, +CrewAI, MCP — can expose a knowledge graph as tools. Purely additive; +no breaking changes. + +### Added + +#### Agentic GraphRAG toolkit (`graphrag_sdk.tools`) + +- **`GraphRAGToolkit`** — framework-neutral async agent surface over + `GraphRAG`: `search()`, `answer()`, `schema()`, `entity()`, guarded + `cypher_read()`, `remember()`, `flush()`; knobs `finalize_policy=` + (`"manual"`/`"on_write"`/`"never"`), `read_only=`, `include=`; generic + `call(name, arguments)` dispatch for adapters; `for_tenant(...)` for + tenant-scoped graphs with owned-connection lifecycle (`async with`). +- **`GraphRAGToolkit.tool_specs()`** — machine-readable tool definitions + (name, LLM-ready description, JSON-Schema input, output hint); the + single source of truth for agent-framework adapters and MCP servers. +- **Typed tool results** — pydantic models (`SearchResult`, + `AnswerResult`, `SchemaResult`, `CypherResult`, `EntityResult`, + `RememberResult`) with deterministic, budget-bounded + `to_llm_text(max_chars=...)` rendering: truncation only at item + boundaries with an explicit `…(N more)` marker, control characters + stripped. Answers carry `document_id`/`chunk_id` citations. +- **`ReadOnlyViolation`** — raised by the fail-closed Cypher guard + (write keywords, non-allowlisted procedures, multi-statement; lexer- + based string/comment handling, NFKC-normalized scan) and by write + tools on `read_only=True` toolkits. +- **`MultiPathRetrieval` provenance** — retrieval results now expose + `metadata["provenance"]` (entity ids, kept chunk ids/texts/document + paths in rerank order, fact and relationship strings); section + content is byte-unchanged. +- New docs page `docs/agentic.md`; example + `examples/11_agent_toolkit.py`; `jsonschema` added to the dev extra. + +### Fixed + +- **`graphrag_sdk.__version__`** — was stale at `"1.2.0"`; now matches + the packaged version and is guarded by a pyproject drift test (the + value is stamped into graphs as `__GraphRAGConfig__.sdk_version`). + ## [1.3.0] - 2026-06-04 Ontology discovery (#271): bootstrap an ontology straight from a diff --git a/README.md b/README.md index 9fd2ff48..fb0dc08f 100644 --- a/README.md +++ b/README.md @@ -192,7 +192,24 @@ guards the default. | Retrieval | Relationship expansion | DB | | Retrieval | Cosine reranking | Local | -> 💡 Every answer is traceable to its source chunks via `MENTIONS` edges. Pass `return_context=True` to `completion()` to get the retrieval trail alongside the answer. +> 💡 Every answer is traceable to its source chunks via `MENTIONED_IN` edges. Pass `return_context=True` to `completion()` to get the retrieval trail alongside the answer. + +--- + +## Use with Agents + +`graphrag_sdk.tools` turns any `GraphRAG` into a set of agent tools — typed results with `document_id`/`chunk_id` citations, a guarded read-only Cypher escape hatch, and machine-readable `tool_specs()` that adapters (pydantic-ai, LangGraph, MCP) generate their tool definitions from: + +```python +from graphrag_sdk.tools import GraphRAGToolkit + +toolkit = GraphRAGToolkit(rag) # wrap any GraphRAG +result = await toolkit.search("Who works at Acme?") # typed, citation-ready +print(result.to_llm_text()) # prompt-ready rendering +specs = toolkit.tool_specs() # JSON-Schema tool definitions +``` + +See the [Agentic GraphRAG guide](docs/agentic.md) and [`examples/11_agent_toolkit.py`](graphrag_sdk/examples/11_agent_toolkit.py). --- @@ -208,6 +225,7 @@ guards the default. | 4 | [Custom Provider](graphrag_sdk/examples/04_custom_provider.py) | Plug in any LLM or embedder behind a clean interface | | 5 | [Notebook Demo](graphrag_sdk/examples/05_notebook_demo.ipynb) | An interactive walkthrough that shows the provenance trail | | 7 | [Incremental Updates](graphrag_sdk/examples/07_incremental_updates.py) | `update`, `delete_document`, and `apply_changes` for CI-driven graph syncs | +| 11 | [Agent Toolkit](graphrag_sdk/examples/11_agent_toolkit.py) | Expose your graph as agent tools with citations and guarded Cypher | --- @@ -233,7 +251,7 @@ guards the default. - 🎉 2026-04: Version 1.0 is released with a new set of benchmarks based on a year's worth of research and customer PoCs - 📦 Still on the v0.x API? Pin the legacy release: `pip install graphrag-sdk==0.8.2` - 2026-Q2: Production observability; expand ingestion support — tables, structured data -- 2026-Q3: Introduce Agentic GraphRAG; complete PDF ingestion +- 2026-Q3: Introduce Agentic GraphRAG (✅ `graphrag_sdk.tools` agent toolkit); complete PDF ingestion - 2026-Q4: Smarter retrieval — dynamic traversal, temporal graph --- diff --git a/docs/agentic.md b/docs/agentic.md new file mode 100644 index 00000000..199462e6 --- /dev/null +++ b/docs/agentic.md @@ -0,0 +1,168 @@ +# Agentic GraphRAG + +`graphrag_sdk.tools` gives any agent framework — pydantic-ai, LangGraph, CrewAI, +MCP clients — a small, stable set of async operations over a knowledge graph. +Every operation returns a typed pydantic model that renders itself into +compact, deterministic plain text for LLM consumption, and the whole surface +is described by machine-readable [`tool_specs()`](#tool_specs-for-adapter-authors) +so adapters never hand-copy names, descriptions, or schemas. + +## Quickstart + +```python +import asyncio + +from graphrag_sdk import ConnectionConfig, GraphRAG, LiteLLM, LiteLLMEmbedder +from graphrag_sdk.tools import GraphRAGToolkit + + +async def main(): + rag = GraphRAG( + connection=ConnectionConfig(host="localhost", graph_name="agent_demo"), + llm=LiteLLM(model="openai/gpt-5.5"), + embedder=LiteLLMEmbedder(model="openai/text-embedding-3-large", dimensions=256), + embedding_dimension=256, + ) + toolkit = GraphRAGToolkit(rag) # finalize_policy="manual" by default + + await toolkit.remember("Alice is a software engineer at Acme Corp.") + await toolkit.flush() # dedup + embeddings + indexes (expensive — see policies) + + result = await toolkit.search("Who works at Acme?") + print(result.to_llm_text()) # prompt-ready text with ids for citations + print(result.model_dump_json()) # or full structured output + + answer = await toolkit.answer("Who works at Acme?") + print(answer.answer, answer.citations) + + await rag.close() + + +asyncio.run(main()) +``` + +## Tools + +| Tool name | Method | When to use | +|---|---|---| +| `graph_search` | `search(query, *, top_k=8, expand_hops=1, include_chunks=True)` | Retrieval only — ranked entities, relations, facts, and source chunks the host LLM composes from. The default agent mode; cite with the returned `document_id`/`chunk_id`. | +| `graph_answer` | `answer(question, *, top_k=8)` | Full RAG: the SDK's `completion()` pipeline plus provenance citations. One-shot Q&A. | +| `graph_schema` | `schema()` | Entity labels, relation types, directional patterns, and live counts. Call once up front to plan queries. | +| `graph_entity` | `entity(name, *, hops=1)` | Entity card: best name match, properties, neighbors up to `hops`, source documents, and "nearby" alternative matches. | +| `cypher_read` | `cypher_read(query, params=None, *, limit=100, timeout_ms=5000)` | Guarded read-only Cypher for aggregations and precise filters search cannot express. | +| `graph_remember` | `remember(text, *, document_id=None)` | Store new text into the graph (agent memory, fact capture). | +| `graph_flush` | `flush()` | Run finalization after a write session (only advertised under the `"manual"` policy). | + +All methods are async and validate their arguments through the same pydantic +input models that generate the JSON Schemas in `tool_specs()`. Every result +model supports `model_dump()` / `model_dump_json()` for structured consumers +and `to_llm_text(max_chars=4000)` for prompt building — deterministic, +truncated only at item boundaries with an explicit `…(N more)` marker, and +control-character-stripped. + +## `tool_specs()` for adapter authors + +`tool_specs()` is the single source of truth. Adapters and MCP servers must +generate their tool definitions from it — names, LLM-facing descriptions, and +input JSON Schemas are never duplicated downstream. + +```python +for spec in toolkit.tool_specs(): + print(spec.name) # "graph_search" + print(spec.description) # when-to-use guidance written for the model + print(spec.input_schema) # JSON Schema (draft 2020-12, additionalProperties: false) + print(spec.output_hint) # one-line shape hint for the result +``` + +Dispatch generically with `call()` — it validates arguments against the +tool's input model and enforces the same gates as `tool_specs()`: + +```python +result = await toolkit.call("graph_search", {"query": "Who works at Acme?", "top_k": 5}) +``` + +Two knobs shape the advertised surface: + +- `read_only=True` — `graph_remember`/`graph_flush` are removed from + `tool_specs()` and raise `ReadOnlyViolation` if invoked anyway. +- `include=["graph_search", "graph_schema"]` — advertise a subset. `include` + filters `tool_specs()` and `call()`; direct method calls still work. + +## Finalize policies + +`GraphRAG.finalize()` runs cross-document entity deduplication, entity and +relationship embeddings, and index creation. + +!!! warning "finalize() is O(graph size)" + Finalization cost grows with the whole graph, not with the size of the + last write. `finalize_policy="on_write"` (finalize after **every** + `remember`) is for demos and tiny graphs only. + +| Policy | Behavior | +|---|---| +| `"manual"` (default) | Writes accumulate; call `flush()` once at the end of a write session. `graph_flush` is advertised in `tool_specs()`. | +| `"on_write"` | Every `remember()` finalizes (`RememberResult.finalized=True`). `flush()` is a no-op and is not advertised. | +| `"never"` | The toolkit never finalizes; `flush()` raises `ConfigError`. You own the `rag.finalize()` lifecycle. | + +## Security: treat agent Cypher as untrusted input + +Natural-language-to-Cypher — whether generated by your model or typed by an +agent — is **model-generated, untrusted input**. `cypher_read` is guarded: + +- Write clauses (`CREATE`, `MERGE`, `DELETE`, `DETACH`, `SET`, `REMOVE`, + `DROP`, `FOREACH`, `LOAD CSV`) are rejected with `ReadOnlyViolation` + naming the offending token. Detection is case-insensitive, tolerant of + comments/whitespace, immune to string-literal smuggling (a single-pass + lexer masks literals and strips comments), and runs on both the raw and + the NFKC-normalized query text. +- `CALL` is allowed only for read-safe procedures: + `db.labels`, `db.relationshipTypes`, `db.propertyKeys`, `db.indexes`, + `db.idx.fulltext.queryNodes`, `db.idx.fulltext.queryRelationships`, + `db.idx.vector.queryNodes`, `db.idx.vector.queryRelationships` + (full-name match — `db.idx.fulltext.createNodeIndex` is a write and is + rejected). `CALL { … }` subqueries are permitted; their contents are + scanned like everything else. +- A `LIMIT {limit}` is appended when the query has none, and `timeout_ms` + is enforced server-side per query. The connection layer retries transient + failures (`ConnectionConfig.retry_count`, default 3), so worst-case wall + time is about `retry_count × timeout_ms`. +- Node/edge values returned by `cypher_read` are converted to JSON-safe + data with bulky internals (`embedding`, `source_chunk_ids`) stripped. + +!!! tip "Defense in depth" + The guard is one layer. In production, point agent toolkits at read-only + replicas or construct them with `read_only=True`; give write access only + to toolkits that need it. + +All strings rendered by `to_llm_text()` pass the SDK's control-character +sanitizer, so tool output cannot smuggle terminal escapes into prompts. + +## Multi-tenancy + +Bind a toolkit to a tenant-scoped graph with `for_tenant`. It derives +`graph_name = f"{base.graph_name}__{tenant_id}"` (tenant ids are validated +against `^[A-Za-z0-9_-]{1,64}$`), builds a dedicated `GraphRAG` the toolkit +owns, and closes it on `aclose()` / `async with`: + +```python +async with GraphRAGToolkit.for_tenant( + base_config, "acme", llm=llm, embedder=embedder +) as toolkit: + await toolkit.search("...") +``` + +## Limitations + +- `AnswerResult.cypher_used` is currently always `None`; the experimental + text-to-Cypher retrieval path does not yet surface its generated query. + The field exists for forward compatibility. +- `search()`/`answer()` drive a toolkit-tuned `MultiPathRetrieval` so that + `top_k` and citations behave predictably. A custom `retrieval_strategy` + configured on the `GraphRAG` instance is **not** used by the toolkit — + call `rag.completion()` / `rag.retrieve()` directly when you need it. +- Citations come from the retrieval pipeline's provenance + (`RetrieverResult.metadata["provenance"]`), which lists the context the + model actually saw — in relevance order, capped at `top_k`. + +See [`examples/11_agent_toolkit.py`](https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/examples/11_agent_toolkit.py) +for a runnable end-to-end script. diff --git a/docs/incremental-updates.md b/docs/incremental-updates.md index 896a843a..96e70e84 100644 --- a/docs/incremental-updates.md +++ b/docs/incremental-updates.md @@ -290,7 +290,7 @@ Steps 2 and 3 are why `apply_changes()` does not call it automatically: they sca ## End-to-End Example -A full runnable example lives in [`graphrag_sdk/examples/07_incremental_updates.py`](../graphrag_sdk/examples/07_incremental_updates.py). It exercises: +A full runnable example lives in [`graphrag_sdk/examples/07_incremental_updates.py`](https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/examples/07_incremental_updates.py). It exercises: - Initial ingest with a stable `document_id` - No-op update (content-hash short-circuit) diff --git a/graphrag_sdk/examples/11_agent_toolkit.py b/graphrag_sdk/examples/11_agent_toolkit.py new file mode 100644 index 00000000..f4b3c380 --- /dev/null +++ b/graphrag_sdk/examples/11_agent_toolkit.py @@ -0,0 +1,91 @@ +""" +GraphRAG SDK -- Agent Toolkit +============================= +The framework-neutral agent surface: wrap a GraphRAG in GraphRAGToolkit, +store a few facts, then exercise every tool an agent would call -- +graph_schema, graph_search, graph_entity, graph_answer, cypher_read -- +and print the machine-readable tool_specs() adapters consume. + +Prerequisites: + docker run -p 6379:6379 falkordb/falkordb + pip install graphrag-sdk[litellm] + export OPENAI_API_KEY="sk-..." + +More providers: see docs/providers.md (mirror 01_quickstart.py's setup). +""" + +import asyncio +import json +import os + +from graphrag_sdk import ConnectionConfig, GraphRAG, LiteLLM, LiteLLMEmbedder +from graphrag_sdk.tools import GraphRAGToolkit, ReadOnlyViolation + +if not os.getenv("OPENAI_API_KEY"): + raise SystemExit("Set OPENAI_API_KEY before running this example.") + +FACTS = [ + "Alice Johnson is a software engineer at Acme Corp in London.", + "Bob Smith is the CTO of Acme Corp and Alice's manager.", + "Acme Corp is headquartered in Berlin and builds cloud infrastructure.", +] + + +async def main(): + llm = LiteLLM(model="openai/gpt-5.5") + embedder = LiteLLMEmbedder(model="openai/text-embedding-3-large", dimensions=256) + + rag = GraphRAG( + connection=ConnectionConfig(host="localhost", graph_name="agent_toolkit_demo"), + llm=llm, + embedder=embedder, + embedding_dimension=256, + ) + toolkit = GraphRAGToolkit(rag, finalize_policy="manual") + + try: + # 1. Write path: remember facts, then flush once (finalize is O(graph size)) + for i, fact in enumerate(FACTS): + stored = await toolkit.remember(fact, document_id=f"fact-{i}") + print(stored.to_llm_text(), "\n") + await toolkit.flush() + + # 2. graph_schema -- what does the graph contain? + schema = await toolkit.schema() + print("=== graph_schema ===\n" + schema.to_llm_text(max_chars=800), "\n") + + # 3. graph_search -- typed retrieval, no generation (the default agent mode) + search = await toolkit.search("Who works at Acme Corp?", top_k=5) + print("=== graph_search ===\n" + search.to_llm_text(max_chars=1200), "\n") + + # 4. graph_entity -- one entity's card + entity = await toolkit.entity("Alice Johnson", hops=2) + print("=== graph_entity ===\n" + entity.to_llm_text(max_chars=800), "\n") + + # 5. graph_answer -- full RAG with citations + answer = await toolkit.answer("Who is the CTO of Acme Corp?") + print("=== graph_answer ===\n" + answer.to_llm_text(max_chars=800), "\n") + + # 6. cypher_read -- guarded read-only escape hatch + rows = await toolkit.cypher_read( + "MATCH (e:__Entity__) RETURN e.name ORDER BY e.name LIMIT 5" + ) + print("=== cypher_read ===\n" + rows.to_llm_text(max_chars=600), "\n") + + # ... and the guard rejecting a write attempt: + try: + await toolkit.cypher_read("CREATE (n:Hack) RETURN n") + except ReadOnlyViolation as exc: + print(f"Guard blocked write: {exc} (token={exc.offending_token})\n") + + # 7. tool_specs() -- what adapters (pydantic-ai, MCP, ...) consume + spec = toolkit.tool_specs()[0] + print("=== tool_specs()[0] ===") + print(json.dumps(spec.model_dump(), indent=2)[:600], "...") + finally: + await rag.delete_all() # demo cleanup: drop the example graph + await rag.close() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/mkdocs.yml b/mkdocs.yml index aea56ad6..f84e3483 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -49,6 +49,7 @@ nav: - Ingestion: ingestion.md - Extraction: extraction.md - Retrieval: retrieval.md + - Agentic GraphRAG: agentic.md - Storage: storage.md - Graph Schema: graph-schema.md - Ontology Evolution: ontology-evolution.md