graphrag-core 0.2.0__py3-none-any.whl

graphrag_core/__init__.py

@@ -0,0 +1,138 @@
+"""graphrag-core: Domain-agnostic Graph RAG framework."""
+
+__version__ = "0.2.0"
+
+from graphrag_core.interfaces import (
+    Agent,
+    ApprovalGateway,
+    Chunker,
+    DetectionLayer,
+    DocumentParser,
+    EmbeddingModel,
+    EntityRegistry,
+    ExtractionEngine,
+    GraphStore,
+    IngestionPipeline,
+    LLMClient,
+    LLMCurationLayer,
+    Orchestrator,
+    ReportRenderer,
+    SearchEngine,
+)
+from graphrag_core.ingestion import (
+    DocxParser,
+    MarkdownParser,
+    PdfParser,
+    TextParser,
+    TokenChunker,
+)
+from graphrag_core.extraction import LLMExtractionEngine
+from graphrag_core.graph import InMemoryGraphStore
+from graphrag_core.search import InMemorySearchEngine
+from graphrag_core.registry import InMemoryEntityRegistry
+from graphrag_core.curation import DeterministicDetectionLayer, CurationPipeline
+from graphrag_core.tools import Tool, ToolLibrary, register_core_tools
+from graphrag_core.agents import AgentContext, SequentialOrchestrator
+from graphrag_core.models import (
+    AgentResult,
+    CurationIssue,
+    CurationReport,
+    DocumentChunk,
+    ExtractionResult,
+    GraphNode,
+    ImportRun,
+    KnownEntity,
+    NodeTypeDefinition,
+    OntologySchema,
+    PropertyDefinition,
+    RegistryMatch,
+    RelationshipTypeDefinition,
+    RenderConfig,
+    ReportData,
+    SearchResult,
+    ToolParameter,
+    ToolResult,
+    WorkflowResult,
+)
+
+__all__ = [
+    # Protocols
+    "Agent",
+    "ApprovalGateway",
+    "Chunker",
+    "DetectionLayer",
+    "DocumentParser",
+    "EmbeddingModel",
+    "EntityRegistry",
+    "ExtractionEngine",
+    "GraphStore",
+    "IngestionPipeline",
+    "LLMClient",
+    "LLMCurationLayer",
+    "Orchestrator",
+    "ReportRenderer",
+    "SearchEngine",
+    # BB1 implementations
+    "DocxParser",
+    "MarkdownParser",
+    "PdfParser",
+    "TextParser",
+    "TokenChunker",
+    # BB2 implementations
+    "LLMExtractionEngine",
+    # BB3 implementations
+    "InMemoryGraphStore",
+    # BB4 implementations
+    "InMemorySearchEngine",
+    # BB5 implementations
+    "CurationPipeline",
+    "DeterministicDetectionLayer",
+    # BB6 implementations
+    "InMemoryEntityRegistry",
+    # BB7 implementations
+    "Tool",
+    "ToolLibrary",
+    "register_core_tools",
+    # BB8 implementations
+    "AgentContext",
+    "SequentialOrchestrator",
+    # Models
+    "AgentResult",
+    "CurationIssue",
+    "CurationReport",
+    "DocumentChunk",
+    "ExtractionResult",
+    "GraphNode",
+    "ImportRun",
+    "KnownEntity",
+    "NodeTypeDefinition",
+    "OntologySchema",
+    "PropertyDefinition",
+    "RegistryMatch",
+    "RelationshipTypeDefinition",
+    "RenderConfig",
+    "ReportData",
+    "SearchResult",
+    "ToolParameter",
+    "ToolResult",
+    "WorkflowResult",
+]
+
+# Optional: Neo4j and Anthropic (require extras)
+try:
+    from graphrag_core.graph import Neo4jGraphStore
+    __all__.append("Neo4jGraphStore")
+except ImportError:
+    pass
+
+try:
+    from graphrag_core.search import Neo4jHybridSearch
+    __all__.append("Neo4jHybridSearch")
+except ImportError:
+    pass
+
+try:
+    from graphrag_core.llm import AnthropicLLMClient
+    __all__.append("AnthropicLLMClient")
+except ImportError:
+    pass

