memkernel 0.2.0__py3-none-any.whl

memkernel/retrieval/search.py

@@ -0,0 +1,235 @@
+# Copyright (c) 2026 EGOGE
+# SPDX-License-Identifier: Apache-2.0
+"""Hybrid search engine: dense (ChromaDB) + sparse (BM25) fused via RRF."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from memkernel.core.models import ArtifactRef, Memory, SearchResult
+from memkernel.memory.lifecycle import recency_score
+
+if TYPE_CHECKING:
+    from memkernel.config import MemKernelSettings
+    from memkernel.core.models import SearchQuery
+    from memkernel.core.ports import EmbeddingPort, SparseIndexPort, VectorStorePort
+
+
+# Map ChromaDB metadata field names to Memory attribute names.
+_FIELD_MAP: dict[str, str] = {
+    "type": "type",
+    "tier": "tier",
+    "session_id": "session_id",
+    "checkpoint": "checkpoint",
+    "project_id": "project_id",
+}
+
+
+class HybridSearchService:
+    """Fuses dense and sparse search results via Reciprocal Rank Fusion (RRF).
+
+    Dense search: ChromaDB cosine-similarity query via VectorStorePort.
+    Sparse search: BM25Plus over the full corpus via SparseIndexPort, then
+                   post-filtered against the ChromaDB metadata via get_by_id.
+    Fusion: RRF with configurable k (default 60).
+    Scoring: the fused relevance rank is re-weighted by recency (per-tier
+             exponential decay) and importance, following the
+             relevance + recency + importance retrieval function of
+             Generative Agents (Park et al., 2023). Weights come from settings.
+    """
+
+    def __init__(
+        self,
+        vector_store: VectorStorePort,
+        embedding: EmbeddingPort,
+        sparse_index: SparseIndexPort,
+        settings: MemKernelSettings,
+    ) -> None:
+        self._vector_store = vector_store
+        self._embedding = embedding
+        self._sparse_index = sparse_index
+        self._settings = settings
+
+    def search(self, query: SearchQuery) -> list[ArtifactRef]:
+        """Run hybrid search and return up to query.top_k ArtifactRefs."""
+        where = _build_where_filter(query)
+        over_fetch = max(1, query.top_k * 3)
+
+        # Dense pass
+        query_embedding = self._embedding.embed([query.query_text])[0]
+        dense_results = self._vector_store.search_dense(query_embedding, over_fetch, where)
+
+        # Sparse pass — score all, then post-filter to honour metadata filters
+        sparse_all = self._sparse_index.search_all(query.query_text)
+        sparse_results = filter_sparse_results(sparse_all, where, self._vector_store)
+        sparse_results = sparse_results[:over_fetch]
+
+        # RRF fusion
+        fused = rrf_fuse(self._settings.rrf_k, dense_results, sparse_results)
+
+        # Hydrate every fused candidate (not just top_k): re-weighting can
+        # promote a lower-ranked but recent/important memory into the top_k.
+        refs = build_artifact_refs(fused, len(fused), self._vector_store, self._settings.preview_length)
+        refs = rescore(refs, self._settings, self._vector_store)
+        return refs[: query.top_k]
+
+
+# ------------------------------------------------------------------
+# Module-level helpers (pure functions, easier to unit-test independently)
+# ------------------------------------------------------------------
+
+
+def _build_where_filter(query: SearchQuery) -> dict[str, Any] | None:
+    """Construct a ChromaDB ``where`` filter dict from a SearchQuery.
+
+    Returns None when no filters are active (avoids passing an empty dict
+    to ChromaDB which would cause a validation error).
+    """
+    conditions: list[dict[str, Any]] = []
+    if query.type_filter is not None:
+        conditions.append({"type": query.type_filter.value})
+    if query.tier_filter is not None:
+        conditions.append({"tier": query.tier_filter.value})
+    if query.session_id_filter is not None:
+        conditions.append({"session_id": query.session_id_filter})
+    if query.checkpoint_only:
+        conditions.append({"checkpoint": True})
+
+    if not conditions:
+        return None
+    if len(conditions) == 1:
+        return conditions[0]
+    return {"$and": conditions}
+
+
+def rrf_fuse(
+    k: int,
+    *result_lists: list[SearchResult],
+) -> list[tuple[str, float]]:
+    """Reciprocal Rank Fusion across arbitrarily many ranked lists.
+
+    score(d) = sum of  1 / (k + rank_i(d))  for each list that contains d.
+
+    Returns (ref_id, fused_score) pairs sorted descending by score.
+    """
+    scores: dict[str, float] = {}
+    for results in result_lists:
+        for result in results:
+            rrf_score = 1.0 / (k + result.rank)
+            scores[result.ref_id] = scores.get(result.ref_id, 0.0) + rrf_score
+    return sorted(scores.items(), key=lambda x: x[1], reverse=True)
+
+
+def rescore(
+    refs: list[ArtifactRef],
+    settings: MemKernelSettings,
+    vector_store: VectorStorePort,
+) -> list[ArtifactRef]:
+    """Re-rank fused results by relevance + recency + importance.
+
+    The fused RRF score is min-max normalised to [0, 1] so the three axes
+    are commensurable; then::
+
+        final = w_rel * relevance + w_rec * recency + w_imp * importance
+
+    Recency uses the per-tier half-life from settings (long-term memories do
+    not decay). The ``last_accessed`` timestamp is consulted so that a memory
+    the agent keeps recalling stays fresh (rehearsal).
+    """
+    if not refs:
+        return refs
+    scores = [r.score for r in refs]
+    lo, hi = min(scores), max(scores)
+    span = hi - lo
+    rescored: list[ArtifactRef] = []
+    for ref in refs:
+        relevance = (ref.score - lo) / span if span > 0 else 1.0
+        memory = vector_store.get_by_id(ref.ref_id)
+        last_accessed = memory.last_accessed if memory is not None else None
+        recency = recency_score(ref.type, ref.timestamp, settings, last_accessed=last_accessed)
+        final = (
+            settings.recall_weight_relevance * relevance
+            + settings.recall_weight_recency * recency
+            + settings.recall_weight_importance * ref.importance
+        )
+        rescored.append(ref.model_copy(update={"score": final}))
+    rescored.sort(key=lambda r: r.score, reverse=True)
+    return rescored
+
+
+def _metadata_matches(memory: Memory, where: dict[str, Any]) -> bool:
+    """Check whether a Memory satisfies a ChromaDB-style where filter.
+
+    Supports the subset of ChromaDB filter syntax used by this module:
+      - Simple equality: {"field": value}
+      - $and: [cond, ...]
+    """
+    if "$and" in where:
+        return all(_metadata_matches(memory, cond) for cond in where["$and"])
+
+    for key, value in where.items():
+        actual = getattr(memory, _FIELD_MAP.get(key, key), None)
+        if actual is None:
+            return False
+        # MemoryType is a StrEnum — compare via string value
+        if hasattr(actual, "value"):
+            actual = actual.value
+        if actual != value:
+            return False
+    return True
+
+
+def filter_sparse_results(
+    results: list[SearchResult],
+    where: dict[str, Any] | None,
+    vector_store: VectorStorePort,
+) -> list[SearchResult]:
+    """Post-filter BM25 results against ChromaDB metadata.
+
+    For each BM25 result, fetch the Memory from the vector store.
+    Stale entries (get_by_id returns None) are skipped silently.
+    Surviving results are re-ranked 1-based.
+    """
+    kept: list[SearchResult] = []
+    for result in results:
+        memory = vector_store.get_by_id(result.ref_id)
+        if memory is None:
+            continue
+        if where is not None and not _metadata_matches(memory, where):
+            continue
+        kept.append(result)
+
+    return [SearchResult(ref_id=r.ref_id, score=r.score, rank=new_rank) for new_rank, r in enumerate(kept, start=1)]
+
+
+def build_artifact_refs(
+    fused: list[tuple[str, float]],
+    top_k: int,
+    vector_store: VectorStorePort,
+    preview_length: int,
+) -> list[ArtifactRef]:
+    """Fetch full Memory for each fused result and build ArtifactRef list.
+
+    Stale entries (get_by_id returns None) are skipped silently.
+    Stops after collecting top_k valid results.
+    """
+    refs: list[ArtifactRef] = []
+    for ref_id, score in fused:
+        if len(refs) >= top_k:
+            break
+        memory = vector_store.get_by_id(ref_id)
+        if memory is None:
+            continue
+        refs.append(
+            ArtifactRef(
+                ref_id=memory.ref_id,
+                preview=memory.content[:preview_length],
+                score=score,
+                token_count=memory.token_count,
+                type=memory.type,
+                timestamp=memory.timestamp,
+                tier=memory.tier,
+                importance=memory.importance,
+            )
+        )
+    return refs

memkernel/retrieval/sparse.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2026 EGOGE
+# SPDX-License-Identifier: Apache-2.0
+"""BM25 sparse search index implementing SparseIndexPort."""
+
+from __future__ import annotations
+
+import logging
+import re
+
+from rank_bm25 import BM25Plus  # type: ignore[import-untyped]
+
+from memkernel.core.models import SearchResult
+
+logger = logging.getLogger(__name__)
+
+
+class BM25Index:
+    """In-memory BM25 index backed by rank_bm25.BM25Plus.
+
+    BM25Plus uses a lower-bounded IDF that always yields non-negative scores,
+    avoiding the negative-score pathology of BM25Okapi on small corpora.
+    BM25Plus is immutable after construction, so the index is rebuilt lazily
+    whenever the document corpus changes (dirty flag pattern).
+    """
+
+    def __init__(self) -> None:
+        self._documents: dict[str, list[str]] = {}  # ref_id -> token list
+        self._ref_ids: list[str] = []
+        self._index: BM25Plus | None = None
+        self._dirty: bool = True
+        self._needs_rebuild_flag: bool = False
+
+    # ------------------------------------------------------------------
+    # Mutation
+    # ------------------------------------------------------------------
+
+    def add(self, ref_id: str, content: str) -> None:
+        """Index a document by ref_id. Replaces existing entry if present."""
+        self._documents[ref_id] = self._tokenize(content)
+        self._dirty = True
+
+    def remove(self, ref_id: str) -> None:
+        """Remove a document from the index. No-op if ref_id is absent."""
+        self._documents.pop(ref_id, None)
+        self._dirty = True
+
+    def rebuild(self, documents: list[tuple[str, str]]) -> None:
+        """Replace the entire corpus and eagerly rebuild the BM25 index."""
+        self._documents.clear()
+        for ref_id, content in documents:
+            self._documents[ref_id] = self._tokenize(content)
+        self._needs_rebuild_flag = False
+        self._dirty = True
+        self._ensure_index()
+
+    # ------------------------------------------------------------------
+    # Search
+    # ------------------------------------------------------------------
+
+    def search(self, query: str, top_k: int) -> list[SearchResult]:
+        """Return up to top_k results ranked by BM25 score."""
+        return self.search_all(query)[:top_k]
+
+    def search_all(self, query: str) -> list[SearchResult]:
+        """Score the entire corpus and return all documents with score > 0.
+
+        Results are sorted by descending score and assigned 1-based ranks.
+        """
+        self._ensure_index()
+        if self._index is None or not self._ref_ids:
+            return []
+
+        query_tokens = self._tokenize(query)
+        raw_scores = self._index.get_scores(query_tokens)
+
+        # BM25Plus uses a lower-bounded IDF that gives every document a
+        # positive baseline score, even those that don't contain the query
+        # term. Guard with a token-overlap check so only documents that
+        # actually share at least one token with the query are returned.
+        query_token_set = set(query_tokens)
+        scored: list[tuple[int, float]] = sorted(
+            (
+                (i, float(s))
+                for i, (s, rid) in enumerate(zip(raw_scores, self._ref_ids, strict=True))
+                if query_token_set & set(self._documents[rid])
+            ),
+            key=lambda t: t[1],
+            reverse=True,
+        )
+
+        return [
+            SearchResult(ref_id=self._ref_ids[i], score=score, rank=rank + 1) for rank, (i, score) in enumerate(scored)
+        ]
+
+    # ------------------------------------------------------------------
+    # Rebuild marker
+    # ------------------------------------------------------------------
+
+    @property
+    def needs_rebuild(self) -> bool:
+        """True if an external caller has signalled that the index needs a full rebuild."""
+        return self._needs_rebuild_flag
+
+    def mark_needs_rebuild(self) -> None:
+        """Signal that the caller should supply fresh documents via rebuild()."""
+        self._needs_rebuild_flag = True
+
+    # ------------------------------------------------------------------
+    # Dunder
+    # ------------------------------------------------------------------
+
+    def __len__(self) -> int:
+        return len(self._documents)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _ensure_index(self) -> None:
+        """Rebuild BM25Plus from the current corpus if dirty."""
+        if not self._dirty:
+            return
+        self._ref_ids = list(self._documents.keys())
+        corpus = [self._documents[rid] for rid in self._ref_ids]
+        self._index = BM25Plus(corpus) if corpus else None
+        self._dirty = False
+
+    _TOKEN_RE = re.compile(r"[a-zA-Z0-9_]+")
+
+    @staticmethod
+    def _tokenize(text: str) -> list[str]:
+        """Split on non-alphanumeric boundaries so code identifiers are isolated.
+
+        'foo()' -> ['foo'], 'Class.method' -> ['class', 'method'],
+        'path/file.py' -> ['path', 'file', 'py']
+        """
+        return [t.lower() for t in BM25Index._TOKEN_RE.findall(text)]

memkernel/storage/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) 2026 EGOGE
+# SPDX-License-Identifier: Apache-2.0
+"""Storage adapters for MemKernel."""

memkernel/storage/chroma.py

@@ -0,0 +1,256 @@
+# Copyright (c) 2026 EGOGE
+# SPDX-License-Identifier: Apache-2.0
+"""ChromaDB adapter implementing VectorStorePort."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from typing import TYPE_CHECKING, Any
+
+import chromadb
+
+from memkernel.core.models import DEFAULT_IMPORTANCE, Memory, MemoryTier, MemoryType, SearchResult, tier_of
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+
+    from chromadb.api import ClientAPI
+    from chromadb.api.models.Collection import Collection
+
+    from memkernel.config import MemKernelSettings
+
+
+_REQUIRED_METADATA_KEYS: frozenset[str] = frozenset(
+    {"project_id", "session_id", "timestamp", "type", "source", "hash", "token_count"}
+)
+
+
+def _to_chroma_metadata(metadata: dict[str, Any]) -> dict[str, str | int | float | bool]:
+    """Normalise a metadata dict to ChromaDB-acceptable scalar types."""
+    result: dict[str, str | int | float | bool] = {}
+    for key, value in metadata.items():
+        if isinstance(value, bool | int | float):
+            result[key] = value
+        elif isinstance(value, datetime):
+            result[key] = value.isoformat()
+        else:
+            result[key] = str(value)
+    return result
+
+
+def _memory_from_chroma(
+    ids: list[str],
+    documents: list[str],
+    metadatas: list[Mapping[str, Any]],
+    index: int,
+) -> Memory:
+    """Reconstruct a Memory from ChromaDB get/query result lists."""
+    meta = metadatas[index]
+    raw_ts = meta.get("timestamp", "")
+    timestamp = datetime.fromisoformat(str(raw_ts)) if raw_ts else datetime.now(tz=UTC)
+    chunk_refs = _split_csv(meta.get("chunk_refs"))
+    memory_type = MemoryType(str(meta.get("type", "episode")))
+    # Records written by memkernel 0.1.x carry no tier/importance metadata;
+    # derive defaults from the type so old stores keep working unchanged.
+    raw_tier = meta.get("tier")
+    tier = MemoryTier(str(raw_tier)) if raw_tier else tier_of(memory_type)
+    raw_importance = meta.get("importance")
+    importance = float(raw_importance) if raw_importance is not None else DEFAULT_IMPORTANCE[memory_type]
+    raw_last = meta.get("last_accessed")
+    last_accessed = datetime.fromisoformat(str(raw_last)) if raw_last else None
+    return Memory(
+        ref_id=ids[index],
+        content=documents[index],
+        type=memory_type,
+        project_id=str(meta.get("project_id", "")),
+        session_id=str(meta.get("session_id", "")),
+        timestamp=timestamp,
+        source=str(meta.get("source", "")),
+        content_hash=str(meta.get("hash", "")),
+        token_count=int(meta.get("token_count", 0)),
+        checkpoint=bool(meta.get("checkpoint", False)),
+        chunk_refs=chunk_refs,
+        tier=tier,
+        importance=importance,
+        access_count=int(meta.get("access_count", 0) or 0),
+        last_accessed=last_accessed,
+        tags=_split_csv(meta.get("tags")),
+        supersedes=str(meta["supersedes"]) if meta.get("supersedes") else None,
+        promoted_from=str(meta["promoted_from"]) if meta.get("promoted_from") else None,
+    )
+
+
+def _split_csv(raw: Any) -> list[str]:
+    """Split a comma-separated metadata value into a list (ChromaDB has no list type)."""
+    if not raw:
+        return []
+    return [part for part in str(raw).split(",") if part]
+
+
+class ChromaVectorStore:
+    """ChromaDB-backed implementation of VectorStorePort.
+
+    Production usage:
+        store = ChromaVectorStore(settings)
+
+    Test usage (in-memory, no persistence):
+        store = ChromaVectorStore(settings, client=chromadb.EphemeralClient())
+    """
+
+    def __init__(
+        self,
+        settings: MemKernelSettings,
+        client: ClientAPI | None = None,
+    ) -> None:
+        if client is not None:
+            self._client: ClientAPI = client
+        else:
+            persist_path = settings.chroma_persist_path.strip()
+            if persist_path:
+                self._client = chromadb.PersistentClient(path=persist_path)
+            else:
+                self._client = chromadb.EphemeralClient()
+
+        collection_name = f"{settings.chroma_collection_prefix}_{settings.project_id}"
+        self._collection: Collection = self._client.get_or_create_collection(
+            name=collection_name,
+            metadata={"hnsw:space": "cosine"},
+        )
+
+    # ------------------------------------------------------------------
+    # VectorStorePort methods
+    # ------------------------------------------------------------------
+
+    def upsert(
+        self,
+        ref_id: str,
+        content: str,
+        embedding: list[float],
+        metadata: dict[str, Any],
+    ) -> None:
+        """Insert or update a memory in the collection."""
+        missing = _REQUIRED_METADATA_KEYS - metadata.keys()
+        if missing:
+            raise ValueError(f"upsert missing required metadata keys: {sorted(missing)}")
+        embeddings: list[Sequence[float]] = [embedding]
+        self._collection.upsert(
+            ids=[ref_id],
+            embeddings=embeddings,
+            documents=[content],
+            metadatas=[_to_chroma_metadata(metadata)],
+        )
+
+    def search_dense(
+        self,
+        query_embedding: list[float],
+        top_k: int,
+        where: dict[str, Any] | None = None,
+    ) -> list[SearchResult]:
+        """Query by vector similarity. Returns results ordered by score descending."""
+        query_embeddings: list[Sequence[float]] = [query_embedding]
+        kwargs: dict[str, Any] = {
+            "query_embeddings": query_embeddings,
+            "n_results": top_k,
+            "include": ["metadatas", "documents", "distances"],
+        }
+        if where is not None:
+            kwargs["where"] = where
+
+        results = self._collection.query(**kwargs)
+
+        # ChromaDB returns nested lists; unpack the first (only) query row.
+        ids: list[str] = (results.get("ids") or [[]])[0]
+        distances: list[float] = (results.get("distances") or [[]])[0]
+
+        search_results: list[SearchResult] = []
+        for rank_idx, (ref_id, distance) in enumerate(zip(ids, distances, strict=True), start=1):
+            score = 1.0 - distance
+            search_results.append(SearchResult(ref_id=ref_id, score=score, rank=rank_idx))
+
+        return search_results
+
+    def get_by_id(self, ref_id: str) -> Memory | None:
+        """Fetch a single memory by its ref_id. Returns None if not found."""
+        result = self._collection.get(
+            ids=[ref_id],
+            include=["metadatas", "documents"],
+        )
+        ids: list[str] = result.get("ids") or []
+        raw_docs = result.get("documents")
+        documents: list[str] = list(raw_docs) if raw_docs is not None else []
+        raw_metas = result.get("metadatas")
+        metadatas: list[Mapping[str, Any]] = list(raw_metas) if raw_metas is not None else []
+
+        if not ids:
+            return None
+
+        return _memory_from_chroma(ids, documents, metadatas, 0)
+
+    def delete(self, ref_id: str) -> bool:
+        """Delete a memory by ref_id. Returns True if it existed, False otherwise."""
+        existing = self.get_by_id(ref_id)
+        if existing is None:
+            return False
+        self._collection.delete(ids=[ref_id])
+        return True
+
+    def list_all(self, where: dict[str, Any] | None = None) -> list[Memory]:
+        """Return all memories, optionally filtered by metadata."""
+        kwargs: dict[str, Any] = {"include": ["metadatas", "documents"]}
+        if where is not None:
+            kwargs["where"] = where
+
+        result = self._collection.get(**kwargs)
+        ids: list[str] = result.get("ids") or []
+        raw_docs = result.get("documents")
+        documents: list[str] = list(raw_docs) if raw_docs is not None else []
+        raw_metas = result.get("metadatas")
+        metadatas: list[Mapping[str, Any]] = list(raw_metas) if raw_metas is not None else []
+
+        return [_memory_from_chroma(ids, documents, metadatas, i) for i in range(len(ids))]
+
+    def find_by_hash(self, content_hash: str, memory_type: str) -> Memory | None:
+        """Find a memory by content hash + type. Returns first match or None."""
+        result = self._collection.get(
+            where={"$and": [{"hash": content_hash}, {"type": memory_type}]},
+            include=["metadatas", "documents"],
+        )
+        ids: list[str] = result.get("ids") or []
+        raw_docs = result.get("documents")
+        documents: list[str] = list(raw_docs) if raw_docs is not None else []
+        raw_metas = result.get("metadatas")
+        metadatas: list[Mapping[str, Any]] = list(raw_metas) if raw_metas is not None else []
+
+        if not ids:
+            return None
+
+        return _memory_from_chroma(ids, documents, metadatas, 0)
+
+    def update_metadata(self, ref_id: str, metadata: dict[str, Any]) -> None:
+        """Fetch-and-merge: caller-specified keys overwrite existing; all other keys preserved."""
+        existing_result = self._collection.get(
+            ids=[ref_id],
+            include=["metadatas"],
+        )
+        existing_ids: list[str] = existing_result.get("ids") or []
+        raw_existing_metas = existing_result.get("metadatas")
+        existing_metadatas: list[Mapping[str, Any]] = list(raw_existing_metas) if raw_existing_metas is not None else []
+
+        if not existing_ids:
+            return
+
+        merged: dict[str, Any] = dict(existing_metadatas[0])
+        merged.update(metadata)
+
+        self._collection.update(
+            ids=[ref_id],
+            metadatas=[_to_chroma_metadata(merged)],
+        )
+
+    def count(self) -> int:
+        """Return the total number of memories in the collection."""
+        return self._collection.count()
+
+    def list_collections(self) -> list[dict[str, Any]]:
+        """List all collections accessible via this client with their record counts."""
+        return [{"name": coll.name, "count": coll.count()} for coll in self._client.list_collections()]

memkernel/wiki/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2026 EGOGE
+# SPDX-License-Identifier: Apache-2.0
+"""LLM-wiki bridge: compile long-term memories into markdown pages and index wiki pages as memories.
+
+The wiki (Karpathy's "LLM Wiki" pattern) is the human-readable, git-versioned
+form of the long-term semantic tier. MemKernel's vector store is the index
+over it plus the home of everything that is *not* wiki-worthy (episodes,
+observations, working state). See docs/architecture/memory-tiers.md.
+"""
+
+from memkernel.wiki.pages import WikiPage, parse_page, render_page
+from memkernel.wiki.sync import WikiExporter, WikiImporter
+
+__all__ = ["WikiExporter", "WikiImporter", "WikiPage", "parse_page", "render_page"]