pretensor 0.1.0__py3-none-any.whl

pretensor/graph_models/node.py

@@ -0,0 +1,79 @@
+"""Graph node payload used when writing to Kuzu."""
+
+from __future__ import annotations
+
+from pydantic import Field
+
+from pretensor.graph_models.base import GraphModel
+
+__all__ = ["GraphNode"]
+
+
+class GraphNode(GraphModel):
+    """A table (or future entity type) represented as a node in the graph."""
+
+    node_id: str = Field(
+        description="Stable primary key in the Kuzu SchemaTable node table."
+    )
+    connection_name: str
+    database: str
+    schema_name: str
+    table_name: str
+    row_count: int | None = None
+    row_count_source: str | None = Field(
+        default=None,
+        description=(
+            "Provenance: 'stat', 'view_count', or 'view_timeout' (row_count = -1)."
+        ),
+    )
+    comment: str | None = None
+    description: str | None = Field(
+        default=None,
+        description="Human-authored description (e.g. dbt); distinct from DB comment.",
+    )
+    tags: list[str] | None = Field(
+        default=None,
+        description="Tag names merged onto the graph node (e.g. dbt tags).",
+    )
+    table_type: str | None = None
+    entity_type: str | None = Field(
+        default=None,
+        description="Business entity label from extraction (e.g. Customer).",
+    )
+    seq_scan_count: int | None = None
+    idx_scan_count: int | None = None
+    insert_count: int | None = None
+    update_count: int | None = None
+    delete_count: int | None = None
+    is_partitioned: bool | None = None
+    partition_key: str | None = None
+    grants_json: str | None = None
+    access_read_count: int | None = None
+    access_write_count: int | None = None
+    days_since_last_access: int | None = None
+    potentially_unused: bool | None = None
+    table_bytes: int | None = None
+    clustering_key: str | None = None
+    role: str | None = Field(
+        default=None,
+        description="Data modeling role from schema classifier (fact, dimension, …).",
+    )
+    role_confidence: float | None = Field(
+        default=None, description="Classifier confidence in ``[0, 1]``."
+    )
+    classification_signals: str | None = Field(
+        default=None,
+        description="JSON array string of human-readable classification signals.",
+    )
+    has_external_consumers: bool | None = Field(
+        default=None,
+        description="True when dbt lineage marks downstream exposure outside the project.",
+    )
+    test_count: int | None = Field(
+        default=None,
+        description="Number of dbt tests attached to this model (when known).",
+    )
+    embedding: list[float] | None = Field(
+        default=None,
+        description="Dense vector embedding of table identity text (dim=384).",
+    )

pretensor/graph_models/relationship.py

