You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@carbondata.apache.org by ch...@apache.org on 2026/08/04 06:10:01 UTC
(carbondata) branch master updated (84268138b4 -> 7ac63b2992)
This is an automated email from the ASF dual-hosted git repository.
asf-gitbox-commits pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/carbondata.git
from 84268138b4 fix agent module issues (#4396)
new 58eaa8550b update Agent module
new 7ac63b2992 rework
The 2 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.gitignore | 3 +
Agent_module/QUICKSTART.md | 37 +-
Agent_module/README.md | 48 +-
Agent_module/__init__.py | 13 -
Agent_module/carbon_data/__init__.py | 78 -
Agent_module/carbon_data/chunkers/__init__.py | 89 -
Agent_module/carbon_data/embedders/__init__.py | 95 -
Agent_module/carbon_data/index/__init__.py | 21 -
Agent_module/carbon_data/index/vector_brute.py | 98 -
Agent_module/carbon_data/index/vector_hnsw.py | 168 --
Agent_module/carbon_data/models.py | 116 --
Agent_module/carbon_data/schema.py | 172 --
Agent_module/carbon_data/store.py | 2064 ----------------------
Agent_module/examples/__init__.py | 0
Agent_module/examples/carbondata_quickstart.py | 342 ----
Agent_module/tests/__init__.py | 0
Agent_module/tests/carbon_data/__init__.py | 0
Agent_module/tests/carbon_data/test_m1_schema.py | 184 --
Agent_module/tests/carbon_data/test_m2_crud.py | 359 ----
Agent_module/tests/carbon_data/test_m3_search.py | 375 ----
Agent_module/tests/carbon_data/test_m4_fts.py | 363 ----
Agent_module/tests/carbon_data/test_m5_hnsw.py | 331 ----
Agent_module/tests/carbon_data/test_m6_graph.py | 354 ----
Agent_module/tests/carbon_data/test_m7_memory.py | 377 ----
Agent_module/tests/carbon_data/test_m8_ingest.py | 278 ---
Agent_module/tests/carbon_data/test_m9_admin.py | 236 ---
Agent_module/tests/conftest.py | 10 -
27 files changed, 59 insertions(+), 6152 deletions(-)
delete mode 100644 Agent_module/__init__.py
delete mode 100644 Agent_module/carbon_data/__init__.py
delete mode 100644 Agent_module/carbon_data/chunkers/__init__.py
delete mode 100644 Agent_module/carbon_data/embedders/__init__.py
delete mode 100644 Agent_module/carbon_data/index/__init__.py
delete mode 100644 Agent_module/carbon_data/index/vector_brute.py
delete mode 100644 Agent_module/carbon_data/index/vector_hnsw.py
delete mode 100644 Agent_module/carbon_data/models.py
delete mode 100644 Agent_module/carbon_data/schema.py
delete mode 100644 Agent_module/carbon_data/store.py
delete mode 100644 Agent_module/examples/__init__.py
delete mode 100644 Agent_module/examples/carbondata_quickstart.py
delete mode 100644 Agent_module/tests/__init__.py
delete mode 100644 Agent_module/tests/carbon_data/__init__.py
delete mode 100644 Agent_module/tests/carbon_data/test_m1_schema.py
delete mode 100644 Agent_module/tests/carbon_data/test_m2_crud.py
delete mode 100644 Agent_module/tests/carbon_data/test_m3_search.py
delete mode 100644 Agent_module/tests/carbon_data/test_m4_fts.py
delete mode 100644 Agent_module/tests/carbon_data/test_m5_hnsw.py
delete mode 100644 Agent_module/tests/carbon_data/test_m6_graph.py
delete mode 100644 Agent_module/tests/carbon_data/test_m7_memory.py
delete mode 100644 Agent_module/tests/carbon_data/test_m8_ingest.py
delete mode 100644 Agent_module/tests/carbon_data/test_m9_admin.py
delete mode 100644 Agent_module/tests/conftest.py
(carbondata) 02/02: rework
Posted by ch...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
asf-gitbox-commits pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/carbondata.git
commit 7ac63b2992decf829e31f1d07a0d08411d90fbf2
Author: chenliang613 <ch...@apache.org>
AuthorDate: Tue Aug 4 14:09:38 2026 +0800
rework
---
Agent_module/__init__.py | 13 -
Agent_module/carbon_data/__init__.py | 78 -
Agent_module/carbon_data/chunkers/__init__.py | 89 -
Agent_module/carbon_data/embedders/__init__.py | 95 -
Agent_module/carbon_data/index/__init__.py | 21 -
Agent_module/carbon_data/index/vector_brute.py | 98 -
Agent_module/carbon_data/index/vector_hnsw.py | 168 --
Agent_module/carbon_data/models.py | 116 --
Agent_module/carbon_data/schema.py | 172 --
Agent_module/carbon_data/store.py | 2125 ----------------------
Agent_module/examples/__init__.py | 0
Agent_module/examples/carbondata_quickstart.py | 342 ----
Agent_module/pyproject.toml | 40 -
Agent_module/tests/__init__.py | 0
Agent_module/tests/carbon_data/__init__.py | 0
Agent_module/tests/carbon_data/test_m1_schema.py | 184 --
Agent_module/tests/carbon_data/test_m2_crud.py | 376 ----
Agent_module/tests/carbon_data/test_m3_search.py | 375 ----
Agent_module/tests/carbon_data/test_m4_fts.py | 363 ----
Agent_module/tests/carbon_data/test_m5_hnsw.py | 426 -----
Agent_module/tests/carbon_data/test_m6_graph.py | 354 ----
Agent_module/tests/carbon_data/test_m7_memory.py | 377 ----
Agent_module/tests/carbon_data/test_m8_ingest.py | 278 ---
Agent_module/tests/carbon_data/test_m9_admin.py | 258 ---
Agent_module/tests/conftest.py | 10 -
25 files changed, 6358 deletions(-)
diff --git a/Agent_module/__init__.py b/Agent_module/__init__.py
deleted file mode 100644
index 69860838af..0000000000
--- a/Agent_module/__init__.py
+++ /dev/null
@@ -1,13 +0,0 @@
-"""
-Agent_module
-============
-Agent Data Infra. The primary product is :mod:`Agent_module.carbon_data`.
-
- from Agent_module.carbon_data import create, open
-
-See ``carbon_data/__init__.py`` for the full public API.
-"""
-
-from . import carbon_data
-
-__all__ = ["carbon_data"]
diff --git a/Agent_module/carbon_data/__init__.py b/Agent_module/carbon_data/__init__.py
deleted file mode 100644
index 5ab78fe863..0000000000
--- a/Agent_module/carbon_data/__init__.py
+++ /dev/null
@@ -1,78 +0,0 @@
-"""
-Agent Data Infra — .carbondata file format and query API.
-
-Quick start (M1 surface):
-
- from Agent_module.carbon_data import open, create
-
- store = create("kb.carbondata")
- store.stats() # {'entities': 0, ...}
- store.close()
-
- store = open("kb.carbondata")
- store.close()
-"""
-from __future__ import annotations
-
-from . import chunkers
-from .chunkers import Chunker, by_paragraph, by_sentence, by_tokens
-from .embedders import BaseEmbedder, Embedder, LambdaEmbedder, NullEmbedder
-from .models import (
- Chunk,
- Entity,
- MemoryHit,
- MemoryItem,
- Neighbor,
- Relation,
- SearchHit,
- Subgraph,
- TraversalHit,
- ValidationReport,
-)
-from .schema import (
- APPLICATION_ID,
- SCHEMA_VERSION,
- NotACarbondataFile,
- UnsupportedSchemaVersion,
-)
-from .store import CarbonStore, PathLike
-
-
-def open(path: PathLike, *, mode: str = "rw") -> CarbonStore: # noqa: A001
- """Open an existing .carbondata file (or create it if missing in rw mode)."""
- return CarbonStore.open(path, mode=mode) # type: ignore[arg-type]
-
-
-def create(path: PathLike, *, exist_ok: bool = False) -> CarbonStore:
- """Create a new .carbondata file. Raises FileExistsError unless exist_ok=True."""
- return CarbonStore.create(path, exist_ok=exist_ok)
-
-
-__all__ = [
- "open",
- "create",
- "CarbonStore",
- "Entity",
- "Chunk",
- "Relation",
- "MemoryItem",
- "MemoryHit",
- "SearchHit",
- "Neighbor",
- "TraversalHit",
- "Subgraph",
- "ValidationReport",
- "chunkers",
- "Chunker",
- "by_tokens",
- "by_paragraph",
- "by_sentence",
- "Embedder",
- "BaseEmbedder",
- "NullEmbedder",
- "LambdaEmbedder",
- "SCHEMA_VERSION",
- "APPLICATION_ID",
- "NotACarbondataFile",
- "UnsupportedSchemaVersion",
-]
diff --git a/Agent_module/carbon_data/chunkers/__init__.py b/Agent_module/carbon_data/chunkers/__init__.py
deleted file mode 100644
index b2120ceba7..0000000000
--- a/Agent_module/carbon_data/chunkers/__init__.py
+++ /dev/null
@@ -1,89 +0,0 @@
-"""
-Pluggable chunker factories.
-
-A ``Chunker`` is just a function ``str -> list[str]``. Factories return
-configured chunkers so the call site stays declarative:
-
- chunker = by_tokens(max_tokens=512, overlap=64)
- store.ingest_text(text, chunker=chunker)
-"""
-from __future__ import annotations
-
-import re
-from typing import Callable
-
-Chunker = Callable[[str], list[str]]
-
-
-def by_tokens(
- *,
- max_tokens: int = 512,
- overlap: int = 0,
- sep: str = " ",
-) -> Chunker:
- """
- Whitespace-token sliding window. Each chunk holds at most ``max_tokens``
- tokens; consecutive chunks share ``overlap`` tokens at the boundary.
-
- Tokens are produced by :py:meth:`str.split` — there is no model-specific
- tokenizer here. Use it as a coarse proxy for real token budgets, not
- as a substitute.
- """
- if max_tokens <= 0:
- raise ValueError("max_tokens must be > 0")
- if overlap < 0:
- raise ValueError("overlap must be >= 0")
- if overlap >= max_tokens:
- raise ValueError("overlap must be < max_tokens")
-
- step = max_tokens - overlap
-
- def _chunk(text: str) -> list[str]:
- words = text.split()
- if not words:
- return []
- out: list[str] = []
- n = len(words)
- start = 0
- while start < n:
- end = start + max_tokens
- out.append(sep.join(words[start:end]))
- if end >= n:
- break
- start += step
- return out
-
- return _chunk
-
-
-def _regex_chunker(splitter: "re.Pattern[str]", min_chars: int) -> Chunker:
- """Split on ``splitter``; keep stripped pieces of at least ``min_chars``."""
-
- def _chunk(text: str) -> list[str]:
- return [
- p.strip()
- for p in splitter.split(text)
- if p.strip() and len(p.strip()) >= min_chars
- ]
-
- return _chunk
-
-
-def by_paragraph(*, min_chars: int = 0) -> Chunker:
- """
- Split on runs of blank lines. Drops any chunk shorter than
- ``min_chars`` after stripping; set to 0 to keep every paragraph.
- """
- return _regex_chunker(re.compile(r"\n\s*\n"), min_chars)
-
-
-def by_sentence(*, min_chars: int = 0) -> Chunker:
- """
- Naive sentence splitter: end-of-sentence punctuation followed by
- whitespace and a capital letter. Good enough for prose; for legal/
- technical text use ``by_tokens`` or a real tokenizer.
- """
- return _regex_chunker(re.compile(r"(?<=[.!?])\s+(?=[A-Z])"), min_chars)
-
-
-__all__ = ["Chunker", "by_tokens", "by_paragraph", "by_sentence"]
diff --git a/Agent_module/carbon_data/embedders/__init__.py b/Agent_module/carbon_data/embedders/__init__.py
deleted file mode 100644
index e9c95941c4..0000000000
--- a/Agent_module/carbon_data/embedders/__init__.py
+++ /dev/null
@@ -1,95 +0,0 @@
-"""
-Embedders.
-
-An `Embedder` is any object with ``model: str``, ``dim: int``, and an
-``encode(texts) -> np.ndarray`` method returning shape ``(len(texts), dim)``
-as float32.
-
-We do NOT bundle embedding models. Users bring their own (OpenAI,
-sentence-transformers, local llama, ...) and wrap them with
-`LambdaEmbedder`, or they subclass `BaseEmbedder`.
-"""
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Callable, Protocol, Sequence, runtime_checkable
-
-if TYPE_CHECKING:
- import numpy as np
-
-
-@runtime_checkable
-class Embedder(Protocol):
- """
- Duck-typed embedder interface.
-
- Attributes:
- model: a stable identifier written alongside each vector; queries
- must use the same ``model`` to be comparable.
- dim: vector dimensionality.
-
- Method:
- encode(texts): return float32 numpy array shape ``(len(texts), dim)``.
- """
-
- model: str
- dim: int
-
- def encode(self, texts: Sequence[str]) -> "np.ndarray": # pragma: no cover
- ...
-
-
-class BaseEmbedder:
- """Convenience base class that validates output shape."""
-
- model: str
- dim: int
-
- def __init__(self, model: str, dim: int) -> None:
- self.model = model
- self.dim = dim
-
- def encode(self, texts: Sequence[str]) -> "np.ndarray": # pragma: no cover
- raise NotImplementedError
-
- def _validate(self, arr, n_texts: int) -> "np.ndarray":
- import numpy as np
-
- out = np.ascontiguousarray(arr, dtype=np.float32)
- if out.ndim != 2 or out.shape != (n_texts, self.dim):
- raise ValueError(
- f"embedder {self.model!r} returned shape {out.shape}, "
- f"expected ({n_texts}, {self.dim})"
- )
- return out
-
-
-class NullEmbedder(BaseEmbedder):
- """All-zero embedder. Intended for tests and placeholders only."""
-
- def __init__(self, model: str = "null", dim: int = 8) -> None:
- super().__init__(model=model, dim=dim)
-
- def encode(self, texts: Sequence[str]) -> "np.ndarray":
- import numpy as np
-
- return np.zeros((len(texts), self.dim), dtype=np.float32)
-
-
-class LambdaEmbedder(BaseEmbedder):
- """Wrap any callable ``fn(texts) -> array-like`` as an Embedder."""
-
- def __init__(
- self,
- fn: Callable[[Sequence[str]], "np.ndarray"],
- *,
- model: str,
- dim: int,
- ) -> None:
- super().__init__(model=model, dim=dim)
- self._fn = fn
-
- def encode(self, texts: Sequence[str]) -> "np.ndarray":
- return self._validate(self._fn(texts), len(texts))
-
-
-__all__ = ["Embedder", "BaseEmbedder", "NullEmbedder", "LambdaEmbedder"]
diff --git a/Agent_module/carbon_data/index/__init__.py b/Agent_module/carbon_data/index/__init__.py
deleted file mode 100644
index c789a2fb79..0000000000
--- a/Agent_module/carbon_data/index/__init__.py
+++ /dev/null
@@ -1,21 +0,0 @@
-"""
-Index layer.
-
-- M3: brute-force vector search (always available; pure numpy).
-- M5: HNSW via hnswlib as an optional dependency. Falls back to brute
- when hnswlib is not installed or when filters/metrics rule it out.
-"""
-from .vector_brute import (
- bytes_to_vector,
- search_brute,
- vector_to_bytes,
-)
-from .vector_hnsw import HnswIndex, hnsw_available
-
-__all__ = [
- "bytes_to_vector",
- "vector_to_bytes",
- "search_brute",
- "HnswIndex",
- "hnsw_available",
-]
diff --git a/Agent_module/carbon_data/index/vector_brute.py b/Agent_module/carbon_data/index/vector_brute.py
deleted file mode 100644
index bb037f605a..0000000000
--- a/Agent_module/carbon_data/index/vector_brute.py
+++ /dev/null
@@ -1,98 +0,0 @@
-"""
-Brute-force nearest-neighbour search over in-memory float32 matrices.
-
-Fine for < ~10k chunks. HNSW (M5) swaps in for larger corpora behind the
-same interface on ``CarbonStore.search``.
-"""
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Literal, Sequence
-
-if TYPE_CHECKING:
- import numpy as np
-
-Metric = Literal["cosine", "dot", "l2"]
-
-_EPS = 1e-12 # guard against divide-by-zero on zero-vectors
-
-
-def vector_to_bytes(vec) -> bytes:
- """Serialize a 1-D float32 vector to bytes for BLOB storage."""
- import numpy as np
-
- arr = np.ascontiguousarray(vec, dtype=np.float32)
- if arr.ndim != 1:
- raise ValueError(f"expected 1-D vector, got shape {arr.shape}")
- return arr.tobytes()
-
-
-def bytes_to_vector(blob: bytes, dim: int) -> "np.ndarray":
- """Deserialize a BLOB back into a 1-D float32 numpy array."""
- import numpy as np
-
- arr = np.frombuffer(blob, dtype=np.float32)
- if arr.size != dim:
- raise ValueError(
- f"vector blob has {arr.size} floats, expected {dim}"
- )
- # frombuffer is read-only and zero-copy; caller may want to mutate
- return arr.copy()
-
-
-def search_brute(
- query: "np.ndarray",
- vectors: "np.ndarray",
- ids: Sequence[str],
- *,
- top_k: int,
- metric: Metric = "cosine",
-) -> list[tuple[str, float]]:
- """
- Brute-force top-k over an in-memory matrix.
-
- Args:
- query: shape (dim,) float32
- vectors: shape (N, dim) float32
- ids: length-N iterable of chunk ids aligned with vectors
- top_k: max hits to return
- metric: 'cosine' | 'dot' | 'l2' (l2 is returned as negative distance
- so that larger = better, same convention as cosine/dot)
-
- Returns:
- list of (chunk_id, score) sorted descending by score.
- """
- import numpy as np
-
- if vectors.size == 0:
- return []
-
- q = np.ascontiguousarray(query, dtype=np.float32).reshape(-1)
- if q.shape[0] != vectors.shape[1]:
- raise ValueError(
- f"query dim {q.shape[0]} != vector dim {vectors.shape[1]}"
- )
-
- if metric == "cosine":
- v_norms = np.linalg.norm(vectors, axis=1)
- q_norm = float(np.linalg.norm(q))
- scores = (vectors @ q) / (v_norms * q_norm + _EPS)
- elif metric == "dot":
- scores = vectors @ q
- elif metric == "l2":
- # negated distance so argsort-descending yields nearest first
- diffs = vectors - q
- scores = -np.linalg.norm(diffs, axis=1)
- else:
- raise ValueError(f"unknown metric: {metric!r}")
-
- if top_k >= scores.shape[0]:
- order = np.argsort(-scores)
- else:
- # argpartition is O(N); sort only the top-k slice
- part = np.argpartition(-scores, top_k)[:top_k]
- order = part[np.argsort(-scores[part])]
-
- return [(ids[int(i)], float(scores[int(i)])) for i in order]
-
-
-__all__ = ["vector_to_bytes", "bytes_to_vector", "search_brute"]
diff --git a/Agent_module/carbon_data/index/vector_hnsw.py b/Agent_module/carbon_data/index/vector_hnsw.py
deleted file mode 100644
index 356e8e944d..0000000000
--- a/Agent_module/carbon_data/index/vector_hnsw.py
+++ /dev/null
@@ -1,168 +0,0 @@
-"""
-HNSW (Hierarchical Navigable Small World) wrapper around hnswlib.
-
-Optional dependency. Use ``hnsw_available()`` before constructing an index;
-when hnswlib is missing the carbondata search path falls back to brute force.
-"""
-from __future__ import annotations
-
-import os
-import tempfile
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
- import numpy as np
-
-
-_DEFAULT_M = 16
-_DEFAULT_EF_CONSTRUCTION = 200
-_DEFAULT_EF_SEARCH = 50
-
-
-def hnsw_available() -> bool:
- """True iff hnswlib can be imported."""
- try:
- import hnswlib # noqa: F401
- except ImportError:
- return False
- return True
-
-
-class HnswIndex:
- """
- Cosine-space HNSW with integer labels and tempfile blob roundtripping.
-
- Labels are caller-assigned ints (typically chunk.rowid). hnswlib's
- save_index / load_index APIs only know how to talk to filesystem
- paths, so serialization stages through a tempfile.
- """
-
- def __init__(
- self,
- dim: int,
- max_elements: int,
- *,
- M: int = _DEFAULT_M,
- ef_construction: int = _DEFAULT_EF_CONSTRUCTION,
- ef_search: int = _DEFAULT_EF_SEARCH,
- ) -> None:
- import hnswlib
-
- self.dim = dim
- self.M = M
- self.ef_construction = ef_construction
- self.ef_search = ef_search
- self._idx = hnswlib.Index(space="cosine", dim=dim)
- self._idx.init_index(
- max_elements=max(max_elements, 1),
- ef_construction=ef_construction,
- M=M,
- )
- self._idx.set_ef(ef_search)
-
- # ------------------------------------------------------------------
- # construction
- # ------------------------------------------------------------------
-
- @classmethod
- def build(
- cls,
- vectors: "np.ndarray",
- labels: list[int],
- **kwargs,
- ) -> "HnswIndex":
- """Build an index from a contiguous (N, dim) array and N labels."""
- import numpy as np
-
- if vectors.ndim != 2:
- raise ValueError(
- f"expected 2-D vectors, got shape {vectors.shape}"
- )
- if len(labels) != vectors.shape[0]:
- raise ValueError(
- f"labels/vectors size mismatch: "
- f"{len(labels)} vs {vectors.shape[0]}"
- )
- idx = cls(
- dim=vectors.shape[1], max_elements=len(labels), **kwargs
- )
- idx._idx.add_items(
- np.ascontiguousarray(vectors, dtype=np.float32), labels
- )
- return idx
-
- # ------------------------------------------------------------------
- # query
- # ------------------------------------------------------------------
-
- def search(
- self, q: "np.ndarray", top_k: int
- ) -> list[tuple[int, float]]:
- """Return ``[(label, similarity), ...]`` ranked best-first."""
- import numpy as np
-
- cur = self._idx.get_current_count()
- if cur == 0 or top_k <= 0:
- return []
- k = min(top_k, cur)
- q2 = np.ascontiguousarray(q, dtype=np.float32).reshape(1, -1)
- labels, distances = self._idx.knn_query(q2, k=k)
- # cosine distance ∈ [0, 2]; similarity = 1 - distance.
- return [
- (int(lbl), float(1.0 - dist))
- for lbl, dist in zip(labels[0], distances[0])
- ]
-
- @property
- def count(self) -> int:
- return self._idx.get_current_count()
-
- # ------------------------------------------------------------------
- # serialization
- # ------------------------------------------------------------------
-
- def serialize(self) -> bytes:
- """Save the index to a tempfile and read it back as bytes."""
- fd, path = tempfile.mkstemp(suffix=".hnsw")
- os.close(fd)
- try:
- self._idx.save_index(path)
- with open(path, "rb") as f:
- return f.read()
- finally:
- try:
- os.unlink(path)
- except FileNotFoundError:
- pass
-
- @classmethod
- def deserialize(
- cls,
- blob: bytes,
- *,
- dim: int,
- ef_search: int = _DEFAULT_EF_SEARCH,
- ) -> "HnswIndex":
- import hnswlib
-
- fd, path = tempfile.mkstemp(suffix=".hnsw")
- try:
- with os.fdopen(fd, "wb") as f:
- f.write(blob)
- inst = cls.__new__(cls)
- inst.dim = dim
- inst.M = _DEFAULT_M
- inst.ef_construction = _DEFAULT_EF_CONSTRUCTION
- inst.ef_search = ef_search
- inst._idx = hnswlib.Index(space="cosine", dim=dim)
- inst._idx.load_index(path)
- inst._idx.set_ef(ef_search)
- return inst
- finally:
- try:
- os.unlink(path)
- except FileNotFoundError:
- pass
-
-
-__all__ = ["HnswIndex", "hnsw_available"]
diff --git a/Agent_module/carbon_data/models.py b/Agent_module/carbon_data/models.py
deleted file mode 100644
index 2098244ad4..0000000000
--- a/Agent_module/carbon_data/models.py
+++ /dev/null
@@ -1,116 +0,0 @@
-"""
-Data model for Agent Data Infra.
-"""
-from __future__ import annotations
-
-import time
-from dataclasses import dataclass, field
-from typing import Any, Optional
-
-
-@dataclass
-class Entity:
- id: str
- kind: str
- namespace: str = "default"
- content: Optional[str] = None
- metadata: dict[str, Any] = field(default_factory=dict)
- created_at: float = field(default_factory=time.time)
- updated_at: float = field(default_factory=time.time)
- version: int = 1
-
-
-@dataclass
-class Chunk:
- id: str
- entity_id: str
- seq: int
- content: str
- kind: str = "text"
- metadata: dict[str, Any] = field(default_factory=dict)
- token_count: Optional[int] = None
-
-
-@dataclass
-class Relation:
- src: str
- dst: str
- kind: str
- weight: float = 1.0
- metadata: dict[str, Any] = field(default_factory=dict)
-
-
-@dataclass
-class MemoryItem(Entity):
- kind: str = "memory"
- session_id: Optional[str] = None
- actor: Optional[str] = None
- salience: float = 0.5
- expires_at: Optional[float] = None
-
-
-@dataclass
-class SearchHit:
- """Result row from `CarbonStore.search`."""
-
- chunk: Chunk
- entity: Entity
- score: float
-
-
-@dataclass
-class MemoryHit:
- """Result row from `CarbonStore.recall`."""
-
- memory: MemoryItem
- score: float
-
-
-@dataclass
-class Neighbor:
- """A single-hop graph neighbor with the edge that connects it."""
-
- entity: Entity
- relation_kind: str # the edge's kind
- weight: float
- direction: str # "out" | "in" relative to the query node
- metadata: dict[str, Any] = field(default_factory=dict)
-
-
-@dataclass
-class TraversalHit:
- """A node visited during a multi-hop traversal."""
-
- entity: Entity
- hop: int # shortest distance from the start node
-
-
-@dataclass
-class Subgraph:
- """An induced subgraph: entities plus the relations among them."""
-
- entities: list[Entity] = field(default_factory=list)
- relations: list[Relation] = field(default_factory=list)
-
-
-@dataclass
-class ValidationReport:
- """Outcome of :py:meth:`CarbonStore.validate`."""
-
- ok: bool
- issues: list[str] = field(default_factory=list)
- counts: dict[str, int] = field(default_factory=dict)
-
-
-__all__ = [
- "Entity",
- "Chunk",
- "Relation",
- "MemoryItem",
- "SearchHit",
- "MemoryHit",
- "Neighbor",
- "TraversalHit",
- "Subgraph",
- "ValidationReport",
-]
diff --git a/Agent_module/carbon_data/schema.py b/Agent_module/carbon_data/schema.py
deleted file mode 100644
index 9e0ec8c06f..0000000000
--- a/Agent_module/carbon_data/schema.py
+++ /dev/null
@@ -1,172 +0,0 @@
-"""
-SQLite schema for .carbondata files.
-"""
-from __future__ import annotations
-
-import sqlite3
-import time
-
-SCHEMA_VERSION = 1
-
-# "CBND" = 0x43 42 4E 44. Stored via PRAGMA application_id so `file` recognizes
-# the carbondata signature and other carbondata libraries can distinguish it
-# from arbitrary SQLite files.
-APPLICATION_ID = 0x4342_4E44
-
-
-_DDL = [
- """CREATE TABLE IF NOT EXISTS _meta (
- key TEXT PRIMARY KEY,
- value TEXT
- )""",
- """CREATE TABLE IF NOT EXISTS entity (
- id TEXT PRIMARY KEY,
- kind TEXT NOT NULL,
- namespace TEXT NOT NULL DEFAULT 'default',
- content TEXT,
- metadata TEXT,
- created_at REAL NOT NULL,
- updated_at REAL NOT NULL,
- version INTEGER NOT NULL DEFAULT 1
- )""",
- "CREATE INDEX IF NOT EXISTS idx_entity_kind_ns ON entity(kind, namespace)",
- "CREATE INDEX IF NOT EXISTS idx_entity_updated ON entity(updated_at)",
- """CREATE TABLE IF NOT EXISTS chunk (
- id TEXT PRIMARY KEY,
- entity_id TEXT NOT NULL REFERENCES entity(id) ON DELETE CASCADE,
- seq INTEGER NOT NULL,
- content TEXT NOT NULL,
- kind TEXT NOT NULL DEFAULT 'text',
- metadata TEXT,
- token_count INTEGER
- )""",
- "CREATE INDEX IF NOT EXISTS idx_chunk_entity ON chunk(entity_id, seq)",
- """CREATE VIRTUAL TABLE IF NOT EXISTS chunk_fts USING fts5(
- content,
- content='chunk',
- content_rowid='rowid'
- )""",
- """CREATE TABLE IF NOT EXISTS embedding (
- chunk_id TEXT NOT NULL REFERENCES chunk(id) ON DELETE CASCADE,
- model TEXT NOT NULL,
- dim INTEGER NOT NULL,
- vector BLOB NOT NULL,
- PRIMARY KEY (chunk_id, model)
- )""",
- "CREATE INDEX IF NOT EXISTS idx_embedding_model ON embedding(model)",
- """CREATE TABLE IF NOT EXISTS relation (
- src TEXT NOT NULL REFERENCES entity(id) ON DELETE CASCADE,
- dst TEXT NOT NULL REFERENCES entity(id) ON DELETE CASCADE,
- kind TEXT NOT NULL,
- weight REAL NOT NULL DEFAULT 1.0,
- metadata TEXT,
- PRIMARY KEY (src, dst, kind)
- )""",
- "CREATE INDEX IF NOT EXISTS idx_relation_dst ON relation(dst, kind)",
- "CREATE INDEX IF NOT EXISTS idx_relation_kind ON relation(kind)",
- """CREATE TABLE IF NOT EXISTS memory_ext (
- entity_id TEXT PRIMARY KEY REFERENCES entity(id) ON DELETE CASCADE,
- session_id TEXT,
- actor TEXT,
- salience REAL NOT NULL DEFAULT 0.5,
- expires_at REAL
- )""",
- "CREATE INDEX IF NOT EXISTS idx_memory_session ON memory_ext(session_id)",
- "CREATE INDEX IF NOT EXISTS idx_memory_expires ON memory_ext(expires_at)",
- """CREATE TABLE IF NOT EXISTS vector_index (
- model TEXT NOT NULL,
- namespace TEXT NOT NULL DEFAULT 'default',
- index_blob BLOB NOT NULL,
- index_meta TEXT,
- updated_at REAL NOT NULL,
- PRIMARY KEY (model, namespace)
- )""",
-]
-
-
-# FTS5 external-content tables don't auto-sync from their source table; we
-# install row triggers so chunk writes propagate into chunk_fts. Defined
-# separately so `open()` can re-apply them idempotently on legacy files.
-_TRIGGERS = [
- """CREATE TRIGGER IF NOT EXISTS chunk_ai AFTER INSERT ON chunk BEGIN
- INSERT INTO chunk_fts(rowid, content) VALUES (new.rowid, new.content);
- END""",
- """CREATE TRIGGER IF NOT EXISTS chunk_ad AFTER DELETE ON chunk BEGIN
- INSERT INTO chunk_fts(chunk_fts, rowid, content)
- VALUES ('delete', old.rowid, old.content);
- END""",
- """CREATE TRIGGER IF NOT EXISTS chunk_au AFTER UPDATE ON chunk BEGIN
- INSERT INTO chunk_fts(chunk_fts, rowid, content)
- VALUES ('delete', old.rowid, old.content);
- INSERT INTO chunk_fts(rowid, content) VALUES (new.rowid, new.content);
- END""",
-]
-
-
-def init_schema(conn: sqlite3.Connection) -> None:
- """Apply carbondata schema to a fresh SQLite connection."""
- conn.execute(f"PRAGMA application_id = {APPLICATION_ID}")
- conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
- conn.execute("PRAGMA foreign_keys = ON")
- conn.execute("PRAGMA journal_mode = WAL")
- with conn:
- for stmt in _DDL:
- conn.execute(stmt)
- for stmt in _TRIGGERS:
- conn.execute(stmt)
- conn.execute(
- "INSERT OR IGNORE INTO _meta (key, value) VALUES (?, ?)",
- ("schema_version", str(SCHEMA_VERSION)),
- )
- conn.execute(
- "INSERT OR IGNORE INTO _meta (key, value) VALUES (?, ?)",
- ("created_at", str(time.time())),
- )
-
-
-def ensure_triggers(conn: sqlite3.Connection) -> None:
- """Idempotently install FTS sync triggers. Safe to call on every open."""
- with conn:
- for stmt in _TRIGGERS:
- conn.execute(stmt)
-
-
-class NotACarbondataFile(ValueError):
- pass
-
-
-class UnsupportedSchemaVersion(ValueError):
- pass
-
-
-def verify_schema(conn: sqlite3.Connection) -> int:
- """
- Validate that `conn` points to a carbondata file.
-
- Returns the on-disk schema version. Raises if the file is not carbondata
- or is written by a future version this library cannot read.
- """
- (app_id,) = conn.execute("PRAGMA application_id").fetchone()
- if app_id != APPLICATION_ID:
- raise NotACarbondataFile(
- f"Not a carbondata file: application_id={app_id:#010x} "
- f"(expected {APPLICATION_ID:#010x})"
- )
- (ver,) = conn.execute("PRAGMA user_version").fetchone()
- if ver > SCHEMA_VERSION:
- raise UnsupportedSchemaVersion(
- f"carbondata file schema version {ver} is newer than "
- f"library version {SCHEMA_VERSION}"
- )
- return ver
-
-
-__all__ = [
- "SCHEMA_VERSION",
- "APPLICATION_ID",
- "init_schema",
- "ensure_triggers",
- "verify_schema",
- "NotACarbondataFile",
- "UnsupportedSchemaVersion",
-]
diff --git a/Agent_module/carbon_data/store.py b/Agent_module/carbon_data/store.py
deleted file mode 100644
index 8f70ddfc50..0000000000
--- a/Agent_module/carbon_data/store.py
+++ /dev/null
@@ -1,2125 +0,0 @@
-"""
-CarbonStore: primary handle to a .carbondata file.
-
-Milestones implemented here:
-- M1: open/create/close + schema verification + stats
-- M2: Entity/Chunk CRUD + structured query + transactions
-"""
-from __future__ import annotations
-
-import json
-import sqlite3
-import time
-import uuid
-from contextlib import contextmanager
-from pathlib import Path
-from typing import Any, Iterable, Iterator, Literal, Optional, Union
-
-from . import schema as _schema
-from .chunkers import Chunker, by_paragraph
-from .embedders import Embedder
-from .index import (
- HnswIndex,
- bytes_to_vector,
- hnsw_available,
- search_brute,
- vector_to_bytes,
-)
-from .models import (
- Chunk,
- Entity,
- MemoryHit,
- MemoryItem,
- Neighbor,
- Relation,
- SearchHit,
- Subgraph,
- TraversalHit,
- ValidationReport,
-)
-
-PathLike = Union[str, Path]
-Mode = Literal["r", "rw"]
-
-# Whitelist of entity columns users may reference directly in `where`
-# and `order_by`. Guards against SQL injection: anything outside this
-# set gets rejected before it hits the query string.
-_ENTITY_COLUMNS = frozenset(
- {"id", "kind", "namespace", "content", "created_at", "updated_at", "version"}
-)
-
-_OPERATORS = frozenset({"=", "!=", "<", "<=", ">", ">=", "LIKE", "IN"})
-
-
-def _uuid() -> str:
- return str(uuid.uuid4())
-
-
-def _dumps(obj: Any) -> str:
- return json.dumps(obj or {}, sort_keys=True, ensure_ascii=False)
-
-
-def _loads(s: Optional[str]) -> dict[str, Any]:
- if s is None or s == "":
- return {}
- return json.loads(s)
-
-
-def _render_row(row: dict[str, Any]) -> str:
- """Render a table row as ``key: value`` lines for indexing."""
- return "\n".join(f"{k}: {v}" for k, v in row.items())
-
-
-def _read_csv_rows(path: Path) -> list[dict[str, Any]]:
- """Read a CSV via DictReader; non-string values stay as strings here."""
- import csv
-
- with path.open(newline="", encoding="utf-8") as f:
- return [dict(r) for r in csv.DictReader(f)]
-
-
-def _fts_escape_query(q: str) -> str:
- """
- Convert an arbitrary user string into a safe FTS5 query.
-
- Strips embedded double-quotes, splits on whitespace, wraps each token
- as an FTS5 phrase. Tokens are AND-combined (FTS5 default). Returns
- "" (caller must short-circuit) when the input has no usable tokens.
- """
- tokens = q.replace('"', " ").split()
- if not tokens:
- return ""
- return " ".join(f'"{t}"' for t in tokens)
-
-
-class CarbonStore:
- """
- A single .carbondata file.
-
- Use `CarbonStore.open(path)` or the package-level `carbon_data.open(path)`.
- """
-
- def __init__(
- self,
- path: Path,
- conn: sqlite3.Connection,
- mode: Mode,
- ) -> None:
- self.path = path
- self._conn = conn
- self._mode: Mode = mode
- self._tx_depth = 0 # 0 = autocommit-per-write; >0 = user transaction
- # Per-model HNSW index cache. Populated lazily on first vector
- # search; mutations explicitly invalidate affected models while
- # count checks protect against out-of-band adds/deletes.
- self._hnsw_cache: dict[str, HnswIndex] = {}
-
- # ------------------------------------------------------------------
- # lifecycle
- # ------------------------------------------------------------------
-
- @classmethod
- def create(cls, path: PathLike, *, exist_ok: bool = False) -> "CarbonStore":
- p = Path(path)
- if p.exists():
- if not exist_ok:
- raise FileExistsError(str(p))
- return cls.open(p, mode="rw")
- conn = sqlite3.connect(str(p))
- _schema.init_schema(conn)
- return cls(p, conn, mode="rw")
-
- @classmethod
- def open(cls, path: PathLike, *, mode: Mode = "rw") -> "CarbonStore":
- p = Path(path)
- if not p.exists():
- if mode == "r":
- raise FileNotFoundError(str(p))
- return cls.create(p)
-
- if mode == "r":
- conn = sqlite3.connect(f"file:{p}?mode=ro", uri=True)
- else:
- conn = sqlite3.connect(str(p))
- conn.execute("PRAGMA foreign_keys = ON")
-
- try:
- _schema.verify_schema(conn)
- if mode != "r":
- _schema.ensure_triggers(conn)
- except Exception:
- conn.close()
- raise
- return cls(p, conn, mode)
-
- def close(self) -> None:
- if self._conn is not None:
- self._conn.close()
- self._conn = None # type: ignore[assignment]
-
- def __enter__(self) -> "CarbonStore":
- return self
-
- def __exit__(self, exc_type, exc_val, exc_tb) -> None:
- self.close()
-
- # ------------------------------------------------------------------
- # introspection
- # ------------------------------------------------------------------
-
- @property
- def mode(self) -> Mode:
- return self._mode
-
- @property
- def closed(self) -> bool:
- return self._conn is None
-
- def _require_open(self) -> sqlite3.Connection:
- if self._conn is None:
- raise RuntimeError("CarbonStore is closed")
- return self._conn
-
- def _require_rw(self) -> sqlite3.Connection:
- conn = self._require_open()
- if self._mode == "r":
- raise PermissionError("CarbonStore opened read-only")
- return conn
-
- def stats(self) -> dict[str, object]:
- c = self._require_open()
-
- def count(table: str) -> int:
- return c.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
-
- (ver,) = c.execute("PRAGMA user_version").fetchone()
- (app_id,) = c.execute("PRAGMA application_id").fetchone()
- return {
- "path": str(self.path),
- "mode": self._mode,
- "schema_version": ver,
- "application_id": app_id,
- "entities": count("entity"),
- "chunks": count("chunk"),
- "embeddings": count("embedding"),
- "relations": count("relation"),
- "memories": count("memory_ext"),
- }
-
- # ------------------------------------------------------------------
- # transactions
- # ------------------------------------------------------------------
-
- @contextmanager
- def transaction(self) -> Iterator[None]:
- """
- Group writes atomically.
-
- Nested calls use SQLite savepoints. An exception rolls back the
- current nesting level and is re-raised; if the caller catches an
- inner exception, the outer transaction may continue safely.
- """
- conn = self._require_rw()
- outer = self._tx_depth == 0
- savepoint = f"carbondata_tx_{self._tx_depth}"
- if outer:
- conn.execute("BEGIN")
- else:
- conn.execute(f"SAVEPOINT {savepoint}")
- self._tx_depth += 1
- try:
- yield
- if outer:
- conn.commit()
- else:
- conn.execute(f"RELEASE SAVEPOINT {savepoint}")
- except Exception:
- if outer:
- conn.rollback()
- else:
- conn.execute(f"ROLLBACK TO SAVEPOINT {savepoint}")
- conn.execute(f"RELEASE SAVEPOINT {savepoint}")
- # A search inside the failed scope may have built an in-memory
- # index from data that no longer exists after rollback.
- self._hnsw_cache.clear()
- raise
- finally:
- self._tx_depth -= 1
-
- @contextmanager
- def _write_ctx(self) -> Iterator[sqlite3.Connection]:
- """Internal per-write wrapper: commits iff no user transaction is active."""
- conn = self._require_rw()
- if self._tx_depth > 0:
- yield conn
- return
- try:
- yield conn
- conn.commit()
- except Exception:
- conn.rollback()
- raise
-
- # ------------------------------------------------------------------
- # Entity CRUD
- # ------------------------------------------------------------------
-
- def put_entity(
- self,
- entity: Optional[Entity] = None,
- *,
- id: Optional[str] = None, # noqa: A002
- kind: Optional[str] = None,
- namespace: str = "default",
- content: Optional[str] = None,
- metadata: Optional[dict[str, Any]] = None,
- ) -> str:
- """
- Insert or update an entity. On upsert, `created_at` is preserved and
- `version` is bumped by one. Returns the entity id.
-
- Call either with an Entity:
- store.put_entity(Entity(id="x", kind="doc"))
- or with kwargs:
- store.put_entity(id="x", kind="doc", content="...")
- """
- if entity is not None:
- mixed = any(
- v is not None for v in (id, kind, content, metadata)
- ) or namespace != "default"
- if mixed:
- raise TypeError("pass either an Entity or kwargs, not both")
- e = entity
- else:
- if id is None or kind is None:
- raise TypeError(
- "put_entity requires either an Entity positional "
- "or (id=, kind=, ...) kwargs"
- )
- e = Entity(
- id=id,
- kind=kind,
- namespace=namespace,
- content=content,
- metadata=metadata or {},
- )
-
- now = time.time()
- meta_json = _dumps(e.metadata)
-
- with self._write_ctx() as conn:
- existing = conn.execute(
- "SELECT version FROM entity WHERE id=?", (e.id,)
- ).fetchone()
- if existing is None:
- conn.execute(
- """INSERT INTO entity
- (id, kind, namespace, content, metadata,
- created_at, updated_at, version)
- VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
- (
- e.id, e.kind, e.namespace, e.content, meta_json,
- e.created_at, e.updated_at, e.version,
- ),
- )
- else:
- (old_version,) = existing
- conn.execute(
- """UPDATE entity
- SET kind=?, namespace=?, content=?, metadata=?,
- updated_at=?, version=?
- WHERE id=?""",
- (
- e.kind, e.namespace, e.content, meta_json,
- now, old_version + 1, e.id,
- ),
- )
- return e.id
-
- def get_entity(self, entity_id: str) -> Optional[Entity]:
- conn = self._require_open()
- row = conn.execute(
- """SELECT id, kind, namespace, content, metadata,
- created_at, updated_at, version
- FROM entity WHERE id=?""",
- (entity_id,),
- ).fetchone()
- return _row_to_entity(row) if row else None
-
- def delete_entity(self, entity_id: str) -> bool:
- """Remove an entity and (via ON DELETE CASCADE) its chunks,
- embeddings, relations, and memory_ext row. Returns True iff
- something was deleted."""
- deleted = False
- with self._write_ctx() as conn:
- cur = conn.execute("DELETE FROM entity WHERE id=?", (entity_id,))
- deleted = cur.rowcount > 0
- if deleted:
- conn.execute("DELETE FROM vector_index")
- if deleted:
- self._hnsw_cache.clear()
- return deleted
-
- # ------------------------------------------------------------------
- # Chunk CRUD
- # ------------------------------------------------------------------
-
- def add_chunks(
- self,
- entity_id: str,
- chunks: Iterable[Union[Chunk, str, dict[str, Any]]],
- *,
- base_seq: Optional[int] = None,
- ) -> list[str]:
- """
- Append chunks to an entity.
-
- Each item may be:
- - a ``Chunk`` dataclass
- - a ``str`` (treated as content; id/seq auto-assigned)
- - a ``dict`` with keys matching Chunk fields
-
- ``seq`` auto-assigns continuing from the entity's current max seq
- unless ``base_seq`` is provided explicitly.
-
- Returns the list of chunk ids.
- """
- conn = self._require_rw()
- if not conn.execute(
- "SELECT 1 FROM entity WHERE id=?", (entity_id,)
- ).fetchone():
- raise KeyError(f"entity not found: {entity_id!r}")
-
- if base_seq is None:
- row = conn.execute(
- "SELECT COALESCE(MAX(seq), -1) FROM chunk WHERE entity_id=?",
- (entity_id,),
- ).fetchone()
- base_seq = (row[0] if row else -1) + 1
-
- materialized: list[Chunk] = []
- for offset, raw in enumerate(chunks):
- seq = base_seq + offset
- materialized.append(_coerce_chunk(raw, entity_id=entity_id, default_seq=seq))
-
- rows = [
- (
- c.id, c.entity_id, c.seq, c.content, c.kind,
- _dumps(c.metadata), c.token_count,
- )
- for c in materialized
- ]
- with self._write_ctx() as conn:
- conn.executemany(
- """INSERT INTO chunk
- (id, entity_id, seq, content, kind, metadata, token_count)
- VALUES (?, ?, ?, ?, ?, ?, ?)""",
- rows,
- )
- return [c.id for c in materialized]
-
- def get_chunk(self, chunk_id: str) -> Optional[Chunk]:
- conn = self._require_open()
- row = conn.execute(
- """SELECT id, entity_id, seq, content, kind, metadata, token_count
- FROM chunk WHERE id=?""",
- (chunk_id,),
- ).fetchone()
- return _row_to_chunk(row) if row else None
-
- def list_chunks(self, entity_id: str) -> list[Chunk]:
- conn = self._require_open()
- rows = conn.execute(
- """SELECT id, entity_id, seq, content, kind, metadata, token_count
- FROM chunk WHERE entity_id=? ORDER BY seq""",
- (entity_id,),
- ).fetchall()
- return [_row_to_chunk(r) for r in rows]
-
- def delete_chunks(self, entity_id: str) -> int:
- """Delete all chunks for an entity. Returns the count removed."""
- with self._write_ctx() as conn:
- cur = conn.execute("DELETE FROM chunk WHERE entity_id=?", (entity_id,))
- deleted = cur.rowcount
- if deleted:
- conn.execute("DELETE FROM vector_index")
- if deleted:
- self._hnsw_cache.clear()
- return deleted
-
- # ------------------------------------------------------------------
- # Structured query
- # ------------------------------------------------------------------
-
- def query_entities(
- self,
- *,
- kind: Optional[str] = None,
- namespace: Optional[str] = None,
- where: Optional[dict[str, Any]] = None,
- limit: Optional[int] = None,
- offset: int = 0,
- order_by: str = "updated_at DESC",
- ) -> list[Entity]:
- """
- Structured query over the entity table.
-
- `where` keys:
- - direct column names from _ENTITY_COLUMNS (exact match or tuple (op, val))
- - ``metadata.<path>`` → json_extract(metadata, '$.<path>')
-
- Operator tuples: ``{"created_at": (">=", 1700000000)}``,
- ``{"kind": ("IN", ["document", "code"])}``.
- """
- conn = self._require_open()
- clauses: list[str] = []
- params: list[Any] = []
-
- if kind is not None:
- clauses.append("kind = ?")
- params.append(kind)
- if namespace is not None:
- clauses.append("namespace = ?")
- params.append(namespace)
-
- if where:
- extra_c, extra_p = _build_entity_where(where, table_prefix="")
- clauses.extend(extra_c)
- params.extend(extra_p)
-
- sql = (
- "SELECT id, kind, namespace, content, metadata, "
- "created_at, updated_at, version FROM entity"
- )
- if clauses:
- sql += " WHERE " + " AND ".join(clauses)
- sql += f" ORDER BY {_sanitize_order_by(order_by)}"
- if limit is not None:
- sql += " LIMIT ? OFFSET ?"
- params.append(int(limit))
- params.append(int(offset))
-
- rows = conn.execute(sql, params).fetchall()
- return [_row_to_entity(r) for r in rows]
-
- # ------------------------------------------------------------------
- # Embeddings (M3)
- # ------------------------------------------------------------------
-
- def set_embedding(
- self,
- chunk_id: str,
- vector: Any,
- *,
- model: str,
- ) -> None:
- """Store (or replace) the embedding of a chunk under ``model``."""
- import numpy as np
-
- arr = np.ascontiguousarray(vector, dtype=np.float32).reshape(-1)
- dim = int(arr.shape[0])
- blob = vector_to_bytes(arr)
- with self._write_ctx() as conn:
- exists = conn.execute(
- "SELECT 1 FROM chunk WHERE id=?", (chunk_id,)
- ).fetchone()
- if not exists:
- raise KeyError(f"chunk not found: {chunk_id!r}")
- conn.execute(
- """INSERT OR REPLACE INTO embedding
- (chunk_id, model, dim, vector) VALUES (?, ?, ?, ?)""",
- (chunk_id, model, dim, blob),
- )
- conn.execute("DELETE FROM vector_index WHERE model=?", (model,))
- self._hnsw_cache.pop(model, None)
-
- def get_embedding(self, chunk_id: str, *, model: str) -> Optional[Any]:
- """Return the stored vector (numpy array) or None."""
- conn = self._require_open()
- row = conn.execute(
- "SELECT vector, dim FROM embedding WHERE chunk_id=? AND model=?",
- (chunk_id, model),
- ).fetchone()
- if row is None:
- return None
- blob, dim = row
- return bytes_to_vector(blob, dim)
-
- def embed_chunks(
- self,
- embedder: Embedder,
- *,
- entity_id: Optional[str] = None,
- missing_only: bool = True,
- batch_size: int = 64,
- ) -> int:
- """
- Encode chunks with ``embedder`` and store the resulting vectors.
-
- Args:
- embedder: anything with ``.model``, ``.dim``, ``.encode``.
- entity_id: restrict to chunks of this entity. None = all.
- missing_only: skip chunks that already have an embedding under
- ``embedder.model``. Set False to re-encode.
- batch_size: chunks per ``encode`` call.
-
- Returns:
- number of chunks embedded.
- """
- import numpy as np
-
- conn = self._require_rw()
-
- clauses = []
- params: list[Any] = []
- if entity_id is not None:
- clauses.append("chunk.entity_id = ?")
- params.append(entity_id)
- if missing_only:
- clauses.append(
- "NOT EXISTS (SELECT 1 FROM embedding e "
- "WHERE e.chunk_id = chunk.id AND e.model = ?)"
- )
- params.append(embedder.model)
-
- sql = "SELECT chunk.id, chunk.content FROM chunk"
- if clauses:
- sql += " WHERE " + " AND ".join(clauses)
- sql += " ORDER BY chunk.rowid"
- rows = conn.execute(sql, params).fetchall()
- if not rows:
- return 0
-
- total = 0
- with self._write_ctx() as write_conn:
- for i in range(0, len(rows), batch_size):
- batch = rows[i : i + batch_size]
- texts = [r[1] for r in batch]
- vecs = embedder.encode(texts)
- arr = np.ascontiguousarray(vecs, dtype=np.float32)
- if arr.shape != (len(batch), embedder.dim):
- raise ValueError(
- f"embedder {embedder.model!r} returned shape "
- f"{arr.shape}, expected ({len(batch)}, {embedder.dim})"
- )
- payload = [
- (batch[j][0], embedder.model, embedder.dim, arr[j].tobytes())
- for j in range(len(batch))
- ]
- write_conn.executemany(
- """INSERT OR REPLACE INTO embedding
- (chunk_id, model, dim, vector) VALUES (?, ?, ?, ?)""",
- payload,
- )
- total += len(batch)
- if total:
- write_conn.execute(
- "DELETE FROM vector_index WHERE model=?",
- (embedder.model,),
- )
- if total:
- self._hnsw_cache.pop(embedder.model, None)
- return total
-
- # ------------------------------------------------------------------
- # Semantic / keyword / hybrid search
- # ------------------------------------------------------------------
-
- def search(
- self,
- query: Union[str, Any],
- *,
- mode: Literal["vector", "keyword", "hybrid"] = "vector",
- top_k: int = 5,
- model: Optional[str] = None,
- embedder: Optional[Embedder] = None,
- namespace: Optional[str] = None,
- filters: Optional[dict[str, Any]] = None,
- metric: Literal["cosine", "dot", "l2"] = "cosine",
- rrf_k: int = 60,
- fusion_oversample: int = 2,
- raw_fts: bool = False,
- use_hnsw: Literal["auto", "on", "off"] = "auto",
- ) -> list[SearchHit]:
- """
- Search chunks. Three modes:
-
- - ``mode="vector"`` (default): embedding cosine/dot/l2 ranking.
- - ``mode="keyword"``: SQLite FTS5 BM25.
- - ``mode="hybrid"``: vector + keyword fused by Reciprocal Rank Fusion.
-
- Args:
- query: str (any mode) OR 1-D vector (vector mode only).
- mode: 'vector' | 'keyword' | 'hybrid'.
- top_k: max hits.
- model: embedding model name (vector / hybrid).
- embedder: encodes string queries; provides ``model`` if missing.
- namespace: restrict to entities in this namespace.
- filters: same shape as ``query_entities`` where=, applied
- before ranking.
- metric: vector metric ('cosine' | 'dot' | 'l2').
- rrf_k: Reciprocal Rank Fusion constant (hybrid only).
- fusion_oversample: each modality fetches top_k*N before fusion.
- raw_fts: pass keyword query through to FTS5 unescaped (advanced).
- use_hnsw: 'auto' (default) picks HNSW when applicable;
- 'on' forces HNSW (errors if filters/non-cosine/no
- hnswlib); 'off' always uses brute force.
- """
- if top_k <= 0:
- return []
- if mode == "vector":
- ranked = self._run_vector(
- query, top_k=top_k, model=model, embedder=embedder,
- namespace=namespace, filters=filters, metric=metric,
- use_hnsw=use_hnsw,
- )
- return self._materialize_hits(ranked)
- if mode == "keyword":
- if not isinstance(query, str):
- raise TypeError("keyword mode requires a string query")
- ranked = self._run_keyword(
- query, top_k=top_k, namespace=namespace,
- filters=filters, raw_fts=raw_fts,
- )
- return self._materialize_hits(ranked)
- if mode == "hybrid":
- if not isinstance(query, str):
- raise TypeError("hybrid mode requires a string query")
- return self._search_hybrid(
- query, top_k=top_k, model=model, embedder=embedder,
- namespace=namespace, filters=filters, metric=metric,
- rrf_k=rrf_k, fusion_oversample=fusion_oversample,
- raw_fts=raw_fts, use_hnsw=use_hnsw,
- )
- raise ValueError(
- f"unknown search mode {mode!r}: expected 'vector', "
- f"'keyword', or 'hybrid'"
- )
-
- # ---- vector ------------------------------------------------------
-
- def _run_vector(
- self,
- query: Union[str, Any],
- *,
- top_k: int,
- model: Optional[str],
- embedder: Optional[Embedder],
- namespace: Optional[str],
- filters: Optional[dict[str, Any]],
- metric: Literal["cosine", "dot", "l2"],
- use_hnsw: Literal["auto", "on", "off"] = "auto",
- ) -> list[tuple[str, float]]:
- q_vec, model = _resolve_query_vector(
- query, model=model, embedder=embedder, caller="vector search"
- )
-
- # Pick HNSW vs brute. HNSW only handles cosine + no JSON filters;
- # namespace is fine because we post-filter on the result set.
- hnsw_eligible = (
- metric == "cosine"
- and not filters
- and hnsw_available()
- )
- if use_hnsw == "on":
- if metric != "cosine":
- raise ValueError(
- f"use_hnsw='on' requires metric='cosine', got {metric!r}"
- )
- if filters:
- raise ValueError(
- "use_hnsw='on' is incompatible with filters= "
- "(HNSW cannot pre-filter; use 'auto' or 'off')"
- )
- if not hnsw_available():
- raise RuntimeError(
- "use_hnsw='on' but hnswlib is not installed"
- )
- return self._run_vector_hnsw(
- q_vec, model=model, namespace=namespace, top_k=top_k
- )
- if use_hnsw == "auto" and hnsw_eligible:
- return self._run_vector_hnsw(
- q_vec, model=model, namespace=namespace, top_k=top_k
- )
- return self._run_vector_brute(
- q_vec, model=model, namespace=namespace,
- filters=filters, top_k=top_k, metric=metric,
- )
-
- def _run_vector_brute(
- self,
- q_vec: "Any",
- *,
- model: str,
- namespace: Optional[str],
- filters: Optional[dict[str, Any]],
- top_k: int,
- metric: Literal["cosine", "dot", "l2"],
- ) -> list[tuple[str, float]]:
- conn = self._require_open()
- clauses = ["embedding.model = ?"]
- params: list[Any] = [model]
- if namespace is not None:
- clauses.append("entity.namespace = ?")
- params.append(namespace)
- if filters:
- fc, fp = _build_entity_where(filters, table_prefix="entity.")
- clauses.extend(fc)
- params.extend(fp)
-
- sql = (
- "SELECT chunk.id, embedding.vector, embedding.dim "
- "FROM embedding "
- "JOIN chunk ON chunk.id = embedding.chunk_id "
- "JOIN entity ON entity.id = chunk.entity_id "
- f"WHERE {' AND '.join(clauses)}"
- )
- rows = conn.execute(sql, params).fetchall()
- if not rows:
- return []
-
- ids = [r[0] for r in rows]
- vectors, _dim = _stack_vectors(rows, model=model, q_vec=q_vec)
- return search_brute(q_vec, vectors, ids, top_k=top_k, metric=metric)
-
- def _run_vector_hnsw(
- self,
- q_vec: "Any",
- *,
- model: str,
- namespace: Optional[str],
- top_k: int,
- ) -> list[tuple[str, float]]:
- idx = self._get_hnsw(model)
- if idx is None:
- return []
- if q_vec.shape[0] != idx.dim:
- raise ValueError(
- f"query dim {q_vec.shape[0]} != index dim {idx.dim} "
- f"for model {model!r}"
- )
-
- conn = self._require_open()
- fetch_k = min(
- idx.count,
- max(top_k, top_k * 3 if namespace is not None else top_k),
- )
-
- while fetch_k > 0:
- raw = idx.search(q_vec, top_k=fetch_k)
- if not raw:
- return []
-
- rowids = [r[0] for r in raw]
- placeholders = ",".join("?" * len(rowids))
- params: list[Any] = list(rowids)
- sql = (
- f"SELECT chunk.rowid, chunk.id "
- f"FROM chunk JOIN entity ON entity.id = chunk.entity_id "
- f"WHERE chunk.rowid IN ({placeholders})"
- )
- if namespace is not None:
- sql += " AND entity.namespace = ?"
- params.append(namespace)
-
- rid_to_cid = {
- r[0]: r[1] for r in conn.execute(sql, params).fetchall()
- }
- out = [
- (rid_to_cid[rid], score)
- for rid, score in raw
- if rid in rid_to_cid
- ]
- if len(out) >= top_k or fetch_k >= idx.count:
- return out[:top_k]
-
- # A sparse namespace may not occur in the first 3× candidates.
- # Expand progressively until top_k is satisfied or the complete
- # index has been considered.
- fetch_k = min(idx.count, max(fetch_k * 2, top_k))
-
- return []
-
- # ---- HNSW cache + persistence ------------------------------------
-
- def _embedding_count(self, model: str) -> int:
- conn = self._require_open()
- (n,) = conn.execute(
- "SELECT COUNT(*) FROM embedding WHERE model=?", (model,)
- ).fetchone()
- return int(n)
-
- def _get_hnsw(self, model: str) -> Optional[HnswIndex]:
- """
- Return a usable HNSW for ``model``, building or reloading as needed.
-
- Cache invariant: ``HnswIndex.count`` equals the live embedding
- count. Adds/deletes are caught by the count check; in-process vector
- replacements explicitly invalidate the affected model.
- """
- cur_count = self._embedding_count(model)
- if cur_count == 0:
- self._hnsw_cache.pop(model, None)
- return None
-
- cached = self._hnsw_cache.get(model)
- if cached is not None and cached.count == cur_count:
- return cached
-
- # Try the persisted blob — only useful on first access after open.
- if cached is None:
- loaded = self._load_hnsw_blob(model, expected_count=cur_count)
- if loaded is not None:
- self._hnsw_cache[model] = loaded
- return loaded
-
- return self._rebuild_hnsw(model)
-
- def _load_hnsw_blob(
- self, model: str, *, expected_count: int
- ) -> Optional[HnswIndex]:
- conn = self._require_open()
- row = conn.execute(
- "SELECT index_blob, index_meta FROM vector_index "
- "WHERE model=? AND namespace='*'",
- (model,),
- ).fetchone()
- if row is None:
- return None
- blob, meta_json = row
- try:
- meta = json.loads(meta_json) if meta_json else {}
- except json.JSONDecodeError:
- return None
- if meta.get("count") != expected_count:
- return None
- dim = meta.get("dim")
- if not isinstance(dim, int):
- return None
- try:
- return HnswIndex.deserialize(blob, dim=dim)
- except Exception:
- # Corrupt blob shouldn't block search — fall through to rebuild.
- return None
-
- def _rebuild_hnsw(self, model: str) -> Optional[HnswIndex]:
- conn = self._require_open()
- rows = conn.execute(
- "SELECT chunk.rowid, embedding.vector, embedding.dim "
- "FROM embedding JOIN chunk ON chunk.id = embedding.chunk_id "
- "WHERE embedding.model = ?",
- (model,),
- ).fetchall()
- if not rows:
- return None
-
- vectors, dim = _stack_vectors(rows, model=model)
- labels = [int(r[0]) for r in rows]
- idx = HnswIndex.build(vectors, labels)
- self._hnsw_cache[model] = idx
- if self._mode != "r":
- self._persist_hnsw(model, idx, dim=dim)
- return idx
-
- def _persist_hnsw(self, model: str, idx: HnswIndex, *, dim: int) -> None:
- if self._mode == "r":
- return
- blob = idx.serialize()
- meta = json.dumps(
- {"count": idx.count, "dim": dim, "M": idx.M},
- sort_keys=True,
- )
- with self._write_ctx() as conn:
- conn.execute(
- """INSERT INTO vector_index
- (model, namespace, index_blob, index_meta, updated_at)
- VALUES (?, '*', ?, ?, ?)
- ON CONFLICT(model, namespace) DO UPDATE SET
- index_blob = excluded.index_blob,
- index_meta = excluded.index_meta,
- updated_at = excluded.updated_at""",
- (model, blob, meta, time.time()),
- )
-
- # ---- keyword -----------------------------------------------------
-
- def _search_keyword(
- self,
- query: str,
- *,
- top_k: int,
- namespace: Optional[str],
- filters: Optional[dict[str, Any]],
- raw_fts: bool,
- ) -> list[SearchHit]:
- ranked = self._run_keyword(
- query, top_k=top_k, namespace=namespace,
- filters=filters, raw_fts=raw_fts,
- )
- return self._materialize_hits(ranked)
-
- def _run_keyword(
- self,
- query: str,
- *,
- top_k: int,
- namespace: Optional[str],
- filters: Optional[dict[str, Any]],
- raw_fts: bool,
- ) -> list[tuple[str, float]]:
- fts_q = query if raw_fts else _fts_escape_query(query)
- if not fts_q:
- return []
-
- conn = self._require_open()
- clauses = ["chunk_fts MATCH ?"]
- params: list[Any] = [fts_q]
- if namespace is not None:
- clauses.append("entity.namespace = ?")
- params.append(namespace)
- if filters:
- fc, fp = _build_entity_where(filters, table_prefix="entity.")
- clauses.extend(fc)
- params.extend(fp)
- params.append(top_k)
-
- # bm25() returns negative; smaller (more negative) = better.
- # Negate to expose a "higher is better" score.
- sql = (
- "SELECT chunk.id, -bm25(chunk_fts) AS score "
- "FROM chunk_fts "
- "JOIN chunk ON chunk.rowid = chunk_fts.rowid "
- "JOIN entity ON entity.id = chunk.entity_id "
- f"WHERE {' AND '.join(clauses)} "
- "ORDER BY bm25(chunk_fts) ASC LIMIT ?"
- )
- try:
- rows = conn.execute(sql, params).fetchall()
- except sqlite3.OperationalError as exc:
- # FTS5 syntax errors when raw_fts=True or escape misses an edge case.
- raise ValueError(f"FTS5 query failed: {exc} (query={fts_q!r})") from exc
- return [(r[0], float(r[1])) for r in rows]
-
- # ---- hybrid (RRF) ------------------------------------------------
-
- def _search_hybrid(
- self,
- query: str,
- *,
- top_k: int,
- model: Optional[str],
- embedder: Optional[Embedder],
- namespace: Optional[str],
- filters: Optional[dict[str, Any]],
- metric: Literal["cosine", "dot", "l2"],
- rrf_k: int,
- fusion_oversample: int,
- raw_fts: bool,
- use_hnsw: Literal["auto", "on", "off"] = "auto",
- ) -> list[SearchHit]:
- if fusion_oversample < 1:
- raise ValueError("fusion_oversample must be >= 1")
- per_modality_k = max(top_k * fusion_oversample, top_k)
-
- vec_ranked = self._run_vector(
- query, top_k=per_modality_k, model=model, embedder=embedder,
- namespace=namespace, filters=filters, metric=metric,
- use_hnsw=use_hnsw,
- )
- kw_ranked = self._run_keyword(
- query, top_k=per_modality_k, namespace=namespace,
- filters=filters, raw_fts=raw_fts,
- )
-
- scores: dict[str, float] = {}
- for rank, (cid, _s) in enumerate(vec_ranked, start=1):
- scores[cid] = scores.get(cid, 0.0) + 1.0 / (rrf_k + rank)
- for rank, (cid, _s) in enumerate(kw_ranked, start=1):
- scores[cid] = scores.get(cid, 0.0) + 1.0 / (rrf_k + rank)
-
- fused = sorted(scores.items(), key=lambda kv: kv[1], reverse=True)[:top_k]
- return self._materialize_hits(fused)
-
- # ---- shared materialization -------------------------------------
-
- def _materialize_hits(
- self, ranked: list[tuple[str, float]]
- ) -> list[SearchHit]:
- hits: list[SearchHit] = []
- for chunk_id, score in ranked:
- chunk = self.get_chunk(chunk_id)
- if chunk is None:
- continue # raced against delete
- entity = self.get_entity(chunk.entity_id)
- if entity is None:
- continue
- hits.append(SearchHit(chunk=chunk, entity=entity, score=score))
- return hits
-
- # ------------------------------------------------------------------
- # Memory: remember / recall / forget
- # ------------------------------------------------------------------
-
- def remember(
- self,
- content: str,
- *,
- session_id: Optional[str] = None,
- actor: Optional[str] = None,
- salience: float = 0.5,
- ttl: Optional[float] = None,
- expires_at: Optional[float] = None,
- namespace: str = "default",
- metadata: Optional[dict[str, Any]] = None,
- id: Optional[str] = None, # noqa: A002
- embedder: Optional[Embedder] = None,
- ) -> str:
- """
- Persist a memory item. Atomic across entity + chunk + memory_ext
- (and embedding when ``embedder`` is provided).
-
- Args:
- content: the memory text.
- session_id: which conversation/task this belongs to.
- actor: which agent recorded it.
- salience: importance weight in [0, 1].
- ttl: seconds-from-now until expiry. Mutually exclusive
- with ``expires_at``.
- expires_at: absolute unix timestamp of expiry.
- namespace: tenant/dataset bucket.
- metadata: arbitrary JSON metadata on the underlying entity.
- id: pre-assigned entity id; auto-uuid otherwise.
- embedder: if given, embed the chunk eagerly so ``recall``
- can do semantic retrieval immediately.
-
- Returns the memory's entity id.
- """
- if ttl is not None and expires_at is not None:
- raise TypeError("provide ttl OR expires_at, not both")
- if not 0.0 <= salience <= 1.0:
- raise ValueError(f"salience must be in [0, 1], got {salience}")
-
- eid = id if id is not None else _uuid()
- now = time.time()
- eat: Optional[float]
- if expires_at is not None:
- eat = expires_at
- elif ttl is not None:
- eat = now + ttl
- else:
- eat = None
-
- with self.transaction():
- self.put_entity(
- id=eid,
- kind="memory",
- namespace=namespace,
- content=content,
- metadata=metadata or {},
- )
- self.add_chunks(eid, [content])
- with self._write_ctx() as conn:
- conn.execute(
- """INSERT INTO memory_ext
- (entity_id, session_id, actor, salience, expires_at)
- VALUES (?, ?, ?, ?, ?)""",
- (eid, session_id, actor, salience, eat),
- )
- if embedder is not None:
- self.embed_chunks(embedder, entity_id=eid)
- return eid
-
- def recall(
- self,
- query: Union[str, Any],
- *,
- top_k: int = 5,
- session_id: Optional[str] = None,
- actor: Optional[str] = None,
- namespace: Optional[str] = None,
- min_salience: Optional[float] = None,
- include_expired: bool = False,
- model: Optional[str] = None,
- embedder: Optional[Embedder] = None,
- metric: Literal["cosine", "dot", "l2"] = "cosine",
- ) -> list[MemoryHit]:
- """
- Semantic recall over stored memories.
-
- Filters compose with AND. Expired memories are excluded by default.
- Returns ranked ``MemoryHit`` (highest score first).
- """
- q_vec, model = _resolve_query_vector(
- query, model=model, embedder=embedder, caller="recall"
- )
-
- conn = self._require_open()
- clauses = ["embedding.model = ?", "entity.kind = 'memory'"]
- params: list[Any] = [model]
- if namespace is not None:
- clauses.append("entity.namespace = ?")
- params.append(namespace)
- if session_id is not None:
- clauses.append("memory_ext.session_id = ?")
- params.append(session_id)
- if actor is not None:
- clauses.append("memory_ext.actor = ?")
- params.append(actor)
- if min_salience is not None:
- clauses.append("memory_ext.salience >= ?")
- params.append(min_salience)
- if not include_expired:
- clauses.append(
- "(memory_ext.expires_at IS NULL OR memory_ext.expires_at > ?)"
- )
- params.append(time.time())
-
- sql = (
- "SELECT entity.id, embedding.vector, embedding.dim "
- "FROM embedding "
- "JOIN chunk ON chunk.id = embedding.chunk_id "
- "JOIN entity ON entity.id = chunk.entity_id "
- "JOIN memory_ext ON memory_ext.entity_id = entity.id "
- f"WHERE {' AND '.join(clauses)}"
- )
- rows = conn.execute(sql, params).fetchall()
- if not rows:
- return []
-
- ids = [r[0] for r in rows]
- vectors, _dim = _stack_vectors(rows, model=model, q_vec=q_vec)
- ranked = search_brute(q_vec, vectors, ids, top_k=top_k, metric=metric)
- return [
- MemoryHit(memory=m, score=score)
- for entity_id, score in ranked
- for m in [self._load_memory(entity_id)]
- if m is not None
- ]
-
- def forget(
- self,
- *,
- entity_id: Optional[str] = None,
- session_id: Optional[str] = None,
- actor: Optional[str] = None,
- namespace: Optional[str] = None,
- before: Optional[float] = None,
- ) -> int:
- """
- Delete memories matching ALL provided criteria. At least one
- criterion is required (no-arg ``forget()`` raises). Returns the
- number of memories removed.
- """
- if all(
- x is None
- for x in (entity_id, session_id, actor, namespace, before)
- ):
- raise TypeError("forget() requires at least one criterion")
-
- conn = self._require_rw()
- clauses = ["entity.kind = 'memory'"]
- params: list[Any] = []
- if entity_id is not None:
- clauses.append("entity.id = ?")
- params.append(entity_id)
- if namespace is not None:
- clauses.append("entity.namespace = ?")
- params.append(namespace)
- if before is not None:
- clauses.append("entity.created_at < ?")
- params.append(before)
- if session_id is not None:
- clauses.append("memory_ext.session_id = ?")
- params.append(session_id)
- if actor is not None:
- clauses.append("memory_ext.actor = ?")
- params.append(actor)
-
- sql = (
- "SELECT entity.id FROM entity "
- "JOIN memory_ext ON memory_ext.entity_id = entity.id "
- f"WHERE {' AND '.join(clauses)}"
- )
- ids = [r[0] for r in conn.execute(sql, params).fetchall()]
- return self._delete_entities(ids)
-
- def forget_expired(self, *, now: Optional[float] = None) -> int:
- """Delete every memory whose ``expires_at`` is in the past."""
- t = now if now is not None else time.time()
- conn = self._require_rw()
- ids = [
- r[0]
- for r in conn.execute(
- """SELECT entity_id FROM memory_ext
- WHERE expires_at IS NOT NULL AND expires_at <= ?""",
- (t,),
- ).fetchall()
- ]
- return self._delete_entities(ids)
-
- def _delete_entities(self, ids: list[str]) -> int:
- if not ids:
- return 0
- with self._write_ctx() as conn:
- conn.executemany(
- "DELETE FROM entity WHERE id=?", [(i,) for i in ids]
- )
- conn.execute("DELETE FROM vector_index")
- self._hnsw_cache.clear()
- return len(ids)
-
- def _load_memory(self, entity_id: str) -> Optional[MemoryItem]:
- conn = self._require_open()
- row = conn.execute(
- """SELECT entity.id, entity.kind, entity.namespace, entity.content,
- entity.metadata, entity.created_at, entity.updated_at,
- entity.version,
- memory_ext.session_id, memory_ext.actor,
- memory_ext.salience, memory_ext.expires_at
- FROM entity
- JOIN memory_ext ON memory_ext.entity_id = entity.id
- WHERE entity.id = ?""",
- (entity_id,),
- ).fetchone()
- if row is None:
- return None
- return MemoryItem(
- id=row[0],
- kind=row[1],
- namespace=row[2],
- content=row[3],
- metadata=_loads(row[4]),
- created_at=row[5],
- updated_at=row[6],
- version=row[7],
- session_id=row[8],
- actor=row[9],
- salience=row[10],
- expires_at=row[11],
- )
-
- # ------------------------------------------------------------------
- # Graph: relations + neighbors + traverse + subgraph
- # ------------------------------------------------------------------
-
- def add_relation(
- self,
- src: str,
- dst: str,
- kind: str,
- *,
- weight: float = 1.0,
- metadata: Optional[dict[str, Any]] = None,
- ) -> None:
- """
- Insert or replace a directed edge (src -> dst, labeled ``kind``).
- Both endpoints must already exist as entities; missing endpoints
- raise ``KeyError`` rather than relying on the FK error.
- """
- conn = self._require_rw()
- for name, eid in (("src", src), ("dst", dst)):
- if not conn.execute(
- "SELECT 1 FROM entity WHERE id=?", (eid,)
- ).fetchone():
- raise KeyError(f"{name} entity not found: {eid!r}")
-
- meta_json = _dumps(metadata or {})
- with self._write_ctx() as wconn:
- wconn.execute(
- """INSERT INTO relation (src, dst, kind, weight, metadata)
- VALUES (?, ?, ?, ?, ?)
- ON CONFLICT(src, dst, kind) DO UPDATE SET
- weight=excluded.weight,
- metadata=excluded.metadata""",
- (src, dst, kind, weight, meta_json),
- )
-
- def get_relation(
- self, src: str, dst: str, kind: str
- ) -> Optional[Relation]:
- conn = self._require_open()
- row = conn.execute(
- """SELECT src, dst, kind, weight, metadata
- FROM relation WHERE src=? AND dst=? AND kind=?""",
- (src, dst, kind),
- ).fetchone()
- return _row_to_relation(row) if row else None
-
- def delete_relation(self, src: str, dst: str, kind: str) -> bool:
- """Returns True iff a row was deleted."""
- with self._write_ctx() as conn:
- cur = conn.execute(
- "DELETE FROM relation WHERE src=? AND dst=? AND kind=?",
- (src, dst, kind),
- )
- return cur.rowcount > 0
-
- def list_relations(
- self,
- *,
- src: Optional[str] = None,
- dst: Optional[str] = None,
- kind: Optional[str] = None,
- limit: Optional[int] = None,
- ) -> list[Relation]:
- clauses: list[str] = []
- params: list[Any] = []
- if src is not None:
- clauses.append("src=?")
- params.append(src)
- if dst is not None:
- clauses.append("dst=?")
- params.append(dst)
- if kind is not None:
- clauses.append("kind=?")
- params.append(kind)
- sql = "SELECT src, dst, kind, weight, metadata FROM relation"
- if clauses:
- sql += " WHERE " + " AND ".join(clauses)
- sql += " ORDER BY src, dst, kind"
- if limit is not None:
- sql += " LIMIT ?"
- params.append(limit)
- conn = self._require_open()
- return [_row_to_relation(r) for r in conn.execute(sql, params).fetchall()]
-
- def neighbors(
- self,
- entity_id: str,
- *,
- kind: Optional[str] = None,
- direction: Literal["out", "in", "both"] = "out",
- limit: Optional[int] = None,
- ) -> list[Neighbor]:
- """
- Single-hop neighbors of ``entity_id``.
-
- ``direction``:
- - "out": edges where this node is the src; returns dst nodes.
- - "in": edges where this node is the dst; returns src nodes.
- - "both": both, with each result tagged accordingly.
- """
- if direction not in ("out", "in", "both"):
- raise ValueError(
- f"direction must be 'out' | 'in' | 'both', got {direction!r}"
- )
- conn = self._require_open()
- out: list[Neighbor] = []
-
- def _query_one(
- this_col: str, other_col: str, label: str
- ) -> list[Neighbor]:
- clauses = [f"relation.{this_col} = ?"]
- params: list[Any] = [entity_id]
- if kind is not None:
- clauses.append("relation.kind = ?")
- params.append(kind)
- sql = (
- "SELECT entity.id, entity.kind, entity.namespace, "
- "entity.content, entity.metadata, entity.created_at, "
- "entity.updated_at, entity.version, "
- "relation.kind, relation.weight, relation.metadata "
- f"FROM relation JOIN entity ON entity.id = relation.{other_col} "
- f"WHERE {' AND '.join(clauses)}"
- )
- rows = conn.execute(sql, params).fetchall()
- return [
- Neighbor(
- entity=_row_to_entity(r[0:8]),
- relation_kind=r[8],
- weight=r[9],
- direction=label,
- metadata=_loads(r[10]),
- )
- for r in rows
- ]
-
- if direction in ("out", "both"):
- out.extend(_query_one("src", "dst", "out"))
- if direction in ("in", "both"):
- out.extend(_query_one("dst", "src", "in"))
-
- if limit is not None:
- out = out[:limit]
- return out
-
- def traverse(
- self,
- start: str,
- *,
- kind: Optional[str] = None,
- direction: Literal["out", "in", "both"] = "out",
- max_hops: int = 3,
- namespace: Optional[str] = None,
- ) -> list[TraversalHit]:
- """
- Multi-hop BFS from ``start``. Returns one row per reachable node
- (excluding ``start``) with its shortest hop count.
-
- Cycle-safe: bounded by ``max_hops``, hops dedup via ``UNION``,
- final selection takes ``MIN(hop)``.
- """
- if max_hops < 1:
- raise ValueError("max_hops must be >= 1")
- join_cond, next_node = _direction_sql(direction)
- kind_clause = "AND r.kind = ?" if kind is not None else ""
-
- cte_sql = (
- f"WITH RECURSIVE walk(id, hop) AS ("
- f" SELECT ?, 0"
- f" UNION"
- f" SELECT {next_node}, walk.hop + 1 "
- f" FROM walk JOIN relation r ON {join_cond} "
- f" WHERE walk.hop < ? {kind_clause}"
- f") "
- f"SELECT entity.id, entity.kind, entity.namespace, entity.content, "
- f" entity.metadata, entity.created_at, entity.updated_at, "
- f" entity.version, MIN(walk.hop) as min_hop "
- f"FROM walk JOIN entity ON entity.id = walk.id "
- f"WHERE walk.id != ? "
- )
- params: list[Any] = [start, max_hops]
- if kind is not None:
- params.append(kind)
- params.append(start)
- if namespace is not None:
- cte_sql += "AND entity.namespace = ? "
- params.append(namespace)
- cte_sql += "GROUP BY entity.id ORDER BY min_hop, entity.id"
-
- conn = self._require_open()
- rows = conn.execute(cte_sql, params).fetchall()
- return [
- TraversalHit(entity=_row_to_entity(r[0:8]), hop=int(r[8]))
- for r in rows
- ]
-
- def subgraph(
- self,
- seeds: Iterable[str],
- *,
- kind: Optional[str] = None,
- direction: Literal["out", "in", "both"] = "both",
- max_hops: int = 2,
- ) -> Subgraph:
- """
- Induced subgraph reachable from ``seeds``. Includes the seeds
- themselves plus every node within ``max_hops``, and every relation
- where both endpoints land in the visited set.
- """
- seed_list = list(seeds)
- if not seed_list:
- return Subgraph()
- if max_hops < 0:
- raise ValueError("max_hops must be >= 0")
- join_cond, next_node = _direction_sql(direction)
- kind_clause = "AND r.kind = ?" if kind is not None else ""
-
- # Multi-source seed set is fed to the CTE via json_each — works
- # for any seed count without dynamic placeholders.
- seeds_json = json.dumps(seed_list)
- params: list[Any] = [seeds_json]
-
- # When max_hops == 0, just collect the seeds themselves.
- if max_hops == 0:
- cte_sql = (
- "SELECT DISTINCT value AS id FROM json_each(?)"
- )
- else:
- params.append(max_hops)
- if kind is not None:
- params.append(kind)
- cte_sql = (
- f"WITH RECURSIVE walk(id, hop) AS ("
- f" SELECT value, 0 FROM json_each(?)"
- f" UNION"
- f" SELECT {next_node}, walk.hop + 1 "
- f" FROM walk JOIN relation r ON {join_cond} "
- f" WHERE walk.hop < ? {kind_clause}"
- f") "
- f"SELECT DISTINCT id FROM walk"
- )
-
- conn = self._require_open()
- ids = [r[0] for r in conn.execute(cte_sql, params).fetchall()]
- if not ids:
- return Subgraph()
-
- placeholders = ",".join("?" * len(ids))
- ent_rows = conn.execute(
- f"""SELECT id, kind, namespace, content, metadata,
- created_at, updated_at, version
- FROM entity WHERE id IN ({placeholders})
- ORDER BY id""",
- ids,
- ).fetchall()
- entities = [_row_to_entity(r) for r in ent_rows]
-
- rel_clauses = [f"src IN ({placeholders})", f"dst IN ({placeholders})"]
- rel_params = list(ids) + list(ids)
- if kind is not None:
- rel_clauses.append("kind = ?")
- rel_params.append(kind)
- rel_rows = conn.execute(
- "SELECT src, dst, kind, weight, metadata FROM relation "
- f"WHERE {' AND '.join(rel_clauses)} ORDER BY src, dst, kind",
- rel_params,
- ).fetchall()
- relations = [_row_to_relation(r) for r in rel_rows]
- return Subgraph(entities=entities, relations=relations)
-
- # ------------------------------------------------------------------
- # Ingestion helpers
- # ------------------------------------------------------------------
-
- def ingest_text(
- self,
- text: str,
- *,
- id: Optional[str] = None, # noqa: A002
- kind: str = "document",
- namespace: str = "default",
- chunker: Optional[Chunker] = None,
- metadata: Optional[dict[str, Any]] = None,
- embedder: Optional[Embedder] = None,
- ) -> str:
- """
- Atomic text ingestion: entity + chunks (+ embeddings).
-
- ``chunker`` defaults to :py:func:`chunkers.by_paragraph`. If the
- chosen chunker yields no chunks (e.g. the text has no blank-line
- boundaries), the whole text is stored as one chunk so the entity
- is never silently empty.
- """
- eid = id if id is not None else _uuid()
- active_chunker = chunker if chunker is not None else by_paragraph()
- pieces = active_chunker(text)
- if not pieces:
- stripped = text.strip()
- if stripped:
- pieces = [stripped]
-
- with self.transaction():
- self.put_entity(
- id=eid, kind=kind, namespace=namespace,
- content=text, metadata=metadata or {},
- )
- if pieces:
- self.add_chunks(eid, pieces)
- if embedder is not None and pieces:
- self.embed_chunks(embedder, entity_id=eid)
- return eid
-
- def ingest_table(
- self,
- rows: Union[Iterable[dict[str, Any]], str, Path],
- *,
- table_name: str,
- id_column: Optional[str] = None,
- namespace: str = "default",
- embedder: Optional[Embedder] = None,
- ) -> list[str]:
- """
- Atomic table ingestion. Each row becomes one entity (kind=
- ``"table_row"``) holding the row in ``metadata.columns``, plus
- one chunk rendered as ``"k: v\\n..."`` for full-text + embedding.
-
- ``rows`` may be:
- - an iterable of dicts (in-memory rows), or
- - a path / string path to a CSV file (read via ``csv.DictReader``).
-
- ``id_column`` picks the entity id from a row column when present;
- otherwise an auto-uuid is assigned.
- """
- if isinstance(rows, (str, Path)):
- rows = _read_csv_rows(Path(rows))
- materialized = list(rows)
- if not materialized:
- return []
-
- ids: list[str] = []
- with self.transaction():
- for row in materialized:
- if id_column is not None and id_column in row:
- eid = str(row[id_column])
- else:
- eid = _uuid()
- ids.append(eid)
- content = _render_row(row)
- self.put_entity(
- id=eid, kind="table_row", namespace=namespace,
- content=content,
- metadata={"table": table_name, "columns": dict(row)},
- )
- self.add_chunks(eid, [content])
- if embedder is not None:
- for eid in ids:
- self.embed_chunks(embedder, entity_id=eid)
- return ids
-
- # ------------------------------------------------------------------
- # Management: validate / compact / export
- # ------------------------------------------------------------------
-
- def validate(self) -> ValidationReport:
- """
- Run consistency checks. Safe on a read-only handle.
-
- Catches:
- - SQLite page-level corruption (PRAGMA integrity_check)
- - Foreign-key violations (PRAGMA foreign_key_check)
- - Bad application_id / unsupported schema_version
- - FTS5 index corruption (deep check on writable handles)
- - Stale per-model HNSW blob (count != live embedding count)
- """
- conn = self._require_open()
- issues: list[str] = []
- counts: dict[str, int] = {}
-
- rows = conn.execute("PRAGMA integrity_check").fetchall()
- if not (len(rows) == 1 and rows[0][0] == "ok"):
- issues.append(
- "integrity_check: " + "; ".join(r[0] for r in rows)
- )
-
- fk_rows = conn.execute("PRAGMA foreign_key_check").fetchall()
- if fk_rows:
- issues.append(
- f"foreign_key_check: {len(fk_rows)} violation(s) "
- f"(first: {fk_rows[0]})"
- )
-
- (app_id,) = conn.execute("PRAGMA application_id").fetchone()
- if app_id != _schema.APPLICATION_ID:
- issues.append(
- f"application_id={app_id:#010x} "
- f"(expected {_schema.APPLICATION_ID:#010x})"
- )
- (sv,) = conn.execute("PRAGMA user_version").fetchone()
- if sv > _schema.SCHEMA_VERSION:
- issues.append(
- f"schema_version={sv} > library {_schema.SCHEMA_VERSION}"
- )
-
- for table in ("entity", "chunk", "embedding", "relation", "memory_ext"):
- counts[table] = conn.execute(
- f"SELECT COUNT(*) FROM {table}"
- ).fetchone()[0]
-
- if self._mode == "r":
- # FTS5 exposes its deep integrity check as a special INSERT,
- # which SQLite rejects on a read-only connection. Still verify
- # that the virtual table is readable without producing a false
- # corruption report for a healthy read-only store.
- try:
- conn.execute("SELECT rowid FROM chunk_fts LIMIT 1").fetchall()
- except sqlite3.DatabaseError as exc:
- issues.append(f"chunk_fts read-check: {exc}")
- else:
- # FTS5 'integrity-check' compares index entries against the
- # external content table. The control command opens a SQLite
- # transaction even though it does not change user data, so end
- # that transaction when validate() started it.
- was_in_transaction = conn.in_transaction
- try:
- conn.execute(
- "INSERT INTO chunk_fts(chunk_fts) VALUES('integrity-check')"
- )
- except sqlite3.DatabaseError as exc:
- issues.append(f"chunk_fts integrity-check: {exc}")
- finally:
- if not was_in_transaction and conn.in_transaction:
- conn.rollback()
-
- # Per-model HNSW staleness — only complain when a blob exists.
- idx_rows = conn.execute(
- "SELECT model, index_meta FROM vector_index"
- ).fetchall()
- for model, meta_json in idx_rows:
- try:
- meta = json.loads(meta_json) if meta_json else {}
- except json.JSONDecodeError:
- issues.append(f"vector_index[{model}]: corrupt index_meta JSON")
- continue
- (live,) = conn.execute(
- "SELECT COUNT(*) FROM embedding WHERE model=?", (model,)
- ).fetchone()
- if meta.get("count") != live:
- issues.append(
- f"vector_index[{model}] stale: "
- f"meta.count={meta.get('count')} != live={live}"
- )
-
- return ValidationReport(ok=not issues, issues=issues, counts=counts)
-
- def compact(self) -> dict[str, int]:
- """
- Reclaim space and rebuild auxiliary indexes.
-
- Steps:
- 1. Rebuild the FTS5 index from ``chunk``.
- 2. Drop persisted HNSW blobs (next search rebuilds them).
- 3. ``VACUUM`` the database.
-
- Returns ``{"size_before": N, "size_after": M}`` in bytes.
- """
- conn = self._require_rw()
- if self._tx_depth > 0:
- raise RuntimeError("compact() cannot run inside a transaction()")
-
- size_before = self.path.stat().st_size if self.path.exists() else 0
-
- with self._write_ctx() as wconn:
- wconn.execute(
- "INSERT INTO chunk_fts(chunk_fts) VALUES('rebuild')"
- )
- wconn.execute("DELETE FROM vector_index")
- self._hnsw_cache.clear()
-
- # VACUUM cannot run inside a transaction; sqlite3 implicitly opens
- # one for any DML, so issue a COMMIT first to be safe, then VACUUM.
- conn.commit()
- conn.execute("VACUUM")
-
- size_after = self.path.stat().st_size if self.path.exists() else 0
- return {"size_before": size_before, "size_after": size_after}
-
- def export(
- self,
- path: PathLike,
- *,
- include_embeddings: bool = True,
- ) -> dict[str, int]:
- """
- Dump the store to a single JSON file.
-
- Embeddings are emitted as base64-encoded float32 bytes so the
- output stays text-only and round-trippable. Set
- ``include_embeddings=False`` to omit them — useful when you only
- want a content backup and plan to re-embed at import time.
- """
- import base64
-
- conn = self._require_open()
- out_path = Path(path)
-
- entities = []
- for r in conn.execute(
- "SELECT id, kind, namespace, content, metadata, "
- "created_at, updated_at, version FROM entity"
- ):
- entities.append({
- "id": r[0], "kind": r[1], "namespace": r[2],
- "content": r[3], "metadata": _loads(r[4]),
- "created_at": r[5], "updated_at": r[6], "version": r[7],
- })
-
- chunks = []
- for r in conn.execute(
- "SELECT id, entity_id, seq, content, kind, metadata, "
- "token_count FROM chunk"
- ):
- chunks.append({
- "id": r[0], "entity_id": r[1], "seq": r[2],
- "content": r[3], "kind": r[4],
- "metadata": _loads(r[5]), "token_count": r[6],
- })
-
- relations = []
- for r in conn.execute(
- "SELECT src, dst, kind, weight, metadata FROM relation"
- ):
- relations.append({
- "src": r[0], "dst": r[1], "kind": r[2],
- "weight": r[3], "metadata": _loads(r[4]),
- })
-
- memories = []
- for r in conn.execute(
- "SELECT entity_id, session_id, actor, salience, expires_at "
- "FROM memory_ext"
- ):
- memories.append({
- "entity_id": r[0], "session_id": r[1], "actor": r[2],
- "salience": r[3], "expires_at": r[4],
- })
-
- embeddings: list[dict] = []
- if include_embeddings:
- for r in conn.execute(
- "SELECT chunk_id, model, dim, vector FROM embedding"
- ):
- embeddings.append({
- "chunk_id": r[0],
- "model": r[1],
- "dim": r[2],
- "vector_b64": base64.b64encode(r[3]).decode("ascii"),
- })
-
- payload = {
- "schema_version": _schema.SCHEMA_VERSION,
- "exported_at": time.time(),
- "entities": entities,
- "chunks": chunks,
- "embeddings": embeddings,
- "relations": relations,
- "memory_ext": memories,
- }
- out_path.write_text(
- json.dumps(payload, ensure_ascii=False, sort_keys=True),
- encoding="utf-8",
- )
- return {
- "entities": len(entities),
- "chunks": len(chunks),
- "embeddings": len(embeddings),
- "relations": len(relations),
- "memories": len(memories),
- "bytes_written": out_path.stat().st_size,
- }
-
-
-# ----------------------------------------------------------------------
-# shared query helpers
-# ----------------------------------------------------------------------
-
-def _resolve_query_vector(
- query: Union[str, Any],
- *,
- model: Optional[str],
- embedder: Optional[Embedder],
- caller: str,
-) -> tuple[Any, str]:
- """
- Normalize a search/recall query into (float32 query vector, model name).
-
- ``model`` falls back to ``embedder.model``; string queries require an
- embedder to encode them.
- """
- import numpy as np
-
- if embedder is not None and model is None:
- model = embedder.model
- if model is None:
- raise TypeError(
- f"{caller} requires either `model=` or an `embedder=` "
- f"(to infer the model)"
- )
-
- if isinstance(query, str):
- if embedder is None:
- raise TypeError("string query requires an `embedder=`")
- q_vec = embedder.encode([query])[0]
- else:
- q_vec = np.ascontiguousarray(query, dtype=np.float32).reshape(-1)
- return q_vec, model
-
-
-def _stack_vectors(
- rows: list[tuple],
- *,
- model: str,
- q_vec: Optional[Any] = None,
-) -> tuple[Any, int]:
- """
- Stack ``(key, vector_blob, dim)`` rows into an (N, dim) float32 matrix.
-
- Enforces a single consistent dim across rows and, when ``q_vec`` is
- given, that the query vector matches it. Returns (matrix, dim).
- """
- import numpy as np
-
- dims = {r[2] for r in rows}
- if len(dims) != 1:
- raise RuntimeError(
- f"inconsistent dims for model {model!r}: {sorted(dims)}"
- )
- dim = dims.pop()
- if q_vec is not None and q_vec.shape[0] != dim:
- raise ValueError(
- f"query dim {q_vec.shape[0]} != embedding dim {dim} "
- f"for model {model!r}"
- )
-
- vectors = np.empty((len(rows), dim), dtype=np.float32)
- for i, r in enumerate(rows):
- vectors[i] = bytes_to_vector(r[1], dim)
- return vectors, dim
-
-
-def _direction_sql(direction: str) -> tuple[str, str]:
- """
- Map a traversal direction to the recursive-CTE (join condition,
- next-node expression) pair shared by ``traverse`` and ``subgraph``.
- """
- if direction == "out":
- return "r.src = walk.id", "r.dst"
- if direction == "in":
- return "r.dst = walk.id", "r.src"
- if direction == "both":
- return (
- "(r.src = walk.id OR r.dst = walk.id)",
- "CASE WHEN r.src = walk.id THEN r.dst ELSE r.src END",
- )
- raise ValueError(
- f"direction must be 'out' | 'in' | 'both', got {direction!r}"
- )
-
-
-# ----------------------------------------------------------------------
-# row <-> dataclass helpers
-# ----------------------------------------------------------------------
-
-def _row_to_entity(row: tuple) -> Entity:
- return Entity(
- id=row[0],
- kind=row[1],
- namespace=row[2],
- content=row[3],
- metadata=_loads(row[4]),
- created_at=row[5],
- updated_at=row[6],
- version=row[7],
- )
-
-
-def _row_to_chunk(row: tuple) -> Chunk:
- return Chunk(
- id=row[0],
- entity_id=row[1],
- seq=row[2],
- content=row[3],
- kind=row[4],
- metadata=_loads(row[5]),
- token_count=row[6],
- )
-
-
-def _row_to_relation(row: tuple) -> Relation:
- return Relation(
- src=row[0],
- dst=row[1],
- kind=row[2],
- weight=row[3],
- metadata=_loads(row[4]),
- )
-
-
-def _coerce_chunk(
- raw: Union[Chunk, str, dict[str, Any]],
- *,
- entity_id: str,
- default_seq: int,
-) -> Chunk:
- if isinstance(raw, Chunk):
- # Force entity_id to the one we're inserting under; seq defaults
- # only if the caller passed seq=-1 as a sentinel (unusual).
- if raw.entity_id != entity_id:
- raw.entity_id = entity_id
- if not raw.id:
- raw.id = _uuid()
- return raw
- if isinstance(raw, str):
- return Chunk(
- id=_uuid(),
- entity_id=entity_id,
- seq=default_seq,
- content=raw,
- )
- if isinstance(raw, dict):
- d = dict(raw)
- return Chunk(
- id=d.get("id") or _uuid(),
- entity_id=entity_id,
- seq=d.get("seq", default_seq),
- content=d["content"],
- kind=d.get("kind", "text"),
- metadata=d.get("metadata", {}),
- token_count=d.get("token_count"),
- )
- raise TypeError(f"unsupported chunk type: {type(raw).__name__}")
-
-
-def _build_entity_where(
- where: dict[str, Any],
- *,
- table_prefix: str = "",
-) -> tuple[list[str], list[Any]]:
- """
- Shared builder for entity `where=` clauses used by both
- `query_entities` and `search`.
-
- ``table_prefix`` is prepended to column names (e.g. ``"entity."``) so the
- same builder works inside multi-table joins.
- """
- clauses: list[str] = []
- params: list[Any] = []
- for key, raw_val in where.items():
- op, val = _unpack_op(raw_val)
- if key.startswith("metadata."):
- path = "$." + key[len("metadata."):]
- col_expr = f"json_extract({table_prefix}metadata, ?)"
- if op == "IN":
- placeholders = ",".join("?" * len(val))
- clauses.append(f"{col_expr} IN ({placeholders})")
- params.append(path)
- params.extend(val)
- else:
- clauses.append(f"{col_expr} {op} ?")
- params.append(path)
- params.append(val)
- else:
- if key not in _ENTITY_COLUMNS:
- raise ValueError(
- f"unknown entity column: {key!r}. "
- f"Use 'metadata.<path>' for JSON fields."
- )
- col = f"{table_prefix}{key}"
- if op == "IN":
- placeholders = ",".join("?" * len(val))
- clauses.append(f"{col} IN ({placeholders})")
- params.extend(val)
- else:
- clauses.append(f"{col} {op} ?")
- params.append(val)
- return clauses, params
-
-
-def _unpack_op(raw_val: Any) -> tuple[str, Any]:
- """
- Accept either a bare value (=) or a (op, value) tuple.
- Validate op against the allow-list to block SQL injection
- via the operator slot.
- """
- if isinstance(raw_val, tuple) and len(raw_val) == 2 and isinstance(raw_val[0], str):
- op, val = raw_val
- op = op.upper()
- if op not in _OPERATORS:
- raise ValueError(f"unsupported operator: {raw_val[0]!r}")
- if op == "IN" and not isinstance(val, (list, tuple)):
- raise ValueError("IN operator requires a list/tuple value")
- return op, val
- return "=", raw_val
-
-
-def _sanitize_order_by(clause: str) -> str:
- """Accept only `<col> [ASC|DESC]` pairs, comma-separated."""
- out = []
- for segment in clause.split(","):
- tokens = segment.strip().split()
- if len(tokens) == 1:
- col, direction = tokens[0], "ASC"
- elif len(tokens) == 2:
- col, direction = tokens[0], tokens[1].upper()
- else:
- raise ValueError(f"invalid order_by clause: {segment!r}")
- if col not in _ENTITY_COLUMNS:
- raise ValueError(f"unknown order_by column: {col!r}")
- if direction not in {"ASC", "DESC"}:
- raise ValueError(f"invalid direction: {direction!r}")
- out.append(f"{col} {direction}")
- return ", ".join(out)
-
-
-__all__ = ["CarbonStore"]
diff --git a/Agent_module/examples/__init__.py b/Agent_module/examples/__init__.py
deleted file mode 100644
index e69de29bb2..0000000000
diff --git a/Agent_module/examples/carbondata_quickstart.py b/Agent_module/examples/carbondata_quickstart.py
deleted file mode 100644
index f3374b9551..0000000000
--- a/Agent_module/examples/carbondata_quickstart.py
+++ /dev/null
@@ -1,342 +0,0 @@
-"""
-05_carbondata_quickstart.py
-===========================
-End-to-end demo of carbon_data — Agent Data Infra.
-
-Walks through all four core scenarios on a tiny in-memory corpus:
-
- 1. RAG semantic / keyword / hybrid search
- 2. Long-term memory (remember / recall / forget)
- 3. Structured queries over an ingested table
- 4. Knowledge-graph traversal (relations / neighbors / subgraph)
-
-Plus admin: stats, validate, compact, export, persistence across reopen.
-
-The "embedder" used here is a hand-rolled vocabulary bag (no external
-model). It's deterministic and gives semantically meaningful clustering
-on the demo corpus, so you can see real top-k ranking — but it's only
-useful for didactic data of this size. Swap in any
-``embedder.encode(texts) -> np.ndarray`` for real workloads.
-
-Run:
- cd Agent_module
- .venv/bin/python examples/05_carbondata_quickstart.py
-"""
-from __future__ import annotations
-
-import os
-import sys
-import tempfile
-from pathlib import Path
-
-import numpy as np
-
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
-
-from Agent_module.carbon_data import (
- LambdaEmbedder,
- by_paragraph,
- by_tokens,
- create,
- open as cd_open,
-)
-
-
-# ---------------------------------------------------------------------------
-# Tiny vocabulary embedder — gives meaningful cosine clustering for the demo
-# without needing a real embedding model.
-# ---------------------------------------------------------------------------
-
-VOCAB = [
- # web
- "django", "flask", "fastapi", "web", "framework", "http", "route",
- # ml
- "tensorflow", "pytorch", "neural", "network", "model", "training",
- "learning", "embedding", "transformer",
- # cooking
- "recipe", "cooking", "ingredient", "kitchen", "flavor", "dish",
- # generic
- "python", "code", "data", "system",
-]
-DIM = len(VOCAB)
-
-
-def vocab_embed(texts):
- """Bag-of-vocabulary one-hot, L2-normalized — cosine acts like Jaccard."""
- out = np.zeros((len(texts), DIM), dtype=np.float32)
- for i, text in enumerate(texts):
- words = set(text.lower().split())
- for j, term in enumerate(VOCAB):
- if term in words:
- out[i, j] = 1.0
- norm = np.linalg.norm(out[i])
- if norm > 0:
- out[i] /= norm
- return out
-
-
-EMBEDDER = LambdaEmbedder(vocab_embed, model="vocab-demo", dim=DIM)
-
-
-# ---------------------------------------------------------------------------
-# helpers
-# ---------------------------------------------------------------------------
-
-def _h(title: str) -> None:
- print(f"\n{'=' * 60}\n {title}\n{'=' * 60}")
-
-
-def _row(label: str, value) -> None:
- print(f" {label:<24} {value}")
-
-
-# ---------------------------------------------------------------------------
-# the demo
-# ---------------------------------------------------------------------------
-
-def run() -> bool:
- workdir = Path(tempfile.mkdtemp(prefix="carbondata_demo_"))
- db_path = workdir / "kb.carbondata"
- print(f"Working in {workdir}")
-
- store = create(db_path)
-
- # -----------------------------------------------------------------
- # 1) Ingest text documents
- # -----------------------------------------------------------------
- _h("1. ingest_text — three short documents")
-
- docs = {
- "doc:django": (
- "Django is a high-level Python web framework that encourages "
- "rapid development.\n\n"
- "It ships with an ORM, an admin panel, and a templating system "
- "for building HTTP routes.\n\n"
- "Django is opinionated and convention-heavy."
- ),
- "doc:flask": (
- "Flask is a lightweight Python web framework with a minimal core.\n\n"
- "Routes, request parsing, and templating are unbundled, letting "
- "you choose the components yourself.\n\n"
- "Flask is unopinionated and small."
- ),
- "doc:transformer": (
- "A transformer is a neural network architecture built around "
- "self-attention.\n\n"
- "Modern embedding models and large language models use the "
- "transformer as the core training substrate.\n\n"
- "It largely replaced earlier recurrent neural network designs."
- ),
- "doc:recipe": (
- "Pasta carbonara is an Italian dish from Rome.\n\n"
- "The recipe uses egg, cheese, pepper, and cured pork — "
- "simple ingredients combined in the kitchen for rich flavor."
- ),
- }
-
- for eid, text in docs.items():
- store.ingest_text(
- text,
- id=eid,
- chunker=by_paragraph(),
- embedder=EMBEDDER,
- metadata={"source": "demo"},
- )
-
- stats = store.stats()
- _row("entities", stats["entities"])
- _row("chunks", stats["chunks"])
- _row("embeddings", stats["embeddings"])
-
- # -----------------------------------------------------------------
- # 2) Vector search — semantic ranking
- # -----------------------------------------------------------------
- _h("2. semantic search — query: 'python web framework'")
-
- hits = store.search(
- "python web framework", embedder=EMBEDDER, top_k=3
- )
- for h in hits:
- _row(f"{h.entity.id} (score={h.score:.3f})", h.chunk.content[:60] + "...")
- assert hits[0].entity.id in {"doc:django", "doc:flask"}, hits
-
- # -----------------------------------------------------------------
- # 3) Keyword search — exact-token retrieval via FTS5/BM25
- # -----------------------------------------------------------------
- _h("3. keyword search — query: 'transformer'")
-
- hits = store.search("transformer", mode="keyword", top_k=3)
- for h in hits:
- _row(f"{h.entity.id} (BM25={h.score:.3f})", h.chunk.content[:60] + "...")
- assert any(h.entity.id == "doc:transformer" for h in hits)
-
- # -----------------------------------------------------------------
- # 4) Hybrid search — vector ⊕ keyword via Reciprocal Rank Fusion
- # -----------------------------------------------------------------
- _h("4. hybrid search — query: 'neural training'")
-
- hits = store.search(
- "neural training", mode="hybrid", embedder=EMBEDDER, top_k=3
- )
- for h in hits:
- _row(f"{h.entity.id} (RRF={h.score:.4f})", h.chunk.content[:60] + "...")
-
- # -----------------------------------------------------------------
- # 5) Ingest a table — each row becomes a queryable entity
- # -----------------------------------------------------------------
- _h("5. ingest_table — users")
-
- users = [
- {"id": "u1", "name": "Alice", "team": "infra", "lang": "Go"},
- {"id": "u2", "name": "Bob", "team": "ml", "lang": "Python"},
- {"id": "u3", "name": "Carol", "team": "ml", "lang": "Python"},
- {"id": "u4", "name": "Dan", "team": "product", "lang": "Rust"},
- ]
- store.ingest_table(users, table_name="users", id_column="id",
- embedder=EMBEDDER)
-
- # Table column filter — note the JSON path goes through `columns.`
- # because ingest_table stores the row as metadata = {table, columns}.
- rows = store.query_entities(
- kind="table_row", where={"metadata.columns.team": "ml"}
- )
- print(f" team=ml rows: {len(rows)}")
- for e in rows:
- _row(e.id, e.metadata["columns"])
-
- # -----------------------------------------------------------------
- # 6) Long-term memory
- # -----------------------------------------------------------------
- _h("6. memory — remember + recall")
-
- store.remember(
- "user prefers django over flask",
- session_id="sess-1", actor="CodingAgent",
- salience=0.9, embedder=EMBEDDER,
- )
- store.remember(
- "user works mainly with python web framework code",
- session_id="sess-1", actor="CodingAgent",
- salience=0.7, embedder=EMBEDDER,
- )
- store.remember(
- "user once asked about a neural transformer model",
- session_id="sess-1", actor="CodingAgent",
- salience=0.4, embedder=EMBEDDER,
- )
-
- memories = store.recall(
- "django framework",
- session_id="sess-1",
- embedder=EMBEDDER, top_k=3,
- )
- for m in memories:
- _row(f"score={m.score:.3f} sal={m.memory.salience}", m.memory.content)
-
- # min_salience filters out the third memory
- important = store.recall(
- "framework", session_id="sess-1", min_salience=0.6,
- embedder=EMBEDDER, top_k=5,
- )
- _row("with min_salience=0.6", f"{len(important)} memories")
-
- # -----------------------------------------------------------------
- # 7) Knowledge graph
- # -----------------------------------------------------------------
- _h("7. graph — relations + traversal")
-
- store.add_relation("doc:django", "doc:flask", "compares_with",
- weight=0.8)
- store.add_relation("doc:django", "doc:transformer", "tangentially_about",
- weight=0.2)
- store.add_relation("doc:flask", "doc:transformer", "tangentially_about",
- weight=0.2)
- store.add_relation("doc:transformer", "doc:recipe", "unrelated",
- weight=0.0)
-
- print(" neighbors of doc:django (out):")
- for n in store.neighbors("doc:django", direction="out"):
- _row(f" -> {n.entity.id}", f"{n.relation_kind} (w={n.weight})")
-
- print(" traverse from doc:django, max_hops=2:")
- for t in store.traverse("doc:django", direction="both", max_hops=2):
- _row(f" hop={t.hop}", t.entity.id)
-
- print(" subgraph from {doc:django, doc:flask}, max_hops=1:")
- sg = store.subgraph(["doc:django", "doc:flask"], max_hops=1)
- _row(" entities", [e.id for e in sg.entities])
- _row(" relations", [(r.src, r.dst, r.kind) for r in sg.relations])
-
- # -----------------------------------------------------------------
- # 8) Filter pushdown — restrict to a single entity kind
- # -----------------------------------------------------------------
- _h("8. filter pushdown — kind=document only")
-
- # Without the filter, table rows and memories also rank in.
- # With filters set, the search routes through brute force (HNSW
- # cannot pre-filter on JSON metadata).
- hits = store.search(
- "python web framework", embedder=EMBEDDER,
- filters={"kind": "document"}, top_k=3,
- )
- for h in hits:
- _row(f"{h.entity.id} ({h.entity.kind})",
- h.chunk.content[:60] + "...")
- assert all(h.entity.kind == "document" for h in hits)
-
- # -----------------------------------------------------------------
- # 9) Admin: validate, compact, export
- # -----------------------------------------------------------------
- _h("9. admin")
-
- # First validate may flag a stale HNSW blob: section 2 persisted the
- # index at the original 11 embeddings, but later sections added
- # table rows + memories. Filter-bearing searches in section 8 used
- # brute force, so the persisted blob was never refreshed. validate()
- # catches the divergence; compact() resolves it (next search lazily
- # rebuilds from the live embedding table).
- report = store.validate()
- _row("validate.ok (before compact)", report.ok)
- for i in report.issues:
- print(f" ! {i}")
-
- sizes = store.compact()
- _row("compact size_before", sizes["size_before"])
- _row("compact size_after", sizes["size_after"])
-
- report2 = store.validate()
- _row("validate.ok (after compact)", report2.ok)
- assert report2.ok, report2.issues
-
- export_path = workdir / "backup.json"
- summary = store.export(export_path)
- _row("export bytes", summary["bytes_written"])
- _row("export entities", summary["entities"])
- _row("export embeddings", summary["embeddings"])
-
- # -----------------------------------------------------------------
- # 10) Persistence — close, reopen, search still works
- # -----------------------------------------------------------------
- _h("10. persistence — reopen the file")
-
- store.close()
- store2 = cd_open(db_path)
- try:
- hits = store2.search(
- "python web framework", embedder=EMBEDDER, top_k=2
- )
- _row("post-reopen top-1", hits[0].entity.id)
- # FTS persisted via triggers, recall persisted via memory_ext, etc.
- ms = store2.recall("django", session_id="sess-1",
- embedder=EMBEDDER, top_k=1)
- _row("post-reopen recall", ms[0].memory.content if ms else "(none)")
- finally:
- store2.close()
-
- print("\nDone. Demo data lives at:", db_path)
- return True
-
-
-if __name__ == "__main__":
- ok = run()
- sys.exit(0 if ok else 1)
diff --git a/Agent_module/pyproject.toml b/Agent_module/pyproject.toml
deleted file mode 100644
index 5940d6a5ae..0000000000
--- a/Agent_module/pyproject.toml
+++ /dev/null
@@ -1,40 +0,0 @@
-[build-system]
-requires = ["setuptools>=68"]
-build-backend = "setuptools.build_meta"
-
-[project]
-name = "apache-carbondata-agent"
-version = "0.1.0"
-description = "Agent-native context data format and local store for cross-agent collaboration"
-readme = "README.md"
-requires-python = ">=3.10"
-license = { text = "Apache-2.0" }
-dependencies = [
- "numpy>=1.24",
-]
-
-[project.optional-dependencies]
-hnsw = [
- "hnswlib>=0.8.0",
-]
-test = [
- "pytest>=7.4",
-]
-all = [
- "hnswlib>=0.8.0",
- "pytest>=7.4",
-]
-
-[tool.setuptools]
-packages = [
- "Agent_module",
- "Agent_module.carbon_data",
- "Agent_module.carbon_data.chunkers",
- "Agent_module.carbon_data.embedders",
- "Agent_module.carbon_data.index",
-]
-package-dir = { Agent_module = "." }
-
-[tool.pytest.ini_options]
-testpaths = ["tests"]
-addopts = "-ra"
diff --git a/Agent_module/tests/__init__.py b/Agent_module/tests/__init__.py
deleted file mode 100644
index e69de29bb2..0000000000
diff --git a/Agent_module/tests/carbon_data/__init__.py b/Agent_module/tests/carbon_data/__init__.py
deleted file mode 100644
index e69de29bb2..0000000000
diff --git a/Agent_module/tests/carbon_data/test_m1_schema.py b/Agent_module/tests/carbon_data/test_m1_schema.py
deleted file mode 100644
index edb5d78ef2..0000000000
--- a/Agent_module/tests/carbon_data/test_m1_schema.py
+++ /dev/null
@@ -1,184 +0,0 @@
-"""M1: schema + open/create/close tests."""
-from __future__ import annotations
-
-import sqlite3
-from pathlib import Path
-
-import pytest
-
-from Agent_module.carbon_data import (
- APPLICATION_ID,
- SCHEMA_VERSION,
- CarbonStore,
- NotACarbondataFile,
- UnsupportedSchemaVersion,
- create,
-)
-from Agent_module.carbon_data import open as cd_open
-
-
-EXPECTED_TABLES = {
- "_meta",
- "entity",
- "chunk",
- "embedding",
- "relation",
- "memory_ext",
- "vector_index",
- "chunk_fts",
-}
-
-
-def _tables(path: Path) -> set[str]:
- conn = sqlite3.connect(str(path))
- try:
- rows = conn.execute(
- "SELECT name FROM sqlite_master WHERE type IN ('table','virtual')"
- ).fetchall()
- finally:
- conn.close()
- return {r[0] for r in rows}
-
-
-def test_create_and_reopen_empty(tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- store = create(p)
- assert p.exists()
- stats = store.stats()
- assert stats["entities"] == 0
- assert stats["chunks"] == 0
- assert stats["embeddings"] == 0
- assert stats["relations"] == 0
- assert stats["memories"] == 0
- assert stats["schema_version"] == SCHEMA_VERSION
- assert stats["application_id"] == APPLICATION_ID
- store.close()
-
- reopened = cd_open(p)
- assert reopened.stats()["entities"] == 0
- reopened.close()
-
-
-def test_application_id_and_user_version_on_disk(tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- create(p).close()
- conn = sqlite3.connect(str(p))
- try:
- (app_id,) = conn.execute("PRAGMA application_id").fetchone()
- (ver,) = conn.execute("PRAGMA user_version").fetchone()
- finally:
- conn.close()
- assert app_id == APPLICATION_ID
- assert ver == SCHEMA_VERSION
-
-
-def test_schema_has_all_expected_tables(tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- create(p).close()
- tables = _tables(p)
- missing = EXPECTED_TABLES - tables
- assert not missing, f"missing tables: {missing}"
-
-
-def test_reject_non_carbondata_sqlite(tmp_path: Path) -> None:
- p = tmp_path / "random.db"
- conn = sqlite3.connect(str(p))
- conn.execute("CREATE TABLE foo (x INTEGER)")
- conn.commit()
- conn.close()
- with pytest.raises(NotACarbondataFile):
- cd_open(p)
-
-
-def test_create_refuses_existing_by_default(tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- create(p).close()
- with pytest.raises(FileExistsError):
- create(p)
-
-
-def test_create_exist_ok_reopens(tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- create(p).close()
- store = create(p, exist_ok=True)
- assert store.stats()["schema_version"] == SCHEMA_VERSION
- store.close()
-
-
-def test_open_missing_readonly_fails(tmp_path: Path) -> None:
- p = tmp_path / "nope.carbondata"
- with pytest.raises(FileNotFoundError):
- cd_open(p, mode="r")
-
-
-def test_open_missing_rw_creates(tmp_path: Path) -> None:
- p = tmp_path / "new.carbondata"
- store = cd_open(p)
- assert p.exists()
- assert store.stats()["entities"] == 0
- store.close()
-
-
-def test_context_manager_closes(tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- with create(p) as store:
- assert not store.closed
- assert store.stats()["entities"] == 0
- assert store.closed
-
-
-def test_stats_after_close_errors(tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- store = create(p)
- store.close()
- with pytest.raises(RuntimeError):
- store.stats()
-
-
-def test_future_schema_version_rejected(tmp_path: Path) -> None:
- """Simulate a file written by a newer library version."""
- p = tmp_path / "future.carbondata"
- create(p).close()
- # Bump user_version past SCHEMA_VERSION
- conn = sqlite3.connect(str(p))
- conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION + 1}")
- conn.close()
- with pytest.raises(UnsupportedSchemaVersion):
- cd_open(p)
-
-
-def test_readonly_open_rejects_writes(tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- create(p).close()
- ro = cd_open(p, mode="r")
- assert ro.mode == "r"
- # Attempting a write through the underlying connection must fail.
- with pytest.raises(sqlite3.OperationalError):
- ro._require_open().execute(
- "INSERT INTO _meta (key, value) VALUES ('x', 'y')"
- )
- ro.close()
-
-
-def test_dataclasses_constructable() -> None:
- """Sanity: data model dataclasses can be instantiated without numpy."""
- from Agent_module.carbon_data.models import (
- Chunk,
- Entity,
- MemoryItem,
- Relation,
- )
-
- e = Entity(id="e1", kind="document", content="hi")
- assert e.namespace == "default"
- assert e.version == 1
-
- c = Chunk(id="c1", entity_id="e1", seq=0, content="hello")
- assert c.kind == "text"
-
- r = Relation(src="e1", dst="e2", kind="references")
- assert r.weight == 1.0
-
- m = MemoryItem(id="m1", session_id="s1", actor="A", salience=0.9)
- assert m.kind == "memory"
- assert m.session_id == "s1"
diff --git a/Agent_module/tests/carbon_data/test_m2_crud.py b/Agent_module/tests/carbon_data/test_m2_crud.py
deleted file mode 100644
index 82d0dd6473..0000000000
--- a/Agent_module/tests/carbon_data/test_m2_crud.py
+++ /dev/null
@@ -1,376 +0,0 @@
-"""M2: Entity/Chunk CRUD + structured query + transactions."""
-from __future__ import annotations
-
-from pathlib import Path
-
-import pytest
-
-from Agent_module.carbon_data import CarbonStore, Chunk, Entity, create
-from Agent_module.carbon_data import open as cd_open
-
-
-# ---------------------------------------------------------------------------
-# fixtures
-# ---------------------------------------------------------------------------
-
-@pytest.fixture
-def store(tmp_path: Path):
- p = tmp_path / "kb.carbondata"
- s = create(p)
- yield s
- if not s.closed:
- s.close()
-
-
-# ---------------------------------------------------------------------------
-# Entity CRUD
-# ---------------------------------------------------------------------------
-
-class TestEntityCrud:
- def test_put_then_get_kwargs(self, store: CarbonStore) -> None:
- eid = store.put_entity(
- id="doc1", kind="document", content="hello",
- metadata={"author": "X", "tags": ["a", "b"]},
- )
- assert eid == "doc1"
- got = store.get_entity("doc1")
- assert got is not None
- assert got.id == "doc1"
- assert got.kind == "document"
- assert got.namespace == "default"
- assert got.content == "hello"
- assert got.metadata == {"author": "X", "tags": ["a", "b"]}
- assert got.version == 1
-
- def test_put_with_entity_object(self, store: CarbonStore) -> None:
- e = Entity(id="doc2", kind="document", namespace="ns1", content="x")
- store.put_entity(e)
- got = store.get_entity("doc2")
- assert got is not None
- assert got.namespace == "ns1"
-
- def test_put_requires_id_and_kind(self, store: CarbonStore) -> None:
- with pytest.raises(TypeError):
- store.put_entity()
- with pytest.raises(TypeError):
- store.put_entity(id="x")
- with pytest.raises(TypeError):
- store.put_entity(kind="doc")
-
- def test_put_mixing_entity_and_kwargs_rejected(self, store: CarbonStore) -> None:
- e = Entity(id="x", kind="doc")
- with pytest.raises(TypeError):
- store.put_entity(e, content="hi")
-
- def test_get_missing_returns_none(self, store: CarbonStore) -> None:
- assert store.get_entity("nope") is None
-
- def test_upsert_bumps_version_preserves_created_at(
- self, store: CarbonStore
- ) -> None:
- store.put_entity(id="x", kind="doc", content="v1")
- first = store.get_entity("x")
- assert first is not None
- store.put_entity(id="x", kind="doc", content="v2", metadata={"y": 1})
- second = store.get_entity("x")
- assert second is not None
- assert second.content == "v2"
- assert second.metadata == {"y": 1}
- assert second.version == first.version + 1
- assert second.created_at == first.created_at
- assert second.updated_at >= first.updated_at
-
- def test_delete_entity_returns_true(self, store: CarbonStore) -> None:
- store.put_entity(id="x", kind="doc")
- assert store.delete_entity("x") is True
- assert store.get_entity("x") is None
-
- def test_delete_missing_returns_false(self, store: CarbonStore) -> None:
- assert store.delete_entity("ghost") is False
-
- def test_delete_cascades_to_chunks(self, store: CarbonStore) -> None:
- store.put_entity(id="x", kind="doc")
- store.add_chunks("x", ["a", "b", "c"])
- assert len(store.list_chunks("x")) == 3
- store.delete_entity("x")
- assert store.list_chunks("x") == []
-
- def test_metadata_roundtrip_nested(self, store: CarbonStore) -> None:
- meta = {
- "nested": {"a": [1, 2, {"b": True}]},
- "unicode": "你好",
- "none": None,
- }
- store.put_entity(id="x", kind="doc", metadata=meta)
- got = store.get_entity("x")
- assert got is not None
- assert got.metadata == meta
-
-
-# ---------------------------------------------------------------------------
-# Chunk CRUD
-# ---------------------------------------------------------------------------
-
-class TestChunkCrud:
- def test_add_chunks_from_strings(self, store: CarbonStore) -> None:
- store.put_entity(id="x", kind="doc")
- ids = store.add_chunks("x", ["first", "second", "third"])
- assert len(ids) == 3
- chunks = store.list_chunks("x")
- assert [c.content for c in chunks] == ["first", "second", "third"]
- assert [c.seq for c in chunks] == [0, 1, 2]
-
- def test_add_chunks_mixed_types(self, store: CarbonStore) -> None:
- store.put_entity(id="x", kind="doc")
- ids = store.add_chunks("x", [
- "plain string",
- Chunk(id="c2", entity_id="x", seq=0, content="from dataclass", kind="code"),
- {"content": "from dict", "kind": "row", "metadata": {"col": "name"}},
- ])
- assert len(ids) == 3
- chunks = store.list_chunks("x")
- assert chunks[0].kind == "text"
- assert chunks[1].id == "c2"
- assert chunks[1].kind == "code"
- assert chunks[2].kind == "row"
- assert chunks[2].metadata == {"col": "name"}
-
- def test_add_chunks_seq_continues(self, store: CarbonStore) -> None:
- store.put_entity(id="x", kind="doc")
- store.add_chunks("x", ["a", "b"])
- store.add_chunks("x", ["c", "d"])
- chunks = store.list_chunks("x")
- assert [c.seq for c in chunks] == [0, 1, 2, 3]
- assert [c.content for c in chunks] == ["a", "b", "c", "d"]
-
- def test_add_chunks_base_seq_override(self, store: CarbonStore) -> None:
- store.put_entity(id="x", kind="doc")
- store.add_chunks("x", ["a"])
- store.add_chunks("x", ["b", "c"], base_seq=100)
- chunks = store.list_chunks("x")
- assert [c.seq for c in chunks] == [0, 100, 101]
-
- def test_add_chunks_entity_id_overridden(self, store: CarbonStore) -> None:
- store.put_entity(id="x", kind="doc")
- # Chunk constructed with a different entity_id gets corrected.
- store.add_chunks("x", [
- Chunk(id="c1", entity_id="WRONG", seq=0, content="hi"),
- ])
- assert store.get_chunk("c1").entity_id == "x" # type: ignore[union-attr]
-
- def test_add_chunks_missing_entity_raises(self, store: CarbonStore) -> None:
- with pytest.raises(KeyError):
- store.add_chunks("nope", ["hi"])
-
- def test_add_chunks_bad_type_raises(self, store: CarbonStore) -> None:
- store.put_entity(id="x", kind="doc")
- with pytest.raises(TypeError):
- store.add_chunks("x", [123]) # type: ignore[list-item]
-
- def test_get_chunk_missing(self, store: CarbonStore) -> None:
- assert store.get_chunk("ghost") is None
-
- def test_delete_chunks(self, store: CarbonStore) -> None:
- store.put_entity(id="x", kind="doc")
- store.add_chunks("x", ["a", "b"])
- assert store.delete_chunks("x") == 2
- assert store.list_chunks("x") == []
-
-
-# ---------------------------------------------------------------------------
-# Structured query
-# ---------------------------------------------------------------------------
-
-class TestQueryEntities:
- @pytest.fixture
- def populated(self, store: CarbonStore) -> CarbonStore:
- store.put_entity(
- id="d1", kind="document",
- metadata={"author": "Alice", "year": 2020},
- )
- store.put_entity(
- id="d2", kind="document",
- metadata={"author": "Bob", "year": 2023},
- )
- store.put_entity(
- id="c1", kind="code",
- metadata={"author": "Alice", "lang": "python"},
- )
- store.put_entity(
- id="n1", kind="note", namespace="archive",
- metadata={"author": "Charlie"},
- )
- return store
-
- def test_filter_by_kind(self, populated: CarbonStore) -> None:
- results = populated.query_entities(kind="document")
- assert {e.id for e in results} == {"d1", "d2"}
-
- def test_filter_by_namespace(self, populated: CarbonStore) -> None:
- results = populated.query_entities(namespace="archive")
- assert {e.id for e in results} == {"n1"}
-
- def test_filter_by_metadata_equality(self, populated: CarbonStore) -> None:
- results = populated.query_entities(where={"metadata.author": "Alice"})
- assert {e.id for e in results} == {"d1", "c1"}
-
- def test_filter_by_operator_tuple(self, populated: CarbonStore) -> None:
- results = populated.query_entities(
- where={"metadata.year": (">=", 2023)},
- )
- assert {e.id for e in results} == {"d2"}
-
- def test_filter_by_in(self, populated: CarbonStore) -> None:
- results = populated.query_entities(
- where={"kind": ("IN", ["document", "code"])},
- )
- assert {e.id for e in results} == {"d1", "d2", "c1"}
-
- def test_combined_filters(self, populated: CarbonStore) -> None:
- results = populated.query_entities(
- kind="document",
- where={"metadata.author": "Alice"},
- )
- assert {e.id for e in results} == {"d1"}
-
- def test_limit_and_offset(self, populated: CarbonStore) -> None:
- all_ids = [e.id for e in populated.query_entities()]
- assert len(all_ids) == 4
- page1 = populated.query_entities(limit=2, offset=0)
- page2 = populated.query_entities(limit=2, offset=2)
- combined = [e.id for e in page1] + [e.id for e in page2]
- assert combined == all_ids
-
- def test_order_by_custom(self, populated: CarbonStore) -> None:
- results = populated.query_entities(order_by="id ASC")
- assert [e.id for e in results] == ["c1", "d1", "d2", "n1"]
-
- def test_reject_unknown_column_in_where(self, populated: CarbonStore) -> None:
- with pytest.raises(ValueError, match="unknown entity column"):
- populated.query_entities(where={"evil; DROP TABLE entity": "x"})
-
- def test_reject_unknown_order_by_column(self, populated: CarbonStore) -> None:
- with pytest.raises(ValueError, match="unknown order_by column"):
- populated.query_entities(order_by="password ASC")
-
- def test_reject_bad_order_by_direction(self, populated: CarbonStore) -> None:
- with pytest.raises(ValueError, match="invalid direction"):
- populated.query_entities(order_by="id SIDEWAYS")
-
- def test_reject_bad_operator(self, populated: CarbonStore) -> None:
- with pytest.raises(ValueError, match="unsupported operator"):
- populated.query_entities(where={"kind": ("DROP", "document")})
-
- def test_in_requires_list(self, populated: CarbonStore) -> None:
- with pytest.raises(ValueError, match="IN operator"):
- populated.query_entities(where={"kind": ("IN", "document")})
-
-
-# ---------------------------------------------------------------------------
-# Transactions
-# ---------------------------------------------------------------------------
-
-class TestTransactions:
- def test_commit_on_success(self, store: CarbonStore) -> None:
- with store.transaction():
- store.put_entity(id="a", kind="doc")
- store.put_entity(id="b", kind="doc")
- assert store.get_entity("a") is not None
- assert store.get_entity("b") is not None
-
- def test_rollback_on_exception(self, store: CarbonStore) -> None:
- with pytest.raises(RuntimeError):
- with store.transaction():
- store.put_entity(id="a", kind="doc")
- store.put_entity(id="b", kind="doc")
- raise RuntimeError("boom")
- assert store.get_entity("a") is None
- assert store.get_entity("b") is None
-
- def test_nested_transactions_commit_once(self, store: CarbonStore) -> None:
- with store.transaction():
- store.put_entity(id="a", kind="doc")
- with store.transaction():
- store.put_entity(id="b", kind="doc")
- store.put_entity(id="c", kind="doc")
- assert store.get_entity("a") is not None
- assert store.get_entity("b") is not None
- assert store.get_entity("c") is not None
-
- def test_nested_inner_exception_rolls_back_all(
- self, store: CarbonStore
- ) -> None:
- with pytest.raises(RuntimeError):
- with store.transaction():
- store.put_entity(id="a", kind="doc")
- with store.transaction():
- store.put_entity(id="b", kind="doc")
- raise RuntimeError("boom")
- assert store.get_entity("a") is None
- assert store.get_entity("b") is None
-
- def test_caught_inner_exception_rolls_back_to_savepoint(
- self, store: CarbonStore
- ) -> None:
- with store.transaction():
- store.put_entity(id="a", kind="doc")
- try:
- with store.transaction():
- store.put_entity(id="b", kind="doc")
- raise RuntimeError("inner")
- except RuntimeError:
- pass
- store.put_entity(id="c", kind="doc")
-
- assert store.get_entity("a") is not None
- assert store.get_entity("b") is None
- assert store.get_entity("c") is not None
-
- def test_chunks_added_in_transaction_all_or_nothing(
- self, store: CarbonStore
- ) -> None:
- store.put_entity(id="x", kind="doc")
- with pytest.raises(RuntimeError):
- with store.transaction():
- store.add_chunks("x", ["one"])
- store.add_chunks("x", ["two"])
- raise RuntimeError("boom")
- assert store.list_chunks("x") == []
-
-
-# ---------------------------------------------------------------------------
-# Read-only mode
-# ---------------------------------------------------------------------------
-
-class TestReadOnly:
- def test_readonly_rejects_put(self, tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- create(p).close()
- ro = cd_open(p, mode="r")
- try:
- with pytest.raises(PermissionError):
- ro.put_entity(id="x", kind="doc")
- finally:
- ro.close()
-
- def test_readonly_rejects_delete(self, tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- s = create(p)
- s.put_entity(id="x", kind="doc")
- s.close()
- ro = cd_open(p, mode="r")
- try:
- with pytest.raises(PermissionError):
- ro.delete_entity("x")
- finally:
- ro.close()
-
- def test_readonly_get_and_query_work(self, tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- with create(p) as s:
- s.put_entity(id="x", kind="doc", content="hi")
- with cd_open(p, mode="r") as ro:
- got = ro.get_entity("x")
- assert got is not None
- assert got.content == "hi"
- assert len(ro.query_entities(kind="doc")) == 1
diff --git a/Agent_module/tests/carbon_data/test_m3_search.py b/Agent_module/tests/carbon_data/test_m3_search.py
deleted file mode 100644
index 40a1fa7816..0000000000
--- a/Agent_module/tests/carbon_data/test_m3_search.py
+++ /dev/null
@@ -1,375 +0,0 @@
-"""M3: embeddings + semantic search."""
-from __future__ import annotations
-
-from pathlib import Path
-
-import numpy as np
-import pytest
-
-from Agent_module.carbon_data import (
- CarbonStore,
- LambdaEmbedder,
- NullEmbedder,
- SearchHit,
- create,
-)
-from Agent_module.carbon_data.index import (
- bytes_to_vector,
- search_brute,
- vector_to_bytes,
-)
-
-
-# ---------------------------------------------------------------------------
-# fixtures
-# ---------------------------------------------------------------------------
-
-@pytest.fixture
-def store(tmp_path: Path):
- p = tmp_path / "kb.carbondata"
- s = create(p)
- yield s
- if not s.closed:
- s.close()
-
-
-def _fixed_embedder(mapping: dict[str, list[float]], *, model="test", dim=3):
- """Embedder that returns pre-defined vectors for each text, else zeros."""
- arr_cache = {
- k: np.asarray(v, dtype=np.float32).reshape(-1) for k, v in mapping.items()
- }
- for k, v in arr_cache.items():
- if v.shape[0] != dim:
- raise ValueError(f"mapping[{k!r}] has dim {v.shape[0]} != {dim}")
-
- def _fn(texts):
- out = np.zeros((len(texts), dim), dtype=np.float32)
- for i, t in enumerate(texts):
- if t in arr_cache:
- out[i] = arr_cache[t]
- return out
-
- return LambdaEmbedder(_fn, model=model, dim=dim)
-
-
-# ---------------------------------------------------------------------------
-# pure vector utils
-# ---------------------------------------------------------------------------
-
-class TestVectorPacking:
- def test_roundtrip(self) -> None:
- v = np.array([1.0, 2.0, 3.0, 4.5], dtype=np.float32)
- blob = vector_to_bytes(v)
- back = bytes_to_vector(blob, 4)
- assert np.array_equal(back, v)
-
- def test_rejects_2d(self) -> None:
- with pytest.raises(ValueError):
- vector_to_bytes(np.zeros((2, 3), dtype=np.float32))
-
- def test_dim_mismatch_raises(self) -> None:
- blob = vector_to_bytes(np.array([1, 2, 3], dtype=np.float32))
- with pytest.raises(ValueError):
- bytes_to_vector(blob, 4)
-
-
-class TestSearchBrute:
- def test_cosine_ranking(self) -> None:
- q = np.array([1.0, 0.0, 0.0], dtype=np.float32)
- vectors = np.array(
- [[1.0, 0.0, 0.0], [0.9, 0.1, 0.0], [0.0, 1.0, 0.0]],
- dtype=np.float32,
- )
- ids = ["match", "close", "far"]
- hits = search_brute(q, vectors, ids, top_k=3)
- assert [h[0] for h in hits] == ["match", "close", "far"]
- assert hits[0][1] > hits[1][1] > hits[2][1]
-
- def test_top_k_smaller_than_corpus(self) -> None:
- q = np.array([1.0, 0.0], dtype=np.float32)
- vectors = np.stack(
- [np.array([i, 0.0], dtype=np.float32) for i in range(1, 6)]
- )
- ids = [str(i) for i in range(1, 6)]
- hits = search_brute(q, vectors, ids, top_k=2)
- assert len(hits) == 2
- assert hits[0][0] in ids # all cosines are 1.0 — order just stable top-k
-
- def test_empty_corpus(self) -> None:
- assert (
- search_brute(
- np.array([1.0], dtype=np.float32),
- np.zeros((0, 1), dtype=np.float32),
- [],
- top_k=5,
- )
- == []
- )
-
- def test_dim_mismatch(self) -> None:
- q = np.array([1.0, 2.0, 3.0], dtype=np.float32)
- vectors = np.array([[1.0, 2.0]], dtype=np.float32)
- with pytest.raises(ValueError):
- search_brute(q, vectors, ["a"], top_k=1)
-
-
-# ---------------------------------------------------------------------------
-# Embedders
-# ---------------------------------------------------------------------------
-
-class TestEmbedders:
- def test_null_embedder_shape(self) -> None:
- emb = NullEmbedder(dim=16)
- out = emb.encode(["a", "b", "c"])
- assert out.shape == (3, 16)
- assert out.dtype == np.float32
- assert np.all(out == 0.0)
-
- def test_lambda_validates_shape(self) -> None:
- bad = LambdaEmbedder(
- lambda texts: np.zeros((len(texts), 5), dtype=np.float32),
- model="bad", dim=3,
- )
- with pytest.raises(ValueError, match="expected"):
- bad.encode(["x"])
-
- def test_lambda_casts_to_float32(self) -> None:
- emb = LambdaEmbedder(
- lambda texts: [[1, 2, 3] for _ in texts],
- model="ints", dim=3,
- )
- out = emb.encode(["x", "y"])
- assert out.dtype == np.float32
- assert out.shape == (2, 3)
-
-
-# ---------------------------------------------------------------------------
-# set_embedding / get_embedding
-# ---------------------------------------------------------------------------
-
-class TestEmbeddingCrud:
- def test_set_then_get(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="doc")
- store.add_chunks("e1", ["hello"])
- c = store.list_chunks("e1")[0]
- v = np.array([0.1, 0.2, 0.3], dtype=np.float32)
- store.set_embedding(c.id, v, model="m1")
- got = store.get_embedding(c.id, model="m1")
- assert got is not None
- np.testing.assert_array_equal(got, v)
-
- def test_set_replaces_existing(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="doc")
- store.add_chunks("e1", ["hello"])
- c = store.list_chunks("e1")[0]
- store.set_embedding(c.id, [0.1, 0.2], model="m1")
- store.set_embedding(c.id, [0.9, 0.8], model="m1")
- got = store.get_embedding(c.id, model="m1")
- np.testing.assert_array_almost_equal(got, [0.9, 0.8])
-
- def test_multiple_models_per_chunk(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="doc")
- store.add_chunks("e1", ["hello"])
- c = store.list_chunks("e1")[0]
- store.set_embedding(c.id, [1.0, 0.0], model="small")
- store.set_embedding(c.id, [0.0, 0.0, 1.0], model="big")
- np.testing.assert_array_equal(
- store.get_embedding(c.id, model="small"), [1.0, 0.0]
- )
- np.testing.assert_array_equal(
- store.get_embedding(c.id, model="big"), [0.0, 0.0, 1.0]
- )
-
- def test_set_on_missing_chunk(self, store: CarbonStore) -> None:
- with pytest.raises(KeyError):
- store.set_embedding("ghost", [1.0, 2.0], model="m1")
-
- def test_get_missing_returns_none(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="doc")
- store.add_chunks("e1", ["hello"])
- c = store.list_chunks("e1")[0]
- assert store.get_embedding(c.id, model="nothere") is None
-
-
-# ---------------------------------------------------------------------------
-# embed_chunks
-# ---------------------------------------------------------------------------
-
-class TestEmbedChunks:
- def test_embeds_all_pending(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="doc")
- store.add_chunks("e1", ["a", "b", "c"])
- emb = NullEmbedder(dim=4, model="m1")
- count = store.embed_chunks(emb)
- assert count == 3
- for c in store.list_chunks("e1"):
- v = store.get_embedding(c.id, model="m1")
- assert v is not None and v.shape == (4,)
-
- def test_missing_only_skips_existing(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="doc")
- store.add_chunks("e1", ["a", "b"])
- emb = NullEmbedder(dim=2, model="m1")
- assert store.embed_chunks(emb) == 2
- # second call should skip all
- assert store.embed_chunks(emb) == 0
-
- def test_missing_only_false_reencodes(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="doc")
- store.add_chunks("e1", ["a"])
- emb = NullEmbedder(dim=2, model="m1")
- assert store.embed_chunks(emb) == 1
- assert store.embed_chunks(emb, missing_only=False) == 1
-
- def test_scoped_to_entity(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="doc")
- store.put_entity(id="e2", kind="doc")
- store.add_chunks("e1", ["a", "b"])
- store.add_chunks("e2", ["c", "d"])
- emb = NullEmbedder(dim=2, model="m1")
- assert store.embed_chunks(emb, entity_id="e1") == 2
- # e2 chunks still pending
- assert store.embed_chunks(emb) == 2
-
-
-# ---------------------------------------------------------------------------
-# search
-# ---------------------------------------------------------------------------
-
-class TestSearch:
- def _seed_docs(self, store: CarbonStore):
- store.put_entity(id="d1", kind="doc", metadata={"topic": "cats"})
- store.put_entity(id="d2", kind="doc", metadata={"topic": "dogs"})
- store.add_chunks("d1", ["a cat sat"])
- store.add_chunks("d2", ["a dog ran"])
- emb = _fixed_embedder({
- "a cat sat": [1.0, 0.0, 0.0],
- "a dog ran": [0.0, 1.0, 0.0],
- })
- store.embed_chunks(emb)
- return emb
-
- def test_vector_query(self, store: CarbonStore) -> None:
- emb = self._seed_docs(store)
- hits = store.search(
- np.array([1.0, 0.0, 0.0], dtype=np.float32),
- model=emb.model,
- top_k=2,
- )
- assert len(hits) == 2
- assert isinstance(hits[0], SearchHit)
- assert hits[0].entity.id == "d1"
- assert hits[0].score > hits[1].score
-
- def test_string_query_uses_embedder(self, store: CarbonStore) -> None:
- emb = self._seed_docs(store)
- # embedder also maps the query string:
- emb_with_query = _fixed_embedder({
- "a cat sat": [1.0, 0.0, 0.0],
- "a dog ran": [0.0, 1.0, 0.0],
- "kitten": [1.0, 0.0, 0.0],
- })
- store.embed_chunks(emb_with_query, missing_only=False)
- hits = store.search("kitten", embedder=emb_with_query, top_k=1)
- assert len(hits) == 1
- assert hits[0].entity.id == "d1"
-
- def test_string_query_without_embedder_errors(
- self, store: CarbonStore
- ) -> None:
- self._seed_docs(store)
- with pytest.raises(TypeError, match="embedder"):
- store.search("hi", model="test", top_k=1)
-
- def test_search_requires_model_or_embedder(
- self, store: CarbonStore
- ) -> None:
- self._seed_docs(store)
- with pytest.raises(TypeError, match="model"):
- store.search(
- np.array([1.0, 0.0, 0.0], dtype=np.float32), top_k=1
- )
-
- def test_filters_restrict_results(self, store: CarbonStore) -> None:
- emb = self._seed_docs(store)
- hits = store.search(
- np.array([1.0, 0.0, 0.0], dtype=np.float32),
- model=emb.model,
- filters={"metadata.topic": "dogs"},
- top_k=5,
- )
- # Only d2 survives filter
- assert {h.entity.id for h in hits} == {"d2"}
-
- def test_namespace_isolation(self, store: CarbonStore) -> None:
- store.put_entity(id="a", kind="doc", namespace="ns_a")
- store.put_entity(id="b", kind="doc", namespace="ns_b")
- store.add_chunks("a", ["same text"])
- store.add_chunks("b", ["same text"])
- emb = _fixed_embedder({"same text": [1.0, 0.0, 0.0]})
- store.embed_chunks(emb)
-
- hits_a = store.search(
- np.array([1.0, 0.0, 0.0], dtype=np.float32),
- model=emb.model,
- namespace="ns_a",
- top_k=10,
- )
- assert {h.entity.id for h in hits_a} == {"a"}
-
- hits_b = store.search(
- np.array([1.0, 0.0, 0.0], dtype=np.float32),
- model=emb.model,
- namespace="ns_b",
- top_k=10,
- )
- assert {h.entity.id for h in hits_b} == {"b"}
-
- def test_model_isolation(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="doc")
- store.add_chunks("e1", ["x"])
- c = store.list_chunks("e1")[0]
- store.set_embedding(c.id, [1.0, 0.0], model="m1")
- store.set_embedding(c.id, [0.0, 1.0], model="m2")
- # Query under m1: found
- hits_m1 = store.search(
- np.array([1.0, 0.0], dtype=np.float32), model="m1", top_k=1
- )
- assert len(hits_m1) == 1
- # Query under m2: different vector, still finds the chunk (only 1)
- hits_m2 = store.search(
- np.array([0.0, 1.0], dtype=np.float32), model="m2", top_k=1
- )
- assert len(hits_m2) == 1
-
- def test_search_empty_store(self, store: CarbonStore) -> None:
- assert (
- store.search(
- np.array([1.0, 0.0, 0.0], dtype=np.float32),
- model="nothing",
- top_k=5,
- )
- == []
- )
-
- def test_search_query_dim_mismatch(self, store: CarbonStore) -> None:
- emb = self._seed_docs(store)
- with pytest.raises(ValueError, match="query dim"):
- store.search(
- np.array([1.0, 0.0], dtype=np.float32), # 2-D query vs 3-D corpus
- model=emb.model,
- top_k=1,
- )
-
- def test_top_k_caps_results(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="doc")
- store.add_chunks("e1", [f"chunk {i}" for i in range(5)])
- emb = NullEmbedder(dim=3, model="null")
- # With zero vectors the ranks will all tie at 0 cosine after eps —
- # still need to verify top_k truncates.
- store.embed_chunks(emb)
- hits = store.search(
- np.array([1.0, 0.0, 0.0], dtype=np.float32),
- model="null", top_k=3,
- )
- assert len(hits) == 3
diff --git a/Agent_module/tests/carbon_data/test_m4_fts.py b/Agent_module/tests/carbon_data/test_m4_fts.py
deleted file mode 100644
index 2baa6bf41f..0000000000
--- a/Agent_module/tests/carbon_data/test_m4_fts.py
+++ /dev/null
@@ -1,363 +0,0 @@
-"""M4: FTS5 keyword search + RRF hybrid fusion."""
-from __future__ import annotations
-
-from pathlib import Path
-
-import numpy as np
-import pytest
-
-from Agent_module.carbon_data import (
- CarbonStore,
- LambdaEmbedder,
- SearchHit,
- create,
- open as cd_open,
-)
-from Agent_module.carbon_data.store import _fts_escape_query
-
-
-# ---------------------------------------------------------------------------
-# fixtures
-# ---------------------------------------------------------------------------
-
-@pytest.fixture
-def store(tmp_path: Path):
- p = tmp_path / "kb.carbondata"
- s = create(p)
- yield s
- if not s.closed:
- s.close()
-
-
-def _fixed_embedder(mapping: dict[str, list[float]], *, model="test", dim=3):
- """Returns pre-defined vectors for known texts, zero vector otherwise."""
- arr = {k: np.asarray(v, dtype=np.float32).reshape(-1) for k, v in mapping.items()}
-
- def _fn(texts):
- out = np.zeros((len(texts), dim), dtype=np.float32)
- for i, t in enumerate(texts):
- if t in arr:
- out[i] = arr[t]
- return out
-
- return LambdaEmbedder(_fn, model=model, dim=dim)
-
-
-def _fts_count(store: CarbonStore) -> int:
- conn = store._require_open()
- return conn.execute("SELECT COUNT(*) FROM chunk_fts").fetchone()[0]
-
-
-# ---------------------------------------------------------------------------
-# FTS escape helper
-# ---------------------------------------------------------------------------
-
-class TestFtsEscape:
- def test_simple_tokens(self) -> None:
- assert _fts_escape_query("hello world") == '"hello" "world"'
-
- def test_empty(self) -> None:
- assert _fts_escape_query("") == ""
- assert _fts_escape_query(" ") == ""
-
- def test_strips_quotes(self) -> None:
- # Internal double-quotes get stripped (replaced with whitespace),
- # so "exact phrase" never produces an unbalanced FTS5 string.
- assert _fts_escape_query('say "hi"') == '"say" "hi"'
-
- def test_special_chars_become_phrases(self) -> None:
- # FTS5 reserves AND/OR/NOT/NEAR/(/)/* etc — quoting neutralizes them.
- out = _fts_escape_query("foo AND bar* (baz)")
- assert out == '"foo" "AND" "bar*" "(baz)"'
-
-
-# ---------------------------------------------------------------------------
-# Trigger sync: chunk → chunk_fts
-# ---------------------------------------------------------------------------
-
-class TestFtsTriggers:
- def test_insert_propagates(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["the quick brown fox", "lazy dog"])
- assert _fts_count(store) == 2
-
- def test_delete_chunks_propagates(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["alpha", "beta"])
- assert _fts_count(store) == 2
- store.delete_chunks("e1")
- assert _fts_count(store) == 0
-
- def test_delete_entity_cascades(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["alpha", "beta", "gamma"])
- assert _fts_count(store) == 3
- store.delete_entity("e1")
- assert _fts_count(store) == 0
-
- def test_update_propagates(self, store: CarbonStore) -> None:
- # Only chunk.content updates need to flow into FTS. We don't have
- # a public update_chunk yet; exercise the trigger via direct UPDATE
- # so the sync invariant is locked in for future writers.
- store.put_entity(id="e1", kind="document")
- ids = store.add_chunks("e1", ["original text"])
- with store._write_ctx() as conn:
- conn.execute("UPDATE chunk SET content=? WHERE id=?", ("rewritten", ids[0]))
- # Old token gone, new token findable.
- hits_old = store.search("original", mode="keyword")
- hits_new = store.search("rewritten", mode="keyword")
- assert hits_old == []
- assert len(hits_new) == 1
-
-
-# ---------------------------------------------------------------------------
-# Keyword (FTS5 + BM25) search
-# ---------------------------------------------------------------------------
-
-class TestKeywordSearch:
- def test_basic_match(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", [
- "the quick brown fox jumps",
- "totally unrelated content",
- "fox hunting in the woods",
- ])
- hits = store.search("fox", mode="keyword", top_k=5)
- contents = {h.chunk.content for h in hits}
- assert any("quick brown fox" in c for c in contents)
- assert any("fox hunting" in c for c in contents)
- assert not any("unrelated" in c for c in contents)
-
- def test_bm25_ranking_prefers_rarer_terms(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- # "common" appears everywhere; "unique" appears in only one chunk.
- store.add_chunks("e1", [
- "common common common common",
- "common rare and unique words",
- "common words again here",
- ])
- hits = store.search("unique", mode="keyword", top_k=3)
- assert len(hits) == 1
- assert "unique" in hits[0].chunk.content
-
- def test_score_higher_is_better(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", [
- "fox cat dog bird fish", # 1 hit on "fox"
- "fox fox fox terrier fox", # repeated → higher BM25
- ])
- hits = store.search("fox", mode="keyword", top_k=2)
- assert len(hits) == 2
- # Our score = -bm25(): higher means better match.
- assert hits[0].score >= hits[1].score
-
- def test_namespace_isolation(self, store: CarbonStore) -> None:
- store.put_entity(id="a", kind="document", namespace="ns1")
- store.put_entity(id="b", kind="document", namespace="ns2")
- store.add_chunks("a", ["unique_term in namespace one"])
- store.add_chunks("b", ["unique_term in namespace two"])
- hits = store.search("unique_term", mode="keyword", namespace="ns1")
- assert len(hits) == 1
- assert hits[0].entity.id == "a"
-
- def test_filters_pushdown(self, store: CarbonStore) -> None:
- store.put_entity(id="a", kind="document", metadata={"lang": "en"})
- store.put_entity(id="b", kind="document", metadata={"lang": "fr"})
- store.add_chunks("a", ["the keyword appears here"])
- store.add_chunks("b", ["the keyword appears here too"])
- hits = store.search(
- "keyword", mode="keyword", filters={"metadata.lang": "fr"}
- )
- assert len(hits) == 1
- assert hits[0].entity.id == "b"
-
- def test_top_k_caps(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", [f"token chunk-{i}" for i in range(10)])
- hits = store.search("token", mode="keyword", top_k=3)
- assert len(hits) == 3
-
- def test_empty_query_short_circuits(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["something"])
- assert store.search("", mode="keyword") == []
- assert store.search(" ", mode="keyword") == []
-
- def test_special_chars_dont_crash(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["match this content"])
- # FTS5 reserved words/syntax must not raise via the escape path.
- for q in ["AND", "OR (", "match*", '"', "(foo) AND NOT bar"]:
- store.search(q, mode="keyword") # any of these may legitimately
- # return zero hits — we only assert no exception.
-
- def test_raw_fts_passes_operators(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["fox jumps over the lazy dog"])
- # Without raw_fts, "fox AND jumps" would be quoted into nonsense.
- # With raw_fts, FTS5 sees the AND operator.
- hits = store.search("fox AND jumps", mode="keyword", raw_fts=True)
- assert len(hits) == 1
-
- def test_raw_fts_invalid_raises(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["whatever"])
- with pytest.raises(ValueError, match="FTS5 query failed"):
- # Unbalanced quote — only reachable via raw_fts=True.
- store.search('"unclosed', mode="keyword", raw_fts=True)
-
- def test_keyword_returns_search_hit(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["needle in haystack"])
- hits = store.search("needle", mode="keyword")
- assert len(hits) == 1
- assert isinstance(hits[0], SearchHit)
- assert hits[0].entity.id == "e1"
- assert hits[0].chunk.content == "needle in haystack"
-
- def test_keyword_rejects_vector_query(self, store: CarbonStore) -> None:
- with pytest.raises(TypeError, match="keyword mode"):
- store.search(np.array([1.0, 0.0, 0.0]), mode="keyword")
-
-
-# ---------------------------------------------------------------------------
-# Hybrid (RRF)
-# ---------------------------------------------------------------------------
-
-class TestHybridSearch:
- def test_combines_both_modalities(self, store: CarbonStore) -> None:
- # Build a corpus where vector and keyword disagree, and check
- # that hybrid surfaces both winners.
- emb = _fixed_embedder(
- {
- "vec_winner": [1.0, 0.0, 0.0], # vec match for q
- "shared": [0.5, 0.5, 0.0], # weak vec, no kw
- "kw_winner": [0.0, 0.0, 1.0], # weak vec, strong kw
- "irrelevant": [0.0, 1.0, 0.0],
- "needle": [0.9, 0.1, 0.0], # query vector (~vec_winner)
- },
- dim=3,
- )
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["vec_winner", "shared", "kw_winner", "irrelevant"])
- store.embed_chunks(embedder=emb)
-
- hits = store.search(
- "kw_winner", mode="hybrid", top_k=2, embedder=emb,
- )
- ids = [h.chunk.content for h in hits]
- # kw_winner: in keyword top, also in vector tail → strong RRF
- # vec_winner: not matching keyword, but #1 in vector → also strong
- assert "kw_winner" in ids
- assert "vec_winner" in ids
-
- def test_rrf_dedupes_overlap(self, store: CarbonStore) -> None:
- # A chunk that wins both modalities should rank above either alone.
- emb = _fixed_embedder(
- {
- "needle": [1.0, 0.0, 0.0],
- "decoy_a": [0.0, 1.0, 0.0],
- "decoy_b": [0.0, 0.0, 1.0],
- },
- dim=3,
- )
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["needle", "decoy_a", "decoy_b"])
- store.embed_chunks(embedder=emb)
-
- hits = store.search("needle", mode="hybrid", top_k=3, embedder=emb)
- assert hits[0].chunk.content == "needle"
- # Should appear once, not twice.
- contents = [h.chunk.content for h in hits]
- assert contents.count("needle") == 1
-
- def test_rrf_score_formula(self, store: CarbonStore) -> None:
- # Single-source case: chunk only ranks in one modality, score should
- # equal 1/(rrf_k + 1).
- emb = _fixed_embedder({"only_vec": [1.0, 0.0, 0.0]}, dim=3)
- store.put_entity(id="e1", kind="document")
- # Use a content string that won't match the vector query token.
- store.add_chunks("e1", ["only_vec"])
- store.embed_chunks(embedder=emb)
-
- # Query with a vector-side term that does NOT appear as a keyword.
- # We construct a query string that the embedder maps to [1,0,0]
- # but whose tokens won't FTS-match the chunk content.
- # The embedder maps unknown strings to zero, so we map a synthetic.
- emb2 = _fixed_embedder(
- {"zzzqueryzzz": [1.0, 0.0, 0.0], "only_vec": [1.0, 0.0, 0.0]},
- dim=3,
- )
- hits = store.search(
- "zzzqueryzzz", mode="hybrid", top_k=1,
- embedder=emb2, rrf_k=60,
- )
- assert len(hits) == 1
- assert abs(hits[0].score - 1.0 / (60 + 1)) < 1e-9
-
- def test_hybrid_requires_string(self, store: CarbonStore) -> None:
- with pytest.raises(TypeError, match="hybrid mode"):
- store.search(np.array([1.0, 0.0, 0.0]), mode="hybrid")
-
- def test_hybrid_requires_embedder_or_model(self, store: CarbonStore) -> None:
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["something"])
- with pytest.raises(TypeError, match="model="):
- store.search("something", mode="hybrid")
-
- def test_hybrid_degrades_when_keyword_empty(self, store: CarbonStore) -> None:
- # Whitespace query → keyword side returns nothing; hybrid should
- # still return vector-side results without error.
- emb = _fixed_embedder({"a": [1.0, 0.0, 0.0]}, dim=3)
- store.put_entity(id="e1", kind="document")
- store.add_chunks("e1", ["a"])
- store.embed_chunks(embedder=emb)
- # Encoder returns zero-vec for whitespace; that's fine, just no
- # vector match either. Confirm it doesn't crash.
- result = store.search(" ", mode="hybrid", top_k=3, embedder=emb)
- assert isinstance(result, list)
-
- def test_oversample_validation(self, store: CarbonStore) -> None:
- emb = _fixed_embedder({"x": [1.0, 0.0, 0.0]}, dim=3)
- with pytest.raises(ValueError, match="fusion_oversample"):
- store.search("x", mode="hybrid", embedder=emb, fusion_oversample=0)
-
-
-# ---------------------------------------------------------------------------
-# Mode dispatch
-# ---------------------------------------------------------------------------
-
-class TestSearchModeDispatch:
- def test_default_is_vector(self, store: CarbonStore) -> None:
- # No mode= → vector path. Should require model/embedder.
- with pytest.raises(TypeError, match="model="):
- store.search("anything")
-
- def test_unknown_mode_raises(self, store: CarbonStore) -> None:
- with pytest.raises(ValueError, match="unknown search mode"):
- store.search("foo", mode="laser") # type: ignore[arg-type]
-
-
-# ---------------------------------------------------------------------------
-# ensure_triggers idempotency on reopen
-# ---------------------------------------------------------------------------
-
-class TestTriggerIdempotency:
- def test_reopen_keeps_fts_in_sync(self, tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- s = create(p)
- s.put_entity(id="e1", kind="document")
- s.add_chunks("e1", ["alpha beta gamma"])
- s.close()
-
- s2 = cd_open(p)
- try:
- hits = s2.search("beta", mode="keyword")
- assert len(hits) == 1
- # Writes after reopen still flow into FTS.
- s2.put_entity(id="e2", kind="document")
- s2.add_chunks("e2", ["delta epsilon"])
- hits2 = s2.search("epsilon", mode="keyword")
- assert len(hits2) == 1
- finally:
- s2.close()
diff --git a/Agent_module/tests/carbon_data/test_m5_hnsw.py b/Agent_module/tests/carbon_data/test_m5_hnsw.py
deleted file mode 100644
index 60bcb57030..0000000000
--- a/Agent_module/tests/carbon_data/test_m5_hnsw.py
+++ /dev/null
@@ -1,426 +0,0 @@
-"""M5: HNSW (hnswlib) ANN index — wrapper, dispatch, cache, persistence."""
-from __future__ import annotations
-
-from pathlib import Path
-
-import numpy as np
-import pytest
-
-from Agent_module.carbon_data import (
- CarbonStore,
- LambdaEmbedder,
- create,
- open as cd_open,
-)
-from Agent_module.carbon_data.index import (
- HnswIndex,
- hnsw_available,
-)
-
-
-pytestmark = pytest.mark.skipif(
- not hnsw_available(), reason="hnswlib not installed"
-)
-
-
-# ---------------------------------------------------------------------------
-# fixtures
-# ---------------------------------------------------------------------------
-
-@pytest.fixture
-def store(tmp_path: Path):
- p = tmp_path / "kb.carbondata"
- s = create(p)
- yield s
- if not s.closed:
- s.close()
-
-
-def _seed(store: CarbonStore, *, namespace="default", n=20, dim=8, seed=0):
- """Add n random embeddings under one entity. Returns the embedder."""
- rng = np.random.default_rng(seed)
- matrix = rng.standard_normal((n, dim)).astype(np.float32)
- # Map text -> row in the matrix so the embedder is deterministic.
- texts = [f"chunk-{i}" for i in range(n)]
- text_to_vec = {t: matrix[i] for i, t in enumerate(texts)}
-
- def _fn(ts):
- return np.stack([text_to_vec.get(t, np.zeros(dim, dtype=np.float32))
- for t in ts])
-
- emb = LambdaEmbedder(_fn, model="m5-test", dim=dim)
-
- eid = store.put_entity(id=f"e-{namespace}", kind="document", namespace=namespace)
- store.add_chunks(eid, texts)
- store.embed_chunks(emb)
- return emb, texts, matrix
-
-
-def _vector_index_count(store: CarbonStore) -> int:
- conn = store._require_open()
- return conn.execute("SELECT COUNT(*) FROM vector_index").fetchone()[0]
-
-
-# ---------------------------------------------------------------------------
-# bare wrapper — independent of CarbonStore
-# ---------------------------------------------------------------------------
-
-class TestHnswWrapper:
- def test_available(self) -> None:
- assert hnsw_available() is True
-
- def test_build_and_search_self_hit(self) -> None:
- rng = np.random.default_rng(42)
- v = rng.standard_normal((20, 6)).astype(np.float32)
- labels = list(range(1000, 1020))
- idx = HnswIndex.build(v, labels)
- # Query equals row 7 → should top-rank label 1007.
- hits = idx.search(v[7], top_k=3)
- assert hits[0][0] == 1007
- # Cosine sim with self ~= 1.0
- assert hits[0][1] > 0.99
-
- def test_search_caps_at_count(self) -> None:
- v = np.eye(3, dtype=np.float32)
- idx = HnswIndex.build(v, [10, 20, 30])
- hits = idx.search(v[0], top_k=100)
- assert len(hits) == 3
-
- def test_serialize_roundtrip_preserves_ranking(self) -> None:
- rng = np.random.default_rng(7)
- v = rng.standard_normal((30, 4)).astype(np.float32)
- labels = list(range(30))
- idx = HnswIndex.build(v, labels)
-
- before = idx.search(v[5], top_k=5)
- blob = idx.serialize()
- assert isinstance(blob, bytes) and len(blob) > 0
-
- idx2 = HnswIndex.deserialize(blob, dim=4)
- after = idx2.search(v[5], top_k=5)
- assert before == after
-
- def test_dimension_mismatch(self) -> None:
- v = np.zeros((5, 3), dtype=np.float32)
- with pytest.raises(ValueError, match="size mismatch"):
- HnswIndex.build(v, labels=[1, 2])
-
- def test_2d_required(self) -> None:
- with pytest.raises(ValueError, match="2-D"):
- HnswIndex.build(np.zeros(5, dtype=np.float32), labels=[1])
-
-
-# ---------------------------------------------------------------------------
-# dispatch — when does HNSW get used?
-# ---------------------------------------------------------------------------
-
-class TestDispatch:
- def test_auto_uses_hnsw_when_eligible(self, store: CarbonStore) -> None:
- emb, _, _ = _seed(store)
- # Default search → should populate the HNSW cache.
- store.search(np.array([1.0] + [0.0] * 7, dtype=np.float32),
- model=emb.model, top_k=3)
- assert emb.model in store._hnsw_cache
-
- def test_off_skips_hnsw(self, store: CarbonStore) -> None:
- emb, _, _ = _seed(store)
- store.search(np.array([1.0] + [0.0] * 7, dtype=np.float32),
- model=emb.model, top_k=3, use_hnsw="off")
- assert emb.model not in store._hnsw_cache
-
- def test_filters_force_brute(self, store: CarbonStore) -> None:
- emb, _, _ = _seed(store)
- # filters present → auto falls back to brute, cache untouched.
- store.search(
- np.array([1.0] + [0.0] * 7, dtype=np.float32),
- model=emb.model, top_k=3,
- filters={"kind": "document"},
- )
- assert emb.model not in store._hnsw_cache
-
- def test_non_cosine_forces_brute(self, store: CarbonStore) -> None:
- emb, _, _ = _seed(store)
- store.search(
- np.array([1.0] + [0.0] * 7, dtype=np.float32),
- model=emb.model, top_k=3, metric="dot",
- )
- assert emb.model not in store._hnsw_cache
-
- def test_on_with_filters_raises(self, store: CarbonStore) -> None:
- emb, _, _ = _seed(store)
- with pytest.raises(ValueError, match="use_hnsw='on' is incompatible"):
- store.search(
- np.array([1.0] + [0.0] * 7, dtype=np.float32),
- model=emb.model, use_hnsw="on",
- filters={"kind": "document"},
- )
-
- def test_on_with_non_cosine_raises(self, store: CarbonStore) -> None:
- emb, _, _ = _seed(store)
- with pytest.raises(ValueError, match="metric='cosine'"):
- store.search(
- np.array([1.0] + [0.0] * 7, dtype=np.float32),
- model=emb.model, use_hnsw="on", metric="l2",
- )
-
- def test_on_without_embeddings_returns_empty(self, store: CarbonStore) -> None:
- # Force HNSW with a model that has zero embeddings — should not
- # raise; just produce no hits.
- out = store.search(
- np.array([1.0, 0.0, 0.0], dtype=np.float32),
- model="ghost-model", use_hnsw="on", top_k=5,
- )
- assert out == []
-
-
-# ---------------------------------------------------------------------------
-# correctness — HNSW vs brute should agree at top-1 for clean data
-# ---------------------------------------------------------------------------
-
-class TestParity:
- def test_top1_matches_brute(self, store: CarbonStore) -> None:
- emb, texts, matrix = _seed(store, n=50, dim=16, seed=1)
- rng = np.random.default_rng(99)
- for _ in range(5):
- q = rng.standard_normal(16).astype(np.float32)
- hnsw_hits = store.search(q, model=emb.model, top_k=1, use_hnsw="on")
- brute_hits = store.search(q, model=emb.model, top_k=1, use_hnsw="off")
- assert hnsw_hits[0].chunk.id == brute_hits[0].chunk.id
-
-
-# ---------------------------------------------------------------------------
-# persistence — blob is written and reloaded across open()
-# ---------------------------------------------------------------------------
-
-class TestPersistence:
- def test_blob_persisted_after_first_search(self, store: CarbonStore) -> None:
- emb, _, _ = _seed(store)
- assert _vector_index_count(store) == 0
- store.search(np.array([1.0] + [0.0] * 7, dtype=np.float32),
- model=emb.model, top_k=3)
- assert _vector_index_count(store) == 1
-
- def test_reopen_loads_from_blob_not_rebuilt(self, tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- s = create(p)
- emb, _, matrix = _seed(s, n=12, dim=4, seed=2)
- # Trigger initial build + persist.
- s.search(matrix[0], model=emb.model, top_k=1)
- s.close()
-
- s2 = cd_open(p)
- try:
- # Cache empty, but blob is current → load path, no rebuild.
- assert emb.model not in s2._hnsw_cache
- hits = s2.search(matrix[0], model=emb.model, top_k=1)
- assert hits[0].chunk.content == "chunk-0"
- # Cache populated via the load branch.
- assert emb.model in s2._hnsw_cache
- finally:
- s2.close()
-
- def test_blob_with_stale_count_triggers_rebuild(self, tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- s = create(p)
- emb, _, matrix = _seed(s, n=8, dim=4, seed=3)
- s.search(matrix[0], model=emb.model, top_k=1)
- s.close()
-
- # Tamper with persisted meta count → reopen should treat blob as
- # stale and rebuild from the embedding table.
- import sqlite3
- with sqlite3.connect(str(p)) as raw:
- raw.execute(
- "UPDATE vector_index SET index_meta = ? WHERE model = ?",
- ('{"count": 0, "dim": 4}', emb.model),
- )
-
- s2 = cd_open(p)
- try:
- hits = s2.search(matrix[0], model=emb.model, top_k=1)
- assert hits[0].chunk.content == "chunk-0"
- # After rebuild, cache count == real embedding count.
- assert s2._hnsw_cache[emb.model].count == 8
- finally:
- s2.close()
-
-
-# ---------------------------------------------------------------------------
-# cache invalidation — adds/deletes change the embedding count
-# ---------------------------------------------------------------------------
-
-class TestCacheInvalidation:
- def test_add_invalidates_cache(self, store: CarbonStore) -> None:
- emb, _, _ = _seed(store, n=5, dim=4, seed=10)
- store.search(np.zeros(4, dtype=np.float32),
- model=emb.model, top_k=1)
- cached_before = store._hnsw_cache[emb.model]
- assert cached_before.count == 5
-
- # Add another chunk + embedding under the same entity.
- eid = store.get_entity("e-default").id
- new_ids = store.add_chunks(eid, ["chunk-extra"])
-
- def _fn(texts):
- arr = np.zeros((len(texts), 4), dtype=np.float32)
- for i, t in enumerate(texts):
- if t == "chunk-extra":
- arr[i] = np.array([0.5, 0.5, 0.5, 0.5], dtype=np.float32)
- return arr
- ext_emb = LambdaEmbedder(_fn, model=emb.model, dim=4)
- store.embed_chunks(ext_emb, entity_id=eid, missing_only=True)
-
- store.search(np.zeros(4, dtype=np.float32),
- model=emb.model, top_k=1)
- cached_after = store._hnsw_cache[emb.model]
- # Either a new instance (rebuild) or the same instance with bumped
- # count — both are acceptable. The invariant is: count is current.
- assert cached_after.count == 6
-
- def test_delete_invalidates_cache(self, store: CarbonStore) -> None:
- emb, _, _ = _seed(store, n=6, dim=4, seed=11)
- store.search(np.zeros(4, dtype=np.float32),
- model=emb.model, top_k=1)
- assert store._hnsw_cache[emb.model].count == 6
-
- # Delete a chunk → cascade also drops its embedding row.
- first_chunk = store.list_chunks("e-default")[0]
- with store._write_ctx() as conn:
- conn.execute("DELETE FROM chunk WHERE id=?", (first_chunk.id,))
-
- store.search(np.zeros(4, dtype=np.float32),
- model=emb.model, top_k=1)
- assert store._hnsw_cache[emb.model].count == 5
-
- def test_replace_embedding_invalidates_cache_and_blob(
- self, store: CarbonStore
- ) -> None:
- emb, _, _ = _seed(store, n=5, dim=4, seed=13)
- store.search(np.ones(4, dtype=np.float32),
- model=emb.model, top_k=1)
- assert emb.model in store._hnsw_cache
- assert _vector_index_count(store) == 1
-
- first_chunk = store.list_chunks("e-default")[0]
- store.set_embedding(
- first_chunk.id,
- np.array([9.0, 0.0, 0.0, 0.0], dtype=np.float32),
- model=emb.model,
- )
-
- assert emb.model not in store._hnsw_cache
- assert _vector_index_count(store) == 0
-
- def test_reembed_invalidates_cache_and_blob(
- self, store: CarbonStore
- ) -> None:
- emb, _, _ = _seed(store, n=5, dim=4, seed=14)
- store.search(np.ones(4, dtype=np.float32),
- model=emb.model, top_k=1)
- assert emb.model in store._hnsw_cache
- assert _vector_index_count(store) == 1
-
- store.embed_chunks(emb, missing_only=False)
-
- assert emb.model not in store._hnsw_cache
- assert _vector_index_count(store) == 0
-
- def test_rollback_clears_index_built_from_uncommitted_vectors(
- self, store: CarbonStore
- ) -> None:
- emb, _, _ = _seed(store, n=5, dim=4, seed=15)
- first_chunk = store.list_chunks("e-default")[0]
-
- with pytest.raises(RuntimeError):
- with store.transaction():
- store.set_embedding(
- first_chunk.id,
- np.array([9.0, 0.0, 0.0, 0.0], dtype=np.float32),
- model=emb.model,
- )
- store.search(
- np.ones(4, dtype=np.float32),
- model=emb.model,
- top_k=1,
- use_hnsw="on",
- )
- assert emb.model in store._hnsw_cache
- raise RuntimeError("rollback")
-
- assert emb.model not in store._hnsw_cache
-
- def test_full_cleanup_drops_cache_entry(self, store: CarbonStore) -> None:
- emb, _, _ = _seed(store, n=3, dim=4, seed=12)
- store.search(np.zeros(4, dtype=np.float32),
- model=emb.model, top_k=1)
- assert emb.model in store._hnsw_cache
-
- store.delete_entity("e-default")
- # Next search → no embeddings → cache entry should be removed.
- out = store.search(np.zeros(4, dtype=np.float32),
- model=emb.model, top_k=1)
- assert out == []
- assert emb.model not in store._hnsw_cache
-
-
-# ---------------------------------------------------------------------------
-# namespace post-filter
-# ---------------------------------------------------------------------------
-
-class TestNamespacePostFilter:
- def test_namespace_filter_applied(self, tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- s = create(p)
- try:
- # Two namespaces, same model; HNSW index spans both.
- _seed(s, namespace="ns-a", n=4, dim=4, seed=20)
- _seed(s, namespace="ns-b", n=4, dim=4, seed=21)
- # Sanity: 8 embeddings now across one model.
- assert s.stats()["embeddings"] == 8
-
- # Searching with namespace=ns-a must only return ns-a hits.
- hits = s.search(
- np.zeros(4, dtype=np.float32),
- model="m5-test", top_k=4, namespace="ns-a",
- )
- assert all(h.entity.namespace == "ns-a" for h in hits)
- finally:
- s.close()
-
- def test_sparse_namespace_expands_until_top_k(self, tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- s = create(p)
- try:
- s.put_entity(id="target", kind="document", namespace="target")
- target_chunks = s.add_chunks(
- "target", [f"target-{i}" for i in range(4)]
- )
- s.put_entity(id="other", kind="document", namespace="other")
- other_chunks = s.add_chunks(
- "other", [f"other-{i}" for i in range(12)]
- )
- for chunk_id in target_chunks:
- s.set_embedding(
- chunk_id,
- np.array([-1.0, 0.0], dtype=np.float32),
- model="sparse-ns",
- )
- for chunk_id in other_chunks:
- s.set_embedding(
- chunk_id,
- np.array([1.0, 0.0], dtype=np.float32),
- model="sparse-ns",
- )
-
- hits = s.search(
- np.array([1.0, 0.0], dtype=np.float32),
- model="sparse-ns",
- top_k=4,
- namespace="target",
- use_hnsw="on",
- )
-
- assert len(hits) == 4
- assert all(hit.entity.namespace == "target" for hit in hits)
- finally:
- s.close()
diff --git a/Agent_module/tests/carbon_data/test_m6_graph.py b/Agent_module/tests/carbon_data/test_m6_graph.py
deleted file mode 100644
index f925bfcf4e..0000000000
--- a/Agent_module/tests/carbon_data/test_m6_graph.py
+++ /dev/null
@@ -1,354 +0,0 @@
-"""M6: knowledge graph — relations, neighbors, traverse, subgraph."""
-from __future__ import annotations
-
-from pathlib import Path
-
-import pytest
-
-from Agent_module.carbon_data import (
- CarbonStore,
- Neighbor,
- Relation,
- Subgraph,
- TraversalHit,
- create,
-)
-
-
-# ---------------------------------------------------------------------------
-# fixtures
-# ---------------------------------------------------------------------------
-
-@pytest.fixture
-def store(tmp_path: Path):
- p = tmp_path / "kb.carbondata"
- s = create(p)
- yield s
- if not s.closed:
- s.close()
-
-
-def _node(s: CarbonStore, id: str, *, kind="document", namespace="default") -> str:
- return s.put_entity(id=id, kind=kind, namespace=namespace)
-
-
-# ---------------------------------------------------------------------------
-# CRUD
-# ---------------------------------------------------------------------------
-
-class TestRelationCRUD:
- def test_add_and_get(self, store: CarbonStore) -> None:
- _node(store, "a")
- _node(store, "b")
- store.add_relation("a", "b", "references", weight=0.7,
- metadata={"reason": "cite"})
- r = store.get_relation("a", "b", "references")
- assert isinstance(r, Relation)
- assert r.weight == 0.7
- assert r.metadata == {"reason": "cite"}
-
- def test_missing_endpoints_raise(self, store: CarbonStore) -> None:
- _node(store, "a")
- with pytest.raises(KeyError, match="dst entity not found"):
- store.add_relation("a", "ghost", "references")
- with pytest.raises(KeyError, match="src entity not found"):
- store.add_relation("ghost", "a", "references")
-
- def test_upsert_replaces_weight_and_metadata(self, store: CarbonStore) -> None:
- _node(store, "a")
- _node(store, "b")
- store.add_relation("a", "b", "references", weight=0.5)
- store.add_relation("a", "b", "references", weight=0.9,
- metadata={"updated": True})
- r = store.get_relation("a", "b", "references")
- assert r.weight == 0.9
- assert r.metadata == {"updated": True}
-
- def test_kind_is_part_of_pk(self, store: CarbonStore) -> None:
- _node(store, "a")
- _node(store, "b")
- store.add_relation("a", "b", "references")
- store.add_relation("a", "b", "supersedes")
- # Two distinct edges between the same endpoints with different kinds.
- assert len(store.list_relations(src="a", dst="b")) == 2
-
- def test_delete(self, store: CarbonStore) -> None:
- _node(store, "a")
- _node(store, "b")
- store.add_relation("a", "b", "references")
- assert store.delete_relation("a", "b", "references") is True
- assert store.delete_relation("a", "b", "references") is False
- assert store.get_relation("a", "b", "references") is None
-
- def test_list_filters(self, store: CarbonStore) -> None:
- for n in ("a", "b", "c"):
- _node(store, n)
- store.add_relation("a", "b", "references")
- store.add_relation("a", "c", "references")
- store.add_relation("b", "c", "supersedes")
-
- assert len(store.list_relations(src="a")) == 2
- assert len(store.list_relations(kind="supersedes")) == 1
- assert len(store.list_relations(dst="c")) == 2
- assert len(store.list_relations(src="a", kind="references")) == 2
-
- def test_list_limit(self, store: CarbonStore) -> None:
- for n in ("a", "b", "c", "d"):
- _node(store, n)
- for d in ("b", "c", "d"):
- store.add_relation("a", d, "ref")
- rels = store.list_relations(src="a", limit=2)
- assert len(rels) == 2
-
- def test_entity_delete_cascades_relations(self, store: CarbonStore) -> None:
- _node(store, "a")
- _node(store, "b")
- store.add_relation("a", "b", "references")
- store.delete_entity("b")
- assert store.list_relations() == []
-
-
-# ---------------------------------------------------------------------------
-# neighbors (1-hop)
-# ---------------------------------------------------------------------------
-
-class TestNeighbors:
- def _build_star(self, s: CarbonStore) -> None:
- # a -> b, a -> c (out of a)
- # d -> a (in to a)
- for n in ("a", "b", "c", "d"):
- _node(s, n)
- s.add_relation("a", "b", "references")
- s.add_relation("a", "c", "references", weight=0.5)
- s.add_relation("d", "a", "supersedes")
-
- def test_out(self, store: CarbonStore) -> None:
- self._build_star(store)
- ns = store.neighbors("a", direction="out")
- assert all(isinstance(n, Neighbor) for n in ns)
- assert {n.entity.id for n in ns} == {"b", "c"}
- assert all(n.direction == "out" for n in ns)
-
- def test_in(self, store: CarbonStore) -> None:
- self._build_star(store)
- ns = store.neighbors("a", direction="in")
- assert {n.entity.id for n in ns} == {"d"}
- assert all(n.direction == "in" for n in ns)
-
- def test_both(self, store: CarbonStore) -> None:
- self._build_star(store)
- ns = store.neighbors("a", direction="both")
- assert {n.entity.id for n in ns} == {"b", "c", "d"}
- per_dir = {n.entity.id: n.direction for n in ns}
- assert per_dir["b"] == "out"
- assert per_dir["d"] == "in"
-
- def test_kind_filter(self, store: CarbonStore) -> None:
- self._build_star(store)
- ns = store.neighbors("a", direction="both", kind="supersedes")
- assert {n.entity.id for n in ns} == {"d"}
-
- def test_weight_and_metadata(self, store: CarbonStore) -> None:
- _node(store, "a")
- _node(store, "b")
- store.add_relation("a", "b", "references", weight=0.42,
- metadata={"src_page": 7})
- ns = store.neighbors("a")
- assert ns[0].weight == 0.42
- assert ns[0].metadata == {"src_page": 7}
- assert ns[0].relation_kind == "references"
-
- def test_isolated_node(self, store: CarbonStore) -> None:
- _node(store, "lonely")
- assert store.neighbors("lonely") == []
-
- def test_limit(self, store: CarbonStore) -> None:
- _node(store, "a")
- for d in ("b", "c", "d"):
- _node(store, d)
- store.add_relation("a", d, "ref")
- ns = store.neighbors("a", limit=2)
- assert len(ns) == 2
-
- def test_invalid_direction(self, store: CarbonStore) -> None:
- _node(store, "a")
- with pytest.raises(ValueError, match="direction"):
- store.neighbors("a", direction="sideways") # type: ignore[arg-type]
-
-
-# ---------------------------------------------------------------------------
-# traverse (multi-hop)
-# ---------------------------------------------------------------------------
-
-class TestTraverse:
- def _build_chain(self, s: CarbonStore) -> None:
- # a -> b -> c -> d -> e
- for n in ("a", "b", "c", "d", "e"):
- _node(s, n)
- for u, v in [("a", "b"), ("b", "c"), ("c", "d"), ("d", "e")]:
- s.add_relation(u, v, "next")
-
- def test_max_hops_bounds_walk(self, store: CarbonStore) -> None:
- self._build_chain(store)
- hits = store.traverse("a", max_hops=2, direction="out")
- ids_to_hop = {h.entity.id: h.hop for h in hits}
- assert ids_to_hop == {"b": 1, "c": 2}
-
- def test_max_hops_three(self, store: CarbonStore) -> None:
- self._build_chain(store)
- hits = store.traverse("a", max_hops=3, direction="out")
- assert {h.entity.id: h.hop for h in hits} == {"b": 1, "c": 2, "d": 3}
-
- def test_excludes_start(self, store: CarbonStore) -> None:
- self._build_chain(store)
- hits = store.traverse("a", max_hops=5)
- assert "a" not in {h.entity.id for h in hits}
-
- def test_returns_min_hop(self, store: CarbonStore) -> None:
- # a -> b -> c, a -> c (direct). c reachable at hop 1 and hop 2.
- for n in ("a", "b", "c"):
- _node(store, n)
- store.add_relation("a", "b", "ref")
- store.add_relation("b", "c", "ref")
- store.add_relation("a", "c", "ref")
- hits = store.traverse("a", max_hops=3)
- c_hop = next(h.hop for h in hits if h.entity.id == "c")
- assert c_hop == 1 # picks shortest path
-
- def test_cycle_safety(self, store: CarbonStore) -> None:
- # a -> b -> a (cycle). Without bounding this would loop forever.
- _node(store, "a")
- _node(store, "b")
- store.add_relation("a", "b", "x")
- store.add_relation("b", "a", "x")
- hits = store.traverse("a", max_hops=5, direction="out")
- # b reachable; a is the start, excluded.
- assert {h.entity.id for h in hits} == {"b"}
-
- def test_direction_in(self, store: CarbonStore) -> None:
- self._build_chain(store)
- hits = store.traverse("e", max_hops=2, direction="in")
- assert {h.entity.id: h.hop for h in hits} == {"d": 1, "c": 2}
-
- def test_direction_both(self, store: CarbonStore) -> None:
- # a -> b, c -> b (b has both incoming and outgoing options if both used)
- _node(store, "a")
- _node(store, "b")
- _node(store, "c")
- store.add_relation("a", "b", "x")
- store.add_relation("c", "b", "x")
- hits = store.traverse("b", max_hops=1, direction="both")
- assert {h.entity.id for h in hits} == {"a", "c"}
-
- def test_kind_filter(self, store: CarbonStore) -> None:
- # a -ref-> b -ref-> c, a -other-> d
- for n in ("a", "b", "c", "d"):
- _node(store, n)
- store.add_relation("a", "b", "ref")
- store.add_relation("b", "c", "ref")
- store.add_relation("a", "d", "other")
- hits = store.traverse("a", kind="ref", max_hops=3)
- assert {h.entity.id for h in hits} == {"b", "c"}
-
- def test_namespace_filter(self, store: CarbonStore) -> None:
- _node(store, "a", namespace="ns1")
- _node(store, "b", namespace="ns1")
- _node(store, "c", namespace="ns2")
- store.add_relation("a", "b", "x")
- store.add_relation("b", "c", "x")
- hits = store.traverse("a", max_hops=2, namespace="ns1")
- assert {h.entity.id for h in hits} == {"b"} # c filtered out
-
- def test_invalid_max_hops(self, store: CarbonStore) -> None:
- _node(store, "a")
- with pytest.raises(ValueError, match="max_hops"):
- store.traverse("a", max_hops=0)
-
- def test_invalid_direction(self, store: CarbonStore) -> None:
- _node(store, "a")
- with pytest.raises(ValueError, match="direction"):
- store.traverse("a", direction="weird") # type: ignore[arg-type]
-
-
-# ---------------------------------------------------------------------------
-# subgraph
-# ---------------------------------------------------------------------------
-
-class TestSubgraph:
- def test_includes_seeds_at_hop_zero(self, store: CarbonStore) -> None:
- _node(store, "a")
- _node(store, "b")
- sg = store.subgraph(["a"], max_hops=0)
- assert {e.id for e in sg.entities} == {"a"}
- assert sg.relations == []
-
- def test_one_hop_collects_neighbors(self, store: CarbonStore) -> None:
- for n in ("a", "b", "c", "d"):
- _node(store, n)
- store.add_relation("a", "b", "x")
- store.add_relation("a", "c", "x")
- store.add_relation("c", "d", "x")
- sg = store.subgraph(["a"], max_hops=1, direction="out")
- assert {e.id for e in sg.entities} == {"a", "b", "c"}
- assert {(r.src, r.dst) for r in sg.relations} == {("a", "b"), ("a", "c")}
-
- def test_two_hops(self, store: CarbonStore) -> None:
- for n in ("a", "b", "c", "d"):
- _node(store, n)
- store.add_relation("a", "b", "x")
- store.add_relation("b", "c", "x")
- store.add_relation("c", "d", "x")
- sg = store.subgraph(["a"], max_hops=2, direction="out")
- assert {e.id for e in sg.entities} == {"a", "b", "c"}
- assert len(sg.relations) == 2
-
- def test_relations_only_internal(self, store: CarbonStore) -> None:
- # Subgraph picks up a, b, c (max_hops=1 from a). The edge c -> d
- # has only c inside the visited set; it must NOT appear in the
- # subgraph's relation list.
- for n in ("a", "b", "c", "d"):
- _node(store, n)
- store.add_relation("a", "b", "x")
- store.add_relation("a", "c", "x")
- store.add_relation("c", "d", "x") # leaks outside
- sg = store.subgraph(["a"], max_hops=1, direction="out")
- assert all(
- r.src in {"a", "b", "c"} and r.dst in {"a", "b", "c"}
- for r in sg.relations
- )
-
- def test_multi_seed(self, store: CarbonStore) -> None:
- # a -> b ; c -> d. Two disjoint seeds, two disjoint subgraphs.
- for n in ("a", "b", "c", "d"):
- _node(store, n)
- store.add_relation("a", "b", "x")
- store.add_relation("c", "d", "x")
- sg = store.subgraph(["a", "c"], max_hops=1, direction="out")
- assert {e.id for e in sg.entities} == {"a", "b", "c", "d"}
- assert {(r.src, r.dst) for r in sg.relations} == {
- ("a", "b"), ("c", "d"),
- }
-
- def test_empty_seeds(self, store: CarbonStore) -> None:
- sg = store.subgraph([])
- assert isinstance(sg, Subgraph)
- assert sg.entities == []
- assert sg.relations == []
-
- def test_kind_filter(self, store: CarbonStore) -> None:
- for n in ("a", "b", "c"):
- _node(store, n)
- store.add_relation("a", "b", "ref")
- store.add_relation("a", "c", "other")
- sg = store.subgraph(["a"], kind="ref", max_hops=1, direction="out")
- # Only "ref" walked → b reachable; c not.
- assert {e.id for e in sg.entities} == {"a", "b"}
- assert all(r.kind == "ref" for r in sg.relations)
-
- def test_direction_in_only(self, store: CarbonStore) -> None:
- for n in ("a", "b", "c"):
- _node(store, n)
- store.add_relation("a", "b", "x")
- store.add_relation("c", "b", "x")
- sg = store.subgraph(["b"], direction="in", max_hops=1)
- assert {e.id for e in sg.entities} == {"a", "b", "c"}
diff --git a/Agent_module/tests/carbon_data/test_m7_memory.py b/Agent_module/tests/carbon_data/test_m7_memory.py
deleted file mode 100644
index 9c15b40bf5..0000000000
--- a/Agent_module/tests/carbon_data/test_m7_memory.py
+++ /dev/null
@@ -1,377 +0,0 @@
-"""M7: long-term memory — remember / recall / forget."""
-from __future__ import annotations
-
-import time
-from pathlib import Path
-
-import numpy as np
-import pytest
-
-from Agent_module.carbon_data import (
- CarbonStore,
- LambdaEmbedder,
- MemoryHit,
- MemoryItem,
- create,
-)
-
-
-# ---------------------------------------------------------------------------
-# fixtures
-# ---------------------------------------------------------------------------
-
-@pytest.fixture
-def store(tmp_path: Path):
- p = tmp_path / "kb.carbondata"
- s = create(p)
- yield s
- if not s.closed:
- s.close()
-
-
-def _fixed_embedder(mapping: dict[str, list[float]], *, model="test", dim=3):
- arr = {k: np.asarray(v, dtype=np.float32).reshape(-1) for k, v in mapping.items()}
-
- def _fn(texts):
- out = np.zeros((len(texts), dim), dtype=np.float32)
- for i, t in enumerate(texts):
- if t in arr:
- out[i] = arr[t]
- return out
-
- return LambdaEmbedder(_fn, model=model, dim=dim)
-
-
-# ---------------------------------------------------------------------------
-# remember
-# ---------------------------------------------------------------------------
-
-class TestRemember:
- def test_creates_entity_chunk_and_memory_ext(self, store: CarbonStore) -> None:
- eid = store.remember("user prefers dark mode", session_id="s1", actor="ui-bot")
- ent = store.get_entity(eid)
- assert ent is not None
- assert ent.kind == "memory"
- assert ent.content == "user prefers dark mode"
-
- chunks = store.list_chunks(eid)
- assert len(chunks) == 1
- assert chunks[0].content == "user prefers dark mode"
-
- mem = store._load_memory(eid)
- assert isinstance(mem, MemoryItem)
- assert mem.session_id == "s1"
- assert mem.actor == "ui-bot"
- assert mem.salience == 0.5 # default
-
- def test_returns_assigned_id(self, store: CarbonStore) -> None:
- eid = store.remember("foo", id="mem-001")
- assert eid == "mem-001"
- assert store.get_entity("mem-001") is not None
-
- def test_auto_uuid_when_no_id(self, store: CarbonStore) -> None:
- eid = store.remember("foo")
- assert isinstance(eid, str) and len(eid) >= 16
-
- def test_ttl_sets_expires_at(self, store: CarbonStore) -> None:
- before = time.time()
- eid = store.remember("expires soon", ttl=60)
- mem = store._load_memory(eid)
- assert mem.expires_at is not None
- assert before + 59 <= mem.expires_at <= time.time() + 61
-
- def test_explicit_expires_at(self, store: CarbonStore) -> None:
- eat = time.time() + 3600
- eid = store.remember("X", expires_at=eat)
- mem = store._load_memory(eid)
- assert mem.expires_at == eat
-
- def test_ttl_and_expires_at_mutually_exclusive(self, store: CarbonStore) -> None:
- with pytest.raises(TypeError, match="ttl OR expires_at"):
- store.remember("x", ttl=60, expires_at=time.time() + 60)
-
- def test_salience_bounds_validated(self, store: CarbonStore) -> None:
- with pytest.raises(ValueError, match="salience"):
- store.remember("x", salience=1.5)
- with pytest.raises(ValueError, match="salience"):
- store.remember("x", salience=-0.1)
-
- def test_eager_embedding_via_embedder(self, store: CarbonStore) -> None:
- emb = _fixed_embedder({"likes vim": [1.0, 0.0, 0.0]})
- eid = store.remember("likes vim", embedder=emb)
- # Should have an embedding under the embedder's model.
- chunk = store.list_chunks(eid)[0]
- vec = store.get_embedding(chunk.id, model=emb.model)
- assert vec is not None
- assert vec.shape == (3,)
-
- def test_atomic_failure_rolls_back(self, store: CarbonStore) -> None:
- # If embedding fails, no entity/chunk should remain.
- class BoomEmbedder:
- model = "boom"
- dim = 3
-
- def encode(self, texts):
- raise RuntimeError("boom")
-
- before = store.stats()
- with pytest.raises(RuntimeError, match="boom"):
- store.remember("doomed", embedder=BoomEmbedder())
- after = store.stats()
- assert after["entities"] == before["entities"]
- assert after["chunks"] == before["chunks"]
- assert after["memories"] == before["memories"]
-
- def test_namespace_isolation(self, store: CarbonStore) -> None:
- eid = store.remember("private", namespace="tenant-a")
- ent = store.get_entity(eid)
- assert ent.namespace == "tenant-a"
-
- def test_metadata_passes_through(self, store: CarbonStore) -> None:
- eid = store.remember("note", metadata={"topic": "auth"})
- ent = store.get_entity(eid)
- assert ent.metadata == {"topic": "auth"}
-
-
-# ---------------------------------------------------------------------------
-# recall
-# ---------------------------------------------------------------------------
-
-class TestRecall:
- def test_semantic_ranking(self, store: CarbonStore) -> None:
- emb = _fixed_embedder(
- {
- "loves vim": [1.0, 0.0, 0.0],
- "prefers emacs": [0.0, 1.0, 0.0],
- "likes nano": [0.0, 0.0, 1.0],
- "vim editor": [0.95, 0.05, 0.0], # query
- },
- dim=3,
- )
- store.remember("loves vim", embedder=emb)
- store.remember("prefers emacs", embedder=emb)
- store.remember("likes nano", embedder=emb)
-
- hits = store.recall("vim editor", top_k=3, embedder=emb)
- assert len(hits) == 3
- assert all(isinstance(h, MemoryHit) for h in hits)
- assert hits[0].memory.content == "loves vim"
-
- def test_session_filter(self, store: CarbonStore) -> None:
- emb = _fixed_embedder(
- {"a": [1.0, 0.0, 0.0], "b": [1.0, 0.0, 0.0]}, dim=3
- )
- store.remember("a", session_id="s1", embedder=emb)
- store.remember("b", session_id="s2", embedder=emb)
-
- hits = store.recall("a", session_id="s1", top_k=5, embedder=emb)
- assert len(hits) == 1
- assert hits[0].memory.session_id == "s1"
-
- def test_actor_filter(self, store: CarbonStore) -> None:
- emb = _fixed_embedder(
- {"a": [1.0, 0.0, 0.0], "b": [1.0, 0.0, 0.0]}, dim=3
- )
- store.remember("a", actor="bot1", embedder=emb)
- store.remember("b", actor="bot2", embedder=emb)
-
- hits = store.recall("a", actor="bot1", embedder=emb)
- assert {h.memory.actor for h in hits} == {"bot1"}
-
- def test_min_salience_filter(self, store: CarbonStore) -> None:
- emb = _fixed_embedder(
- {"low": [1.0, 0.0, 0.0], "high": [1.0, 0.0, 0.0]}, dim=3
- )
- store.remember("low", salience=0.2, embedder=emb)
- store.remember("high", salience=0.9, embedder=emb)
-
- hits = store.recall("low", min_salience=0.5, embedder=emb)
- assert len(hits) == 1
- assert hits[0].memory.salience == 0.9
-
- def test_excludes_expired_by_default(self, store: CarbonStore) -> None:
- emb = _fixed_embedder(
- {"old": [1.0, 0.0, 0.0], "fresh": [1.0, 0.0, 0.0]}, dim=3
- )
- store.remember("old", expires_at=time.time() - 1, embedder=emb)
- store.remember("fresh", embedder=emb)
-
- hits = store.recall("old", embedder=emb)
- assert {h.memory.content for h in hits} == {"fresh"}
-
- def test_include_expired_opts_in(self, store: CarbonStore) -> None:
- emb = _fixed_embedder(
- {"old": [1.0, 0.0, 0.0], "fresh": [1.0, 0.0, 0.0]}, dim=3
- )
- store.remember("old", expires_at=time.time() - 1, embedder=emb)
- store.remember("fresh", embedder=emb)
-
- hits = store.recall(
- "old", embedder=emb, include_expired=True, top_k=5
- )
- assert {h.memory.content for h in hits} == {"old", "fresh"}
-
- def test_namespace_isolation(self, store: CarbonStore) -> None:
- emb = _fixed_embedder(
- {"a": [1.0, 0.0, 0.0], "b": [1.0, 0.0, 0.0]}, dim=3
- )
- store.remember("a", namespace="ns1", embedder=emb)
- store.remember("b", namespace="ns2", embedder=emb)
-
- hits = store.recall("a", namespace="ns1", embedder=emb)
- assert len(hits) == 1
- assert hits[0].memory.namespace == "ns1"
-
- def test_does_not_match_non_memory_entities(self, store: CarbonStore) -> None:
- # Even if a regular document-kind entity is in the store, recall
- # restricts to kind='memory'.
- emb = _fixed_embedder(
- {"doc-text": [1.0, 0.0, 0.0], "mem-text": [1.0, 0.0, 0.0]}, dim=3
- )
- store.put_entity(id="doc1", kind="document", content="doc-text")
- store.add_chunks("doc1", ["doc-text"])
- store.embed_chunks(emb, entity_id="doc1")
- store.remember("mem-text", embedder=emb)
-
- hits = store.recall("doc-text", embedder=emb, top_k=5)
- assert all(h.memory.kind == "memory" for h in hits)
- assert len(hits) == 1
-
- def test_empty_store_returns_empty(self, store: CarbonStore) -> None:
- emb = _fixed_embedder({"q": [1.0, 0.0, 0.0]}, dim=3)
- assert store.recall("q", embedder=emb) == []
-
- def test_requires_model_or_embedder(self, store: CarbonStore) -> None:
- with pytest.raises(TypeError, match="model="):
- store.recall("q")
-
- def test_string_query_requires_embedder(self, store: CarbonStore) -> None:
- with pytest.raises(TypeError, match="embedder="):
- store.recall("q", model="m")
-
- def test_vector_query_works_without_embedder(self, store: CarbonStore) -> None:
- emb = _fixed_embedder({"a": [1.0, 0.0, 0.0]}, dim=3)
- store.remember("a", embedder=emb)
- hits = store.recall(
- np.array([1.0, 0.0, 0.0]), model=emb.model, top_k=3
- )
- assert len(hits) == 1
-
-
-# ---------------------------------------------------------------------------
-# forget
-# ---------------------------------------------------------------------------
-
-class TestForget:
- def test_no_criteria_raises(self, store: CarbonStore) -> None:
- with pytest.raises(TypeError, match="at least one criterion"):
- store.forget()
-
- def test_by_entity_id(self, store: CarbonStore) -> None:
- eid = store.remember("x", session_id="s1")
- n = store.forget(entity_id=eid)
- assert n == 1
- assert store.get_entity(eid) is None
-
- def test_by_session(self, store: CarbonStore) -> None:
- store.remember("a", session_id="s1")
- store.remember("b", session_id="s1")
- store.remember("c", session_id="s2")
- n = store.forget(session_id="s1")
- assert n == 2
- # s2 untouched
- remaining_sessions = [
- store._load_memory(eid).session_id
- for eid, in store._require_open().execute(
- "SELECT id FROM entity WHERE kind='memory'"
- )
- ]
- assert remaining_sessions == ["s2"]
-
- def test_by_actor(self, store: CarbonStore) -> None:
- store.remember("a", actor="bot")
- store.remember("b", actor="user")
- n = store.forget(actor="bot")
- assert n == 1
-
- def test_by_namespace(self, store: CarbonStore) -> None:
- store.remember("a", namespace="t1")
- store.remember("b", namespace="t2")
- n = store.forget(namespace="t1")
- assert n == 1
-
- def test_by_before(self, store: CarbonStore) -> None:
- eid_old = store.remember("old")
- # Force created_at to a known past timestamp.
- store._require_rw().execute(
- "UPDATE entity SET created_at=? WHERE id=?",
- (1000.0, eid_old),
- )
- store._require_rw().commit()
- store.remember("new") # created_at = now
- n = store.forget(before=2000.0)
- assert n == 1
-
- def test_only_targets_memories(self, store: CarbonStore) -> None:
- # A document-kind entity in the same namespace should not be
- # touched by forget(namespace=...).
- store.put_entity(id="doc", kind="document", namespace="ns1")
- store.remember("m", namespace="ns1")
- n = store.forget(namespace="ns1")
- assert n == 1
- assert store.get_entity("doc") is not None
-
- def test_combined_criteria_are_anded(self, store: CarbonStore) -> None:
- store.remember("a", session_id="s1", actor="bot")
- store.remember("b", session_id="s1", actor="user")
- store.remember("c", session_id="s2", actor="bot")
- n = store.forget(session_id="s1", actor="bot")
- assert n == 1
-
- def test_no_match_returns_zero(self, store: CarbonStore) -> None:
- store.remember("x", session_id="s1")
- assert store.forget(session_id="nonexistent") == 0
-
- def test_cascade_removes_chunks_and_embeddings(self, store: CarbonStore) -> None:
- emb = _fixed_embedder({"x": [1.0, 0.0, 0.0]})
- eid = store.remember("x", session_id="s1", embedder=emb)
- before = store.stats()
- assert before["chunks"] >= 1
- assert before["embeddings"] >= 1
-
- store.forget(session_id="s1")
- after = store.stats()
- assert after["entities"] == before["entities"] - 1
- assert after["chunks"] == before["chunks"] - 1
- assert after["embeddings"] == before["embeddings"] - 1
- assert after["memories"] == before["memories"] - 1
-
-
-# ---------------------------------------------------------------------------
-# forget_expired
-# ---------------------------------------------------------------------------
-
-class TestForgetExpired:
- def test_clears_past_expirations(self, store: CarbonStore) -> None:
- store.remember("a", expires_at=time.time() - 10)
- store.remember("b", expires_at=time.time() - 1)
- store.remember("c", expires_at=time.time() + 3600)
- store.remember("d") # no expiry
- n = store.forget_expired()
- assert n == 2
-
- def test_explicit_now_boundary(self, store: CarbonStore) -> None:
- # expires_at == now → expired (uses <=).
- eid = store.remember("boundary", expires_at=1500.0)
- n = store.forget_expired(now=1500.0)
- assert n == 1
- assert store.get_entity(eid) is None
-
- def test_explicit_now_strictly_before(self, store: CarbonStore) -> None:
- store.remember("not_yet", expires_at=2000.0)
- n = store.forget_expired(now=1999.999)
- assert n == 0
-
- def test_no_expirations_returns_zero(self, store: CarbonStore) -> None:
- store.remember("permanent")
- assert store.forget_expired() == 0
diff --git a/Agent_module/tests/carbon_data/test_m8_ingest.py b/Agent_module/tests/carbon_data/test_m8_ingest.py
deleted file mode 100644
index 32f0ad5a1d..0000000000
--- a/Agent_module/tests/carbon_data/test_m8_ingest.py
+++ /dev/null
@@ -1,278 +0,0 @@
-"""M8: chunkers + ingest_text + ingest_table."""
-from __future__ import annotations
-
-from pathlib import Path
-
-import numpy as np
-import pytest
-
-from Agent_module.carbon_data import (
- CarbonStore,
- LambdaEmbedder,
- by_paragraph,
- by_sentence,
- by_tokens,
- create,
-)
-
-
-@pytest.fixture
-def store(tmp_path: Path):
- p = tmp_path / "kb.carbondata"
- s = create(p)
- yield s
- if not s.closed:
- s.close()
-
-
-# ---------------------------------------------------------------------------
-# chunker primitives
-# ---------------------------------------------------------------------------
-
-class TestByTokens:
- def test_basic_window(self) -> None:
- chunk = by_tokens(max_tokens=3, overlap=0)
- out = chunk("a b c d e f g")
- assert out == ["a b c", "d e f", "g"]
-
- def test_overlap(self) -> None:
- chunk = by_tokens(max_tokens=3, overlap=1)
- out = chunk("a b c d e")
- # step = 2; windows: [a b c], [c d e]
- assert out == ["a b c", "c d e"]
-
- def test_empty_input(self) -> None:
- assert by_tokens(max_tokens=10)("") == []
- assert by_tokens(max_tokens=10)(" \n ") == []
-
- def test_single_token(self) -> None:
- assert by_tokens(max_tokens=10)("solo") == ["solo"]
-
- def test_validation(self) -> None:
- with pytest.raises(ValueError, match="max_tokens"):
- by_tokens(max_tokens=0)
- with pytest.raises(ValueError, match="overlap must be >= 0"):
- by_tokens(max_tokens=10, overlap=-1)
- with pytest.raises(ValueError, match="overlap must be < max_tokens"):
- by_tokens(max_tokens=10, overlap=10)
-
- def test_no_double_terminal_chunk(self) -> None:
- # Boundary case: when N is a multiple of step, must not emit empty tail.
- chunk = by_tokens(max_tokens=2, overlap=0)
- out = chunk("a b c d")
- assert out == ["a b", "c d"]
-
-
-class TestByParagraph:
- def test_blank_line_split(self) -> None:
- text = "Para one.\n\nPara two.\n\n\nPara three."
- out = by_paragraph()(text)
- assert out == ["Para one.", "Para two.", "Para three."]
-
- def test_min_chars(self) -> None:
- text = "short\n\nlonger paragraph"
- out = by_paragraph(min_chars=10)(text)
- assert out == ["longer paragraph"]
-
- def test_no_paragraphs(self) -> None:
- assert by_paragraph()("") == []
- assert by_paragraph()("just one line") == ["just one line"]
-
-
-class TestBySentence:
- def test_split_on_punctuation(self) -> None:
- out = by_sentence()(
- "Hello world. This is a test! Is it working? Yes."
- )
- assert out == [
- "Hello world.",
- "This is a test!",
- "Is it working?",
- "Yes.",
- ]
-
- def test_lowercase_continuation_kept_together(self) -> None:
- # The naive splitter requires whitespace+capital; "1.5 inches" stays
- # in one sentence.
- out = by_sentence()("It is 1.5 inches long.")
- assert out == ["It is 1.5 inches long."]
-
-
-# ---------------------------------------------------------------------------
-# ingest_text
-# ---------------------------------------------------------------------------
-
-class TestIngestText:
- def test_default_chunker_paragraphs(self, store: CarbonStore) -> None:
- eid = store.ingest_text("Para A.\n\nPara B.\n\nPara C.")
- chunks = store.list_chunks(eid)
- assert [c.content for c in chunks] == ["Para A.", "Para B.", "Para C."]
-
- def test_returns_id(self, store: CarbonStore) -> None:
- eid = store.ingest_text("hi", id="doc-1")
- assert eid == "doc-1"
- assert store.get_entity("doc-1") is not None
-
- def test_auto_id(self, store: CarbonStore) -> None:
- eid = store.ingest_text("hi")
- assert isinstance(eid, str) and len(eid) >= 16
-
- def test_entity_carries_full_text_and_metadata(self, store: CarbonStore) -> None:
- eid = store.ingest_text(
- "Para A.\n\nPara B.",
- id="d", metadata={"author": "x"}, namespace="ns1",
- )
- ent = store.get_entity(eid)
- assert ent.content == "Para A.\n\nPara B."
- assert ent.metadata == {"author": "x"}
- assert ent.namespace == "ns1"
- assert ent.kind == "document"
-
- def test_custom_chunker(self, store: CarbonStore) -> None:
- eid = store.ingest_text(
- "a b c d e f", chunker=by_tokens(max_tokens=2, overlap=0)
- )
- assert [c.content for c in store.list_chunks(eid)] == ["a b", "c d", "e f"]
-
- def test_falls_back_to_whole_text_when_chunker_empty(
- self, store: CarbonStore
- ) -> None:
- # by_paragraph yields nothing on a single line; fallback should
- # store the whole text as one chunk so the entity isn't empty.
- eid = store.ingest_text("just one sentence with no breaks")
- assert [c.content for c in store.list_chunks(eid)] == [
- "just one sentence with no breaks"
- ]
-
- def test_pure_whitespace_yields_no_chunks(self, store: CarbonStore) -> None:
- eid = store.ingest_text(" \n ")
- # Entity exists; chunks is empty (no usable content to index).
- assert store.get_entity(eid) is not None
- assert store.list_chunks(eid) == []
-
- def test_with_embedder(self, store: CarbonStore) -> None:
- def _fn(texts):
- return np.ones((len(texts), 4), dtype=np.float32)
- emb = LambdaEmbedder(_fn, model="m8", dim=4)
-
- eid = store.ingest_text(
- "Hello.\n\nWorld.", embedder=emb,
- )
- chunks = store.list_chunks(eid)
- assert len(chunks) == 2
- for c in chunks:
- assert store.get_embedding(c.id, model="m8") is not None
-
- def test_atomic_rollback(self, store: CarbonStore) -> None:
- class BoomEmbedder:
- model = "boom"
- dim = 3
-
- def encode(self, texts):
- raise RuntimeError("boom")
-
- before = store.stats()
- with pytest.raises(RuntimeError, match="boom"):
- store.ingest_text("a\n\nb", embedder=BoomEmbedder())
- after = store.stats()
- assert after["entities"] == before["entities"]
- assert after["chunks"] == before["chunks"]
-
-
-# ---------------------------------------------------------------------------
-# ingest_table
-# ---------------------------------------------------------------------------
-
-class TestIngestTable:
- def test_list_of_dicts(self, store: CarbonStore) -> None:
- rows = [
- {"id": "u1", "name": "Alice", "team": "infra"},
- {"id": "u2", "name": "Bob", "team": "ml"},
- ]
- ids = store.ingest_table(rows, table_name="users", id_column="id")
- assert ids == ["u1", "u2"]
-
- e = store.get_entity("u1")
- assert e.kind == "table_row"
- assert e.metadata["table"] == "users"
- assert e.metadata["columns"] == rows[0]
- # Rendered chunk should embed the row content for FTS / vector retrieval.
- chunks = store.list_chunks("u1")
- assert len(chunks) == 1
- assert "Alice" in chunks[0].content
- assert "infra" in chunks[0].content
-
- def test_auto_uuid_when_no_id_column(self, store: CarbonStore) -> None:
- rows = [{"name": "Alice"}, {"name": "Bob"}]
- ids = store.ingest_table(rows, table_name="users")
- assert len(ids) == 2
- assert all(len(i) >= 16 for i in ids)
-
- def test_id_column_missing_falls_back_to_uuid(
- self, store: CarbonStore
- ) -> None:
- # First row has id, second row doesn't — second should auto-uuid
- # rather than error or skip.
- rows = [{"id": "u1", "x": 1}, {"x": 2}]
- ids = store.ingest_table(rows, table_name="t", id_column="id")
- assert ids[0] == "u1"
- assert ids[1] != "u1" and len(ids[1]) >= 16
-
- def test_csv_file(self, tmp_path: Path, store: CarbonStore) -> None:
- csv_path = tmp_path / "users.csv"
- csv_path.write_text(
- "id,name,team\nu1,Alice,infra\nu2,Bob,ml\n",
- encoding="utf-8",
- )
- ids = store.ingest_table(csv_path, table_name="users", id_column="id")
- assert ids == ["u1", "u2"]
- e = store.get_entity("u1")
- assert e.metadata["columns"] == {"id": "u1", "name": "Alice", "team": "infra"}
-
- def test_csv_path_as_string(self, tmp_path: Path, store: CarbonStore) -> None:
- csv_path = tmp_path / "users.csv"
- csv_path.write_text("id,name\nu1,Alice\n", encoding="utf-8")
- ids = store.ingest_table(str(csv_path), table_name="users", id_column="id")
- assert ids == ["u1"]
-
- def test_empty_rows(self, store: CarbonStore) -> None:
- assert store.ingest_table([], table_name="t") == []
-
- def test_with_embedder(self, store: CarbonStore) -> None:
- rows = [{"id": "u1", "x": "alpha"}, {"id": "u2", "x": "beta"}]
-
- def _fn(texts):
- return np.ones((len(texts), 4), dtype=np.float32)
- emb = LambdaEmbedder(_fn, model="m8t", dim=4)
-
- ids = store.ingest_table(
- rows, table_name="t", id_column="id", embedder=emb
- )
- for eid in ids:
- chunks = store.list_chunks(eid)
- assert store.get_embedding(chunks[0].id, model="m8t") is not None
-
- def test_namespace_passes_through(self, store: CarbonStore) -> None:
- store.ingest_table(
- [{"id": "u1", "x": 1}], table_name="t", id_column="id",
- namespace="tenant-a",
- )
- assert store.get_entity("u1").namespace == "tenant-a"
-
- def test_atomic_rollback(self, store: CarbonStore) -> None:
- class BoomEmbedder:
- model = "boom"
- dim = 3
-
- def encode(self, texts):
- raise RuntimeError("boom")
-
- before = store.stats()
- with pytest.raises(RuntimeError, match="boom"):
- store.ingest_table(
- [{"id": "u1", "x": 1}], table_name="t", id_column="id",
- embedder=BoomEmbedder(),
- )
- after = store.stats()
- assert after["entities"] == before["entities"]
- assert after["chunks"] == before["chunks"]
diff --git a/Agent_module/tests/carbon_data/test_m9_admin.py b/Agent_module/tests/carbon_data/test_m9_admin.py
deleted file mode 100644
index 33dae5266e..0000000000
--- a/Agent_module/tests/carbon_data/test_m9_admin.py
+++ /dev/null
@@ -1,258 +0,0 @@
-"""M9: management tools — validate, compact, export."""
-from __future__ import annotations
-
-import base64
-import json
-from pathlib import Path
-
-import numpy as np
-import pytest
-
-from Agent_module.carbon_data import (
- APPLICATION_ID,
- CarbonStore,
- LambdaEmbedder,
- ValidationReport,
- create,
- open as cd_open,
-)
-
-
-@pytest.fixture
-def store(tmp_path: Path):
- p = tmp_path / "kb.carbondata"
- s = create(p)
- yield s
- if not s.closed:
- s.close()
-
-
-def _populate(s: CarbonStore) -> None:
- """Drop a small but multi-feature corpus into the store."""
- s.put_entity(id="d1", kind="document", content="hello world")
- s.add_chunks("d1", ["hello world"])
- s.put_entity(id="d2", kind="document", content="another")
- s.add_chunks("d2", ["another"])
- s.add_relation("d1", "d2", "references")
-
- def _fn(texts):
- return np.ones((len(texts), 4), dtype=np.float32)
- emb = LambdaEmbedder(_fn, model="m9", dim=4)
- s.embed_chunks(emb)
- s.remember("user likes vim", session_id="s1")
-
-
-# ---------------------------------------------------------------------------
-# validate
-# ---------------------------------------------------------------------------
-
-class TestValidateClean:
- def test_clean_db_passes(self, store: CarbonStore) -> None:
- _populate(store)
- r = store.validate()
- assert isinstance(r, ValidationReport)
- assert r.ok is True
- assert r.issues == []
- assert r.counts["entity"] == 3 # d1, d2, memory entity
- assert r.counts["chunk"] == 3
- assert r.counts["embedding"] == 2
- assert r.counts["relation"] == 1
- assert r.counts["memory_ext"] == 1
-
- def test_empty_db_passes(self, store: CarbonStore) -> None:
- r = store.validate()
- assert r.ok is True
- assert r.counts["entity"] == 0
-
- def test_validate_does_not_leave_transaction_open(
- self, store: CarbonStore
- ) -> None:
- _populate(store)
- assert store._require_open().in_transaction is False
- assert store.validate().ok is True
- assert store._require_open().in_transaction is False
-
- def test_readonly_validate_passes(self, tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- s = create(p)
- _populate(s)
- s.close()
-
- readonly = cd_open(p, mode="r")
- try:
- report = readonly.validate()
- assert report.ok is True
- assert report.issues == []
- finally:
- readonly.close()
-
-
-class TestValidateDetectsIssues:
- def test_detects_stale_hnsw_meta(self, tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- s = create(p)
- try:
- _populate(s)
- # Trigger HNSW build + persist.
- s.search(np.ones(4, dtype=np.float32), model="m9", top_k=1)
- # Tamper with persisted count.
- s._require_rw().execute(
- "UPDATE vector_index SET index_meta=? WHERE model='m9'",
- ('{"count": 999, "dim": 4}',),
- )
- s._require_rw().commit()
- r = s.validate()
- assert r.ok is False
- assert any("vector_index[m9]" in msg and "stale" in msg
- for msg in r.issues)
- finally:
- s.close()
-
- def test_detects_corrupt_index_meta(self, tmp_path: Path) -> None:
- p = tmp_path / "kb.carbondata"
- s = create(p)
- try:
- _populate(s)
- s.search(np.ones(4, dtype=np.float32), model="m9", top_k=1)
- s._require_rw().execute(
- "UPDATE vector_index SET index_meta=? WHERE model='m9'",
- ("not a json",),
- )
- s._require_rw().commit()
- r = s.validate()
- assert r.ok is False
- assert any("corrupt index_meta" in msg for msg in r.issues)
- finally:
- s.close()
-
-
-# ---------------------------------------------------------------------------
-# compact
-# ---------------------------------------------------------------------------
-
-class TestCompact:
- def test_returns_size_pair(self, store: CarbonStore) -> None:
- _populate(store)
- result = store.compact()
- assert "size_before" in result and "size_after" in result
- assert result["size_before"] >= 0
- assert result["size_after"] >= 0
-
- def test_clears_persisted_hnsw_blob(self, store: CarbonStore) -> None:
- _populate(store)
- # Force persistence.
- store.search(np.ones(4, dtype=np.float32), model="m9", top_k=1)
- assert store._require_open().execute(
- "SELECT COUNT(*) FROM vector_index"
- ).fetchone()[0] == 1
- store.compact()
- assert store._require_open().execute(
- "SELECT COUNT(*) FROM vector_index"
- ).fetchone()[0] == 0
- assert "m9" not in store._hnsw_cache
-
- def test_rebuild_recovers_search(self, store: CarbonStore) -> None:
- _populate(store)
- store.compact()
- # After compact, search should still work — HNSW rebuilt lazily.
- hits = store.search(np.ones(4, dtype=np.float32), model="m9", top_k=2)
- assert len(hits) == 2
-
- def test_fts_still_works_after_compact(self, store: CarbonStore) -> None:
- _populate(store)
- store.compact()
- hits = store.search("hello", mode="keyword")
- assert len(hits) >= 1
- assert any("hello" in h.chunk.content for h in hits)
-
- def test_rejected_inside_transaction(self, store: CarbonStore) -> None:
- with pytest.raises(RuntimeError, match="transaction"):
- with store.transaction():
- store.compact()
-
-
-# ---------------------------------------------------------------------------
-# export
-# ---------------------------------------------------------------------------
-
-class TestExport:
- def test_writes_json_file_with_expected_shape(
- self, store: CarbonStore, tmp_path: Path
- ) -> None:
- _populate(store)
- out = tmp_path / "dump.json"
- report = store.export(out)
- assert out.exists()
- assert report["bytes_written"] == out.stat().st_size
-
- data = json.loads(out.read_text(encoding="utf-8"))
- assert data["schema_version"] == 1
- assert {"entities", "chunks", "embeddings", "relations",
- "memory_ext"} <= set(data.keys())
- assert len(data["entities"]) == report["entities"]
- assert len(data["chunks"]) == report["chunks"]
- assert len(data["relations"]) == report["relations"]
-
- def test_embeddings_round_trip_via_base64(
- self, store: CarbonStore, tmp_path: Path
- ) -> None:
- _populate(store)
- out = tmp_path / "dump.json"
- store.export(out)
- data = json.loads(out.read_text(encoding="utf-8"))
-
- for emb_record in data["embeddings"]:
- raw = base64.b64decode(emb_record["vector_b64"])
- arr = np.frombuffer(raw, dtype=np.float32)
- assert arr.shape == (emb_record["dim"],)
- # _populate seeds all-ones vectors.
- assert np.allclose(arr, 1.0)
-
- def test_include_embeddings_false_omits_them(
- self, store: CarbonStore, tmp_path: Path
- ) -> None:
- _populate(store)
- out = tmp_path / "dump.json"
- store.export(out, include_embeddings=False)
- data = json.loads(out.read_text(encoding="utf-8"))
- assert data["embeddings"] == []
-
- def test_export_then_reload_matches_counts(
- self, store: CarbonStore, tmp_path: Path
- ) -> None:
- _populate(store)
- out = tmp_path / "dump.json"
- store.export(out)
- data = json.loads(out.read_text(encoding="utf-8"))
-
- # Re-open the original; counts should match what we exported.
- path = store.path
- store.close()
- s2 = cd_open(path)
- try:
- stats = s2.stats()
- assert len(data["entities"]) == stats["entities"]
- assert len(data["chunks"]) == stats["chunks"]
- assert len(data["embeddings"]) == stats["embeddings"]
- assert len(data["relations"]) == stats["relations"]
- assert len(data["memory_ext"]) == stats["memories"]
- finally:
- s2.close()
-
- def test_metadata_preserved(
- self, store: CarbonStore, tmp_path: Path
- ) -> None:
- store.put_entity(
- id="d", kind="document",
- metadata={"author": "Alice", "tags": ["x", "y"]},
- )
- out = tmp_path / "dump.json"
- store.export(out)
- data = json.loads(out.read_text(encoding="utf-8"))
- ent = next(e for e in data["entities"] if e["id"] == "d")
- assert ent["metadata"] == {"author": "Alice", "tags": ["x", "y"]}
-
- def test_application_id_constant_exported(self) -> None:
- # Sanity: confirm the carbondata signature still matches what the
- # spec promises (used as a doc/regression anchor).
- assert APPLICATION_ID == 0x4342_4E44
diff --git a/Agent_module/tests/conftest.py b/Agent_module/tests/conftest.py
deleted file mode 100644
index a97804a619..0000000000
--- a/Agent_module/tests/conftest.py
+++ /dev/null
@@ -1,10 +0,0 @@
-"""Make the Agent_module package importable when running pytest from the repo root."""
-
-import os
-import sys
-
-# Add the repo root (parent of Agent_module/) to sys.path so `import Agent_module` works
-# regardless of where pytest is invoked from.
-_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-if _REPO_ROOT not in sys.path:
- sys.path.insert(0, _REPO_ROOT)
(carbondata) 01/02: update Agent module
Posted by ch...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
asf-gitbox-commits pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/carbondata.git
commit 58eaa8550b4ec6a3e6033d7cc0931dfd0afeb186
Author: chenliang613 <ch...@apache.org>
AuthorDate: Thu Jul 30 15:28:16 2026 +0800
update Agent module
---
.gitignore | 3 +
Agent_module/QUICKSTART.md | 37 +++---
Agent_module/README.md | 48 +++++---
Agent_module/carbon_data/store.py | 157 ++++++++++++++++--------
Agent_module/pyproject.toml | 40 ++++++
Agent_module/tests/carbon_data/test_m2_crud.py | 17 +++
Agent_module/tests/carbon_data/test_m5_hnsw.py | 95 ++++++++++++++
Agent_module/tests/carbon_data/test_m9_admin.py | 22 ++++
8 files changed, 342 insertions(+), 77 deletions(-)
diff --git a/.gitignore b/.gitignore
index 363455e06c..0475aceee3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,6 +7,9 @@
*.iws
*.pyc
*.pyo
+.pytest_cache/
+.venv/
+*.egg-info/
.idea/
.idea_modules/
.settings
diff --git a/Agent_module/QUICKSTART.md b/Agent_module/QUICKSTART.md
index f58fbe03b1..c59e441b65 100644
--- a/Agent_module/QUICKSTART.md
+++ b/Agent_module/QUICKSTART.md
@@ -1,10 +1,15 @@
# carbon_data Quickstart
-`Agent_module.carbon_data` is the Agent Data Infra: one `.carbondata` file =
-one self-contained knowledge base, simultaneously serving **RAG semantic
-search / long-term memory / structured queries / knowledge-graph traversal**.
+`Agent_module.carbon_data` is an Agent-native context store: one
+`.carbondata` file persists context that multiple Agents can write, retrieve,
+and hand off. It combines **semantic context retrieval / cross-agent memory /
+structured state queries / context relationship traversal**.
> Companion runnable demo: [`examples/carbondata_quickstart.py`](../examples/carbondata_quickstart.py)
+>
+> The Agent module stores SQLite data. It is not binary-compatible with
+> classic Apache CarbonData columnar fact files, even though both currently
+> use the `.carbondata` suffix.
---
@@ -20,13 +25,16 @@ Optional dependencies are loaded on demand:
| `hnswlib` | fast ANN for 10k+ vectors | optional (falls back to brute force when missing) |
```bash
-pip install numpy
-pip install hnswlib # optional
+# Run from the repository root.
+python3 -m pip install -e "./Agent_module[test]"
+python3 -m pip install -e "./Agent_module[all]" # optional HNSW support too
```
A `.carbondata` file **is just a SQLite database**, so you can crack it open
with any SQLite tool (`sqlite3 kb.carbondata`, DBeaver, DataGrip) for ad-hoc
-debugging.
+debugging. WAL mode can create temporary `-wal` and `-shm` sidecars while
+the database is open, so copy or back up a live store using SQLite-aware
+tools rather than copying only the primary file.
---
@@ -193,7 +201,8 @@ acceleration whenever:
- `mode="vector"` (including the vector leg of `mode="hybrid"`)
- `metric="cosine"`
- no JSON `filters=` are present
-- `namespace=` is applied as a post-filter (with 3× over-sampling)
+- `namespace=` is applied as an adaptive post-filter that expands until
+ `top_k` is satisfied or the full index has been considered
Explicit override:
```python
@@ -203,8 +212,8 @@ store.search(q, embedder=emb, use_hnsw="on") # force HNSW (raises if not elig
```
On the first search, HNSW is built from the entire `embedding` table and
-persisted to a `vector_index` blob; later add/delete calls are detected via
-a row-count comparison and trigger a lazy rebuild.
+persisted to a `vector_index` blob. API mutations explicitly invalidate the
+affected index; row-count checks also detect out-of-band adds and deletes.
### 4.3 Namespace isolation
@@ -226,7 +235,9 @@ with store.transaction():
```
`ingest_text` / `ingest_table` / `remember` are already wrapped in their own
-transactions internally.
+transactions internally. Nested `transaction()` calls use savepoints: an
+inner exception rolls back the inner block, while an outer transaction may
+continue if it catches that exception.
### 4.5 Admin & ops
@@ -260,8 +271,7 @@ print(store.stats())
— each file covers one milestone end-to-end.
- **Run the suite:**
```bash
- cd Agent_module
- .venv/bin/pytest tests/carbon_data/ -q
+ python3 -m pytest Agent_module/tests/carbon_data/ -q
```
---
@@ -272,8 +282,7 @@ print(store.stats())
Run the demo to see it all in action:
```bash
-cd Agent_module
-.venv/bin/python examples/carbondata_quickstart.py
+python3 Agent_module/examples/carbondata_quickstart.py
```
For the full API surface, read the `class CarbonStore` docstring in
diff --git a/Agent_module/README.md b/Agent_module/README.md
index 9ff2bc8e76..957baa34b6 100644
--- a/Agent_module/README.md
+++ b/Agent_module/README.md
@@ -1,20 +1,31 @@
-# Agent Module — Agent Data Infra
+# Agent Module — Agent-Native Context Data
-`Agent_module` is a Python library that gives an LLM agent **a single
-knowledge file** to read from. That file — with the `.carbondata` extension
-— is a self-contained SQLite database that simultaneously serves the four
-data-access patterns agents actually need:
+`Agent_module` provides a data format and local store designed specifically
+for Agent access. Its primary job is to persist and share context across
+Agents — including memories, task artifacts, structured state, semantic
+fragments, and the relationships among them — so another Agent can recover
+the right context efficiently instead of replaying an entire conversation.
+
+The `.carbondata` context file is a self-contained SQLite database serving
+four complementary access patterns:
| Pattern | What an agent calls it for |
|---|---|
-| **RAG semantic search** | "Find chunks relevant to this question." |
-| **Long-term memory** | "What have I learned about this user / session before?" |
-| **Structured query** | "Give me all rows where `team = 'ml'`." |
-| **Knowledge-graph traversal** | "Which documents does this one reference, two hops out?" |
+| **Semantic context retrieval** | "Find the prior context relevant to this task." |
+| **Cross-agent memory** | "What did another Agent learn or decide earlier?" |
+| **Structured state query** | "Give me the state for this task/session/actor." |
+| **Context relationship traversal** | "Which artifacts, decisions, and tasks led here?" |
One file. One Python handle (`CarbonStore`). One schema. No separate vector
DB, KV store, graph DB, or relational DB to wire together.
+> **Format boundary:** the Agent module's `.carbondata` file is an
+> application-identified SQLite database. It is not the same binary format
+> as the classic Apache CarbonData columnar fact files that also use the
+> `.carbondata` suffix, and the two are not interchangeable. Treat the Agent
+> format as an experimental local-store format while its long-term naming
+> and interoperability contract is defined.
+
```
┌─────────────────────────────┐
agent code ───► │ CarbonStore (kb.carbondata)│
@@ -43,7 +54,7 @@ Agent_module/
├── QUICKSTART.md # full hands-on guide (read this next)
├── examples/
│ └── carbondata_quickstart.py # end-to-end runnable demo
-└── tests/carbon_data/ # 249 tests across M1–M9
+└── tests/carbon_data/ # 250+ tests across M1–M9
```
---
@@ -52,13 +63,15 @@ Agent_module/
- Python 3.10+
- `numpy` (required)
-- `hnswlib` (optional — enables ANN acceleration when corpus > ~10k chunks;
- falls back to brute force when missing)
+- `hnswlib` (optional — enables ANN acceleration and is most useful for
+ larger corpora; falls back to brute force when missing)
- `pytest` (optional — for running the test suite)
+From the repository root, install the module and its test dependencies:
+
```bash
-pip install numpy
-pip install hnswlib # optional but recommended
+python3 -m pip install -e "./Agent_module[test]"
+python3 -m pip install -e "./Agent_module[all]" # also install optional HNSW support
```
---
@@ -132,6 +145,11 @@ You now have a working `kb.carbondata` file on disk. Open it again in
another script with `carbon_data.open("kb.carbondata")` — all your
entities, embeddings, relations, and memories will still be there.
+SQLite WAL mode may create temporary `-wal` and `-shm` sidecar files while
+the store is open. A clean close normally checkpoints the database back to
+the primary file; applications must not copy a live store without its WAL
+state.
+
### Step 3 — go deeper
[**QUICKSTART.md**](QUICKSTART.md) covers each scenario in detail:
@@ -147,7 +165,7 @@ For the full API surface, read the `class CarbonStore` docstring in
## Running the tests
```bash
-pytest Agent_module/tests/ -v
+python3 -m pytest Agent_module/tests/ -v
```
Tests are split by milestone (M1 schema → M9 admin); each file is a
diff --git a/Agent_module/carbon_data/store.py b/Agent_module/carbon_data/store.py
index 55cf4729ab..8f70ddfc50 100644
--- a/Agent_module/carbon_data/store.py
+++ b/Agent_module/carbon_data/store.py
@@ -110,7 +110,8 @@ class CarbonStore:
self._mode: Mode = mode
self._tx_depth = 0 # 0 = autocommit-per-write; >0 = user transaction
# Per-model HNSW index cache. Populated lazily on first vector
- # search; invalidated whenever embedding count diverges.
+ # search; mutations explicitly invalidate affected models while
+ # count checks protect against out-of-band adds/deletes.
self._hnsw_cache: dict[str, HnswIndex] = {}
# ------------------------------------------------------------------
@@ -214,20 +215,33 @@ class CarbonStore:
"""
Group writes atomically.
- Nested calls are a no-op at the inner level: only the outermost
- context commits/rolls back. On any exception, the outer frame
- rolls back and re-raises.
+ Nested calls use SQLite savepoints. An exception rolls back the
+ current nesting level and is re-raised; if the caller catches an
+ inner exception, the outer transaction may continue safely.
"""
conn = self._require_rw()
outer = self._tx_depth == 0
+ savepoint = f"carbondata_tx_{self._tx_depth}"
+ if outer:
+ conn.execute("BEGIN")
+ else:
+ conn.execute(f"SAVEPOINT {savepoint}")
self._tx_depth += 1
try:
yield
if outer:
conn.commit()
+ else:
+ conn.execute(f"RELEASE SAVEPOINT {savepoint}")
except Exception:
if outer:
conn.rollback()
+ else:
+ conn.execute(f"ROLLBACK TO SAVEPOINT {savepoint}")
+ conn.execute(f"RELEASE SAVEPOINT {savepoint}")
+ # A search inside the failed scope may have built an in-memory
+ # index from data that no longer exists after rollback.
+ self._hnsw_cache.clear()
raise
finally:
self._tx_depth -= 1
@@ -336,9 +350,15 @@ class CarbonStore:
"""Remove an entity and (via ON DELETE CASCADE) its chunks,
embeddings, relations, and memory_ext row. Returns True iff
something was deleted."""
+ deleted = False
with self._write_ctx() as conn:
cur = conn.execute("DELETE FROM entity WHERE id=?", (entity_id,))
- return cur.rowcount > 0
+ deleted = cur.rowcount > 0
+ if deleted:
+ conn.execute("DELETE FROM vector_index")
+ if deleted:
+ self._hnsw_cache.clear()
+ return deleted
# ------------------------------------------------------------------
# Chunk CRUD
@@ -420,7 +440,12 @@ class CarbonStore:
"""Delete all chunks for an entity. Returns the count removed."""
with self._write_ctx() as conn:
cur = conn.execute("DELETE FROM chunk WHERE entity_id=?", (entity_id,))
- return cur.rowcount
+ deleted = cur.rowcount
+ if deleted:
+ conn.execute("DELETE FROM vector_index")
+ if deleted:
+ self._hnsw_cache.clear()
+ return deleted
# ------------------------------------------------------------------
# Structured query
@@ -505,6 +530,8 @@ class CarbonStore:
(chunk_id, model, dim, vector) VALUES (?, ?, ?, ?)""",
(chunk_id, model, dim, blob),
)
+ conn.execute("DELETE FROM vector_index WHERE model=?", (model,))
+ self._hnsw_cache.pop(model, None)
def get_embedding(self, chunk_id: str, *, model: str) -> Optional[Any]:
"""Return the stored vector (numpy array) or None."""
@@ -585,6 +612,13 @@ class CarbonStore:
payload,
)
total += len(batch)
+ if total:
+ write_conn.execute(
+ "DELETE FROM vector_index WHERE model=?",
+ (embedder.model,),
+ )
+ if total:
+ self._hnsw_cache.pop(embedder.model, None)
return total
# ------------------------------------------------------------------
@@ -631,6 +665,8 @@ class CarbonStore:
'on' forces HNSW (errors if filters/non-cosine/no
hnswlib); 'off' always uses brute force.
"""
+ if top_k <= 0:
+ return []
if mode == "vector":
ranked = self._run_vector(
query, top_k=top_k, model=model, embedder=embedder,
@@ -764,37 +800,46 @@ class CarbonStore:
f"for model {model!r}"
)
- # Over-fetch when post-filtering by namespace so we can survive
- # candidates that get dropped. 3x is plenty for typical workloads.
- oversample = 3 if namespace is not None else 1
- raw = idx.search(q_vec, top_k=top_k * oversample)
- if not raw:
- return []
-
- rowids = [r[0] for r in raw]
- placeholders = ",".join("?" * len(rowids))
- params: list[Any] = list(rowids)
- sql = (
- f"SELECT chunk.rowid, chunk.id "
- f"FROM chunk JOIN entity ON entity.id = chunk.entity_id "
- f"WHERE chunk.rowid IN ({placeholders})"
+ conn = self._require_open()
+ fetch_k = min(
+ idx.count,
+ max(top_k, top_k * 3 if namespace is not None else top_k),
)
- if namespace is not None:
- sql += " AND entity.namespace = ?"
- params.append(namespace)
- conn = self._require_open()
- rid_to_cid = {r[0]: r[1] for r in conn.execute(sql, params).fetchall()}
+ while fetch_k > 0:
+ raw = idx.search(q_vec, top_k=fetch_k)
+ if not raw:
+ return []
- out: list[tuple[str, float]] = []
- for rid, score in raw:
- cid = rid_to_cid.get(rid)
- if cid is None:
- continue
- out.append((cid, score))
- if len(out) >= top_k:
- break
- return out
+ rowids = [r[0] for r in raw]
+ placeholders = ",".join("?" * len(rowids))
+ params: list[Any] = list(rowids)
+ sql = (
+ f"SELECT chunk.rowid, chunk.id "
+ f"FROM chunk JOIN entity ON entity.id = chunk.entity_id "
+ f"WHERE chunk.rowid IN ({placeholders})"
+ )
+ if namespace is not None:
+ sql += " AND entity.namespace = ?"
+ params.append(namespace)
+
+ rid_to_cid = {
+ r[0]: r[1] for r in conn.execute(sql, params).fetchall()
+ }
+ out = [
+ (rid_to_cid[rid], score)
+ for rid, score in raw
+ if rid in rid_to_cid
+ ]
+ if len(out) >= top_k or fetch_k >= idx.count:
+ return out[:top_k]
+
+ # A sparse namespace may not occur in the first 3× candidates.
+ # Expand progressively until top_k is satisfied or the complete
+ # index has been considered.
+ fetch_k = min(idx.count, max(fetch_k * 2, top_k))
+
+ return []
# ---- HNSW cache + persistence ------------------------------------
@@ -810,8 +855,8 @@ class CarbonStore:
Return a usable HNSW for ``model``, building or reloading as needed.
Cache invariant: ``HnswIndex.count`` equals the live embedding
- count. Any divergence — adds, deletes, or first use after a
- reopen — triggers a rebuild from the embedding table.
+ count. Adds/deletes are caught by the count check; in-process vector
+ replacements explicitly invalidate the affected model.
"""
cur_count = self._embedding_count(model)
if cur_count == 0:
@@ -1220,6 +1265,8 @@ class CarbonStore:
conn.executemany(
"DELETE FROM entity WHERE id=?", [(i,) for i in ids]
)
+ conn.execute("DELETE FROM vector_index")
+ self._hnsw_cache.clear()
return len(ids)
def _load_memory(self, entity_id: str) -> Optional[MemoryItem]:
@@ -1617,13 +1664,13 @@ class CarbonStore:
def validate(self) -> ValidationReport:
"""
- Run consistency checks. Read-only; safe on a r-mode handle.
+ Run consistency checks. Safe on a read-only handle.
Catches:
- SQLite page-level corruption (PRAGMA integrity_check)
- Foreign-key violations (PRAGMA foreign_key_check)
- Bad application_id / unsupported schema_version
- - chunk vs chunk_fts row-count divergence
+ - FTS5 index corruption (deep check on writable handles)
- Stale per-model HNSW blob (count != live embedding count)
"""
conn = self._require_open()
@@ -1660,16 +1707,30 @@ class CarbonStore:
f"SELECT COUNT(*) FROM {table}"
).fetchone()[0]
- # FTS5 'integrity-check' compares index entries against the
- # external content table (chunk). It surfaces missing index rows,
- # stale tokens, and corruption — anything a row-count check
- # cannot, since external-content COUNT(*) is served from chunk.
- try:
- conn.execute(
- "INSERT INTO chunk_fts(chunk_fts) VALUES('integrity-check')"
- )
- except sqlite3.DatabaseError as exc:
- issues.append(f"chunk_fts integrity-check: {exc}")
+ if self._mode == "r":
+ # FTS5 exposes its deep integrity check as a special INSERT,
+ # which SQLite rejects on a read-only connection. Still verify
+ # that the virtual table is readable without producing a false
+ # corruption report for a healthy read-only store.
+ try:
+ conn.execute("SELECT rowid FROM chunk_fts LIMIT 1").fetchall()
+ except sqlite3.DatabaseError as exc:
+ issues.append(f"chunk_fts read-check: {exc}")
+ else:
+ # FTS5 'integrity-check' compares index entries against the
+ # external content table. The control command opens a SQLite
+ # transaction even though it does not change user data, so end
+ # that transaction when validate() started it.
+ was_in_transaction = conn.in_transaction
+ try:
+ conn.execute(
+ "INSERT INTO chunk_fts(chunk_fts) VALUES('integrity-check')"
+ )
+ except sqlite3.DatabaseError as exc:
+ issues.append(f"chunk_fts integrity-check: {exc}")
+ finally:
+ if not was_in_transaction and conn.in_transaction:
+ conn.rollback()
# Per-model HNSW staleness — only complain when a blob exists.
idx_rows = conn.execute(
diff --git a/Agent_module/pyproject.toml b/Agent_module/pyproject.toml
new file mode 100644
index 0000000000..5940d6a5ae
--- /dev/null
+++ b/Agent_module/pyproject.toml
@@ -0,0 +1,40 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "apache-carbondata-agent"
+version = "0.1.0"
+description = "Agent-native context data format and local store for cross-agent collaboration"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "Apache-2.0" }
+dependencies = [
+ "numpy>=1.24",
+]
+
+[project.optional-dependencies]
+hnsw = [
+ "hnswlib>=0.8.0",
+]
+test = [
+ "pytest>=7.4",
+]
+all = [
+ "hnswlib>=0.8.0",
+ "pytest>=7.4",
+]
+
+[tool.setuptools]
+packages = [
+ "Agent_module",
+ "Agent_module.carbon_data",
+ "Agent_module.carbon_data.chunkers",
+ "Agent_module.carbon_data.embedders",
+ "Agent_module.carbon_data.index",
+]
+package-dir = { Agent_module = "." }
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-ra"
diff --git a/Agent_module/tests/carbon_data/test_m2_crud.py b/Agent_module/tests/carbon_data/test_m2_crud.py
index d9ce8b5ea0..82d0dd6473 100644
--- a/Agent_module/tests/carbon_data/test_m2_crud.py
+++ b/Agent_module/tests/carbon_data/test_m2_crud.py
@@ -309,6 +309,23 @@ class TestTransactions:
assert store.get_entity("a") is None
assert store.get_entity("b") is None
+ def test_caught_inner_exception_rolls_back_to_savepoint(
+ self, store: CarbonStore
+ ) -> None:
+ with store.transaction():
+ store.put_entity(id="a", kind="doc")
+ try:
+ with store.transaction():
+ store.put_entity(id="b", kind="doc")
+ raise RuntimeError("inner")
+ except RuntimeError:
+ pass
+ store.put_entity(id="c", kind="doc")
+
+ assert store.get_entity("a") is not None
+ assert store.get_entity("b") is None
+ assert store.get_entity("c") is not None
+
def test_chunks_added_in_transaction_all_or_nothing(
self, store: CarbonStore
) -> None:
diff --git a/Agent_module/tests/carbon_data/test_m5_hnsw.py b/Agent_module/tests/carbon_data/test_m5_hnsw.py
index f95788aa15..60bcb57030 100644
--- a/Agent_module/tests/carbon_data/test_m5_hnsw.py
+++ b/Agent_module/tests/carbon_data/test_m5_hnsw.py
@@ -292,6 +292,63 @@ class TestCacheInvalidation:
model=emb.model, top_k=1)
assert store._hnsw_cache[emb.model].count == 5
+ def test_replace_embedding_invalidates_cache_and_blob(
+ self, store: CarbonStore
+ ) -> None:
+ emb, _, _ = _seed(store, n=5, dim=4, seed=13)
+ store.search(np.ones(4, dtype=np.float32),
+ model=emb.model, top_k=1)
+ assert emb.model in store._hnsw_cache
+ assert _vector_index_count(store) == 1
+
+ first_chunk = store.list_chunks("e-default")[0]
+ store.set_embedding(
+ first_chunk.id,
+ np.array([9.0, 0.0, 0.0, 0.0], dtype=np.float32),
+ model=emb.model,
+ )
+
+ assert emb.model not in store._hnsw_cache
+ assert _vector_index_count(store) == 0
+
+ def test_reembed_invalidates_cache_and_blob(
+ self, store: CarbonStore
+ ) -> None:
+ emb, _, _ = _seed(store, n=5, dim=4, seed=14)
+ store.search(np.ones(4, dtype=np.float32),
+ model=emb.model, top_k=1)
+ assert emb.model in store._hnsw_cache
+ assert _vector_index_count(store) == 1
+
+ store.embed_chunks(emb, missing_only=False)
+
+ assert emb.model not in store._hnsw_cache
+ assert _vector_index_count(store) == 0
+
+ def test_rollback_clears_index_built_from_uncommitted_vectors(
+ self, store: CarbonStore
+ ) -> None:
+ emb, _, _ = _seed(store, n=5, dim=4, seed=15)
+ first_chunk = store.list_chunks("e-default")[0]
+
+ with pytest.raises(RuntimeError):
+ with store.transaction():
+ store.set_embedding(
+ first_chunk.id,
+ np.array([9.0, 0.0, 0.0, 0.0], dtype=np.float32),
+ model=emb.model,
+ )
+ store.search(
+ np.ones(4, dtype=np.float32),
+ model=emb.model,
+ top_k=1,
+ use_hnsw="on",
+ )
+ assert emb.model in store._hnsw_cache
+ raise RuntimeError("rollback")
+
+ assert emb.model not in store._hnsw_cache
+
def test_full_cleanup_drops_cache_entry(self, store: CarbonStore) -> None:
emb, _, _ = _seed(store, n=3, dim=4, seed=12)
store.search(np.zeros(4, dtype=np.float32),
@@ -329,3 +386,41 @@ class TestNamespacePostFilter:
assert all(h.entity.namespace == "ns-a" for h in hits)
finally:
s.close()
+
+ def test_sparse_namespace_expands_until_top_k(self, tmp_path: Path) -> None:
+ p = tmp_path / "kb.carbondata"
+ s = create(p)
+ try:
+ s.put_entity(id="target", kind="document", namespace="target")
+ target_chunks = s.add_chunks(
+ "target", [f"target-{i}" for i in range(4)]
+ )
+ s.put_entity(id="other", kind="document", namespace="other")
+ other_chunks = s.add_chunks(
+ "other", [f"other-{i}" for i in range(12)]
+ )
+ for chunk_id in target_chunks:
+ s.set_embedding(
+ chunk_id,
+ np.array([-1.0, 0.0], dtype=np.float32),
+ model="sparse-ns",
+ )
+ for chunk_id in other_chunks:
+ s.set_embedding(
+ chunk_id,
+ np.array([1.0, 0.0], dtype=np.float32),
+ model="sparse-ns",
+ )
+
+ hits = s.search(
+ np.array([1.0, 0.0], dtype=np.float32),
+ model="sparse-ns",
+ top_k=4,
+ namespace="target",
+ use_hnsw="on",
+ )
+
+ assert len(hits) == 4
+ assert all(hit.entity.namespace == "target" for hit in hits)
+ finally:
+ s.close()
diff --git a/Agent_module/tests/carbon_data/test_m9_admin.py b/Agent_module/tests/carbon_data/test_m9_admin.py
index 1f894f24a6..33dae5266e 100644
--- a/Agent_module/tests/carbon_data/test_m9_admin.py
+++ b/Agent_module/tests/carbon_data/test_m9_admin.py
@@ -64,6 +64,28 @@ class TestValidateClean:
assert r.ok is True
assert r.counts["entity"] == 0
+ def test_validate_does_not_leave_transaction_open(
+ self, store: CarbonStore
+ ) -> None:
+ _populate(store)
+ assert store._require_open().in_transaction is False
+ assert store.validate().ok is True
+ assert store._require_open().in_transaction is False
+
+ def test_readonly_validate_passes(self, tmp_path: Path) -> None:
+ p = tmp_path / "kb.carbondata"
+ s = create(p)
+ _populate(s)
+ s.close()
+
+ readonly = cd_open(p, mode="r")
+ try:
+ report = readonly.validate()
+ assert report.ok is True
+ assert report.issues == []
+ finally:
+ readonly.close()
+
class TestValidateDetectsIssues:
def test_detects_stale_hnsw_meta(self, tmp_path: Path) -> None: