embed-tree 0.0.6__py3-none-any.whl

embed_tree/loaders/sqlalchemy.py

@@ -0,0 +1,91 @@
+"""SQLAlchemy-backed tree loader."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from embed_tree.persisters.model import MaterializedTreeState
+from embed_tree.representation import PartialTree
+from embed_tree.representation.default import partial_tree_from_dict, partial_tree_to_dict
+
+
+class SQLAlchemyTreeLoader:
+    """Load/save tree representation from an existing SQL table.
+
+    This base implementation deliberately does not create tables or run
+    migrations. Database shape is an application decision. Subclasses can
+    override ``post_init`` for setup; ``SQLiteTreeLoader`` is the built-in
+    simple implementation that creates its table.
+    """
+
+    def __init__(
+        self,
+        engine_or_url: Any,
+        *,
+        table_name: str = "embed_tree_state",
+        cache_key: str = "default",
+    ) -> None:
+        self.engine_or_url = engine_or_url
+        self.table_name = table_name
+        self.cache_key = cache_key
+        sa = _sqlalchemy()
+        self.engine = _engine(sa, self.engine_or_url)
+        self.post_init()
+
+    def post_init(self) -> None:
+        """Hook for migrations or validation owned by the caller/subclass."""
+        pass
+
+    def load(self) -> PartialTree | MaterializedTreeState | None:
+        sa = _sqlalchemy()
+        table = self._table(sa)
+        stmt = sa.select(table.c.kind, table.c.payload).where(table.c.cache_key == self.cache_key)
+        with self.engine.connect() as conn:
+            row = conn.execute(stmt).mappings().first()
+        if row is None:
+            return None
+        if row["kind"] == "partial_tree":
+            return partial_tree_from_dict(row["payload"])
+        if row["kind"] == "materialized_tree_state":
+            return row["payload"]
+        return row["payload"]
+
+    def save(self, state: PartialTree | MaterializedTreeState) -> None:
+        sa = _sqlalchemy()
+        table = self._table(sa)
+        if isinstance(state, PartialTree):
+            kind = "partial_tree"
+            payload: dict[str, Any] = partial_tree_to_dict(state)
+        else:
+            kind = "materialized_tree_state"
+            payload = state
+
+        delete = table.delete().where(table.c.cache_key == self.cache_key)
+        insert = table.insert().values(cache_key=self.cache_key, kind=kind, payload=payload)
+        with self.engine.begin() as conn:
+            conn.execute(delete)
+            conn.execute(insert)
+
+    def _table(self, sa: Any) -> Any:
+        meta = sa.MetaData()
+        return sa.Table(
+            self.table_name,
+            meta,
+            sa.Column("cache_key", sa.String, primary_key=True),
+            sa.Column("kind", sa.String, nullable=False),
+            sa.Column("payload", sa.JSON, nullable=False),
+        )
+
+
+def _sqlalchemy() -> Any:
+    try:
+        import sqlalchemy as sa
+    except ImportError as e:  # pragma: no cover
+        raise ImportError('SQLAlchemyTreeLoader needs the "sql" extra: pip install "embed-tree[sql]"') from e
+    return sa
+
+
+def _engine(sa: Any, engine_or_url: Any) -> Any:
+    if isinstance(engine_or_url, str):
+        return sa.create_engine(engine_or_url)
+    return engine_or_url

embed_tree/loaders/sqlalchemy_content.py

@@ -0,0 +1,63 @@
+"""SQLAlchemy-backed content loader."""
+
+from __future__ import annotations
+
+from typing import Any, Iterable
+
+from embed_tree.representation import ContentNode, PartialTree
+
+from .sqlalchemy import _engine, _sqlalchemy
+
+
+class SQLAlchemyContentLoader:
+    """Load content nodes from an existing SQL table via SQLAlchemy Core."""
+
+    def __init__(
+        self,
+        engine_or_url: Any,
+        table_name: str,
+        *,
+        id_column: str = "id",
+        content_column: str = "content",
+        text_column: str | None = None,
+        payload_columns: Iterable[str] | None = None,
+        where: Any | None = None,
+    ) -> None:
+        self.engine_or_url = engine_or_url
+        self.table_name = table_name
+        self.id_column = id_column
+        self.content_column = content_column
+        self.text_column = text_column
+        self.payload_columns = list(payload_columns or [])
+        self.where = where
+        self.post_init()
+
+    def post_init(self) -> None:
+        """Hook for migrations or validation owned by the caller/subclass."""
+        pass
+
+    def load(self) -> PartialTree | None:
+        sa = _sqlalchemy()
+        engine = _engine(sa, self.engine_or_url)
+        meta = sa.MetaData()
+        table = sa.Table(self.table_name, meta, autoload_with=engine)
+
+        stmt = sa.select(table)
+        if self.where is not None:
+            stmt = stmt.where(self.where)
+
+        nodes: list[ContentNode] = []
+        with engine.connect() as conn:
+            for row in conn.execute(stmt).mappings():
+                payload = {name: row[name] for name in self.payload_columns}
+                nodes.append(
+                    ContentNode(
+                        id=row[self.id_column],
+                        content=row[self.content_column],
+                        text=None if self.text_column is None else row[self.text_column],
+                        payload=payload or None,
+                    )
+                )
+
+        return PartialTree(content_nodes=nodes, metadata={"source": "sqlalchemy", "table": self.table_name})
+

embed_tree/loaders/sqlite.py

@@ -0,0 +1,21 @@
+"""SQLite-specific tree loader."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from .sqlalchemy import SQLAlchemyTreeLoader, _sqlalchemy
+
+
+class SQLiteTreeLoader(SQLAlchemyTreeLoader):
+    """SQLite tree-state loader with built-in table creation."""
+
+    def __init__(self, path: str | Path, *, table_name: str = "embed_tree_state", cache_key: str = "default") -> None:
+        self.path = Path(path)
+        super().__init__(f"sqlite:///{self.path}", table_name=table_name, cache_key=cache_key)
+
+    def post_init(self) -> None:
+        sa = _sqlalchemy()
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self._table(sa).metadata.create_all(self.engine)
+

