atomicmemory 1.0.0__py3-none-any.whl

atomicmemory/search/chunking.py

@@ -0,0 +1,161 @@
+"""Text chunking helpers.
+
+Port of `atomicmemory-sdk/src/utils/chunking.ts`. Three strategies:
+
+- :func:`chunk_text_with_metadata` — character-window with optional
+  word-boundary preservation and overlap.
+- :func:`chunk_by_sentences` — split on ``.!?`` then re-group with
+  overlap.
+- :func:`chunk_by_paragraphs` — split on blank lines then re-group with
+  overlap.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ChunkOptions:
+    """Configuration for :func:`chunk_text_with_metadata`."""
+
+    chunk_size: int
+    chunk_overlap: int
+    preserve_words: bool = True
+    separator: str = " "
+    min_chunk_size: int | None = None
+
+
+@dataclass(frozen=True)
+class ChunkMetadata:
+    word_count: int
+    char_count: int
+    has_overlap: bool
+
+
+@dataclass(frozen=True)
+class ChunkResult:
+    text: str
+    index: int
+    start_offset: int
+    end_offset: int
+    metadata: ChunkMetadata
+
+
+def _word_count(text: str) -> int:
+    return len([w for w in re.split(r"\s+", text) if w])
+
+
+def _validate(options: ChunkOptions) -> int:
+    if options.chunk_size <= 0:
+        raise ValueError("chunk_size must be positive")
+    if options.chunk_overlap < 0:
+        raise ValueError("chunk_overlap must be >= 0")
+    if options.chunk_overlap >= options.chunk_size:
+        raise ValueError("chunk_overlap must be less than chunk_size")
+    raw_min = options.min_chunk_size if options.min_chunk_size is not None else max(1, options.chunk_size // 10)
+    return max(1, min(raw_min, options.chunk_size))
+
+
+def chunk_text(text: str, options: ChunkOptions) -> list[str]:
+    """Return chunks as plain strings; thin wrapper around metadata variant."""
+    return [c.text for c in chunk_text_with_metadata(text, options)]
+
+
+def chunk_text_with_metadata(text: str, options: ChunkOptions) -> list[ChunkResult]:
+    """Sliding-window chunker with optional word-boundary preservation."""
+    if not isinstance(text, str) or text == "":
+        return []
+    min_chunk_size = _validate(options)
+    if len(text) <= options.chunk_size:
+        trimmed = text.strip()
+        return [
+            ChunkResult(
+                text=trimmed,
+                index=0,
+                start_offset=0,
+                end_offset=len(text),
+                metadata=ChunkMetadata(
+                    word_count=_word_count(trimmed),
+                    char_count=len(text),
+                    has_overlap=False,
+                ),
+            )
+        ]
+    return _chunk_window(text, options, min_chunk_size)
+
+
+def _chunk_window(text: str, options: ChunkOptions, min_chunk_size: int) -> list[ChunkResult]:
+    chunks: list[ChunkResult] = []
+    current_offset = 0
+    chunk_index = 0
+    while current_offset < len(text):
+        end_offset = min(current_offset + options.chunk_size, len(text))
+        chunk_text_slice = text[current_offset:end_offset]
+        if options.preserve_words and end_offset < len(text):
+            last_sep = chunk_text_slice.rfind(options.separator)
+            if last_sep > min_chunk_size:
+                end_offset = current_offset + last_sep
+                chunk_text_slice = text[current_offset:end_offset]
+        trimmed = chunk_text_slice.strip()
+        if len(trimmed) >= min_chunk_size:
+            chunks.append(
+                ChunkResult(
+                    text=trimmed,
+                    index=chunk_index,
+                    start_offset=current_offset,
+                    end_offset=end_offset,
+                    metadata=ChunkMetadata(
+                        word_count=_word_count(trimmed),
+                        char_count=len(trimmed),
+                        has_overlap=chunk_index > 0,
+                    ),
+                )
+            )
+            chunk_index += 1
+        next_offset = end_offset - options.chunk_overlap
+        current_offset = end_offset if next_offset <= current_offset else next_offset
+        if end_offset >= len(text):
+            break
+    return chunks
+
+
+def chunk_by_sentences(text: str, max_sentences: int, overlap_sentences: int = 1) -> list[str]:
+    """Group sentences with overlap; sentence boundary = ``[.!?]+``."""
+    if not isinstance(text, str) or text == "":
+        return []
+    sentences = [s.strip() for s in re.split(r"[.!?]+", text) if s.strip()]
+    if len(sentences) <= max_sentences:
+        return [text]
+    chunks: list[str] = []
+    current = 0
+    while current < len(sentences):
+        end = min(current + max_sentences, len(sentences))
+        slice_ = sentences[current:end]
+        if slice_:
+            chunks.append(". ".join(slice_) + ".")
+        current = end - overlap_sentences
+        if current <= 0 or end >= len(sentences):
+            break
+    return chunks
+
+
+def chunk_by_paragraphs(text: str, max_paragraphs: int, overlap_paragraphs: int = 0) -> list[str]:
+    """Group paragraphs with overlap; paragraph boundary = blank lines."""
+    if not isinstance(text, str) or text == "":
+        return []
+    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
+    if len(paragraphs) <= max_paragraphs:
+        return [text]
+    chunks: list[str] = []
+    current = 0
+    while current < len(paragraphs):
+        end = min(current + max_paragraphs, len(paragraphs))
+        slice_ = paragraphs[current:end]
+        if slice_:
+            chunks.append("\n\n".join(slice_))
+        current = end - overlap_paragraphs
+        if current <= 0 or end >= len(paragraphs):
+            break
+    return chunks

atomicmemory/search/ranking.py

@@ -0,0 +1,94 @@
+"""Heuristic reranking for local semantic search results.
+
+Port of the rerank heuristics in
+`atomicmemory-sdk/src/search/semantic-search.ts:rerankResults`.
+Three signals: short-content boost, long-content penalty, recency
+boost. Pure functions; no I/O.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timezone
+
+
+@dataclass(frozen=True)
+class RankingConfig:
+    """Tunable reranking knobs."""
+
+    short_threshold: int = 500
+    short_boost: float = 1.1
+    long_threshold: int = 2000
+    long_penalty: float = 0.9
+    recency_window_days: int = 7
+    recency_boost: float = 1.05
+    score_cap: float = 1.0
+
+
+@dataclass
+class RankableHit:
+    """Minimal shape rerank operates on; field names match the TS port.
+
+    ``original_index`` lets callers map the reranked output back to their
+    own list without relying on content/score equality (rerank mutates
+    score, and content can collide between hits).
+    """
+
+    content: str
+    score: float
+    timestamp_seconds: float | None = None
+    original_index: int | None = None
+
+
+def _length_factor(content: str, config: RankingConfig) -> float:
+    length = len(content)
+    if length < config.short_threshold:
+        return config.short_boost
+    if length > config.long_threshold:
+        return config.long_penalty
+    return 1.0
+
+
+def _recency_factor(timestamp_seconds: float | None, *, now: datetime, config: RankingConfig) -> float:
+    if timestamp_seconds is None:
+        return 1.0
+    delta = now.timestamp() - timestamp_seconds
+    window_seconds = config.recency_window_days * 24 * 60 * 60
+    if delta < window_seconds:
+        return config.recency_boost
+    return 1.0
+
+
+def rerank(
+    hits: list[RankableHit],
+    *,
+    config: RankingConfig | None = None,
+    now: datetime | None = None,
+) -> list[RankableHit]:
+    """Apply length + recency heuristics; sort by adjusted score descending.
+
+    Args:
+        hits: Hits to rerank. Mutated in place is **not** allowed — a new
+            list is returned.
+        config: Tunable thresholds; defaults to :class:`RankingConfig`.
+        now: Reference time for recency. Defaults to UTC now.
+
+    Returns:
+        New list of hits with adjusted ``score`` values, sorted descending.
+    """
+    cfg = config or RankingConfig()
+    reference = now or datetime.now(tz=timezone.utc)
+    adjusted: list[RankableHit] = []
+    for hit in hits:
+        factor = _length_factor(hit.content, cfg) * _recency_factor(hit.timestamp_seconds, now=reference, config=cfg)
+        new_score = min(hit.score * factor, cfg.score_cap)
+        adjusted.append(
+            RankableHit(
+                content=hit.content,
+                score=new_score,
+                timestamp_seconds=hit.timestamp_seconds,
+                original_index=hit.original_index,
+            )
+        )
+    adjusted.sort(key=lambda h: h.score, reverse=True)
+    return adjusted