graphrag_core/_cypher.py

@@ -0,0 +1,15 @@
+"""Shared Cypher safety utilities."""
+
+from __future__ import annotations
+
+import re
+
+SAFE_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+MAX_DEPTH = 10
+
+
+def validate_identifier(value: str, kind: str) -> str:
+    """Reject identifiers that could cause Cypher injection."""
+    if not SAFE_IDENTIFIER.match(value):
+        raise ValueError(f"Invalid {kind}: {value!r}")
+    return value

graphrag_core/agents/__init__.py

@@ -0,0 +1,6 @@
+"""BB8: Agent orchestration."""
+
+from graphrag_core.agents.context import AgentContext
+from graphrag_core.agents.orchestrator import SequentialOrchestrator
+
+__all__ = ["AgentContext", "SequentialOrchestrator"]

graphrag_core/agents/context.py

@@ -0,0 +1,16 @@
+"""BB8: Shared context for agent workflows."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class AgentContext:
+    """Runtime context passed between agents in a workflow."""
+
+    graph_store: Any
+    tool_library: Any
+    search_engine: Any
+    workflow_state: dict[str, Any] = field(default_factory=dict)

graphrag_core/agents/orchestrator.py

@@ -0,0 +1,34 @@
+"""BB8: Sequential agent orchestrator."""
+
+from __future__ import annotations
+
+from graphrag_core.agents.context import AgentContext
+from graphrag_core.models import WorkflowResult
+
+
+class SequentialOrchestrator:
+    """Runs agents sequentially, stopping on first failure."""
+
+    async def run_workflow(
+        self,
+        workflow_id: str,
+        agents: list,
+        context: AgentContext,
+    ) -> WorkflowResult:
+        agent_results = []
+
+        for agent in agents:
+            result = await agent.execute(context)
+            agent_results.append(result)
+            if not result.success:
+                return WorkflowResult(
+                    workflow_id=workflow_id,
+                    success=False,
+                    agent_results=agent_results,
+                )
+
+        return WorkflowResult(
+            workflow_id=workflow_id,
+            success=True,
+            agent_results=agent_results,
+        )

graphrag_core/curation/__init__.py

@@ -0,0 +1,6 @@
+"""BB5: Governed curation implementations."""
+
+from graphrag_core.curation.detection import DeterministicDetectionLayer
+from graphrag_core.curation.pipeline import CurationPipeline
+
+__all__ = ["DeterministicDetectionLayer", "CurationPipeline"]

graphrag_core/curation/detection.py

@@ -0,0 +1,158 @@
+"""BB5: Deterministic detection layer for graph quality issues."""
+
+from __future__ import annotations
+
+import uuid
+from collections import defaultdict
+
+from graphrag_core.interfaces import EntityRegistry, GraphStore
+from graphrag_core.models import CurationIssue, GraphNode, OntologySchema
+from graphrag_core.registry.matching import fuzzy_score
+
+_PAIRWISE_CAP = 1000
+_FUZZY_THRESHOLD = 0.7
+
+
+class DeterministicDetectionLayer:
+    """Finds duplicates, orphans, and schema violations without LLM calls."""
+
+    def __init__(self, entity_registry: EntityRegistry | None = None) -> None:
+        self._registry = entity_registry
+
+    async def detect(
+        self, graph_store: GraphStore, schema: OntologySchema
+    ) -> list[CurationIssue]:
+        issues: list[CurationIssue] = []
+
+        nodes = await graph_store.list_nodes()
+
+        issues.extend(await self._detect_duplicates(nodes))
+        issues.extend(await self._detect_orphans(nodes, graph_store))
+        issues.extend(await self._detect_schema_violations(graph_store, schema))
+
+        return issues
+
+    async def _detect_duplicates(self, nodes: list[GraphNode]) -> list[CurationIssue]:
+        issues: list[CurationIssue] = []
+
+        groups: dict[str, list[GraphNode]] = defaultdict(list)
+        for node in nodes:
+            groups[node.label].append(node)
+
+        for label, group in groups.items():
+            if self._registry is not None:
+                issues.extend(await self._detect_duplicates_with_registry(group))
+            else:
+                if len(group) > _PAIRWISE_CAP:
+                    issues.append(CurationIssue(
+                        id=str(uuid.uuid4()),
+                        issue_type="skipped_detection",
+                        severity="warning",
+                        affected_nodes=[],
+                        suggested_action=f"Register entities in EntityRegistry for label '{label}' ({len(group)} nodes exceeds pairwise cap of {_PAIRWISE_CAP})",
+                        auto_fixable=False,
+                        source_layer="deterministic",
+                    ))
+                else:
+                    issues.extend(self._detect_duplicates_pairwise(group))
+
+        return issues
+
+    async def _detect_duplicates_with_registry(
+        self, nodes: list[GraphNode]
+    ) -> list[CurationIssue]:
+        issues: list[CurationIssue] = []
+        entity_to_nodes: dict[str, list[str]] = defaultdict(list)
+
+        for node in nodes:
+            name = node.properties.get("name", "")
+            if not name:
+                continue
+            matches = await self._registry.lookup(name, node.label, match_strategy="fuzzy")
+            if matches:
+                entity_to_nodes[matches[0].entity_id].append(node.id)
+
+        for entity_id, node_ids in entity_to_nodes.items():
+            if len(node_ids) > 1:
+                issues.append(CurationIssue(
+                    id=str(uuid.uuid4()),
+                    issue_type="duplicate",
+                    severity="warning",
+                    affected_nodes=node_ids,
+                    suggested_action=f"Merge nodes {node_ids} — they match registry entity '{entity_id}'",
+                    auto_fixable=False,
+                    source_layer="deterministic",
+                ))
+
+        return issues
+
+    def _detect_duplicates_pairwise(
+        self, nodes: list[GraphNode]
+    ) -> list[CurationIssue]:
+        issues: list[CurationIssue] = []
+        seen_pairs: set[tuple[str, str]] = set()
+
+        for i, a in enumerate(nodes):
+            name_a = a.properties.get("name", "")
+            if not name_a:
+                continue
+            for b in nodes[i + 1:]:
+                name_b = b.properties.get("name", "")
+                if not name_b:
+                    continue
+                pair = (min(a.id, b.id), max(a.id, b.id))
+                if pair in seen_pairs:
+                    continue
+                score = fuzzy_score(name_a, name_b)
+                if score >= _FUZZY_THRESHOLD:
+                    seen_pairs.add(pair)
+                    issues.append(CurationIssue(
+                        id=str(uuid.uuid4()),
+                        issue_type="duplicate",
+                        severity="warning",
+                        affected_nodes=[a.id, b.id],
+                        suggested_action=f"Merge '{name_a}' and '{name_b}' (similarity: {score:.2f})",
+                        auto_fixable=False,
+                        source_layer="deterministic",
+                    ))
+
+        return issues
+
+    async def _detect_orphans(
+        self, nodes: list[GraphNode], graph_store: GraphStore
+    ) -> list[CurationIssue]:
+        issues: list[CurationIssue] = []
+
+        for node in nodes:
+            related = await graph_store.get_related(node.id)
+            if not related:
+                issues.append(CurationIssue(
+                    id=str(uuid.uuid4()),
+                    issue_type="orphan",
+                    severity="info",
+                    affected_nodes=[node.id],
+                    suggested_action=f"Node '{node.id}' ({node.label}) has no relationships",
+                    auto_fixable=False,
+                    source_layer="deterministic",
+                ))
+
+        return issues
+
+    async def _detect_schema_violations(
+        self, graph_store: GraphStore, schema: OntologySchema
+    ) -> list[CurationIssue]:
+        await graph_store.apply_schema(schema)
+        violations = await graph_store.validate_schema()
+
+        return [
+            CurationIssue(
+                id=str(uuid.uuid4()),
+                issue_type="schema_violation",
+                severity="error",
+                affected_nodes=[v.node_id],
+                suggested_action=v.message,
+                auto_fixable=False,
+                source_layer="deterministic",
+            )
+            for v in violations
+        ]

graphrag_core/curation/pipeline.py

@@ -0,0 +1,39 @@
+"""BB5: Curation pipeline orchestrator."""
+
+from __future__ import annotations
+
+from graphrag_core.interfaces import ApprovalGateway, DetectionLayer, GraphStore, LLMCurationLayer
+from graphrag_core.models import CurationReport, OntologySchema
+
+
+class CurationPipeline:
+    """Orchestrates the curation flow: detect -> (curate) -> (approve)."""
+
+    def __init__(
+        self,
+        detection: DetectionLayer,
+        llm_curation: LLMCurationLayer | None = None,
+        approval: ApprovalGateway | None = None,
+    ) -> None:
+        self._detection = detection
+        self._llm_curation = llm_curation
+        self._approval = approval
+
+    async def run(
+        self,
+        graph_store: GraphStore,
+        schema: OntologySchema,
+    ) -> CurationReport:
+        issues = await self._detection.detect(graph_store, schema)
+
+        if self._llm_curation is not None:
+            issues = await self._llm_curation.curate(issues)
+
+        nodes = await graph_store.list_nodes()
+        rel_count = await graph_store.count_relationships()
+
+        return CurationReport(
+            issues=issues,
+            nodes_scanned=len(nodes),
+            relationships_scanned=rel_count,
+        )

graphrag_core/extraction/__init__.py

@@ -0,0 +1,5 @@
+"""BB2: Schema-guided entity extraction."""
+
+from graphrag_core.extraction.engine import LLMExtractionEngine
+
+__all__ = ["LLMExtractionEngine"]

graphrag_core/extraction/engine.py

@@ -0,0 +1,154 @@
+"""BB2: LLM-powered schema-guided entity extraction engine."""
+
+from __future__ import annotations
+
+import json
+
+from graphrag_core.interfaces import LLMClient
+from graphrag_core.models import (
+    DocumentChunk,
+    ExtractedNode,
+    ExtractedRelationship,
+    ExtractionResult,
+    ImportRun,
+    OntologySchema,
+    ProvenanceLink,
+)
+
+
+class LLMExtractionEngine:
+    """Extracts entities and relationships from text using an LLM, guided by an ontology schema."""
+
+    def __init__(self, llm_client: LLMClient) -> None:
+        self._llm = llm_client
+
+    async def extract(
+        self,
+        chunks: list[DocumentChunk],
+        schema: OntologySchema,
+        import_run: ImportRun,
+    ) -> ExtractionResult:
+        all_nodes: list[ExtractedNode] = []
+        all_rels: list[ExtractedRelationship] = []
+        all_provenance: list[ProvenanceLink] = []
+
+        system_prompt = self._build_system_prompt(schema)
+
+        for chunk in chunks:
+            nodes, rels = await self._extract_chunk(chunk, system_prompt)
+            nodes, rels = self._validate(nodes, rels, schema)
+
+            for node in nodes:
+                all_provenance.append(
+                    ProvenanceLink(chunk_id=chunk.id, node_id=node.id, confidence=1.0)
+                )
+
+            all_nodes.extend(nodes)
+            all_rels.extend(rels)
+
+        return ExtractionResult(
+            nodes=all_nodes,
+            relationships=all_rels,
+            provenance=all_provenance,
+        )
+
+    async def _extract_chunk(
+        self, chunk: DocumentChunk, system_prompt: str
+    ) -> tuple[list[ExtractedNode], list[ExtractedRelationship]]:
+        response = await self._llm.complete(
+            messages=[{"role": "user", "content": chunk.text}],
+            system=system_prompt,
+            temperature=0.0,
+        )
+        return self._parse_response(response)
+
+    def _build_system_prompt(self, schema: OntologySchema) -> str:
+        node_descriptions = []
+        for nt in schema.node_types:
+            props = ", ".join(
+                f"{p.name} ({p.type}{', required' if p.required else ''})"
+                for p in nt.properties
+            )
+            node_descriptions.append(f"- {nt.label}: properties=[{props}]")
+
+        rel_descriptions = []
+        for rt in schema.relationship_types:
+            rel_descriptions.append(
+                f"- {rt.type}: {rt.source_types} -> {rt.target_types}"
+            )
+
+        return (
+            "You are an entity extraction engine. Extract entities and relationships "
+            "from the provided text according to this schema.\n\n"
+            "ALLOWED NODE TYPES:\n"
+            + "\n".join(node_descriptions)
+            + "\n\nALLOWED RELATIONSHIP TYPES:\n"
+            + "\n".join(rel_descriptions)
+            + "\n\nDo not extract entities or relationships not listed above.\n\n"
+            "Respond with ONLY a JSON object in this exact format:\n"
+            '{"nodes": [{"id": "<unique_id>", "label": "<NodeType>", "properties": {<key>: <value>}}], '
+            '"relationships": [{"source_id": "<node_id>", "target_id": "<node_id>", "type": "<RelType>", "properties": {}}]}\n\n'
+            "Rules:\n"
+            "- Every node id must be unique and descriptive (e.g., 'person-alice', 'company-acme')\n"
+            "- Only use node types and relationship types listed above\n"
+            "- Include all required properties for each node type\n"
+            "- Return empty arrays if no entities are found"
+        )
+
+    def _parse_response(
+        self, response: str
+    ) -> tuple[list[ExtractedNode], list[ExtractedRelationship]]:
+        data = json.loads(response)
+
+        nodes = [
+            ExtractedNode(
+                id=n["id"],
+                label=n["label"],
+                properties=n.get("properties", {}),
+            )
+            for n in data.get("nodes", [])
+        ]
+
+        rels = [
+            ExtractedRelationship(
+                source_id=r["source_id"],
+                target_id=r["target_id"],
+                type=r["type"],
+                properties=r.get("properties", {}),
+            )
+            for r in data.get("relationships", [])
+        ]
+
+        return nodes, rels
+
+    def _validate(
+        self,
+        nodes: list[ExtractedNode],
+        rels: list[ExtractedRelationship],
+        schema: OntologySchema,
+    ) -> tuple[list[ExtractedNode], list[ExtractedRelationship]]:
+        allowed_labels = {nt.label for nt in schema.node_types}
+        allowed_rel_types = {rt.type for rt in schema.relationship_types}
+        rel_constraints = {
+            rt.type: (set(rt.source_types), set(rt.target_types))
+            for rt in schema.relationship_types
+        }
+
+        valid_nodes = [n for n in nodes if n.label in allowed_labels]
+        valid_node_ids = {n.id for n in valid_nodes}
+        node_labels = {n.id: n.label for n in valid_nodes}
+
+        valid_rels = []
+        for rel in rels:
+            if rel.type not in allowed_rel_types:
+                continue
+            if rel.source_id not in valid_node_ids or rel.target_id not in valid_node_ids:
+                continue
+            source_types, target_types = rel_constraints[rel.type]
+            if node_labels[rel.source_id] not in source_types:
+                continue
+            if node_labels[rel.target_id] not in target_types:
+                continue
+            valid_rels.append(rel)
+
+        return valid_nodes, valid_rels

graphrag_core/graph/__init__.py

@@ -0,0 +1,11 @@
+"""BB3: Graph store implementations."""
+
+from graphrag_core.graph.memory import InMemoryGraphStore
+
+__all__ = ["InMemoryGraphStore"]
+
+try:
+    from graphrag_core.graph.neo4j import Neo4jGraphStore
+    __all__.append("Neo4jGraphStore")
+except ImportError:
+    pass