embed_tree/persisters/__init__.py

@@ -0,0 +1,15 @@
+"""Tree persister contracts and implementations."""
+
+from .filesystem import FileSystemTreePersister, FolderTreePersister
+from .json import JsonTreePersister
+from .model import MaterializedTreeState, TreePersister
+from .sqlalchemy import SQLAlchemyTreePersister
+
+__all__ = [
+    "TreePersister",
+    "MaterializedTreeState",
+    "FolderTreePersister",
+    "FileSystemTreePersister",
+    "JsonTreePersister",
+    "SQLAlchemyTreePersister",
+]

embed_tree/persisters/filesystem.py

@@ -0,0 +1,293 @@
+"""Folder-backed external persister."""
+
+from __future__ import annotations
+
+import hashlib
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from embed_tree.representation import ContentNode, KeyNode, NodeId, PartialTree
+
+_UNSAFE_PATH_CHARS = re.compile(r'[\\/:*?"<>|#\[\]^\n\r\t]+')
+
+
+class FolderTreePersister:
+    """Materialize a live tree as folders and files.
+
+    The local folder is treated as mutable ground truth: every save reloads its
+    current file-md5 map, compares it with the target layout from the in-memory
+    tree, moves known files into place, creates missing files only when content
+    is available, prunes empty folders, and leaves unknown files untouched.
+    """
+
+    def __init__(
+        self,
+        root: str | Path,
+        *,
+        include_suffixes: set[str] | list[str] | tuple[str, ...] | None = None,
+        encoding: str = "utf-8",
+        include_hidden: bool = False,
+    ) -> None:
+        self.root = Path(root)
+        self.include_suffixes = None if include_suffixes is None else {s.lower() for s in include_suffixes}
+        self.encoding = encoding
+        self.include_hidden = include_hidden
+
+    def save(self, state: Any) -> None:
+        self.root.mkdir(parents=True, exist_ok=True)
+        current = self._current_files_by_md5()
+        desired = self._desired_files(state, current)
+
+        reserved: set[Path] = set()
+        for file in desired:
+            target = _unique_path(self.root / file.folder / file.filename, reserved)
+            reserved.add(target)
+            existing = current.get(file.md5)
+            if existing is not None:
+                self._move_file(existing, target)
+            elif file.content is not None:
+                self._write_file(target, file.content)
+
+        self._prune_empty_dirs()
+
+    def _desired_files(self, state: Any, current: dict[str, Path]) -> list["_DesiredFile"]:
+        tree = state.get_tree() if hasattr(state, "get_tree") and callable(state.get_tree) else state
+        if isinstance(tree, PartialTree):
+            return self._desired_files_from_partial_tree(tree, current)
+        if _looks_like_live_node(tree):
+            return list(self._desired_files_from_live_node(tree, current))
+        raise TypeError("FolderTreePersister persists an EmbedTree, live Node, or PartialTree")
+
+    def _current_files_by_md5(self) -> dict[str, Path]:
+        files: dict[str, Path] = {}
+        for path in sorted(self.root.rglob("*")):
+            if not path.is_file() or not self._included(path):
+                continue
+            rel_parts = path.relative_to(self.root).parts
+            if not self.include_hidden and any(part.startswith(".") for part in rel_parts):
+                continue
+            files.setdefault(_file_md5(path), path)
+        return files
+
+    def _desired_files_from_partial_tree(
+        self,
+        tree: PartialTree,
+        current: dict[str, Path],
+    ) -> list["_DesiredFile"]:
+        key_nodes = {node.id: node for node in tree.key_nodes}
+        parent_by_child = {edge.child_id: edge.parent_id for edge in tree.edges}
+        desired: list[_DesiredFile] = []
+
+        for node in tree.content_nodes:
+            folder = self._folder_for(node.id, key_nodes, parent_by_child)
+            filename = self._filename_for_node(node, current.get(str(node.id)))
+            content = None if node.content is None else str(node.content)
+            desired.append(_DesiredFile(str(node.id), folder, filename, content))
+        return desired
+
+    def _desired_files_from_live_node(
+        self,
+        root: Any,
+        current: dict[str, Path],
+    ) -> list["_DesiredFile"]:
+        if root.is_leaf:
+            folder = Path(_safe_name(root.label or "topics"))
+            return [
+                file
+                for item in root.items or []
+                if (file := self._desired_file_for_item(item, folder, current)) is not None
+            ]
+
+        desired: list[_DesiredFile] = []
+        self._collect_live_node_files(root, Path(), current, desired, include_self=False)
+        return desired
+
+    def _collect_live_node_files(
+        self,
+        node: Any,
+        prefix: Path,
+        current: dict[str, Path],
+        desired: list["_DesiredFile"],
+        *,
+        include_self: bool,
+    ) -> None:
+        folder = prefix
+        if include_self:
+            folder = prefix / _safe_name(node.label or f"topic-{getattr(node, 'id', 'node')}")
+        if node.is_leaf:
+            for item in node.items or []:
+                file = self._desired_file_for_item(item, folder, current)
+                if file is not None:
+                    desired.append(file)
+            return
+
+        used: set[str] = set()
+        for index, child in enumerate(node.children or []):
+            label = _dedupe_name(_safe_name(child.label or f"topic-{index + 1}"), used)
+            self._collect_live_node_files(child, folder / label, current, desired, include_self=False)
+
+    def _desired_file_for_item(
+        self,
+        item: Any,
+        folder: Path,
+        current: dict[str, Path],
+    ) -> "_DesiredFile | None":
+        md5 = self._md5_for_item(item)
+        if md5 is None:
+            return None
+        existing = current.get(md5)
+        filename = self._filename_for_item(item, existing)
+        content = self._content_for_item(item)
+        return _DesiredFile(md5, folder, filename, content)
+
+    def _folder_for(
+        self,
+        node_id: NodeId,
+        key_nodes: dict[NodeId, KeyNode],
+        parent_by_child: dict[NodeId, NodeId],
+    ) -> Path:
+        parts: list[str] = []
+        current = parent_by_child.get(node_id)
+        seen: set[NodeId] = set()
+        while current is not None and current not in seen:
+            seen.add(current)
+            if str(current) == ".":
+                break
+            key = key_nodes.get(current)
+            raw = key.label if key is not None and key.label else str(current)
+            part = _safe_name(raw)
+            if part:
+                parts.append(part)
+            current = parent_by_child.get(current)
+        return Path(*reversed(parts)) if parts else Path()
+
+    def _filename_for_node(self, node: ContentNode, existing: Path | None) -> str:
+        if existing is not None:
+            return existing.name
+        payload = node.payload if isinstance(node.payload, dict) else {}
+        for key in ("filename", "relative_path", "path"):
+            value = payload.get(key)
+            if value:
+                name = Path(str(value)).name
+                if name:
+                    return _safe_name(name, fallback=f"{node.id}.txt")
+        text = node.text.strip() if isinstance(node.text, str) else ""
+        return _safe_name(text, fallback=f"{node.id}.txt")
+
+    def _md5_for_item(self, item: Any) -> str | None:
+        item_id = str(item.id)
+        if _is_md5(item_id):
+            return item_id
+        payload = item.payload if isinstance(item.payload, dict) else {}
+        for key in ("md5", "file_md5", "content_md5", "content_id", "id"):
+            value = payload.get(key)
+            if value is not None and _is_md5(str(value)):
+                return str(value)
+        return None
+
+    def _filename_for_item(self, item: Any, existing: Path | None) -> str:
+        if existing is not None:
+            return existing.name
+        payload = item.payload if isinstance(item.payload, dict) else {}
+        for key in ("filename", "relative_path", "output_path", "path"):
+            value = payload.get(key)
+            if value:
+                name = Path(str(value)).name
+                if name:
+                    return _safe_name(name, fallback=f"{item.id}.txt")
+        text = item.text.strip() if isinstance(item.text, str) else ""
+        return _safe_name(text, fallback=f"{item.id}.txt")
+
+    def _content_for_item(self, item: Any) -> str | None:
+        payload = item.payload if isinstance(item.payload, dict) else {}
+        for key in ("content", "body", "text"):
+            value = payload.get(key)
+            if value is not None:
+                return str(value)
+        return None
+
+    def _move_file(self, source: Path, target: Path) -> None:
+        if source == target:
+            return
+        target.parent.mkdir(parents=True, exist_ok=True)
+        source.rename(target)
+
+    def _write_file(self, target: Path, content: str) -> None:
+        target.parent.mkdir(parents=True, exist_ok=True)
+        target.write_text(content, encoding=self.encoding)
+
+    def _prune_empty_dirs(self) -> None:
+        for directory in sorted(
+            (path for path in self.root.rglob("*") if path.is_dir()),
+            key=lambda path: len(path.parts),
+            reverse=True,
+        ):
+            rel_parts = directory.relative_to(self.root).parts
+            if not self.include_hidden and any(part.startswith(".") for part in rel_parts):
+                continue
+            try:
+                next(directory.iterdir())
+            except StopIteration:
+                directory.rmdir()
+
+    def _included(self, path: Path) -> bool:
+        return self.include_suffixes is None or path.suffix.lower() in self.include_suffixes
+
+
+FileSystemTreePersister = FolderTreePersister
+
+
+@dataclass(frozen=True)
+class _DesiredFile:
+    md5: str
+    folder: Path
+    filename: str
+    content: str | None
+
+
+def _file_md5(path: Path) -> str:
+    digest = hashlib.md5()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _safe_name(value: str, *, fallback: str = "untitled") -> str:
+    cleaned = _UNSAFE_PATH_CHARS.sub(" ", value)
+    cleaned = " ".join(cleaned.split()).strip(". ")
+    return cleaned or fallback
+
+
+def _dedupe_name(value: str, used: set[str]) -> str:
+    candidate = value
+    i = 2
+    while candidate in used:
+        candidate = f"{value}-{i}"
+        i += 1
+    used.add(candidate)
+    return candidate
+
+
+def _is_md5(value: str) -> bool:
+    return len(value) == 32 and all(c in "0123456789abcdefABCDEF" for c in value)
+
+
+def _looks_like_live_node(value: Any) -> bool:
+    return hasattr(value, "is_leaf") and (hasattr(value, "items") or hasattr(value, "children"))
+
+
+def _unique_path(target: Path, reserved: set[Path]) -> Path:
+    if target not in reserved and not target.exists():
+        return target
+    stem = target.stem
+    suffix = target.suffix
+    parent = target.parent
+    i = 2
+    while True:
+        candidate = parent / f"{stem}-{i}{suffix}"
+        if candidate not in reserved and not candidate.exists():
+            return candidate
+        i += 1

