embed-tree 0.0.6__py3-none-any.whl

embed_tree/__init__.py

@@ -0,0 +1,121 @@
+"""embed-tree: an incremental hierarchical clustering tree over embeddings.
+
+See DESIGN.md for the full design. Minimal usage:
+
+    from embed_tree import EmbedTree, TreeConfig, FileTreeStore
+
+    tree = EmbedTree(embedder=my_embed_fn, store=FileTreeStore("./tree.json"))
+    tree.add("some content")
+    hits = tree.query("similar content", k=5)
+"""
+
+from .config import LLMConfig, RebalanceConfig, TreeConfig
+from .embedders import HuggingFaceTextEmbedder, TextEmbedder, embed_texts
+from .labelers import FunctionLabeler, LabelCandidate, Labeler, LabelRequest, LLMLabeler
+from .loaders import (
+    FileSystemTreeLoader,
+    JsonTreeLoader,
+    SQLAlchemyContentLoader,
+    SQLAlchemyTreeLoader,
+    SQLiteTreeLoader,
+    TreeLoader,
+)
+from .persisters import (
+    FileSystemTreePersister,
+    FolderTreePersister,
+    JsonTreePersister,
+    MaterializedTreeState,
+    SQLAlchemyTreePersister,
+    TreePersister,
+)
+from .projectors import PCAConfig, PCAProjector, VectorProjector
+from .providers import (
+    EmbeddingProvider,
+    FakeEmbeddingProvider,
+    OpenAIEmbeddingProvider,
+    SentenceTransformerProvider,
+)
+from .reducers import (
+    FreezePCAReducer,
+    IdentityReducer,
+    IncrementalPCAReducer,
+    Reducer,
+)
+from .reconcilers import DefaultTreeReconciler, TreeReconciler
+from .representation import (
+    ContentNode,
+    DefaultTreeRepresentation,
+    KeyNode,
+    NodeAggregate,
+    NodeEmbedding,
+    NodeId,
+    PartialTree,
+    TreeEdge,
+    VectorData,
+    partial_tree_from_dict,
+    partial_tree_to_dict,
+)
+from .store import FileTreeStore, NullTreeStore, TreeState, TreeStore
+from .taggers import KeywordTagger, LLMTagger, Tagger, make_tagger
+from .tree import EmbedTree, Item, Node
+
+__all__ = [
+    "EmbedTree",
+    "Item",
+    "Node",
+    "TreeConfig",
+    "RebalanceConfig",
+    "LLMConfig",
+    "TextEmbedder",
+    "HuggingFaceTextEmbedder",
+    "embed_texts",
+    "PCAConfig",
+    "PCAProjector",
+    "VectorProjector",
+    "LabelCandidate",
+    "LabelRequest",
+    "Labeler",
+    "FunctionLabeler",
+    "LLMLabeler",
+    "TreeLoader",
+    "FileSystemTreeLoader",
+    "JsonTreeLoader",
+    "SQLAlchemyContentLoader",
+    "SQLAlchemyTreeLoader",
+    "SQLiteTreeLoader",
+    "MaterializedTreeState",
+    "TreeReconciler",
+    "DefaultTreeReconciler",
+    "TreePersister",
+    "FolderTreePersister",
+    "FileSystemTreePersister",
+    "JsonTreePersister",
+    "SQLAlchemyTreePersister",
+    "DefaultTreeRepresentation",
+    "PartialTree",
+    "ContentNode",
+    "KeyNode",
+    "TreeEdge",
+    "NodeEmbedding",
+    "NodeAggregate",
+    "NodeId",
+    "VectorData",
+    "partial_tree_from_dict",
+    "partial_tree_to_dict",
+    "TreeState",
+    "TreeStore",
+    "FileTreeStore",
+    "NullTreeStore",
+    "EmbeddingProvider",
+    "FakeEmbeddingProvider",
+    "OpenAIEmbeddingProvider",
+    "SentenceTransformerProvider",
+    "Reducer",
+    "IdentityReducer",
+    "FreezePCAReducer",
+    "IncrementalPCAReducer",
+    "Tagger",
+    "KeywordTagger",
+    "LLMTagger",
+    "make_tagger",
+]

embed_tree/cache/__init__.py