atomicmemory/search/semantic_search.py

@@ -0,0 +1,130 @@
+"""Local SemanticSearch orchestrator — embed → score → rank.
+
+Port of `atomicmemory-sdk/src/search/semantic-search.ts`. Brute-force
+linear scan against an in-memory list of :class:`StoredContext`. No
+HNSW / IVF — designed for small-to-medium local stores (<10k items)
+where simplicity beats indexing.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+from atomicmemory.search.ranking import RankableHit, RankingConfig, rerank
+from atomicmemory.search.similarity import batch_cosine_similarity, rank_by_similarity
+
+EmbedFn = Callable[[str], list[float]]
+"""Callable that turns a query string into an embedding vector."""
+
+
+@dataclass(frozen=True)
+class StoredContext:
+    """A single context record indexed for local search."""
+
+    id: str
+    content: str
+    embedding: list[float]
+    metadata: dict[str, Any] | None = None
+    timestamp: float | None = None
+    user_id: str | None = None
+
+
+@dataclass(frozen=True)
+class SemanticSearchResult:
+    context: StoredContext
+    score: float
+
+
+@dataclass(frozen=True)
+class SemanticSearchConfig:
+    """Search-time knobs."""
+
+    default_top_k: int = 10
+    default_threshold: float = 0.1
+    max_results: int = 100
+    reranking_enabled: bool = True
+    ranking: RankingConfig = field(default_factory=RankingConfig)
+
+
+class SemanticSearch:
+    """Brute-force semantic search over an in-memory context list."""
+
+    def __init__(
+        self,
+        embed_fn: EmbedFn,
+        *,
+        config: SemanticSearchConfig | None = None,
+    ) -> None:
+        self._embed = embed_fn
+        self._config = config or SemanticSearchConfig()
+
+    @property
+    def config(self) -> SemanticSearchConfig:
+        return self._config
+
+    def search(
+        self,
+        query: str,
+        contexts: Sequence[StoredContext],
+        *,
+        top_k: int | None = None,
+        threshold: float | None = None,
+        filter_fn: Callable[[StoredContext], bool] | None = None,
+        rerank_results: bool | None = None,
+    ) -> list[SemanticSearchResult]:
+        """Return up to ``top_k`` matches, optionally reranked.
+
+        Args:
+            query: Free-text query.
+            contexts: Candidate pool to score.
+            top_k: Cap on returned results. Defaults to
+                ``config.default_top_k``; clamped to ``config.max_results``.
+            threshold: Minimum cosine similarity. Defaults to
+                ``config.default_threshold``.
+            filter_fn: Optional predicate run before scoring (cheap pre-
+                filter, e.g. by user_id).
+            rerank_results: Override config's ``reranking_enabled``.
+        """
+        if not contexts:
+            return []
+        candidates = [ctx for ctx in contexts if filter_fn(ctx)] if filter_fn is not None else list(contexts)
+        if not candidates:
+            return []
+        query_embedding = self._embed(query)
+        similarities = batch_cosine_similarity(query_embedding, [ctx.embedding for ctx in candidates])
+        effective_threshold = threshold if threshold is not None else self._config.default_threshold
+        ranked_indices = rank_by_similarity(similarities, threshold=effective_threshold)
+        effective_top_k = min(
+            top_k if top_k is not None else self._config.default_top_k,
+            self._config.max_results,
+        )
+        ranked_indices = ranked_indices[:effective_top_k]
+        primary = [SemanticSearchResult(context=candidates[i], score=similarities[i]) for i in ranked_indices]
+        do_rerank = rerank_results if rerank_results is not None else self._config.reranking_enabled
+        if not do_rerank or not primary:
+            return primary
+        return self._apply_rerank(primary)
+
+    def _apply_rerank(self, results: list[SemanticSearchResult]) -> list[SemanticSearchResult]:
+        # Tag each rankable with its original-index so the reranked
+        # output maps back unambiguously even when scores or content
+        # collide.
+        rankable = [
+            RankableHit(
+                content=r.context.content,
+                score=r.score,
+                timestamp_seconds=r.context.timestamp,
+                original_index=i,
+            )
+            for i, r in enumerate(results)
+        ]
+        adjusted = rerank(rankable, config=self._config.ranking, now=datetime.now(tz=timezone.utc))
+        reordered: list[SemanticSearchResult] = []
+        for hit in adjusted:
+            if hit.original_index is None:
+                continue
+            reordered.append(SemanticSearchResult(context=results[hit.original_index].context, score=hit.score))
+        return reordered

atomicmemory/search/similarity.py

@@ -0,0 +1,110 @@
+"""Cosine similarity + top-k helpers.
+
+Port of `atomicmemory-sdk/src/search/similarity-calculator.ts`. Uses
+numpy for the dot/norm math; correctly handles zero-vectors and length
+mismatches.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+
+
+def cosine_similarity(a: Sequence[float], b: Sequence[float]) -> float:
+    """Cosine similarity between two equal-length vectors.
+
+    Args:
+        a: First embedding vector.
+        b: Second embedding vector.
+
+    Returns:
+        The cosine similarity in ``[-1, 1]``. Returns ``0.0`` when either
+        vector has zero L2 norm (avoids NaN).
+
+    Raises:
+        ValueError: When ``a`` and ``b`` have different lengths.
+    """
+    if len(a) != len(b):
+        raise ValueError(f"Vector length mismatch: {len(a)} != {len(b)}")
+    arr_a = np.asarray(a, dtype=np.float64)
+    arr_b = np.asarray(b, dtype=np.float64)
+    norm_a = float(np.linalg.norm(arr_a))
+    norm_b = float(np.linalg.norm(arr_b))
+    if norm_a == 0.0 or norm_b == 0.0:
+        return 0.0
+    return float(np.dot(arr_a, arr_b) / (norm_a * norm_b))
+
+
+def batch_cosine_similarity(
+    query: Sequence[float],
+    candidates: Sequence[Sequence[float]],
+) -> list[float]:
+    """Compute cosine similarity between ``query`` and every candidate.
+
+    Returns a list of similarities in candidate order; empty when
+    ``candidates`` is empty.
+    """
+    if not candidates:
+        return []
+    q = np.asarray(query, dtype=np.float64)
+    if q.size == 0:
+        return [0.0] * len(candidates)
+    matrix = np.asarray(candidates, dtype=np.float64)
+    if matrix.shape[1] != q.shape[0]:
+        raise ValueError(f"Query length {q.shape[0]} does not match candidate width {matrix.shape[1]}")
+    q_norm = float(np.linalg.norm(q))
+    if q_norm == 0.0:
+        return [0.0] * len(candidates)
+    candidate_norms = np.linalg.norm(matrix, axis=1)
+    dots = matrix @ q
+    out: list[float] = []
+    for value, norm in zip(dots, candidate_norms, strict=True):
+        if norm == 0.0:
+            out.append(0.0)
+        else:
+            out.append(float(value / (norm * q_norm)))
+    return out
+
+
+def rank_by_similarity(
+    similarities: Sequence[float],
+    *,
+    threshold: float | None = None,
+) -> list[int]:
+    """Return candidate indices sorted by similarity descending.
+
+    Optional ``threshold`` filters out indices whose similarity is
+    strictly less than the threshold.
+    """
+    indices = list(range(len(similarities)))
+    indices.sort(key=lambda i: similarities[i], reverse=True)
+    if threshold is None:
+        return indices
+    return [i for i in indices if similarities[i] >= threshold]
+
+
+def find_top_k(
+    query_embedding: Sequence[float],
+    candidates: Sequence[Sequence[float]],
+    k: int,
+    *,
+    metadata: Sequence[Any] | None = None,
+    threshold: float | None = None,
+) -> list[tuple[int, float, Any]]:
+    """Return up to ``k`` candidates as ``(index, similarity, metadata?)`` tuples.
+
+    Sorted by similarity descending. ``metadata`` (when supplied) is
+    aligned by index with ``candidates``.
+    """
+    if k <= 0:
+        return []
+    similarities = batch_cosine_similarity(query_embedding, candidates)
+    ranked = rank_by_similarity(similarities, threshold=threshold)[:k]
+    if metadata is None:
+        return [(i, similarities[i], None) for i in ranked]
+    if len(metadata) != len(candidates):
+        raise ValueError(f"metadata length {len(metadata)} does not match candidates length {len(candidates)}")
+    return [(i, similarities[i], metadata[i]) for i in ranked]