embed_tree/persisters/json.py

@@ -0,0 +1,29 @@
+"""JSON external persister."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+from embed_tree.representation import PartialTree
+from embed_tree.representation.default import partial_tree_to_dict
+
+
+class JsonTreePersister:
+    """Persist an exported tree artifact to JSON."""
+
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+
+    def save(self, state: Any) -> None:
+        payload = partial_tree_to_dict(state) if isinstance(state, PartialTree) else state
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = self.path.with_name(f"{self.path.name}.tmp.{os.getpid()}")
+        with tmp.open("w", encoding="utf-8") as f:
+            json.dump(payload, f)
+            f.flush()
+            os.fsync(f.fileno())
+        os.replace(tmp, self.path)
+

embed_tree/persisters/model.py

@@ -0,0 +1,23 @@
+"""Abstract persister contract."""
+
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+from embed_tree.representation import PartialTree
+
+MaterializedTreeState = dict[str, Any]
+
+
+@runtime_checkable
+class TreePersister(Protocol):
+    """Persist a tree representation.
+
+    The same protocol covers durable internal state, reusable state for future
+    builds, and user-facing exports. The role comes from where the persister is
+    passed, not from a separate cache type.
+    """
+
+    def save(self, state: PartialTree | MaterializedTreeState | Any) -> None:
+        """Persist a tree representation or export artifact."""
+        ...

embed_tree/persisters/sqlalchemy.py

@@ -0,0 +1,76 @@
+"""SQLAlchemy external persister."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from embed_tree.representation import PartialTree
+
+
+class SQLAlchemyTreePersister:
+    """Persist PartialTree content nodes to a SQL table via SQLAlchemy Core."""
+
+    def __init__(
+        self,
+        engine_or_url: Any,
+        table_name: str,
+        *,
+        id_column: str = "id",
+        content_column: str = "content",
+        text_column: str = "text",
+        payload_column: str = "payload",
+    ) -> None:
+        self.engine_or_url = engine_or_url
+        self.table_name = table_name
+        self.id_column = id_column
+        self.content_column = content_column
+        self.text_column = text_column
+        self.payload_column = payload_column
+
+    def save(self, state: Any) -> None:
+        if not isinstance(state, PartialTree):
+            raise TypeError("SQLAlchemyTreePersister only persists PartialTree instances")
+
+        sa = _sqlalchemy()
+        engine = _engine(sa, self.engine_or_url)
+        table = self._table(sa)
+        table.metadata.create_all(engine)
+        rows = [
+            {
+                self.id_column: node.id,
+                self.content_column: node.content,
+                self.text_column: node.text,
+                self.payload_column: node.payload,
+            }
+            for node in state.content_nodes
+        ]
+        with engine.begin() as conn:
+            conn.execute(table.delete())
+            if rows:
+                conn.execute(table.insert(), rows)
+
+    def _table(self, sa: Any) -> Any:
+        meta = sa.MetaData()
+        return sa.Table(
+            self.table_name,
+            meta,
+            sa.Column(self.id_column, sa.String, primary_key=True),
+            sa.Column(self.content_column, sa.Text, nullable=False),
+            sa.Column(self.text_column, sa.Text),
+            sa.Column(self.payload_column, sa.JSON),
+        )
+
+
+def _sqlalchemy() -> Any:
+    try:
+        import sqlalchemy as sa
+    except ImportError as e:  # pragma: no cover
+        raise ImportError('SQLAlchemyTreePersister needs the "sql" extra: pip install "embed-tree[sql]"') from e
+    return sa
+
+
+def _engine(sa: Any, engine_or_url: Any) -> Any:
+    if isinstance(engine_or_url, str):
+        return sa.create_engine(engine_or_url)
+    return engine_or_url
+