@@ -0,0 +1,10 @@
+"""Deprecated cache compatibility imports.
+
+Use loaders plus persisters directly for new code.
+"""
+
+from .json import JsonTreeCache
+from .model import MaterializedTreeState, TreeCache
+from .sqlalchemy import SQLAlchemyTreeCache
+
+__all__ = ["MaterializedTreeState", "TreeCache", "JsonTreeCache", "SQLAlchemyTreeCache"]

embed_tree/cache/json.py

@@ -0,0 +1,7 @@
+"""Compatibility import for JSON tree cache."""
+
+from embed_tree.loaders.json import JsonTreeLoader
+
+
+class JsonTreeCache(JsonTreeLoader):
+    """Backward-compatible name for JsonTreeLoader."""

embed_tree/cache/model.py

@@ -0,0 +1,13 @@
+"""Deprecated cache compatibility contract."""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from embed_tree.loaders.model import TreeLoader
+from embed_tree.persisters.model import MaterializedTreeState, TreePersister
+
+
+@runtime_checkable
+class TreeCache(TreeLoader, TreePersister, Protocol):
+    """Deprecated alias for TreeLoader + TreePersister."""

embed_tree/cache/sqlalchemy.py

@@ -0,0 +1,7 @@
+"""Compatibility import for SQLAlchemy tree cache."""
+
+from embed_tree.loaders.sqlalchemy import SQLAlchemyTreeLoader
+
+
+class SQLAlchemyTreeCache(SQLAlchemyTreeLoader):
+    """Backward-compatible name for SQLAlchemyTreeLoader."""

embed_tree/config.py

