reasongraph 0.2.0__py3-none-any.whl

reasongraph/__init__.py

@@ -0,0 +1,21 @@
+__version__ = "0.2.0"
+
+import os as _os
+
+# Suppress noisy transformers/tqdm output (load reports, progress bars, pipeline hints)
+# unless the user has explicitly configured verbosity.
+if "TRANSFORMERS_VERBOSITY" not in _os.environ:
+    _os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+if "TQDM_DISABLE" not in _os.environ:
+    _os.environ["TQDM_DISABLE"] = "1"
+
+from reasongraph._types import Node, Edge
+from reasongraph._extraction import NERExtractor, GLiNER2Extractor
+from reasongraph.graph import ReasonGraph
+from reasongraph.datasets import load_dataset
+
+__all__ = [
+    "ReasonGraph", "Node", "Edge",
+    "NERExtractor", "GLiNER2Extractor",
+    "load_dataset",
+]

reasongraph/_embeddings.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from sentence_transformers import SentenceTransformer, CrossEncoder
+
+
+class EmbeddingManager:
+    """Manages embedding generation and cross-encoder reranking."""
+
+    DEFAULT_EMBED_MODEL = "all-MiniLM-L12-v2"
+    DEFAULT_RERANK_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+
+    def __init__(
+        self,
+        embed_model: str | SentenceTransformer | None = None,
+        rerank_model: str | CrossEncoder | None = None,
+    ) -> None:
+        if isinstance(embed_model, SentenceTransformer):
+            self._embed = embed_model
+        else:
+            self._embed = SentenceTransformer(embed_model or self.DEFAULT_EMBED_MODEL)
+
+        self._rerank: CrossEncoder | None = None
+        self._rerank_name = rerank_model
+
+    def encode(self, text: str) -> list[float]:
+        """Encode a single text string to a float vector."""
+        return self._embed.encode(text).tolist()
+
+    def encode_batch(self, texts: list[str]) -> list[list[float]]:
+        """Encode multiple texts at once."""
+        return self._embed.encode(texts).tolist()
+
+    def rerank(
+        self,
+        query: str,
+        results: list[dict[str, str]],
+        top_k: int,
+    ) -> list[dict[str, str]]:
+        """Rerank results using a cross-encoder. Lazy-loads the model on first call."""
+        if not results:
+            return []
+
+        # Deduplicate by content
+        seen = {}
+        for r in results:
+            if r["content"] not in seen:
+                seen[r["content"]] = r
+        unique = list(seen.values())
+
+        if len(unique) <= 1:
+            return unique[:top_k]
+
+        if self._rerank is None:
+            model_name = (
+                self._rerank_name
+                if isinstance(self._rerank_name, str)
+                else self.DEFAULT_RERANK_MODEL
+            )
+            if isinstance(self._rerank_name, CrossEncoder):
+                self._rerank = self._rerank_name
+            else:
+                self._rerank = CrossEncoder(model_name)
+
+        pairs = [(query, r["content"]) for r in unique]
+        scores = self._rerank.predict(pairs)
+        ranked = [r for _, r in sorted(zip(scores, unique), reverse=True)]
+        return ranked[:top_k]

