embed-tree 0.0.6__tar.gz

embed_tree-0.0.6/PKG-INFO

@@ -0,0 +1,182 @@
+Metadata-Version: 2.4
+Name: embed-tree
+Version: 0.0.6
+Summary: Incremental hierarchical clustering tree over content embeddings
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.23
+Requires-Dist: scikit-learn>=1.3
+Requires-Dist: pydantic>=2
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == "dev"
+Provides-Extra: openai
+Requires-Dist: openai>=1; extra == "openai"
+Provides-Extra: local
+Requires-Dist: sentence-transformers>=2; extra == "local"
+Provides-Extra: sql
+Requires-Dist: sqlalchemy>=2; extra == "sql"
+
+# embed-tree
+
+`embed-tree` turns content embeddings into a browsable, labeled hierarchy.
+It is useful when you have documents, notes, tickets, search results, or any
+other text-like records and want a small taxonomy that a person can inspect.
+
+The library stays model-agnostic: you provide an embedder, and `embed-tree`
+handles clustering, labeling, querying, deletion, and persistence.
+
+## Install
+
+```bash
+pip install embed-tree
+```
+
+Optional adapters:
+
+```bash
+pip install "embed-tree[openai]"  # OpenAI embeddings
+pip install "embed-tree[local]"   # sentence-transformers embeddings
+pip install "embed-tree[sql]"     # SQLAlchemy loaders/persisters
+```
+
+## Quick Start
+
+```python
+from embed_tree import EmbedTree, FakeEmbeddingProvider, TreeConfig
+
+embedder = FakeEmbeddingProvider(dim=32)  # deterministic demo embedder
+tree = EmbedTree(
+    embedder=embedder,
+    config=TreeConfig(max_branches=5, leaf_capacity=10),
+)
+
+tree.add_batch(
+    [
+        "Write import pipeline documentation",
+        "Fix login session refresh",
+        "Reduce report query latency",
+        "Add retry handling to data ingestion",
+    ]
+)
+
+tree.organize()       # rebuild a clean hierarchy and label every node
+print(tree.show())    # human-readable outline
+```
+
+Use a real embedding provider in production:
+
+```python
+from embed_tree import EmbedTree, OpenAIEmbeddingProvider
+
+embedder = OpenAIEmbeddingProvider(
+    model="text-embedding-3-small",
+    api_key="...",
+)
+
+tree = EmbedTree(embedder=embedder)
+tree.add("Some document text", payload={"source": "docs"})
+```
+
+## Core API
+
+`EmbedTree` is the main entry point.
+
+```python
+tree.add(content, item_id=None, payload=None, text=None)
+tree.add_batch(contents, item_ids=None, payloads=None, texts=None)
+tree.add_node(content_node)
+tree.add_nodes(content_nodes)
+tree.add_partial_tree(partial_tree)
+
+tree.organize(tagger=None)
+tree.rebalance()
+tree.label(tagger=None)
+
+tree.query(content, k=10, exhaustive=False)
+tree.remove(item_id)
+tree.remove_batch(item_ids)
+
+tree.show(max_items=3)
+tree.to_dict(max_items=5)
+tree.get_tree()
+len(tree)
+```
+
+`content` is what gets embedded. `text` is the human-readable string used in
+labels and browse output; it defaults to `content` when `content` is a string.
+`payload` is returned in query results and exported browse data.
+
+## Configuration
+
+Configuration is explicit and code-driven through `TreeConfig`. It does not
+read environment variables.
+
+```python
+from embed_tree import LLMConfig, RebalanceConfig, TreeConfig
+
+config = TreeConfig(
+    max_branches=5,
+    leaf_capacity=10,
+    rebalance=RebalanceConfig(enabled=True, every_n_inserts=10_000),
+    llm=LLMConfig(provider="none"),  # default keyword labels, no network
+)
+```
+
+Defaults are tuned for readable taxonomies: small fan-out and small leaves.
+Raise `max_branches` and `leaf_capacity` when using the tree primarily as a
+retrieval index.
+
+## Querying
+
+```python
+hits = tree.query("related content", k=5)
+# [(item_id, distance, payload), ...]
+
+exact_hits = tree.query("related content", k=5, exhaustive=True)
+```
+
+Default queries route to one leaf and rank items there, which is fast but
+approximate. `exhaustive=True` scans every item for exact nearest neighbors.
+
+## Persistence
+
+```python
+from embed_tree import EmbedTree, FileTreeStore
+
+tree = EmbedTree(
+    embedder=embedder,
+    store=FileTreeStore("./tree.json"),
+)
+```
+
+`FileTreeStore` saves an atomic JSON snapshot after writes and reloads it when
+the tree is constructed again.
+
+## Labeling
+
+Without extra configuration, node labels are generated locally from keywords.
+For LLM labels:
+
+```python
+from embed_tree import LLMConfig, TreeConfig
+
+config = TreeConfig(
+    llm=LLMConfig(provider="openai", model="gpt-4o-mini", api_key="...")
+)
+```
+
+You can also pass a custom `tagger: Callable[[list[str]], str]` to
+`EmbedTree(..., tagger=...)`, `tree.label(tagger=...)`, or
+`tree.organize(tagger=...)`.
+
+## More Documentation
+
+See [docs/API.md](docs/API.md) for the fuller API reference, provider details,
+loader/persister abstractions, PCA options, and extension points.
+
+## Development
+
+```bash
+uv sync --extra dev
+uv run pytest -q
+```