embed_tree/projectors/__init__.py

@@ -0,0 +1,7 @@
+"""Vector projection integrations."""
+
+from .model import PCAConfig, VectorProjector
+from .pca import PCAProjector
+
+__all__ = ["PCAConfig", "VectorProjector", "PCAProjector"]
+

embed_tree/projectors/model.py

@@ -0,0 +1,39 @@
+"""Vector projection contracts."""
+
+from __future__ import annotations
+
+from typing import Literal, Protocol, runtime_checkable
+
+import numpy as np
+from pydantic import BaseModel
+
+
+class PCAConfig(BaseModel):
+    """Configuration for PCA-based projection."""
+
+    dims: int | None = None
+    mode: Literal["freeze", "incremental"] = "freeze"
+    warmup: int = 1000
+    batch_size: int = 256
+
+
+@runtime_checkable
+class VectorProjector(Protocol):
+    """Map raw vectors to operational vectors."""
+
+    @property
+    def is_fitted(self) -> bool:
+        """Whether this projector can transform vectors."""
+        ...
+
+    def fit(self, vectors: np.ndarray) -> None:
+        """Fit projector state from vectors."""
+        ...
+
+    def transform(self, vectors: np.ndarray) -> np.ndarray:
+        """Project vectors."""
+        ...
+
+    def __call__(self, vectors: np.ndarray) -> np.ndarray:
+        """Project vectors."""
+        ...