@@ -0,0 +1,33 @@
+"""Relationship discovery candidates (heuristic, LLM, statistical, explicit FK)."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import Field
+
+from pretensor.graph_models.base import GraphModel
+
+__all__ = ["RelationshipCandidate", "RelationshipSource"]
+
+RelationshipSource = Literal[
+    "explicit_fk",
+    "heuristic",
+    "llm_inferred",
+    "statistical",
+    "embedding",
+]
+
+
+class RelationshipCandidate(GraphModel):
+    """A directed join hypothesis between two ``SchemaTable`` nodes."""
+
+    candidate_id: str
+    source_node_id: str
+    target_node_id: str
+    source_column: str
+    target_column: str
+    source: RelationshipSource
+    reasoning: str = ""
+    confidence: float = Field(ge=0.0, le=1.0, default=0.0)
+    status: Literal["suggested", "confirmed", "rejected"] = "suggested"

pretensor/integrations/__init__.py

@@ -0,0 +1,42 @@
+"""Agent-framework SDK adapters for Pretensor graph tools.
+
+Exposes the six graph tools (schema, context, traverse, impact, query,
+validate_sql) as native tool objects for LangChain, LlamaIndex, and Google ADK.
+No MCP server process is required.
+
+Usage::
+
+    from pretensor.integrations import load_langchain_tools
+    from pretensor.integrations import load_llamaindex_tools
+    from pretensor.integrations import load_adk_tools
+"""
+
+# pyright: reportUnsupportedDunderAll=false
+# Names in __all__ are provided via __getattr__; Pyright does not model that.
+
+from __future__ import annotations
+
+from typing import Any
+
+__all__ = [
+    "load_langchain_tools",
+    "load_llamaindex_tools",
+    "load_adk_tools",
+]
+
+
+def __getattr__(name: str) -> Any:
+    if name == "load_langchain_tools":
+        from pretensor.integrations.langchain import load_langchain_tools
+
+        return load_langchain_tools
+    if name == "load_llamaindex_tools":
+        from pretensor.integrations.llamaindex import load_llamaindex_tools
+
+        return load_llamaindex_tools
+    if name == "load_adk_tools":
+        from pretensor.integrations.google_adk import load_adk_tools
+
+        return load_adk_tools
+    msg = f"module {__name__!r} has no attribute {name!r}"
+    raise AttributeError(msg)

pretensor/integrations/_base.py

@@ -0,0 +1,138 @@
+"""Shared tool-function factory for all agent-framework adapters.
+
+Each returned callable has a typed signature and docstring so that
+LangChain, LlamaIndex, and Google ADK can automatically infer the
+tool schema without any manual JSON Schema duplication.
+
+The functions delegate directly to the same payload functions used by
+the MCP server — no logic is duplicated.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Callable
+
+from pretensor.mcp.tools.context import context_payload
+from pretensor.mcp.tools.impact import impact_payload
+from pretensor.mcp.tools.schema import schema_payload
+from pretensor.mcp.tools.search import query_payload
+from pretensor.mcp.tools.traverse import traverse_payload
+from pretensor.mcp.tools.validate_sql import validate_sql_payload
+
+__all__ = ["make_tool_functions"]
+
+
+def _make_schema(graph_dir: Path) -> Callable[..., dict[str, Any]]:
+    def schema(database: str, label: str | None = None) -> dict[str, Any]:
+        """Discover node labels, edge types, and properties in the graph. Call before writing Cypher."""
+        return schema_payload(graph_dir, database=database, label=label)
+
+    return schema
+
+
+def _make_context(graph_dir: Path) -> Callable[..., dict[str, Any]]:
+    def context(
+        table: str,
+        db: str | None = None,
+        detail: str = "standard",
+    ) -> dict[str, Any]:
+        """Full context for one physical table: columns, relationships, lineage, cluster. detail: summary | standard | full."""
+        return context_payload(graph_dir, table=table, db=db, detail=detail)  # type: ignore[arg-type]
+
+    return context
+
+
+def _make_traverse(graph_dir: Path) -> Callable[..., dict[str, Any]]:
+    def traverse(
+        from_table: str,
+        to_table: str,
+        database: str,
+        max_depth: int = 4,
+        top_k: int = 3,
+        edge_types: list[str] | None = None,
+        max_inferred_hops: int = 2,
+    ) -> dict[str, Any]:
+        """Find a join path between two physical tables with SQL JOIN hints. edge_types: fk | inferred (default: both)."""
+        et: tuple[str, ...] | None = None
+        if edge_types:
+            et = tuple(edge_types)
+        return traverse_payload(
+            graph_dir,
+            from_table=from_table,
+            to_table=to_table,
+            database=database,
+            max_depth=max_depth,
+            top_k=max(1, min(10, top_k)),
+            edge_types=et,
+            max_inferred_hops=max(0, min(8, max_inferred_hops)),
+        )
+
+    return traverse
+
+
+def _make_impact(graph_dir: Path) -> Callable[..., dict[str, Any]]:
+    def impact(
+        table: str,
+        database: str,
+        column: str | None = None,
+        max_depth: int = 3,
+    ) -> dict[str, Any]:
+        """Downstream tables reachable via FK and inferred join edges from a table, grouped by hop depth."""
+        return impact_payload(
+            graph_dir,
+            table=table,
+            database=database,
+            column=column,
+            max_depth=max_depth,
+        )
+
+    return impact
+
+
+def _make_query(graph_dir: Path) -> Callable[..., dict[str, Any]]:
+    def query(
+        q: str,
+        db: str | None = None,
+        limit: int = 10,
+    ) -> dict[str, Any]:
+        """BM25 keyword search over table and entity metadata."""
+        return query_payload(graph_dir, q=q, db=db, limit=limit)
+
+    return query
+
+
+def _make_validate_sql(graph_dir: Path) -> Callable[..., dict[str, Any]]:
+    def validate_sql(
+        sql: str,
+        database: str,
+        dialect: str = "postgres",
+    ) -> dict[str, Any]:
+        """Validate a SQL statement against the indexed graph: unknown tables/columns, invalid joins, fuzzy suggestions."""
+        return validate_sql_payload(graph_dir, sql=sql, database=database, dialect=dialect)
+
+    return validate_sql
+
+
+def make_tool_functions(graph_dir: Path) -> list[Callable[..., dict[str, Any]]]:
+    """Return six graph-tool callables bound to ``graph_dir``.
+
+    Each callable has typed parameters and a docstring so that LangChain,
+    LlamaIndex, and Google ADK can automatically derive the tool schema.
+
+    Args:
+        graph_dir: Root directory of the Pretensor graph store (must be
+            pre-indexed via ``pretensor index``).
+
+    Returns:
+        [schema, context, traverse, impact, query, validate_sql]
+    """
+    resolved = Path(graph_dir).resolve()
+    return [
+        _make_schema(resolved),
+        _make_context(resolved),
+        _make_traverse(resolved),
+        _make_impact(resolved),
+        _make_query(resolved),
+        _make_validate_sql(resolved),
+    ]

