feat: add ChromaDB sync hook

Author: scream4ik

memstate/integrations/chroma.py

@@ -0,0 +1,93 @@
+from typing import Any, Callable
+
+from ..constants import Operation
+from ..schemas import Fact
+
+try:
+    from chromadb import EmbeddingFunction
+    from chromadb.api import ClientAPI, Embeddable
+except ImportError:
+    raise ImportError("pip install chromadb")
+
+TextFormatter = Callable[[dict[str, Any]], str]
+MetadataFormatter = Callable[[dict[str, Any]], dict[str, Any]]
+
+
+class ChromaSyncHook:
+    def __init__(
+        self,
+        client: ClientAPI,
+        collection_name: str,
+        embedding_fn: EmbeddingFunction[Embeddable] | None = None,
+        target_types: set[str] | None = None,
+        text_field: str | None = None,
+        text_formatter: TextFormatter | None = None,
+        metadata_fields: list[str] | None = None,
+        metadata_formatter: MetadataFormatter | None = None,
+    ):
+        """
+        :param client: Initialized Chroma Client.
+        :param collection_name: collection name.
+        :param embedding_fn: Function (text -> vector). If None, Chroma uses default.
+        :param target_types: Types of facts for synchronization (to avoid garbage).
+        :param text_field: Field name of text in fact.
+        :param text_formatter: Function.
+        :param metadata_fields: Fields of metadata in fact.
+        :param metadata_formatter: Function.
+        """
+        self.client = client
+        self.collection = client.get_or_create_collection(
+            name=collection_name,
+            embedding_function=embedding_fn,
+        )
+        self.target_types = target_types or set()
+
+        if text_formatter is not None:
+            self._extract_text = text_formatter
+        elif text_field:
+            self._extract_text = lambda data: str(data.get(text_field, ""))
+        else:
+            self._extract_text = lambda data: str(data)
+
+        self.metadata_fields = metadata_fields or []
+        self.metadata_formatter = metadata_formatter
+
+    def _get_metadata(self, data: dict[str, Any]) -> dict[str, Any]:
+        if self.metadata_formatter is not None:
+            return self.metadata_formatter(data)
+
+        if self.metadata_fields:
+            meta = {}
+            for field in self.metadata_fields:
+                val = data.get(field)
+                if val is not None:
+                    if isinstance(val, (str, int, float, bool)):
+                        meta[field] = val
+                    else:
+                        meta[field] = str(val)
+            return meta
+
+        return {}
+
+    def __call__(self, op: Operation, fact_id: str, data: Fact | None) -> None:
+        if op == Operation.DELETE:
+            self.collection.delete(ids=[fact_id])
+            return
+
+        if op == Operation.DISCARD_SESSION:
+            return
+
+        if not data or (self.target_types and data.type not in self.target_types):
+            return
+
+        text = self._extract_text(data.payload)
+
+        if not text.strip():
+            return
+
+        if op in (Operation.COMMIT, Operation.UPDATE, Operation.COMMIT_EPHEMERAL, Operation.PROMOTE):
+            meta = {"type": data.type, "source": data.source or "", "ts": str(data.ts)}
+            metadata = self._get_metadata(data=data.payload)
+            meta.update(metadata)
+
+            self.collection.upsert(ids=[fact_id], documents=[text], metadatas=[meta])

memstate/schemas.py