reasongraph/_extraction.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+
+
+class NERExtractor:
+    """Named entity extractor using a HuggingFace token classification model.
+
+    Default model is dslim/bert-base-NER which recognizes PER, ORG, LOC, MISC.
+    Lazy-loads the pipeline on first call.
+    """
+
+    DEFAULT_MODEL = "dslim/bert-base-NER"
+
+    def __init__(self, model: str | None = None) -> None:
+        self._model_name = model or self.DEFAULT_MODEL
+        self._pipeline = None
+
+    def _load(self):
+        if self._pipeline is None:
+            from transformers import pipeline
+            self._pipeline = pipeline(
+                "ner",
+                model=self._model_name,
+                aggregation_strategy="simple",
+            )
+
+    def __call__(self, text: str) -> list[str]:
+        """Extract entities from text. Returns deduplicated entity strings."""
+        self._load()
+        results = self._pipeline(text)
+        seen = set()
+        entities = []
+        for ent in results:
+            word = ent["word"].strip()
+            # Filter out WordPiece artifacts (##subword tokens) and single chars
+            if not word or word.startswith("##") or len(word) <= 1:
+                continue
+            if word not in seen:
+                seen.add(word)
+                entities.append(word)
+        return entities
+
+
+class GLiNER2Extractor:
+    """Entity and relation extractor using GLiNER2.
+
+    Single model that handles both NER and causal relation extraction.
+    Lazy-loads the model on first call.
+
+    Default model: fastino/gliner2-large-v1 (340M params, DeBERTa-v3-large).
+    Requires: pip install reasongraph[gliner2]
+    """
+
+    DEFAULT_MODEL = "fastino/gliner2-large-v1"
+
+    DEFAULT_ENTITY_TYPES = ["person", "organization", "location", "event"]
+    DEFAULT_RELATION_TYPES = ["causes", "leads_to", "results_in"]
+
+    def __init__(
+        self,
+        model: str | None = None,
+        entity_types: list[str] | None = None,
+        relation_types: list[str] | None = None,
+    ) -> None:
+        self._model_name = model or self.DEFAULT_MODEL
+        self.entity_types = entity_types or self.DEFAULT_ENTITY_TYPES
+        self.relation_types = relation_types or self.DEFAULT_RELATION_TYPES
+        self._model = None
+
+    def _load(self):
+        if self._model is None:
+            try:
+                from gliner2 import GLiNER2
+            except ImportError:
+                raise ImportError(
+                    "GLiNER2 not installed. Install with: pip install reasongraph[gliner2]"
+                )
+            self._model = GLiNER2.from_pretrained(self._model_name)
+
+    def __call__(self, text: str) -> list[str]:
+        """Extract entities from text (compatible with ExtractorFn).
+
+        Returns deduplicated entity strings.
+        """
+        self._load()
+        raw = self._model.extract_entities(text, self.entity_types)
+        # GLiNER2 returns {"entities": {"person": [...], "location": [...]}}
+        entities_by_type = raw.get("entities", raw)
+
+        seen = set()
+        entities = []
+        for type_key, ents in entities_by_type.items():
+            if not isinstance(ents, list):
+                continue
+            for ent in ents:
+                word = ent.strip() if isinstance(ent, str) else ""
+                if word and word not in seen:
+                    seen.add(word)
+                    entities.append(word)
+        return entities
+
+    def extract_causal(self, texts: list[str]) -> list[dict]:
+        """Extract cause-effect relations from texts (compatible with CausalExtractorFn).
+
+        Args:
+            texts: List of text strings to analyze.
+
+        Returns:
+            List of dicts (one per text) with keys:
+                - 'text': original text
+                - 'causal': bool
+                - 'relations': list of {'cause': str, 'effect': str}
+        """
+        self._load()
+        results = []
+        for text in texts:
+            raw = self._model.extract_relations(text, self.relation_types)
+            # GLiNER2 returns {"relation_extraction": {"causes": [("a","b")], ...}}
+            # or directly {"causes": [("a","b")], ...}
+            rel_data = raw.get("relation_extraction", raw)
+
+            relations = []
+            seen_pairs = set()
+            for rel_type, pairs in rel_data.items():
+                if not isinstance(pairs, list):
+                    continue
+                for pair in pairs:
+                    cause, effect = None, None
+                    if isinstance(pair, (list, tuple)) and len(pair) >= 2:
+                        cause, effect = str(pair[0]).strip(), str(pair[1]).strip()
+                    elif isinstance(pair, dict):
+                        h = pair.get("head", {})
+                        t = pair.get("tail", {})
+                        cause = (h.get("text", "") if isinstance(h, dict) else str(h)).strip()
+                        effect = (t.get("text", "") if isinstance(t, dict) else str(t)).strip()
+                    if not cause or not effect or cause == effect:
+                        continue
+                    pair_key = (cause, effect)
+                    if pair_key in seen_pairs:
+                        continue
+                    seen_pairs.add(pair_key)
+                    relations.append({"cause": cause, "effect": effect})
+
+            results.append({
+                "text": text,
+                "causal": len(relations) > 0,
+                "relations": relations,
+            })
+        return results
+
+
+# Type alias for any entity extractor callable: text -> list of entity strings
+ExtractorFn = Callable[[str], list[str]]
+
+# Type alias for causal extractor: list[str] -> list[dict]
+CausalExtractorFn = Callable[[list[str]], list[dict]]