pretensor/integrations/google_adk.py

@@ -0,0 +1,49 @@
+"""Google ADK adapter — wraps Pretensor graph tools as ADK ``FunctionTool`` objects.
+
+Install the optional extra before use::
+
+    pip install 'pretensor[google-adk]'
+
+Usage::
+
+    from pathlib import Path
+    from pretensor.integrations.google_adk import load_adk_tools
+    from google.adk.agents import LlmAgent
+
+    tools = load_adk_tools(Path(".pretensor"))
+    agent = LlmAgent(model="gemini-2.0-flash", tools=tools)
+"""
+
+from __future__ import annotations
+
+import importlib
+from pathlib import Path
+from typing import Any
+
+from pretensor.integrations._base import make_tool_functions
+
+__all__ = ["load_adk_tools"]
+
+
+def load_adk_tools(graph_dir: Path | str) -> list[Any]:
+    """Return Pretensor graph tools as Google ADK ``FunctionTool`` objects.
+
+    Args:
+        graph_dir: Root directory of the Pretensor graph store.
+
+    Returns:
+        List of six ``google.adk.tools.FunctionTool`` instances:
+        schema, context, traverse, impact, query, validate_sql.
+
+    Raises:
+        ImportError: If ``google-adk`` is not installed.
+    """
+    try:
+        adk_tools = importlib.import_module("google.adk.tools")
+    except ImportError as exc:
+        raise ImportError(
+            "Google ADK is not installed. Run: pip install 'pretensor[google-adk]'"
+        ) from exc
+
+    function_tool = adk_tools.FunctionTool
+    return [function_tool(fn) for fn in make_tool_functions(Path(graph_dir))]

pretensor/integrations/langchain.py

