agentic-graphrag 0.0.3__tar.gz → 0.0.4__tar.gz

{agentic_graphrag-0.0.3 → agentic_graphrag-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentic-graphrag
-Version: 0.0.3
+Version: 0.0.4
 Summary: Agentic GraphRAG: graph-based retrieval-augmented generation with agentic reasoning.
 Keywords: graphrag,rag,knowledge-graph,agents,retrieval
 Author: Ashwin Mathur, Varun Mathur
@@ -41,7 +41,7 @@ Requires-Dist: transformers<4.56 ; extra == 'extract'
 Requires-Dist: protobuf ; extra == 'extract'
 Requires-Dist: baml-py>=0.226.1,<0.227 ; extra == 'llm'
 Requires-Dist: pymilvus>=3.0.0,<4 ; extra == 'milvus'
-Requires-Dist: neo4j>=5.0 ; extra == 'neo4j'
+Requires-Dist: neo4j>=5.22 ; extra == 'neo4j'
 Requires-Dist: opentelemetry-sdk>=1.20 ; extra == 'observability'
 Requires-Dist: opentelemetry-exporter-otlp>=1.20 ; extra == 'observability'
 Requires-Dist: opentelemetry-instrumentation-asyncio ; extra == 'observability'

agentic_graphrag-0.0.4/agrag/agents/__init__.py

@@ -0,0 +1,19 @@
+"""Agentic layer: planner/researcher/verifier over SearchEngine.
+
+``build_agent`` lives in ``agrag.agents.build`` and is not re-exported here: it needs
+the ``agents`` extra to import, and this package must import on a base install.
+"""
+
+from agrag.agents.errors import AgentMissingExtraError
+from agrag.agents.ledger import Ledger
+from agrag.agents.result import AgentRunResult
+from agrag.agents.settings import AgentLLMSettings, AgentSettings
+
+
+__all__ = [
+    "AgentLLMSettings",
+    "AgentMissingExtraError",
+    "AgentRunResult",
+    "AgentSettings",
+    "Ledger",
+]

{agentic_graphrag-0.0.3 → agentic_graphrag-0.0.4}/agrag/agents/settings.py

@@ -30,7 +30,10 @@ class AgentLLMSettings(BaseSettings):
     """
 
     model_config = SettingsConfigDict(
-        env_prefix="AGENT_LLM_", env_file=".env", extra="ignore"
+        env_prefix="AGENT_LLM_",
+        env_file=".env",
+        extra="ignore",
+        hide_input_in_errors=True,
     )
 
     clients: list[LLMClientConfig]
@@ -99,7 +102,7 @@ class AgentSettings(BaseSettings):
     """
 
     model_config = SettingsConfigDict(
-        env_prefix="AGENT_", env_file=".env", extra="ignore"
+        env_prefix="AGENT_", env_file=".env", extra="ignore", hide_input_in_errors=True
     )
 
     recursion_limit: int = 50

agentic_graphrag-0.0.4/agrag/common/data_models/__init__.py

@@ -0,0 +1,43 @@
+"""Shared data models used by agrag components."""
+
+from agrag.common.data_models.chunk import Chunk
+from agrag.common.data_models.community import Community
+from agrag.common.data_models.document import Document, DocumentFamily, SourceFormat
+from agrag.common.data_models.entity import Entity
+from agrag.common.data_models.graph_record import NodeRecord, RelationRecord
+from agrag.common.data_models.graph_schema import (
+    GENERIC,
+    EntityType,
+    GraphSchema,
+    RelationType,
+)
+from agrag.common.data_models.normalization import Normalization
+from agrag.common.data_models.provenance import PageProvenance, TextProvenance
+from agrag.common.data_models.relation import Relation
+from agrag.common.data_models.resolved_entity import ResolvedEntity
+from agrag.common.data_models.search_result import SearchResult
+from agrag.common.data_models.vector_record import Distance, VectorRecord
+
+
+__all__ = [
+    "GENERIC",
+    "Chunk",
+    "Community",
+    "Distance",
+    "Document",
+    "DocumentFamily",
+    "Entity",
+    "EntityType",
+    "GraphSchema",
+    "NodeRecord",
+    "Normalization",
+    "PageProvenance",
+    "Relation",
+    "RelationRecord",
+    "RelationType",
+    "ResolvedEntity",
+    "SearchResult",
+    "SourceFormat",
+    "TextProvenance",
+    "VectorRecord",
+]

{agentic_graphrag-0.0.3 → agentic_graphrag-0.0.4}/agrag/cypher/entities.py

@@ -378,8 +378,7 @@ def fetch_entity_neighbors_query() -> str:
     """
     return (
         "UNWIND $ids AS entity_id "
-        "CALL { "
-        "WITH entity_id "
+        "CALL (entity_id) { "
         "MATCH (n:"
         f"{NODE_IDENTITY_LABEL} {{id: entity_id}})-[r]-(m:{NODE_IDENTITY_LABEL}) "
         "WHERE NOT type(r) IN $exclude_types "

{agentic_graphrag-0.0.3 → agentic_graphrag-0.0.4}/agrag/cypher/schema.py

@@ -249,6 +249,26 @@ def vector_index_query(
     )
 
 
+def vector_index_dimensions_query(
+    label: str, vector_property: str
+) -> tuple[str, dict[str, Any]]:
+    """Build a query that reads the dimension of an existing vector index.
+
+    Args:
+        label: The node label. Must already be validated.
+        vector_property: The vector property name. Must already be validated.
+
+    Returns:
+        The query and its parameters. The query returns one ``dimensions`` row
+        when the index exists and no rows when it does not.
+    """
+    return (
+        "SHOW INDEXES YIELD name, options WHERE name = $name "
+        "RETURN options.indexConfig['vector.dimensions'] AS dimensions",
+        {"name": vector_index_name(label, vector_property)},
+    )
+
+
 def vector_search_query(
     index_name: str, filters: dict[str, Any] | None = None
 ) -> tuple[str, dict[str, Any]]:

{agentic_graphrag-0.0.3 → agentic_graphrag-0.0.4}/agrag/embedding/settings.py

@@ -22,7 +22,10 @@ class EmbeddingSettings(BaseSettings):
     """
 
     model_config = SettingsConfigDict(
-        env_prefix="EMBEDDING_", env_file=".env", extra="ignore"
+        env_prefix="EMBEDDING_",
+        env_file=".env",
+        extra="ignore",
+        hide_input_in_errors=True,
     )
 
     model: str = "ibm-granite/granite-embedding-small-english-r2"

{agentic_graphrag-0.0.3 → agentic_graphrag-0.0.4}/agrag/eval/settings.py

@@ -19,7 +19,9 @@ class _JudgeEnvVariables(BaseSettings):
     below.
     """
 
-    model_config = SettingsConfigDict(env_file=".env", extra="ignore")
+    model_config = SettingsConfigDict(
+        env_file=".env", extra="ignore", hide_input_in_errors=True
+    )
 
     eval_judge_base_url: str | None = Field(
         default=None, validation_alias="EVAL_JUDGE_BASE_URL"
@@ -64,7 +66,10 @@ class EvalJudgeSettings(BaseSettings):
     """
 
     model_config = SettingsConfigDict(
-        env_prefix="EVAL_JUDGE_", env_file=".env", extra="ignore"
+        env_prefix="EVAL_JUDGE_",
+        env_file=".env",
+        extra="ignore",
+        hide_input_in_errors=True,
     )
 
     client: LLMClientConfig

{agentic_graphrag-0.0.3 → agentic_graphrag-0.0.4}/agrag/graphdb/base.py

@@ -217,8 +217,13 @@ class GraphStore(ABC):
         Args:
             label: The node label to index.
             vector_property: The embedding property name.
-            dimensions: The embedding dimension.
+            dimensions: The embedding dimension. If the index already exists
+                with a different dimension, this raises.
             distance: The distance metric.
+
+        Raises:
+            EmbeddingDimensionMismatchError: The index already exists with a
+                different dimension than ``dimensions``.
         """
 
     @abstractmethod

{agentic_graphrag-0.0.3 → agentic_graphrag-0.0.4}/agrag/graphdb/neo4j.py

@@ -38,10 +38,12 @@ from agrag.cypher.schema import (
     node_id_constraint_query,
     relation_id_constraint_name,
     relation_id_constraint_query,
+    vector_index_dimensions_query,
     vector_index_name,
     vector_index_query,
     vector_search_query,
 )
+from agrag.embedding.errors import EmbeddingDimensionMismatchError
 from agrag.graphdb.base import GraphStore, GraphStoreTransaction
 from agrag.graphdb.errors import (
     GraphStoreConstraintViolationError,
@@ -357,7 +359,10 @@ class Neo4jGraphStore(GraphStore):
                     # Lazy import: a clean install must raise
                     # GraphStoreMissingExtraError, not ImportError, when neo4j is
                     # absent.
-                    from neo4j import AsyncGraphDatabase  # noqa: PLC0415
+                    from neo4j import (  # noqa: PLC0415
+                        AsyncGraphDatabase,
+                        NotificationDisabledClassification,
+                    )
                 except ImportError as exc:
                     raise GraphStoreMissingExtraError("neo4j") from exc
                 self._driver = AsyncGraphDatabase.driver(
@@ -367,6 +372,14 @@ class Neo4jGraphStore(GraphStore):
                         self._settings.password.get_secret_value(),
                     ),
                     max_connection_lifetime=self._settings.max_connection_lifetime,
+                    # The store probes labels and property keys that a new
+                    # database does not have yet. The server reports each probe
+                    # as UNRECOGNIZED, which is expected noise. Other
+                    # classifications, such as PERFORMANCE and DEPRECATION,
+                    # still reach the neo4j.notifications logger.
+                    notifications_disabled_classifications=[
+                        NotificationDisabledClassification.UNRECOGNIZED
+                    ],
                 )
         return self._driver
 
@@ -861,6 +874,10 @@ class Neo4jGraphStore(GraphStore):
         A concurrent creator can commit the same index after this operation
         starts. Neo4j reports that race as an equivalent-schema error, which
         means the requested index already exists.
+
+        Raises:
+            EmbeddingDimensionMismatchError: The index already exists with a
+                different dimension than ``dimensions``.
         """
         from neo4j.exceptions import ClientError  # noqa: PLC0415
 
@@ -874,6 +891,16 @@ class Neo4jGraphStore(GraphStore):
         except ClientError as exc:
             if exc.code != _EQUIVALENT_SCHEMA_RULE_ERROR:
                 raise
+        # CREATE ... IF NOT EXISTS keeps an existing index of the same name as
+        # it is, so a wrong-size embedder would otherwise write vectors the
+        # index never returns.
+        rows = await self.execute_read(
+            *vector_index_dimensions_query(label, vector_property)
+        )
+        if rows and rows[0]["dimensions"] != dimensions:
+            raise EmbeddingDimensionMismatchError(
+                expected=rows[0]["dimensions"], actual=dimensions
+            )
 
     async def vector_search(
         self,

agentic_graphrag-0.0.4/agrag/graphdb/serialize.py

@@ -0,0 +1,227 @@
+"""Convert graph records to driver parameters and graph node rows to models."""
+
+import contextlib
+from collections.abc import Mapping
+from datetime import datetime
+from typing import Any
+from uuid import UUID
+
+from agrag.common.data_models.chunk import CHUNK_LABEL
+from agrag.common.data_models.entity import Entity
+from agrag.common.data_models.graph_record import (
+    PENDING_JOB_ID_PROPERTY,
+    NodeRecord,
+    RelationRecord,
+)
+from agrag.cypher.entities import NODE_IDENTITY_LABEL
+
+
+def _convert(value: Any) -> Any:
+    """Recursively convert a value to a Neo4j-driver-friendly type.
+
+    Neo4j's driver rejects ``UUID`` objects and other non-primitive types, so any
+    ``UUID`` becomes its string form and nested containers are walked the same
+    way.
+
+    Args:
+        value: The value to convert.
+
+    Returns:
+        The driver-safe equivalent of ``value``.
+    """
+    if isinstance(value, UUID):
+        return str(value)
+    if isinstance(value, Mapping):
+        return {k: _convert(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple, set)):
+        return [_convert(v) for v in value]
+    return value
+
+
+def node_params(record: NodeRecord) -> dict[str, Any]:
+    """Build the ``$records`` entry for a node upsert.
+
+    The Cutover Job tag is split out of ``properties`` into its own key
+    because the upsert queries apply it with ``ON CREATE SET``: a job tags
+    the nodes it creates, never a node it writes over. Left inside the
+    applied property map it would be set on existing nodes too, which
+    would hide a committed node from retrieval for the job's duration and
+    put it in reach of the job's rollback — and rollback deletes tagged
+    rows.
+
+    Args:
+        record: The node record to serialize.
+
+    Returns:
+        A dict with ``id`` (string), ``properties`` (converted, without
+        the pending tag), and ``pending_job_id`` (the tag, or None).
+    """
+    properties = dict(record.properties)
+    pending_job_id = properties.pop(PENDING_JOB_ID_PROPERTY, None)
+    return {
+        "id": str(record.id),
+        "properties": _convert(properties),
+        "pending_job_id": _convert(pending_job_id),
+    }
+
+
+def relation_params(record: RelationRecord) -> dict[str, Any]:
+    """Build the ``$records`` entry for a relationship upsert.
+
+    The Cutover Job tag is split out of ``properties`` for the same reason
+    as in :func:`node_params`: only an edge the job creates carries it.
+
+    Args:
+        record: The relation record to serialize.
+
+    Returns:
+        A dict with ``id``, ``start_id``, ``end_id``, ``properties``
+        (converted, without the pending tag), and ``pending_job_id`` (the
+        tag, or None).
+    """
+    properties = dict(record.properties)
+    pending_job_id = properties.pop(PENDING_JOB_ID_PROPERTY, None)
+    return {
+        "id": str(record.id),
+        "start_id": str(record.start_id),
+        "end_id": str(record.end_id),
+        "properties": _convert(properties),
+        "pending_job_id": _convert(pending_job_id),
+    }
+
+
+def parse_entity_node(node: object) -> Entity | None:  # noqa: PLR0912,PLR0915
+    """Parse a GraphStore node row into an Entity.
+
+    Handles both neo4j Node objects and plain dict mocks used in unit tests.
+    """
+    try:
+        props: dict = {}
+        labels: list[str] = []
+        node_id: object = None
+
+        if isinstance(node, dict) and "labels" in node and "properties" in node:
+            # Mock form: {"id": "...", "labels": [...], "properties": {...}}
+            labels = list(node.get("labels") or [])
+            props = dict(node.get("properties") or {})
+            node_id = node.get("id") or props.get("id")
+        elif isinstance(node, dict) and "id" in node:
+            # Flat mock where properties are top-level alongside id/labels
+            # e.g. {"n": {"id": "...", "name": "...", "merge_key": "..."}}
+            # But here node is that inner dict.
+            props = dict(node)
+            # id may be in props
+            node_id = props.get("id")
+            # labels might be in props or separate
+            maybe_labels = props.pop("labels", None)
+            if isinstance(maybe_labels, list):
+                labels = maybe_labels
+            # Try to get labels from node dict if present alongside
+            if not labels and "labels" in props:
+                labels = props.pop("labels")  # type: ignore[assignment]
+        else:
+            # Attempt neo4j Node: dict(node) gives properties, node.labels gives labels
+            try:
+                props = dict(node)  # ty: ignore[no-matching-overload]  # type: ignore[arg-type]
+            except Exception:
+                props = {}
+            try:
+                maybe_labels = getattr(node, "labels", None)
+                if maybe_labels is not None:
+                    labels = list(maybe_labels)  # type: ignore[arg-type]
+            except Exception:
+                labels = []
+            # Try id from props or node["id"]
+            try:
+                node_id = props.get("id")  # type: ignore[union-attr]
+            except Exception:
+                node_id = None
+            if node_id is None:
+                with contextlib.suppress(Exception):
+                    node_id = node["id"]  # type: ignore[index]  # ty: ignore[not-subscriptable]
+            # Node may also be wrapped as {"n": Node}
+            if isinstance(node, dict) and "n" in node:
+                return parse_entity_node(node["n"])
+
+        if node_id is None:
+            # Fallback: id inside props
+            node_id = props.get("id")
+
+        # Determine label
+        label: str | None = None
+        if labels:
+            for lbl in labels:
+                if lbl not in (NODE_IDENTITY_LABEL, CHUNK_LABEL):
+                    label = str(lbl)
+                    break
+            if label is None:
+                # All labels were system labels, take first
+                label = str(labels[0]) if labels else None
+        if label is None:
+            # Fallback from merge_key
+            mk = props.get("merge_key") or ""
+            if isinstance(mk, str) and ":" in mk:
+                label = mk.split(":", 1)[0]
+
+        if label is None or node_id is None:
+            return None
+
+        # System keys not part of domain properties
+        system_keys = {
+            "name",
+            "merge_key",
+            "merged_from",
+            "merge_count",
+            "source_chunk_ids",
+            "created_at",
+            "embedding",
+            "id",
+            PENDING_JOB_ID_PROPERTY,
+        }
+        entity_props = {k: v for k, v in props.items() if k not in system_keys}
+
+        merged_from_raw = props.get("merged_from") or []
+        merged_from = [UUID(str(x)) for x in merged_from_raw if x]
+
+        try:
+            merge_count = int(props.get("merge_count", 1))
+        except Exception:
+            merge_count = 1
+
+        scids_raw = props.get("source_chunk_ids") or []
+        scids = [UUID(str(x)) for x in scids_raw if x]
+
+        embedding = props.get("embedding")
+
+        created_at_raw = props.get("created_at")
+        created_at = None
+        if isinstance(created_at_raw, str):
+            try:
+                created_at = datetime.fromisoformat(created_at_raw)
+            except Exception:
+                created_at = None
+
+        name_val = props.get("name")
+        if name_val is None:
+            mk = props.get("merge_key", "")
+            if isinstance(mk, str) and ":" in mk:
+                name_val = mk.split(":", 1)[1]
+            else:
+                name_val = mk or ""
+
+        kwargs: dict = {
+            "id": UUID(str(node_id)),
+            "label": label,
+            "name": str(name_val),
+            "properties": entity_props,
+            "merged_from": merged_from,
+            "merge_count": merge_count,
+            "source_chunk_ids": scids,
+        }
+        if embedding is not None:
+            kwargs["embedding"] = list(embedding)  # type: ignore[arg-type]
+        if created_at is not None:
+            kwargs["created_at"] = created_at
+        return Entity(**kwargs)
+    except Exception:
+        return None

{agentic_graphrag-0.0.3 → agentic_graphrag-0.0.4}/agrag/graphdb/settings.py

@@ -27,7 +27,7 @@ class Neo4jSettings(BaseSettings):
     """
 
     model_config = SettingsConfigDict(
-        env_prefix="NEO4J_", env_file=".env", extra="ignore"
+        env_prefix="NEO4J_", env_file=".env", extra="ignore", hide_input_in_errors=True
     )
 
     uri: str = "bolt://localhost:7687"

agentic_graphrag-0.0.4/agrag/ingestion/__init__.py

@@ -0,0 +1,34 @@
+"""The ingestion package."""
+
+from agrag.ingestion.extract import (
+    BAMLExtractor,
+    EscalatingExtractor,
+    ExtractionLLMSettings,
+    Extractor,
+    ExtractorMissingExtraError,
+    GlinerExtractor,
+)
+from agrag.ingestion.graph import Graph
+from agrag.ingestion.reports import (
+    AddResult,
+    CommunityDetectionReport,
+    ConsolidationReport,
+    ReevaluationReport,
+    UpdateResult,
+)
+
+
+__all__ = [
+    "AddResult",
+    "BAMLExtractor",
+    "CommunityDetectionReport",
+    "ConsolidationReport",
+    "EscalatingExtractor",
+    "ExtractionLLMSettings",
+    "Extractor",
+    "ExtractorMissingExtraError",
+    "GlinerExtractor",
+    "Graph",
+    "ReevaluationReport",
+    "UpdateResult",
+]