embed_tree-0.0.6/README.md

@@ -0,0 +1,164 @@
+# embed-tree
+
+`embed-tree` turns content embeddings into a browsable, labeled hierarchy.
+It is useful when you have documents, notes, tickets, search results, or any
+other text-like records and want a small taxonomy that a person can inspect.
+
+The library stays model-agnostic: you provide an embedder, and `embed-tree`
+handles clustering, labeling, querying, deletion, and persistence.
+
+## Install
+
+```bash
+pip install embed-tree
+```
+
+Optional adapters:
+
+```bash
+pip install "embed-tree[openai]"  # OpenAI embeddings
+pip install "embed-tree[local]"   # sentence-transformers embeddings
+pip install "embed-tree[sql]"     # SQLAlchemy loaders/persisters
+```
+
+## Quick Start
+
+```python
+from embed_tree import EmbedTree, FakeEmbeddingProvider, TreeConfig
+
+embedder = FakeEmbeddingProvider(dim=32)  # deterministic demo embedder
+tree = EmbedTree(
+    embedder=embedder,
+    config=TreeConfig(max_branches=5, leaf_capacity=10),
+)
+
+tree.add_batch(
+    [
+        "Write import pipeline documentation",
+        "Fix login session refresh",
+        "Reduce report query latency",
+        "Add retry handling to data ingestion",
+    ]
+)
+
+tree.organize()       # rebuild a clean hierarchy and label every node
+print(tree.show())    # human-readable outline
+```
+
+Use a real embedding provider in production:
+
+```python
+from embed_tree import EmbedTree, OpenAIEmbeddingProvider
+
+embedder = OpenAIEmbeddingProvider(
+    model="text-embedding-3-small",
+    api_key="...",
+)
+
+tree = EmbedTree(embedder=embedder)
+tree.add("Some document text", payload={"source": "docs"})
+```
+
+## Core API
+
+`EmbedTree` is the main entry point.
+
+```python
+tree.add(content, item_id=None, payload=None, text=None)
+tree.add_batch(contents, item_ids=None, payloads=None, texts=None)
+tree.add_node(content_node)
+tree.add_nodes(content_nodes)
+tree.add_partial_tree(partial_tree)
+
+tree.organize(tagger=None)
+tree.rebalance()
+tree.label(tagger=None)
+
+tree.query(content, k=10, exhaustive=False)
+tree.remove(item_id)
+tree.remove_batch(item_ids)
+
+tree.show(max_items=3)
+tree.to_dict(max_items=5)
+tree.get_tree()
+len(tree)
+```
+
+`content` is what gets embedded. `text` is the human-readable string used in
+labels and browse output; it defaults to `content` when `content` is a string.
+`payload` is returned in query results and exported browse data.
+
+## Configuration
+
+Configuration is explicit and code-driven through `TreeConfig`. It does not
+read environment variables.
+
+```python
+from embed_tree import LLMConfig, RebalanceConfig, TreeConfig
+
+config = TreeConfig(
+    max_branches=5,
+    leaf_capacity=10,
+    rebalance=RebalanceConfig(enabled=True, every_n_inserts=10_000),
+    llm=LLMConfig(provider="none"),  # default keyword labels, no network
+)
+```
+
+Defaults are tuned for readable taxonomies: small fan-out and small leaves.
+Raise `max_branches` and `leaf_capacity` when using the tree primarily as a
+retrieval index.
+
+## Querying
+
+```python
+hits = tree.query("related content", k=5)
+# [(item_id, distance, payload), ...]
+
+exact_hits = tree.query("related content", k=5, exhaustive=True)
+```
+
+Default queries route to one leaf and rank items there, which is fast but
+approximate. `exhaustive=True` scans every item for exact nearest neighbors.
+
+## Persistence
+
+```python
+from embed_tree import EmbedTree, FileTreeStore
+
+tree = EmbedTree(
+    embedder=embedder,
+    store=FileTreeStore("./tree.json"),
+)
+```
+
+`FileTreeStore` saves an atomic JSON snapshot after writes and reloads it when
+the tree is constructed again.
+
+## Labeling
+
+Without extra configuration, node labels are generated locally from keywords.
+For LLM labels:
+
+```python
+from embed_tree import LLMConfig, TreeConfig
+
+config = TreeConfig(
+    llm=LLMConfig(provider="openai", model="gpt-4o-mini", api_key="...")
+)
+```
+
+You can also pass a custom `tagger: Callable[[list[str]], str]` to
+`EmbedTree(..., tagger=...)`, `tree.label(tagger=...)`, or
+`tree.organize(tagger=...)`.
+
+## More Documentation
+
+See [docs/API.md](docs/API.md) for the fuller API reference, provider details,
+loader/persister abstractions, PCA options, and extension points.
+
+## Development
+
+```bash
+uv sync --extra dev
+uv run pytest -q
+```

embed_tree-0.0.6/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-0.0.6/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-0.0.6/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-0.0.6/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-0.0.6/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-0.0.6/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-0.0.6/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-0.0.6/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])
+