@@ -0,0 +1,55 @@
+"""LangChain adapter — wraps Pretensor graph tools as ``StructuredTool`` objects.
+
+Install the optional extra before use::
+
+    pip install 'pretensor[langchain]'
+
+Usage::
+
+    from pathlib import Path
+    from pretensor.integrations.langchain import load_langchain_tools
+    from langchain_openai import ChatOpenAI
+    from langchain.agents import create_tool_calling_agent, AgentExecutor
+
+    tools = load_langchain_tools(Path(".pretensor"))
+    llm = ChatOpenAI(model="gpt-4o")
+    agent = AgentExecutor(
+        agent=create_tool_calling_agent(llm, tools, prompt), tools=tools
+    )
+"""
+
+from __future__ import annotations
+
+import importlib
+from pathlib import Path
+from typing import Any
+
+from pretensor.integrations._base import make_tool_functions
+
+__all__ = ["load_langchain_tools"]
+
+
+def load_langchain_tools(graph_dir: Path | str) -> list[Any]:
+    """Return Pretensor graph tools as LangChain ``StructuredTool`` objects.
+
+    Args:
+        graph_dir: Root directory of the Pretensor graph store.
+
+    Returns:
+        List of six ``langchain_core.tools.StructuredTool`` instances:
+        schema, context, traverse, impact, query, validate_sql.
+
+    Raises:
+        ImportError: If ``langchain-core`` is not installed.
+    """
+    try:
+        lc_tools = importlib.import_module("langchain_core.tools")
+    except ImportError as exc:
+        raise ImportError(
+            "LangChain is not installed. Run: pip install 'pretensor[langchain]'"
+        ) from exc
+
+    structured_tool = lc_tools.StructuredTool
+    return [
+        structured_tool.from_function(fn) for fn in make_tool_functions(Path(graph_dir))
+    ]

pretensor/integrations/llamaindex.py

@@ -0,0 +1,53 @@
+"""LlamaIndex adapter — wraps Pretensor graph tools as ``FunctionTool`` objects.
+
+Install the optional extra before use::
+
+    pip install 'pretensor[llama-index]'
+
+Usage::
+
+    from pathlib import Path
+    from pretensor.integrations.llamaindex import load_llamaindex_tools
+    from llama_index.core.agent import ReActAgent
+    from llama_index.llms.openai import OpenAI
+
+    tools = load_llamaindex_tools(Path(".pretensor"))
+    agent = ReActAgent.from_tools(tools, llm=OpenAI(model="gpt-4o"), verbose=True)
+"""
+
+from __future__ import annotations
+
+import importlib
+from pathlib import Path
+from typing import Any
+
+from pretensor.integrations._base import make_tool_functions
+
+__all__ = ["load_llamaindex_tools"]
+
+
+def load_llamaindex_tools(graph_dir: Path | str) -> list[Any]:
+    """Return Pretensor graph tools as LlamaIndex ``FunctionTool`` objects.
+
+    Args:
+        graph_dir: Root directory of the Pretensor graph store.
+
+    Returns:
+        List of six ``llama_index.core.tools.FunctionTool`` instances:
+        schema, context, traverse, impact, query, validate_sql.
+
+    Raises:
+        ImportError: If ``llama-index-core`` is not installed.
+    """
+    try:
+        li_tools = importlib.import_module("llama_index.core.tools")
+    except ImportError as exc:
+        raise ImportError(
+            "LlamaIndex is not installed. Run: pip install 'pretensor[llama-index]'"
+        ) from exc
+
+    function_tool = li_tools.FunctionTool
+    return [
+        function_tool.from_defaults(fn=fn)
+        for fn in make_tool_functions(Path(graph_dir))
+    ]

pretensor/intelligence/__init__.py

@@ -0,0 +1,33 @@
+"""Graph intelligence: relationship discovery, clustering, and join paths."""
+
+from pretensor.intelligence.cluster_labeler import ClusterLabeler, LabeledCluster
+from pretensor.intelligence.clustering import Cluster, ClusteringEngine
+from pretensor.intelligence.graph_export import GraphExporter
+from pretensor.intelligence.join_paths import JoinPathEngine
+from pretensor.intelligence.pipeline import (
+    build_oss_pipeline,
+    run_intelligence_layer,
+    run_intelligence_layer_sync,
+)
+from pretensor.intelligence.steps import (
+    CyclicDependencyError,
+    PipelineContext,
+    PipelineRunner,
+    PipelineStep,
+)
+
+__all__ = [
+    "Cluster",
+    "ClusteringEngine",
+    "ClusterLabeler",
+    "LabeledCluster",
+    "GraphExporter",
+    "JoinPathEngine",
+    "build_oss_pipeline",
+    "run_intelligence_layer",
+    "run_intelligence_layer_sync",
+    "CyclicDependencyError",
+    "PipelineContext",
+    "PipelineRunner",
+    "PipelineStep",
+]