reasongraph/_types.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+
+
+@dataclass
+class Node:
+    """A node in the reason graph."""
+
+    content: str
+    type: str = "text"
+    embedding: list[float] | None = None
+    created_at: datetime = field(default_factory=datetime.now)
+    last_accessed: datetime = field(default_factory=datetime.now)
+
+
+@dataclass
+class Edge:
+    """A directed edge between two nodes."""
+
+    from_content: str
+    to_content: str
+    last_accessed: datetime = field(default_factory=datetime.now)

reasongraph/backends/__init__.py

@@ -0,0 +1,5 @@
+from reasongraph.backends._base import Backend
+from reasongraph.backends._sqlite import SqliteBackend
+from reasongraph.backends._postgres import PostgresBackend
+
+__all__ = ["Backend", "SqliteBackend", "PostgresBackend"]

reasongraph/backends/_base.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from reasongraph._types import Node, Edge
+
+
+class Backend(ABC):
+    """Abstract base class for graph storage backends."""
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        """Create tables/schema if they don't exist."""
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Release resources (connections, pools, etc.)."""
+
+    @abstractmethod
+    async def insert_nodes(self, nodes: list[Node]) -> None:
+        """Insert or upsert a batch of nodes (must have embeddings set)."""
+
+    @abstractmethod
+    async def insert_edges(self, edges: list[Edge]) -> None:
+        """Insert a batch of edges, ignoring duplicates."""
+
+    @abstractmethod
+    async def knn_search(
+        self, embedding: list[float], top_k: int
+    ) -> list[dict[str, str]]:
+        """Return the top_k closest nodes as dicts with 'content' and 'type' keys."""
+
+    @abstractmethod
+    async def get_neighbors(self, content: str) -> list[dict[str, str]]:
+        """Return all direct neighbors (both directions) as dicts with 'content' and 'type'."""
+
+    @abstractmethod
+    async def delete_stale_nodes(self, days: int) -> int:
+        """Delete nodes not accessed within the given number of days. Return count deleted."""
+
+    @abstractmethod
+    async def get_all_nodes(self) -> list[Node]:
+        """Return every node in the graph."""
+
+    @abstractmethod
+    async def hybrid_search(
+        self, embedding: list[float], query_text: str, top_k: int,
+        rrf_k: int = 60, keyword_only: bool = False,
+    ) -> list[dict[str, str]]:
+        """Combined embedding + trigram search using Reciprocal Rank Fusion.
+
+        Each node is ranked independently by cosine similarity and by trigram
+        similarity.  The final score is:
+            rrf_score(d) = 1/(rrf_k + rank_emb(d)) + 1/(rrf_k + rank_kw(d))
+
+        When keyword_only=True, nodes are ranked by trigram similarity alone
+        (no embedding component).
+
+        Args:
+            embedding: Query embedding vector.
+            query_text: Raw query string for trigram matching.
+            top_k: Number of results to return.
+            rrf_k: RRF smoothing constant (default 60).
+            keyword_only: If True, rank by trigram similarity only.
+
+        Returns:
+            Top-k nodes as dicts with 'content' and 'type' keys.
+        """
+
+    @abstractmethod
+    async def get_all_edges(self) -> list[Edge]:
+        """Return every edge in the graph."""

reasongraph/backends/_postgres.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+
+from reasongraph._types import Node, Edge
+from reasongraph.backends._base import Backend
+
+try:
+    from psycopg_pool import AsyncConnectionPool
+    from pgvector.psycopg import register_vector_async
+
+    _HAS_POSTGRES = True
+except ImportError:
+    _HAS_POSTGRES = False
+
+
+class PostgresBackend(Backend):
+    """PostgreSQL + pgvector backend for scalable vector search.
+
+    Requires: pip install reasongraph[postgres]
+    """
+
+    def __init__(self, database_url: str) -> None:
+        if not _HAS_POSTGRES:
+            raise ImportError(
+                "PostgreSQL dependencies not installed. "
+                "Install with: pip install reasongraph[postgres]"
+            )
+        self.database_url = database_url
+        self._pool: AsyncConnectionPool | None = None
+
+    async def _get_pool(self) -> AsyncConnectionPool:
+        if self._pool is None:
+            raise RuntimeError("Backend not initialized. Call initialize() first.")
+        return self._pool
+
+    async def initialize(self) -> None:
+        self._pool = AsyncConnectionPool(self.database_url, open=False)
+        await self._pool.open()
+
+        async with self._pool.connection() as conn:
+            await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+            await register_vector_async(conn)
+
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS nodes (
+                    content TEXT PRIMARY KEY,
+                    embedding VECTOR(384) NOT NULL,
+                    created_at TIMESTAMP DEFAULT NOW(),
+                    last_accessed TIMESTAMP DEFAULT NOW(),
+                    type TEXT NOT NULL CHECK (type IN ('text', 'entity'))
+                )
+            """)
+
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS edges (
+                    from_content TEXT NOT NULL REFERENCES nodes(content) ON DELETE CASCADE,
+                    to_content TEXT NOT NULL REFERENCES nodes(content) ON DELETE CASCADE,
+                    last_accessed TIMESTAMP DEFAULT NOW(),
+                    UNIQUE(from_content, to_content)
+                )
+            """)
+
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS edges_from_idx ON edges (from_content)"
+            )
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS edges_to_idx ON edges (to_content)"
+            )
+
+    async def close(self) -> None:
+        if self._pool is not None:
+            await self._pool.close()
+            self._pool = None
+
+    async def insert_nodes(self, nodes: list[Node]) -> None:
+        pool = await self._get_pool()
+        now = datetime.now()
+        async with pool.connection() as conn:
+            async with conn.cursor() as cur:
+                data = [
+                    (node.content, node.embedding, node.created_at, now, node.type)
+                    for node in nodes
+                ]
+                await cur.executemany(
+                    """
+                    INSERT INTO nodes (content, embedding, created_at, last_accessed, type)
+                    VALUES (%s, %s, %s, %s, %s)
+                    ON CONFLICT (content) DO UPDATE SET last_accessed = NOW()
+                    """,
+                    data,
+                )
+
+    async def insert_edges(self, edges: list[Edge]) -> None:
+        pool = await self._get_pool()
+        async with pool.connection() as conn:
+            async with conn.cursor() as cur:
+                data = [(e.from_content, e.to_content) for e in edges]
+                await cur.executemany(
+                    """
+                    INSERT INTO edges (from_content, to_content)
+                    VALUES (%s, %s)
+                    ON CONFLICT DO NOTHING
+                    """,
+                    data,
+                )
+
+    async def knn_search(
+        self, embedding: list[float], top_k: int
+    ) -> list[dict[str, str]]:
+        pool = await self._get_pool()
+        vec_str = f"[{', '.join(map(str, embedding))}]"
+        async with pool.connection() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(
+                    f"""
+                    SELECT content, type
+                    FROM nodes
+                    ORDER BY embedding <=> '{vec_str}'
+                    LIMIT {top_k}
+                    """
+                )
+                rows = await cur.fetchall()
+                return [{"content": row[0], "type": row[1]} for row in rows]
+
+    async def hybrid_search(
+        self, embedding: list[float], query_text: str, top_k: int,
+        rrf_k: int = 60, keyword_only: bool = False,
+    ) -> list[dict[str, str]]:
+        pool = await self._get_pool()
+        vec_str = f"[{', '.join(map(str, embedding))}]"
+        async with pool.connection() as conn:
+            await conn.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm")
+            async with conn.cursor() as cur:
+                if keyword_only:
+                    await cur.execute(
+                        """
+                        SELECT content, type
+                        FROM nodes
+                        ORDER BY similarity(content, %s) DESC
+                        LIMIT %s
+                        """,
+                        (query_text, top_k),
+                    )
+                    return [
+                        {"content": row[0], "type": row[1]}
+                        for row in await cur.fetchall()
+                    ]
+
+                # RRF entirely in SQL using window functions
+                await cur.execute(
+                    f"""
+                    WITH emb_ranked AS (
+                        SELECT content, type,
+                            ROW_NUMBER() OVER (
+                                ORDER BY embedding <=> '{vec_str}'
+                            ) AS rank
+                        FROM nodes
+                    ),
+                    kw_ranked AS (
+                        SELECT content,
+                            ROW_NUMBER() OVER (
+                                ORDER BY similarity(content, %s) DESC
+                            ) AS rank
+                        FROM nodes
+                    )
+                    SELECT e.content, e.type
+                    FROM emb_ranked e
+                    JOIN kw_ranked k ON e.content = k.content
+                    ORDER BY 1.0 / (%s + e.rank) + 1.0 / (%s + k.rank) DESC
+                    LIMIT %s
+                    """,
+                    (query_text, rrf_k, rrf_k, top_k),
+                )
+                return [
+                    {"content": row[0], "type": row[1]}
+                    for row in await cur.fetchall()
+                ]
+
+    async def get_neighbors(self, content: str) -> list[dict[str, str]]:
+        pool = await self._get_pool()
+        async with pool.connection() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(
+                    """
+                    SELECT DISTINCT n.content, n.type FROM nodes n
+                    INNER JOIN edges e ON (e.to_content = n.content AND e.from_content = %s)
+                                       OR (e.from_content = n.content AND e.to_content = %s)
+                    """,
+                    (content, content),
+                )
+                return [
+                    {"content": row[0], "type": row[1]}
+                    for row in await cur.fetchall()
+                ]
+
+    async def delete_stale_nodes(self, days: int) -> int:
+        pool = await self._get_pool()
+        cutoff = datetime.now() - timedelta(days=days)
+        async with pool.connection() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(
+                    "DELETE FROM nodes WHERE last_accessed < %s", (cutoff,)
+                )
+                return cur.rowcount
+
+    async def get_all_nodes(self) -> list[Node]:
+        pool = await self._get_pool()
+        async with pool.connection() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(
+                    "SELECT content, type, embedding, created_at, last_accessed FROM nodes"
+                )
+                nodes = []
+                for content, node_type, emb, created_at, last_accessed in await cur.fetchall():
+                    nodes.append(Node(
+                        content=content,
+                        type=node_type,
+                        embedding=list(emb) if emb else None,
+                        created_at=created_at,
+                        last_accessed=last_accessed,
+                    ))
+                return nodes
+
+    async def get_all_edges(self) -> list[Edge]:
+        pool = await self._get_pool()
+        async with pool.connection() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(
+                    "SELECT from_content, to_content, last_accessed FROM edges"
+                )
+                return [
+                    Edge(from_content=row[0], to_content=row[1], last_accessed=row[2])
+                    for row in await cur.fetchall()
+                ]