feat: add qdrant hook

scream4ik · scream4ik · commit 269a3ef5ef09 · 2025-12-13T14:46:58.000-05:00
diff --git a/memstate/integrations/qdrant.py b/memstate/integrations/qdrant.py
@@ -0,0 +1,157 @@
+from typing import Any, Callable
+
+from memstate.constants import Operation
+from memstate.schemas import Fact
+
+try:
+    from qdrant_client import QdrantClient, models
+except ImportError:
+    raise ImportError("To use QdrantSyncHook, run: pip install qdrant-client")
+
+TextFormatter = Callable[[dict[str, Any]], str]
+MetadataFormatter = Callable[[dict[str, Any]], dict[str, Any]]
+EmbeddingFunction = Callable[[str], list[float]]
+
+
+class FastEmbedEncoder:
+    """
+    Default embedding implementation using FastEmbed.
+    Used if no custom embedding_fn is provided.
+    """
+
+    def __init__(
+        self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2", options: dict[str, Any] | None = None
+    ):
+        try:
+            from fastembed import TextEmbedding
+        except ImportError:
+            raise ImportError(
+                "FastEmbed is not installed. " "Install it via `pip install fastembed` or pass a custom `embedding_fn`."
+            )
+        self.model = TextEmbedding(model_name, **(options or {}))
+
+    def __call__(self, text: str) -> list[float]:
+        return list(self.model.embed(text))[0].tolist()
+
+
+class QdrantSyncHook:
+    """
+    encoder = FastEmbedEncoder(
+        model_name="BAAI/bge-small-en-v1.5",
+        options={"cuda": True}
+    )
+    hook = QdrantSyncHook(client, "memory", embedding_fn=encoder)
+
+    resp = openai.embeddings.create(input=text, model="text-embedding-3-small")
+    openai_embedder = resp.data[0].embedding
+    hook = QdrantSyncHook(client, "memory", embedding_fn=openai_embedder)
+    """
+
+    def __init__(
+        self,
+        client: QdrantClient,
+        collection_name: str,
+        embedding_fn: EmbeddingFunction | 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,
+        distance: models.Distance = models.Distance.COSINE,
+    ) -> None:
+        self.client = client
+        self.collection_name = collection_name
+
+        self.embedding_fn = embedding_fn or FastEmbedEncoder()
+
+        self.target_types = target_types or set()
+        self.distance = distance
+
+        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
+
+        self._ensure_collection()
+
+    def _ensure_collection(self) -> None:
+        """
+        Auto-detects vector size by running a dummy embedding
+        and ensures the collection exists.
+        """
+        try:
+            dummy_vec = self.embedding_fn("test")
+            vector_size = len(dummy_vec)
+        except Exception as e:
+            raise RuntimeError(f"Failed to initialize embedding function: {e}")
+
+        if not self.client.collection_exists(self.collection_name):
+            self.client.create_collection(
+                collection_name=self.collection_name,
+                vectors_config=models.VectorParams(size=vector_size, distance=self.distance),
+            )
+        else:
+            coll_info = self.client.get_collection(self.collection_name)
+            config = coll_info.config.params.vectors
+
+            existing_size = None
+            if isinstance(config, models.VectorParams):
+                existing_size = config.size
+            elif isinstance(config, dict) and "" in config:  # Default unnamed vector
+                existing_size = config[""].size
+
+            if existing_size and existing_size != vector_size:
+                raise ValueError(
+                    f"Collection '{self.collection_name}' expects vector size {existing_size}, "
+                    f"but your embedding function produces {vector_size}. "
+                    "Mismatch detected."
+                )
+
+    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, list)):
+                        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.client.delete(collection_name=self.collection_name, points_selector=[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
+
+        if op in (Operation.COMMIT, Operation.UPDATE, Operation.COMMIT_EPHEMERAL, Operation.PROMOTE):
+            text = self._extract_text(data.payload)
+            if not text.strip():
+                return
+
+            vector = self.embedding_fn(text)
+
+            meta = {"type": data.type, "source": data.source or "", "ts": str(data.ts), "document": text}
+            user_meta = self._get_metadata(data=data.payload)
+            meta.update(user_meta)
+
+            self.client.upsert(
+                collection_name=self.collection_name,
+                points=[models.PointStruct(id=fact_id, vector=vector, payload=meta)],
+            )
diff --git a/pyproject.toml b/pyproject.toml
@@ -48,13 +48,15 @@ dependencies = [
 redis = ["redis>=7.1.0"]
 langgraph = ["langgraph>=1.0.4"]
 chromadb = ["chromadb>=1.3.5"]
+qdrant = ["qdrant-client>=1.16.2"]
 postgres = ["sqlalchemy>=2.0.0", "psycopg[binary]>=3.3.2"]
 
 [dependency-groups]
 dev = [
     "bandit>=1.9.1",
     "black>=25.11.0",
     "fakeredis>=2.32.1",
+    "fastembed>=0.7.4",
     "isort>=7.0.0",
     "mypy>=1.18.2",
     "pre-commit>=4.4.0",
diff --git a/tests/test_e2e_sync.py b/tests/test_e2e_sync.py
@@ -1,9 +1,11 @@
 import pytest
 
 chromadb = pytest.importorskip("chromadb")
+qdrant_client = pytest.importorskip("qdrant_client")
 
 from memstate import Fact, InMemoryStorage, MemoryStore
 from memstate.integrations.chroma import ChromaSyncHook
+from memstate.integrations.qdrant import QdrantSyncHook
 
 
 def test_e2e_memory_store_syncs_to_chroma():
@@ -25,3 +27,25 @@ def test_e2e_memory_store_syncs_to_chroma():
     assert len(results["ids"]) == 1
     assert results["documents"][0] == "Integration works!"
     assert results["metadatas"][0]["role"] == "system"
+
+
+def test_e2e_memory_store_syncs_to_qdrand():
+    client = qdrant_client.QdrantClient(":memory:")
+    collection_name = "e2e_test"
+
+    hook = QdrantSyncHook(
+        client=client, collection_name=collection_name, text_field="content", metadata_fields=["role"]
+    )
+
+    store = MemoryStore(InMemoryStorage())
+    store.add_hook(hook=hook)
+
+    fact_id = store.commit(fact=Fact(type="test", payload={"content": "Integration works!", "role": "system"}))
+
+    points, _ = client.scroll(collection_name=collection_name, limit=10)
+
+    assert len(points) == 1
+    point = points[0]
+    assert point.id == fact_id
+    assert point.payload["document"] == "Integration works!"
+    assert point.payload["role"] == "system"
diff --git a/tests/test_qdrant_integration.py b/tests/test_qdrant_integration.py
@@ -0,0 +1,134 @@
+import uuid
+
+import pytest
+
+qdrant_client = pytest.importorskip("qdrant_client")
+
+from memstate import Fact, Operation
+from memstate.integrations.qdrant import QdrantSyncHook
+
+
+@pytest.fixture
+def client():
+    return qdrant_client.QdrantClient(":memory:")
+
+
+@pytest.fixture
+def collection_name():
+    return "test_memstate_sync"
+
+
+@pytest.fixture
+def fact_id():
+    return str(uuid.uuid4())
+
+
+def test_initialization_creates_collection(client, collection_name):
+    QdrantSyncHook(client=client, collection_name=collection_name)
+    collections = client.get_collections()
+    assert any(c.name == collection_name for c in collections.collections)
+
+
+def test_commit_upserts_data(client, collection_name, fact_id):
+    hook = QdrantSyncHook(client=client, collection_name=collection_name, text_field="content")
+    hook(op=Operation.COMMIT, fact_id=fact_id, data=Fact(type="memory", payload={"content": "Hello World"}))
+
+    points, _ = client.scroll(collection_name=collection_name, limit=10)
+    point = points[0]
+    assert point.id == fact_id
+    assert point.payload["document"] == "Hello World"
+    assert point.payload["type"] == "memory"
+
+
+def test_promote_updates_data(client, collection_name, fact_id):
+    hook = QdrantSyncHook(client=client, collection_name=collection_name, text_field="text", metadata_fields=["status"])
+
+    # Pre-seed
+    client.upsert(
+        collection_name=collection_name,
+        points=[
+            qdrant_client.models.PointStruct(
+                id=fact_id,
+                vector=qdrant_client.models.Document(
+                    text="Old",
+                    model="sentence-transformers/all-MiniLM-L6-v2",
+                ),
+                payload={"status": "draft"},
+            )
+        ],
+    )
+
+    # Promote
+    hook(
+        op=Operation.PROMOTE, fact_id=fact_id, data=Fact(type="memory", payload={"text": "New", "status": "committed"})
+    )
+
+    points, _ = client.scroll(collection_name=collection_name, limit=10)
+    point = points[0]
+    assert point.id == fact_id
+    assert point.payload["document"] == "New"
+    assert point.payload["status"] == "committed"
+
+
+def test_delete_removes_data(client, collection_name, fact_id):
+    hook = QdrantSyncHook(client=client, collection_name=collection_name)
+
+    client.upsert(
+        collection_name=collection_name,
+        points=[
+            qdrant_client.models.PointStruct(
+                id=fact_id,
+                vector=qdrant_client.models.Document(
+                    text="To delete",
+                    model="sentence-transformers/all-MiniLM-L6-v2",
+                ),
+            )
+        ],
+    )
+
+    hook(op=Operation.DELETE, fact_id=fact_id, data=None)
+
+    points, _ = client.scroll(collection_name=collection_name, limit=10)
+    assert len(points) == 0
+
+
+def test_discard_session_is_ignored(client, collection_name, fact_id):
+    hook = QdrantSyncHook(client=client, collection_name=collection_name)
+
+    client.upsert(
+        collection_name=collection_name,
+        points=[
+            qdrant_client.models.PointStruct(
+                id=fact_id,
+                vector=qdrant_client.models.Document(
+                    text="Stay",
+                    model="sentence-transformers/all-MiniLM-L6-v2",
+                ),
+            )
+        ],
+    )
+
+    hook(op=Operation.DISCARD_SESSION, fact_id=fact_id, data=None)
+
+    points, _ = client.scroll(collection_name=collection_name, limit=10)
+    assert len(points) == 1
+
+
+def test_text_formatter_strategy(client, collection_name, fact_id):
+    hook = QdrantSyncHook(
+        client=client, collection_name=collection_name, text_formatter=lambda d: f"{d['key']}: {d['val']}"
+    )
+    hook(op=Operation.COMMIT, fact_id=fact_id, data=Fact(type="memory", payload={"key": "A", "val": "B"}))
+
+    points, _ = client.scroll(collection_name=collection_name, limit=10)
+    point = points[0]
+    assert point.id == fact_id
+    assert point.payload["document"] == "A: B"
+
+
+def test_fallback_missing_text_skips_upsert(client, collection_name, fact_id):
+    hook = QdrantSyncHook(client=client, collection_name=collection_name, text_field="missing_field")
+    hook(op=Operation.COMMIT, fact_id=fact_id, data=Fact(type="memory", payload={"other": "stuff"}))
+
+    points, _ = client.scroll(collection_name=collection_name, limit=10)
+    assert len(points) == 0
diff --git a/uv.lock b/uv.lock