@@ -0,0 +1,28 @@
+import uuid
+from datetime import datetime, timezone
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from .constants import Operation
+
+
+class Fact(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    type: str
+    payload: dict[str, Any]
+    source: str | None = None
+    session_id: str | None = None
+    ts: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+
+
+class TxEntry(BaseModel):
+    uuid: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    seq: int
+    ts: datetime
+    op: Operation
+    fact_id: str | None
+    fact_before: dict[str, Any] | None = None
+    fact_after: dict[str, Any] | None = None
+    actor: str | None = None
+    reason: str | None = None

memstate/storage.py

@@ -1,38 +1,16 @@
 import copy
 import threading
-import uuid
 from datetime import datetime, timezone
 from typing import Any, Callable
 
-from pydantic import BaseModel, Field, ValidationError
+from pydantic import BaseModel, ValidationError
 
 from .backends.base import StorageBackend
 from .constants import Operation
 from .exceptions import ConflictError, HookError, MemoryStoreError, ValidationFailed
+from .schemas import Fact, TxEntry
 
-
-class Fact(BaseModel):
-    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
-    type: str
-    payload: dict[str, Any]
-    source: str | None = None
-    session_id: str | None = None
-    ts: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
-
-
-class TxEntry(BaseModel):
-    uuid: str = Field(default_factory=lambda: str(uuid.uuid4()))
-    seq: int
-    ts: datetime
-    op: Operation
-    fact_id: str | None
-    fact_before: dict[str, Any] | None = None
-    fact_after: dict[str, Any] | None = None
-    actor: str | None = None
-    reason: str | None = None
-
-
-MemoryHook = Callable[[str, str, dict[str, Any] | None], None]
+MemoryHook = Callable[[Operation, str, Fact | None], None]
 
 
 class SchemaRegistry:
@@ -76,7 +54,7 @@ def register_schema(self, typename: str, model: type[BaseModel], constraint: Con
     def add_hook(self, hook: MemoryHook):
         self._hooks.append(hook)
 
-    def _notify_hooks(self, op: Operation, fact_id: str, data: dict[str, Any] | None) -> None:
+    def _notify_hooks(self, op: Operation, fact_id: str, data: Fact | None) -> None:
         for hook in self._hooks:
             try:
                 hook(op, fact_id, data)
@@ -126,7 +104,7 @@ def commit(
 
             self.storage.save(fact.model_dump())
             self._log_tx(op, fact.id, existing, fact.model_dump(), actor, reason)
-            self._notify_hooks(op, fact.id, fact.model_dump())
+            self._notify_hooks(op, fact.id, fact)
             return fact.id
 
     def update(self, fact_id: str, patch: dict[str, Any], actor: str | None = None, reason: str | None = None) -> str:
@@ -142,7 +120,7 @@ def update(self, fact_id: str, patch: dict[str, Any], actor: str | None = None,
 
             self.storage.save(existing)
             self._log_tx(Operation.UPDATE, fact_id, before, existing, actor, reason)
-            self._notify_hooks(Operation.UPDATE, fact_id, existing)
+            self._notify_hooks(Operation.UPDATE, fact_id, Fact(**existing))
             return fact_id
 
     def delete(self, fact_id: str, actor: str | None = None, reason: str | None = None) -> str:
@@ -153,7 +131,7 @@ def delete(self, fact_id: str, actor: str | None = None, reason: str | None = No
 
             self.storage.delete(fact_id)
             self._log_tx(Operation.DELETE, fact_id, existing, None, actor, reason)
-            self._notify_hooks(Operation.DELETE, fact_id, existing)
+            self._notify_hooks(Operation.DELETE, fact_id, Fact(**existing))
             return fact_id
 
     def get(self, fact_id: str) -> dict[str, Any] | None:
@@ -189,7 +167,7 @@ def promote_session(
 
                 promoted.append(fact_dict["id"])
                 self._log_tx(Operation.PROMOTE, fact_dict["id"], before, fact_dict, actor, reason)
-                self._notify_hooks(Operation.PROMOTE, fact_dict["id"], fact_dict)
+                self._notify_hooks(Operation.PROMOTE, fact_dict["id"], Fact(**fact_dict))
 
             return promoted
 

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
 [project.optional-dependencies]
 redis = ["redis>=7.1.0"]
 langgraph = ["langgraph>=1.0.4"]
+chromadb = ["chromadb>=1.3.5"]
 
 [dependency-groups]
 dev = [

tests/test_chroma_integration.py

@@ -0,0 +1,95 @@
+import pytest
+
+chromadb = pytest.importorskip("chromadb")
+
+from memstate.constants import Operation
+from memstate.integrations.chroma import ChromaSyncHook
+from memstate.schemas import Fact
+
+
+@pytest.fixture
+def chroma_client():
+    return chromadb.Client()
+
+
+@pytest.fixture
+def collection_name():
+    return "test_memstate_sync"
+
+
+def test_initialization_creates_collection(chroma_client, collection_name):
+    ChromaSyncHook(client=chroma_client, collection_name=collection_name)
+    collections = chroma_client.list_collections()
+    assert any(c.name == collection_name for c in collections)
+
+
+def test_commit_upserts_data(chroma_client, collection_name):
+    hook = ChromaSyncHook(client=chroma_client, collection_name=collection_name, text_field="content")
+    hook(op=Operation.COMMIT, fact_id="fact_1", data=Fact(type="memory", payload={"content": "Hello World"}))
+
+    coll = chroma_client.get_collection(collection_name)
+    result = coll.get(ids=["fact_1"])
+    assert result["documents"][0] == "Hello World"
+    assert result["metadatas"][0]["type"] == "memory"
+
+
+def test_promote_updates_data(chroma_client, collection_name):
+    hook = ChromaSyncHook(
+        client=chroma_client, collection_name=collection_name, text_field="text", metadata_fields=["status"]
+    )
+    coll = chroma_client.get_collection(collection_name)
+
+    # Pre-seed
+    coll.add(ids=["fact_1"], documents=["Old"], metadatas=[{"status": "draft"}])
+
+    # Promote
+    hook(
+        op=Operation.PROMOTE, fact_id="fact_1", data=Fact(type="memory", payload={"text": "New", "status": "committed"})
+    )
+
+    result = coll.get(ids=["fact_1"])
+    assert result["documents"][0] == "New"
+    assert result["metadatas"][0]["status"] == "committed"
+
+
+def test_delete_removes_data(chroma_client, collection_name):
+    hook = ChromaSyncHook(client=chroma_client, collection_name=collection_name)
+    coll = chroma_client.get_collection(collection_name)
+
+    coll.add(ids=["del_1"], documents=["To delete"])
+
+    hook(op=Operation.DELETE, fact_id="del_1", data=None)
+
+    result = coll.get(ids=["del_1"])
+    assert len(result["ids"]) == 0
+
+
+def test_discard_session_is_ignored(chroma_client, collection_name):
+    hook = ChromaSyncHook(client=chroma_client, collection_name=collection_name)
+    coll = chroma_client.get_collection(collection_name)
+
+    coll.add(ids=["safe_1"], documents=["Stay"])
+
+    hook(op=Operation.DISCARD_SESSION, fact_id="safe_1", data=None)
+
+    result = coll.get(ids=["safe_1"])
+    assert len(result["ids"]) == 1
+
+
+def test_text_formatter_strategy(chroma_client, collection_name):
+    hook = ChromaSyncHook(
+        client=chroma_client, collection_name=collection_name, text_formatter=lambda d: f"{d['key']}: {d['val']}"
+    )
+    hook(op=Operation.COMMIT, fact_id="fmt_1", data=Fact(type="memory", payload={"key": "A", "val": "B"}))
+
+    coll = chroma_client.get_collection(collection_name)
+    assert coll.get(ids=["fmt_1"])["documents"][0] == "A: B"
+
+
+def test_fallback_missing_text_skips_upsert(chroma_client, collection_name):
+    hook = ChromaSyncHook(client=chroma_client, collection_name=collection_name, text_field="missing_field")
+    hook(op=Operation.COMMIT, fact_id="bad_1", data=Fact(type="memory", payload={"other": "stuff"}))
+
+    coll = chroma_client.get_collection(collection_name)
+    result = coll.get(ids=["bad_1"])
+    assert len(result["ids"]) == 0

tests/test_e2e_sync.py

@@ -0,0 +1,29 @@
+import pytest
+
+chromadb = pytest.importorskip("chromadb")
+
+from memstate.backends.inmemory import InMemoryStorage
+from memstate.integrations.chroma import ChromaSyncHook
+from memstate.schemas import Fact
+from memstate.storage import MemoryStore
+
+
+def test_e2e_memory_store_syncs_to_chroma():
+    chroma_client = chromadb.Client()
+    collection_name = "e2e_test"
+
+    hook = ChromaSyncHook(
+        client=chroma_client, collection_name=collection_name, text_field="content", metadata_fields=["role"]
+    )
+
+    store = MemoryStore(InMemoryStorage())
+    store.add_hook(hook=hook)
+
+    store.commit(fact=Fact(type="test", payload={"content": "Integration works!", "role": "system"}))
+
+    coll = chroma_client.get_collection(collection_name)
+    results = coll.get()
+
+    assert len(results["ids"]) == 1
+    assert results["documents"][0] == "Integration works!"
+    assert results["metadatas"][0]["role"] == "system"

tests/test_storage.py

@@ -5,7 +5,8 @@
 
 from memstate.backends.inmemory import InMemoryStorage
 from memstate.exceptions import ConflictError, HookError, MemoryStoreError, ValidationFailed
-from memstate.storage import Constraint, Fact, MemoryStore
+from memstate.schemas import Fact
+from memstate.storage import Constraint, MemoryStore
 
 
 class User(BaseModel):