@@ -0,0 +1,100 @@
+"""Configuration for embed-tree, as a single pydantic object.
+
+`TreeConfig` is a plain pydantic `BaseModel` (NOT `BaseSettings`): it is
+constructed explicitly and handed whole to `EmbedTree(config=...)`. It does
+**not** read environment variables — every value must be passed in code, so the
+configuration is always explicit and reproducible.
+
+See DESIGN.md §4/§5.3. M1 adds PCA dimensionality reduction (pca_dims) in two
+modes (freeze / incremental); see those sections for the rebalance contract.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+
+
+class RebalanceConfig(BaseModel):
+    """When/whether to rebuild the whole tree from its leaves (DESIGN.md §4)."""
+
+    enabled: bool = True
+    every_n_inserts: int | None = 10_000  # auto-rebuild cadence; None disables
+    on_demand: bool = True  # allow manual tree.rebalance()
+
+
+class LLMConfig(BaseModel):
+    """How to auto-name taxonomy nodes (DESIGN.md §10). Provider "none" uses a
+    no-network keyword tagger; "openai" and "local" generate labels with an LLM.
+    """
+
+    provider: Literal["none", "openai", "local"] = "none"
+    model: str = "gpt-4o-mini"  # OpenAI model id, or HF model id when local
+    api_key: str | None = None  # OpenAI key (explicit; no env)
+    base_url: str | None = None  # OpenAI-compatible endpoint (e.g. a local server)
+    max_samples: int = 15  # member texts shown to the LLM when naming a cluster
+    max_label_words: int = 6  # keep labels short and browsable
+
+
+class TreeConfig(BaseModel):
+    """Top-level knobs. Defaults are tuned for the M0/M1 (<100k items) regime.
+
+    Constructed in code only — no environment-variable loading.
+    """
+
+    model_config = ConfigDict(protected_namespaces=())  # allow `model_args` name
+
+    # Defaults are tuned for a human-browsable taxonomy (DESIGN.md §10): a small
+    # fan-out and small leaves keep every level readable (<=5 sub-topics, <=10
+    # items per leaf). Raise both for a large-scale retrieval index instead.
+    max_branches: int = 5  # max sub-topics per level (k for KMeans)
+    leaf_capacity: int = 10  # max items in a leaf before it subdivides
+    split_algo: str = "kmeans"  # M0: "kmeans" only
+
+    # Distance is always cosine: vectors are L2-normalized and compared with
+    # Euclidean (rank-equivalent on the unit sphere). Embeddings encode meaning
+    # in direction, not magnitude, so there is no separate distance knob.
+
+    # --- PCA dimensionality reduction (M1; see DESIGN.md §5.3) -------------
+    # Off by default: only worth it at scale (thousands+). At tens of items PCA
+    # is meaningless (too few samples) and never reaches pca_warmup anyway.
+    pca_dims: int | None = None  # None = no reduction (operate in raw space)
+    pca_mode: Literal["freeze", "incremental"] = "freeze"
+    pca_warmup: int = 1000  # items buffered before the first PCA fit
+    pca_batch_size: int = 256  # incremental: partial_fit cadence
+
+    rebalance: RebalanceConfig = Field(default_factory=RebalanceConfig)
+    llm: LLMConfig = Field(default_factory=LLMConfig)  # node auto-naming
+    model_args: dict = Field(default_factory=dict)  # passed through to KMeans
+
+    @field_validator("max_branches")
+    @classmethod
+    def _min_branches(cls, v: int) -> int:
+        if v < 2:
+            raise ValueError("max_branches must be >= 2")
+        return v
+
+    @field_validator("split_algo")
+    @classmethod
+    def _supported_split(cls, v: str) -> str:
+        if v != "kmeans":
+            raise NotImplementedError(
+                f"split_algo={v!r} arrives in a later milestone; "
+                "M0 supports 'kmeans' only"
+            )
+        return v
+
+    @model_validator(mode="after")
+    def _cross_field(self) -> "TreeConfig":
+        if self.leaf_capacity < self.max_branches:
+            raise ValueError("leaf_capacity must be >= max_branches")
+        if self.pca_dims is not None:
+            if self.pca_dims < 2:
+                raise ValueError("pca_dims must be >= 2")
+            # PCA needs at least n_components samples to fit / partial_fit.
+            if self.pca_warmup < self.pca_dims:
+                raise ValueError("pca_warmup must be >= pca_dims")
+            if self.pca_batch_size < self.pca_dims:
+                raise ValueError("pca_batch_size must be >= pca_dims")
+        return self

embed_tree/embedders/__init__.py

@@ -0,0 +1,7 @@
+"""Embedding model integrations."""
+
+from .huggingface import HuggingFaceTextEmbedder, embed_texts
+from .model import TextEmbedder, Vector
+
+__all__ = ["TextEmbedder", "Vector", "HuggingFaceTextEmbedder", "embed_texts"]
+

embed_tree/embedders/huggingface.py

@@ -0,0 +1,69 @@
+"""Open-source local embeddings via Hugging Face sentence-transformers."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from embed_tree.providers.local import SentenceTransformerProvider
+
+
+class HuggingFaceTextEmbedder(SentenceTransformerProvider):
+    """Sentence-transformers embedder with Mac-friendly device selection.
+
+    The model is downloaded from Hugging Face by sentence-transformers on first
+    use and then cached by that stack. On Apple Silicon, device="auto" prefers
+    MPS when PyTorch reports it as available; otherwise it falls back to CPU.
+    """
+
+    def __init__(
+        self,
+        model: str = "BAAI/bge-small-en-v1.5",
+        *,
+        device: str | None = "auto",
+        cache_folder: str | Path | None = None,
+        model_obj: Any | None = None,
+        encode_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        resolved_device = _resolve_device(device)
+        self.model_name = model
+        self.device = resolved_device
+        self.cache_folder = None if cache_folder is None else str(cache_folder)
+        if model_obj is None:
+            try:
+                from sentence_transformers import SentenceTransformer
+            except ImportError as e:  # pragma: no cover
+                raise ImportError(
+                    'HuggingFaceTextEmbedder needs the "local" extra: '
+                    'pip install "embed-tree[local]"'
+                ) from e
+            model_obj = SentenceTransformer(model, device=resolved_device, cache_folder=self.cache_folder)
+            super().__init__(model=model, device=resolved_device, model_obj=model_obj, encode_kwargs=encode_kwargs, **kwargs)
+        else:
+            super().__init__(model=model, device=resolved_device, model_obj=model_obj, encode_kwargs=encode_kwargs, **kwargs)
+
+
+def _resolve_device(device: str | None) -> str | None:
+    if device != "auto":
+        return device
+    try:
+        import torch
+    except ImportError:
+        return None
+    if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
+        return "mps"
+    if torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+
+def embed_texts(embedder: Any, texts: list[str]) -> np.ndarray:
+    """Embed a batch through any callable or TextEmbedder-like object."""
+    batch_fn = getattr(embedder, "embed_batch", None)
+    if callable(batch_fn):
+        return np.asarray(batch_fn(texts))
+    return np.asarray([embedder(text) for text in texts])
+

embed_tree/embedders/model.py

@@ -0,0 +1,27 @@
+"""Embedding model contracts."""
+
+from __future__ import annotations
+
+from typing import Protocol, Sequence, runtime_checkable
+
+import numpy as np
+
+Vector = np.ndarray
+
+
+@runtime_checkable
+class TextEmbedder(Protocol):
+    """Turn strings into embedding vectors."""
+
+    def embed(self, text: str) -> Vector:
+        """Embed one string."""
+        ...
+
+    def embed_batch(self, texts: Sequence[str]) -> np.ndarray:
+        """Embed strings in order."""
+        ...
+
+    def __call__(self, text: str) -> Vector:
+        """Embed one string."""
+        ...
+

embed_tree/labelers/__init__.py

@@ -0,0 +1,8 @@
+"""Labeling model integrations."""
+
+from .function import FunctionLabeler
+from .llm import LLMLabeler
+from .model import LabelCandidate, LabelRequest, Labeler
+
+__all__ = ["LabelCandidate", "LabelRequest", "Labeler", "FunctionLabeler", "LLMLabeler"]
+

embed_tree/labelers/function.py

@@ -0,0 +1,25 @@
+"""Function-backed labeler."""
+
+from __future__ import annotations
+
+from typing import Callable, Iterable
+
+from .model import LabelRequest
+
+
+class FunctionLabeler:
+    """Adapt a cheap user function into the streaming labeler protocol."""
+
+    def __init__(self, fn: Callable[[LabelRequest], str | Iterable[str]]) -> None:
+        self.fn = fn
+
+    def stream(self, request: LabelRequest) -> Iterable[str]:
+        out = self.fn(request)
+        if isinstance(out, str):
+            yield out
+        else:
+            yield from out
+
+    def label(self, request: LabelRequest) -> str:
+        return "".join(self.stream(request)).strip()
+

embed_tree/labelers/llm.py

@@ -0,0 +1,26 @@
+"""LLM-backed streaming labeler."""
+
+from __future__ import annotations
+
+from typing import Any, Iterable
+
+from embed_tree.config import LLMConfig
+from embed_tree.taggers import LLMTagger
+
+from .model import LabelRequest
+
+
+class LLMLabeler:
+    """Generate labels from nearby candidates using the existing LLM tagger."""
+
+    def __init__(self, config: LLMConfig | None = None, *, client: Any | None = None, pipeline: Any | None = None) -> None:
+        self.config = config or LLMConfig()
+        self.tagger = LLMTagger(self.config, client=client, pipeline=pipeline)
+
+    def stream(self, request: LabelRequest) -> Iterable[str]:
+        yield self.label(request)
+
+    def label(self, request: LabelRequest) -> str:
+        texts = [candidate.text for candidate in request.candidates]
+        return self.tagger(texts)
+

embed_tree/labelers/model.py

@@ -0,0 +1,38 @@
+"""Labeling model contracts."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Iterable, Protocol, runtime_checkable
+
+
+@dataclass(frozen=True)
+class LabelCandidate:
+    """Nearby node or item used as context for a label."""
+
+    id: Any
+    text: str
+    distance: float | None = None
+    payload: Any = None
+
+
+@dataclass(frozen=True)
+class LabelRequest:
+    """Context for generating a label."""
+
+    candidates: list[LabelCandidate]
+    max_words: int = 6
+
+
+@runtime_checkable
+class Labeler(Protocol):
+    """Generate a label from nearby candidates."""
+
+    def stream(self, request: LabelRequest) -> Iterable[str]:
+        """Yield label chunks."""
+        ...
+
+    def label(self, request: LabelRequest) -> str:
+        """Return the full label."""
+        ...
+

embed_tree/loaders/__init__.py

@@ -0,0 +1,17 @@
+"""Tree loader contracts."""
+
+from .filesystem import FileSystemTreeLoader
+from .json import JsonTreeLoader
+from .model import TreeLoader
+from .sqlalchemy_content import SQLAlchemyContentLoader
+from .sqlalchemy import SQLAlchemyTreeLoader
+from .sqlite import SQLiteTreeLoader
+
+__all__ = [
+    "TreeLoader",
+    "FileSystemTreeLoader",
+    "JsonTreeLoader",
+    "SQLAlchemyContentLoader",
+    "SQLAlchemyTreeLoader",
+    "SQLiteTreeLoader",
+]

embed_tree/loaders/filesystem.py

@@ -0,0 +1,83 @@
+"""Filesystem-backed ground-truth loader."""
+
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+from typing import Iterable
+
+from embed_tree.representation import ContentNode, KeyNode, PartialTree, TreeEdge
+
+
+class FileSystemTreeLoader:
+    """Load files under a directory as content nodes.
+
+    Directory nodes are emitted as ``KeyNode`` records with edges to their child
+    directories/files. File node ids are MD5 hashes of their file bytes, so the
+    same file keeps its identity when it moves locally.
+    """
+
+    def __init__(
+        self,
+        root: str | Path,
+        *,
+        include_suffixes: Iterable[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 load(self) -> PartialTree | None:
+        if not self.root.exists():
+            return None
+
+        tree = PartialTree(metadata={"source": "filesystem", "root": str(self.root)})
+        root_id = "."
+        tree.key_nodes.append(KeyNode(id=root_id, label=self.root.name or str(self.root)))
+
+        for path in sorted(self.root.rglob("*")):
+            rel = path.relative_to(self.root).as_posix()
+            if not self.include_hidden and any(part.startswith(".") for part in path.relative_to(self.root).parts):
+                continue
+            parent = path.parent.relative_to(self.root).as_posix() if path.parent != self.root else root_id
+            if path.is_dir():
+                tree.key_nodes.append(KeyNode(id=rel, label=path.name))
+                tree.edges.append(TreeEdge(parent_id=parent, child_id=rel))
+                continue
+            if not path.is_file() or not self._included(path):
+                continue
+            file_id = _file_md5(path)
+            try:
+                content = path.read_text(encoding=self.encoding)
+            except UnicodeDecodeError:
+                continue
+            tree.content_nodes.append(
+                ContentNode(
+                    id=file_id,
+                    content=content,
+                    text=path.stem,
+                    payload={
+                        "path": str(path),
+                        "relative_path": rel,
+                        "filename": path.name,
+                    },
+                    version=file_id,
+                )
+            )
+            tree.edges.append(TreeEdge(parent_id=parent, child_id=file_id))
+
+        return tree
+
+    def _included(self, path: Path) -> bool:
+        return self.include_suffixes is None or path.suffix.lower() in self.include_suffixes
+
+
+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()

embed_tree/loaders/json.py

@@ -0,0 +1,49 @@
+"""JSON-backed tree loader."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+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 JsonTreeLoader:
+    """Load/save a PartialTree or materialized state as one JSON file."""
+
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+        self.post_init()
+
+    def post_init(self) -> None:
+        """Hook for implementations that need setup after construction."""
+        pass
+
+    def load(self) -> PartialTree | MaterializedTreeState | None:
+        if not self.path.exists():
+            return None
+        with self.path.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+        if data.get("kind") == "partial_tree":
+            return partial_tree_from_dict(data["tree"])
+        if data.get("kind") == "materialized_tree_state":
+            return data["state"]
+        return data
+
+    def save(self, state: PartialTree | MaterializedTreeState) -> None:
+        if isinstance(state, PartialTree):
+            payload: dict[str, Any] = {"kind": "partial_tree", "tree": partial_tree_to_dict(state)}
+        else:
+            payload = {"kind": "materialized_tree_state", "state": 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/loaders/model.py

@@ -0,0 +1,20 @@
+"""Abstract loader contract."""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from embed_tree.representation import PartialTree
+
+
+@runtime_checkable
+class TreeLoader(Protocol):
+    """Load a partial tree from any source.
+
+    Ground truth and reusable-state inputs share this shape. Their semantics
+    come from the argument position where the loader is used.
+    """
+
+    def load(self) -> PartialTree | None:
+        """Return loaded tree data, or None if the source is empty."""
+        ...