atomicmemory/storage/__init__.py

@@ -0,0 +1,63 @@
+"""Backend artifact-storage API.
+
+This package is the Python peer of `atomicmemory-sdk/src/storage`. It
+contains the direct artifact-storage clients and types for
+``/v1/storage/artifacts/*``. Local key/value cache adapters live under
+``atomicmemory.kv_cache``.
+"""
+
+from atomicmemory.storage.async_client import AsyncStorageClient
+from atomicmemory.storage.client import StorageClient
+from atomicmemory.storage.errors import (
+    ArtifactInUseError,
+    ArtifactNotFoundError,
+    FilecoinDirectStorageNotSupportedError,
+    PointerContentNotManagedError,
+    StorageClientError,
+    UnsupportedCapabilityError,
+)
+from atomicmemory.storage.types import (
+    ArtifactHead,
+    ArtifactMetadata,
+    ArtifactRange,
+    ArtifactRef,
+    DeleteArtifactOptions,
+    DeleteArtifactPolicy,
+    DeleteArtifactResult,
+    PutArtifactInput,
+    PutManagedInput,
+    PutPointerInput,
+    StorageArtifactStatus,
+    StorageCapabilities,
+    StorageClientConfig,
+    StoredArtifact,
+    VerificationResult,
+    VerifyArtifactOptions,
+)
+
+__all__ = [
+    "ArtifactHead",
+    "ArtifactInUseError",
+    "ArtifactMetadata",
+    "ArtifactNotFoundError",
+    "ArtifactRange",
+    "ArtifactRef",
+    "AsyncStorageClient",
+    "DeleteArtifactOptions",
+    "DeleteArtifactPolicy",
+    "DeleteArtifactResult",
+    "FilecoinDirectStorageNotSupportedError",
+    "PointerContentNotManagedError",
+    "PutArtifactInput",
+    "PutManagedInput",
+    "PutPointerInput",
+    "StorageArtifactStatus",
+    "StorageCapabilities",
+    "StorageClient",
+    "StorageClientConfig",
+    "StorageClientError",
+    "StoredArtifact",
+    "UnsupportedCapabilityError",
+    "VerificationResult",
+    "VerifyArtifactOptions",
+]