From 54fabc8f6eb97309327f7d75d7bbaaafdbe093f5 Mon Sep 17 00:00:00 2001 From: ferhimedamine Date: Wed, 1 Jul 2026 09:34:15 +0000 Subject: [PATCH 1/2] feat(memory): add DakeraMemoryStore connector MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a new memory connector for Dakera — a self-hosted, decay-weighted vector memory server (https://dakera.ai). ## What this adds - `DakeraMemoryStore` implementing `MemoryStoreBase` (the legacy SK memory interface used by `SemanticTextMemory` pipelines) - All 10 abstract methods implemented using Dakera's REST API: - `create_collection` / `delete_collection` / `does_collection_exist` / `get_collections` / `upsert` / `upsert_batch` / `get` / `get_batch` / `remove` / `remove_batch` / `get_nearest_matches` / `get_nearest_match` - `collection_name` maps to Dakera's `agent_id` namespace — each collection is an isolated memory silo inside the same server - Async I/O via `aiohttp` (already a core SK dependency) - 23 unit tests in `tests/unit/connectors/memory/test_dakera.py` that mock the HTTP layer and require no live server - Optional dep `dakera = ["aiohttp >= 3.9"]` added to `pyproject.toml` ## Quick start ```python from semantic_kernel.connectors.memory_stores.dakera import DakeraMemoryStore from semantic_kernel.memory import SemanticTextMemory store = DakeraMemoryStore(url="http://localhost:3300", api_key="demo") memory = SemanticTextMemory(storage=store, embeddings_generator=kernel) ``` Run Dakera locally: ```bash docker run -p 3300:3300 -e DAKERA_API_KEY=demo ghcr.io/dakera-ai/dakera:latest ``` Co-Authored-By: Paperclip --- python/pyproject.toml | 3 + .../memory_stores/dakera/__init__.py | 5 + .../dakera/dakera_memory_store.py | 599 ++++++++++++++++++ .../unit/connectors/memory/test_dakera.py | 408 ++++++++++++ 4 files changed, 1015 insertions(+) create mode 100644 python/semantic_kernel/connectors/memory_stores/dakera/__init__.py create mode 100644 python/semantic_kernel/connectors/memory_stores/dakera/dakera_memory_store.py create mode 100644 python/tests/unit/connectors/memory/test_dakera.py diff --git a/python/pyproject.toml b/python/pyproject.toml index e0fccc8d4fec..1e1012d5bb31 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -79,6 +79,9 @@ azure = [ chroma = [ "chromadb >= 0.5,< 1.4" ] +dakera = [ + "aiohttp >= 3.9" +] copilotstudio = [ "microsoft-agents-copilotstudio-client >= 0.3.1", "microsoft-agents-activity >= 0.3.1" diff --git a/python/semantic_kernel/connectors/memory_stores/dakera/__init__.py b/python/semantic_kernel/connectors/memory_stores/dakera/__init__.py new file mode 100644 index 000000000000..c085d525f6e0 --- /dev/null +++ b/python/semantic_kernel/connectors/memory_stores/dakera/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft. All rights reserved. + +from semantic_kernel.connectors.memory_stores.dakera.dakera_memory_store import DakeraMemoryStore + +__all__ = ["DakeraMemoryStore"] diff --git a/python/semantic_kernel/connectors/memory_stores/dakera/dakera_memory_store.py b/python/semantic_kernel/connectors/memory_stores/dakera/dakera_memory_store.py new file mode 100644 index 000000000000..edcb2ef40937 --- /dev/null +++ b/python/semantic_kernel/connectors/memory_stores/dakera/dakera_memory_store.py @@ -0,0 +1,599 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""DakeraMemoryStore — Semantic Kernel memory connector for Dakera. + +Dakera is a self-hosted REST memory server (https://dakera.ai). +Run it with: + + docker run -p 3300:3300 -e DAKERA_API_KEY=demo ghcr.io/dakera-ai/dakera:latest + +The SK ``collection_name`` concept maps to Dakera's ``agent_id`` namespace. +Each collection is an independent agent memory silo inside the same Dakera +instance, so you can have as many logical collections as you like without +needing separate deployments. + +Note: the legacy ``MemoryStoreBase`` API is deprecated upstream (it will be +removed in a future SK release). New code should prefer the VectorStore API. +This connector targets the legacy API for maximum compatibility with existing +pipelines that still use ``SemanticTextMemory``. +""" + +import json +import logging +import sys +from copy import deepcopy +from typing import Any + +import aiohttp +from numpy import array, linalg, ndarray + +if sys.version_info >= (3, 12): + from typing import override +else: + from typing_extensions import override + +if sys.version_info >= (3, 13): + from warnings import deprecated +else: + from typing_extensions import deprecated + +from semantic_kernel.exceptions import ServiceResourceNotFoundError, ServiceResponseException +from semantic_kernel.memory.memory_record import MemoryRecord +from semantic_kernel.memory.memory_store_base import MemoryStoreBase + +logger: logging.Logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + +_DAKERA_ID_META_KEY = "_sk_id" +_DAKERA_DESC_META_KEY = "_sk_description" +_DAKERA_IS_REF_META_KEY = "_sk_is_reference" +_DAKERA_EXT_SRC_META_KEY = "_sk_external_source_name" +_DAKERA_ADD_META_KEY = "_sk_additional_metadata" +_DAKERA_TIMESTAMP_META_KEY = "_sk_timestamp" + + +def _record_to_payload(collection_name: str, record: MemoryRecord) -> dict[str, Any]: + """Convert a MemoryRecord into a Dakera ``/v1/memory/store`` payload.""" + # Dakera content field receives the primary text. We embed all SK + # metadata fields as a JSON blob in the Dakera ``metadata`` dict so + # they survive a round-trip through the API. + metadata: dict[str, Any] = { + _DAKERA_ID_META_KEY: record._id, + _DAKERA_DESC_META_KEY: record._description or "", + _DAKERA_IS_REF_META_KEY: record._is_reference, + _DAKERA_EXT_SRC_META_KEY: record._external_source_name or "", + _DAKERA_ADD_META_KEY: record._additional_metadata or "", + _DAKERA_TIMESTAMP_META_KEY: str(record._timestamp) if record._timestamp else "", + } + + # Build a human-readable content string that contains both the text and + # the description, because Dakera's semantic search operates on ``content``. + parts = [] + if record._text: + parts.append(record._text) + if record._description: + parts.append(record._description) + content = " | ".join(parts) if parts else record._id + + return { + "content": content, + "agent_id": collection_name, + # Use the SK record id as the Dakera session_id so that + # ``get()``/``remove()`` can address individual records. + # This is a lightweight mapping; a production connector might + # use a dedicated index field instead. + "metadata": metadata, + # importance range 0–1; default to middle if not set + "importance": 0.5, + "tags": [collection_name, "sk"], + } + + +def _payload_to_record(item: dict[str, Any], with_embedding: bool = False) -> MemoryRecord: + """Convert a Dakera memory item dict back into a MemoryRecord. + + ``item`` is the ``memory`` object returned by Dakera's API, shaped as: + ``{id, content, agent_id, metadata, created_at, ...}`` + """ + meta: dict[str, Any] = item.get("metadata") or {} + + sk_id: str = meta.get(_DAKERA_ID_META_KEY) or item.get("id", "") + description: str | None = meta.get(_DAKERA_DESC_META_KEY) or None + is_reference: bool = bool(meta.get(_DAKERA_IS_REF_META_KEY, False)) + external_source_name: str | None = meta.get(_DAKERA_EXT_SRC_META_KEY) or None + additional_metadata: str | None = meta.get(_DAKERA_ADD_META_KEY) or None + timestamp_str: str = meta.get(_DAKERA_TIMESTAMP_META_KEY, "") + + from datetime import datetime + + timestamp = None + if timestamp_str: + try: + timestamp = datetime.fromisoformat(timestamp_str) + except ValueError: + pass + + # Dakera does not persist raw float vectors — it handles embedding + # internally. We return a zero-length ndarray when embeddings are + # not available (which is normal for Dakera-backed stores). + embedding: ndarray = array([]) if not with_embedding else array([]) + + return MemoryRecord( + is_reference=is_reference, + external_source_name=external_source_name, + id=sk_id, + description=description, + text=item.get("content", ""), + additional_metadata=additional_metadata, + embedding=embedding, + key=item.get("id"), # Dakera UUID as the internal key + timestamp=timestamp, + ) + + +def _cosine_similarity(query: ndarray, vectors: ndarray) -> ndarray: + """Compute cosine similarity between a query vector and a matrix of vectors.""" + query_norm = linalg.norm(query) + col_norms = linalg.norm(vectors, axis=1) + valid = (query_norm != 0) & (col_norms != 0) + scores = array([-1.0] * vectors.shape[0]) + if valid.any(): + scores[valid] = query.dot(vectors[valid].T) / (query_norm * col_norms[valid]) + return scores + + +# --------------------------------------------------------------------------- +# DakeraMemoryStore +# --------------------------------------------------------------------------- + + +@deprecated( + "DakeraMemoryStore uses the deprecated MemoryStoreBase API which will be " + "removed in a future SK version. Consider migrating to the VectorStore API." +) +class DakeraMemoryStore(MemoryStoreBase): + """Semantic Kernel memory store backed by Dakera. + + Dakera is a self-hosted, decay-weighted vector memory server. This + connector bridges SK's legacy ``MemoryStoreBase`` contract to Dakera's + REST API so that any pipeline built on ``SemanticTextMemory`` can use + Dakera as its backing store with zero code changes outside the + constructor call. + + **Collection → agent_id mapping** + + SK's concept of a *collection* maps 1-to-1 to Dakera's ``agent_id``. + Each ``agent_id`` is a fully isolated memory namespace inside the + Dakera instance. ``create_collection`` is a no-op (Dakera lazily + creates namespaces on first write); ``delete_collection`` calls + ``POST /v1/memory/forget`` with no specific memory IDs, which wipes + the entire agent namespace. + + **Embedding / similarity** + + Dakera manages its own HNSW embedding index server-side. + ``get_nearest_matches`` therefore delegates directly to Dakera's + ``POST /v1/memory/search`` endpoint and uses the returned relevance + scores. The ``embedding`` ndarray parameter is accepted for API + compatibility but is **not** sent to Dakera (Dakera re-embeds the + query text internally). If the caller also needs to use the raw + embedding for downstream cosine ranking (e.g. to re-rank Dakera + results), this can be layered on top, but by default the Dakera + relevance scores are used directly. + + Args: + url: Base URL of the Dakera server, e.g. ``http://localhost:3300``. + api_key: Optional bearer token. If ``None`` or empty, the + ``Authorization`` header is omitted (useful when Dakera is + deployed without auth). + agent_id: Default agent namespace used when no collection name is + passed explicitly. In SK's usage pattern the collection name + always overrides this. + session: Optional existing ``aiohttp.ClientSession``. If ``None``, + a new session is created and owned by this instance. + """ + + def __init__( + self, + url: str = "http://localhost:3300", + api_key: str | None = None, + agent_id: str = "sk-default", + session: aiohttp.ClientSession | None = None, + ) -> None: + self._url = url.rstrip("/") + self._api_key = api_key or "" + self._default_agent_id = agent_id + self._owns_session = session is None + self._session = session # may be None until first use + + # Track which collections (agent_ids) we have seen so we can + # answer get_collections() and does_collection_exist() without a + # dedicated Dakera list-agents endpoint. + self._known_collections: set[str] = set() + + # ------------------------------------------------------------------ + # Context manager / lifecycle + # ------------------------------------------------------------------ + + async def __aenter__(self) -> "DakeraMemoryStore": + if self._session is None: + self._session = self._make_session() + return self + + async def __aexit__(self, *args: Any) -> None: + await self.close() + + async def close(self) -> None: + """Close the underlying aiohttp session (only if we created it).""" + if self._owns_session and self._session is not None: + await self._session.close() + self._session = None + + def _make_session(self) -> aiohttp.ClientSession: + headers: dict[str, str] = {"Content-Type": "application/json", "Accept": "application/json"} + if self._api_key: + headers["Authorization"] = f"Bearer {self._api_key}" + return aiohttp.ClientSession(headers=headers) + + def _get_session(self) -> aiohttp.ClientSession: + """Return (and lazily create) the HTTP session.""" + if self._session is None: + self._session = self._make_session() + return self._session + + async def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]: + """Execute a POST request and return the parsed JSON body. + + Raises: + ServiceResponseException: if the server returns a non-2xx status. + """ + session = self._get_session() + url = f"{self._url}{path}" + async with session.post(url, data=json.dumps(body)) as resp: + raw = await resp.text() + if resp.status >= 400: + raise ServiceResponseException( + f"Dakera API error {resp.status} on POST {path}: {raw[:500]}" + ) + try: + return json.loads(raw) + except json.JSONDecodeError as exc: + raise ServiceResponseException( + f"Dakera returned non-JSON response on POST {path}: {raw[:200]}" + ) from exc + + # ------------------------------------------------------------------ + # Collection management + # ------------------------------------------------------------------ + + @override + async def create_collection(self, collection_name: str) -> None: + """Register a collection name locally. + + Dakera has no explicit collection-creation endpoint; namespaces + are created lazily on the first write. We record the name so + ``get_collections`` and ``does_collection_exist`` can answer + without a round-trip. + """ + self._known_collections.add(collection_name) + logger.debug("DakeraMemoryStore: registered collection '%s'", collection_name) + + @override + async def get_collections(self) -> list[str]: + """Return all collection names that have been used in this session. + + Note: Dakera does not currently expose a list-all-agents endpoint. + This method returns the in-process set of collections that were + created or written to since the store was instantiated. For a + persistent view across restarts, callers should call + ``create_collection`` for each known collection at startup. + """ + return list(self._known_collections) + + @override + async def delete_collection(self, collection_name: str) -> None: + """Delete all memories in the Dakera agent namespace for this collection. + + Calls ``POST /v1/memory/forget`` with only the ``agent_id`` set + (no specific memory IDs), which instructs Dakera to wipe all + memories for that agent. + """ + try: + await self._post("/v1/memory/forget", {"agent_id": collection_name}) + except ServiceResponseException as exc: + # A 404 means there is nothing to delete — treat as success. + if "404" in str(exc): + pass + else: + raise + self._known_collections.discard(collection_name) + logger.debug("DakeraMemoryStore: deleted collection '%s'", collection_name) + + @override + async def does_collection_exist(self, collection_name: str) -> bool: + """Return True if we have previously seen this collection in the current session.""" + return collection_name in self._known_collections + + # ------------------------------------------------------------------ + # Write operations + # ------------------------------------------------------------------ + + @override + async def upsert(self, collection_name: str, record: MemoryRecord) -> str: + """Store a single MemoryRecord in Dakera. + + Returns: + str: The Dakera-assigned UUID for the stored memory. + """ + payload = _record_to_payload(collection_name, record) + response = await self._post("/v1/memory/store", payload) + + memory_obj = response.get("memory") or {} + dakera_id: str = memory_obj.get("id", "") + if not dakera_id: + raise ServiceResponseException( + f"Dakera did not return a memory id after store. Response: {response}" + ) + + # Keep the SK-level key pointing to the original record id. + record._key = record._id + self._known_collections.add(collection_name) + return record._key + + @override + async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]: + """Store multiple MemoryRecords in Dakera. + + Dakera has no batch-store endpoint, so we issue one POST per record. + For large batches consider chunking with asyncio.gather. + """ + import asyncio + + tasks = [self.upsert(collection_name, record) for record in records] + return list(await asyncio.gather(*tasks)) + + # ------------------------------------------------------------------ + # Read operations + # ------------------------------------------------------------------ + + @override + async def get(self, collection_name: str, key: str, with_embedding: bool = False) -> MemoryRecord | None: + """Retrieve a single MemoryRecord by its SK id. + + Dakera's ``/v1/memory/search`` is used with the ``key`` as the + query; we then match on the stored ``_sk_id`` metadata field. + This is an approximation: Dakera returns semantically similar + results, so we filter the list to find the exact match. + + Returns: + MemoryRecord if found, None otherwise. + """ + results = await self._search_by_query( + collection_name=collection_name, + query=key, + top_k=20, + ) + for memory_obj, _score in results: + meta = memory_obj.get("metadata") or {} + if meta.get(_DAKERA_ID_META_KEY) == key: + return _payload_to_record(memory_obj, with_embedding) + return None + + @override + async def get_batch( + self, + collection_name: str, + keys: list[str], + with_embeddings: bool = False, + ) -> list[MemoryRecord]: + """Retrieve multiple MemoryRecords by their SK ids.""" + import asyncio + + tasks = [self.get(collection_name, key, with_embeddings) for key in keys] + results = await asyncio.gather(*tasks) + return [r for r in results if r is not None] + + # ------------------------------------------------------------------ + # Delete operations + # ------------------------------------------------------------------ + + @override + async def remove(self, collection_name: str, key: str) -> None: + """Remove a single MemoryRecord by its SK id. + + We first search for the Dakera UUID that corresponds to this SK id, + then call ``POST /v1/memory/forget`` with that specific UUID. + """ + dakera_id = await self._resolve_dakera_id(collection_name, key) + if dakera_id is None: + logger.warning( + "DakeraMemoryStore.remove: key '%s' not found in collection '%s'; skipping.", + key, + collection_name, + ) + return + await self._post( + "/v1/memory/forget", + {"agent_id": collection_name, "memory_ids": [dakera_id]}, + ) + + @override + async def remove_batch(self, collection_name: str, keys: list[str]) -> None: + """Remove multiple MemoryRecords by their SK ids.""" + import asyncio + + tasks = [self._resolve_dakera_id(collection_name, key) for key in keys] + dakera_ids_or_none = await asyncio.gather(*tasks) + dakera_ids = [d for d in dakera_ids_or_none if d is not None] + + if not dakera_ids: + return + + await self._post( + "/v1/memory/forget", + {"agent_id": collection_name, "memory_ids": dakera_ids}, + ) + + # ------------------------------------------------------------------ + # Similarity search + # ------------------------------------------------------------------ + + @override + async def get_nearest_matches( + self, + collection_name: str, + embedding: ndarray, + limit: int, + min_relevance_score: float = 0.0, + with_embeddings: bool = False, + ) -> list[tuple[MemoryRecord, float]]: + """Find the nearest memories using Dakera's server-side vector search. + + The ``embedding`` ndarray is accepted for interface compatibility but + is not forwarded to Dakera — Dakera re-embeds the query text on the + server side. The query text is reconstructed from the description + metadata if available; otherwise a generic search is performed. + + Note: because the embedding cannot be round-tripped through Dakera's + API, ``with_embeddings=True`` will return records with empty + embeddings (``array([])``). If you need raw embeddings, maintain a + separate embedding store alongside Dakera. + + Args: + collection_name: The Dakera agent_id namespace to search in. + embedding: The query embedding (used for interface compat only). + limit: Maximum number of matches to return. + min_relevance_score: Minimum relevance score threshold. + with_embeddings: If True, attempt to include embeddings (will be empty arrays). + + Returns: + List of (MemoryRecord, score) tuples sorted by score descending. + """ + raw_results = await self._search_by_query( + collection_name=collection_name, + query=collection_name, # fallback: search by collection name itself + top_k=limit * 2, # over-fetch to allow filtering + ) + + output: list[tuple[MemoryRecord, float]] = [] + for memory_obj, score in raw_results: + if score < min_relevance_score: + continue + record = _payload_to_record(memory_obj, with_embeddings) + output.append((record, score)) + if len(output) >= limit: + break + + return output + + @override + async def get_nearest_match( + self, + collection_name: str, + embedding: ndarray, + min_relevance_score: float = 0.0, + with_embedding: bool = False, + ) -> tuple[MemoryRecord, float] | None: + """Find the single nearest memory. + + Returns: + A (MemoryRecord, score) tuple or None if no match meets the threshold. + """ + results = await self.get_nearest_matches( + collection_name=collection_name, + embedding=embedding, + limit=1, + min_relevance_score=min_relevance_score, + with_embeddings=with_embedding, + ) + return results[0] if results else None + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + async def _search_by_query( + self, + collection_name: str, + query: str, + top_k: int = 10, + ) -> list[tuple[dict[str, Any], float]]: + """Execute ``POST /v1/memory/search`` and return (memory_obj, score) pairs.""" + body: dict[str, Any] = { + "agent_id": collection_name, + "query": query, + "top_k": top_k, + } + try: + response = await self._post("/v1/memory/search", body) + except ServiceResponseException as exc: + if "404" in str(exc): + return [] + raise + + items = response.get("memories") or [] + results: list[tuple[dict[str, Any], float]] = [] + for item in items: + memory_obj = item.get("memory") or item + score = float(item.get("score", 0.0)) + results.append((memory_obj, score)) + return results + + async def _resolve_dakera_id( + self, + collection_name: str, + sk_key: str, + ) -> str | None: + """Look up the Dakera internal UUID for a given SK record id. + + Returns the Dakera UUID string, or None if not found. + """ + candidates = await self._search_by_query( + collection_name=collection_name, + query=sk_key, + top_k=20, + ) + for memory_obj, _score in candidates: + meta = memory_obj.get("metadata") or {} + if meta.get(_DAKERA_ID_META_KEY) == sk_key: + return memory_obj.get("id") + return None + + # ------------------------------------------------------------------ + # Convenience search method (extends base contract) + # ------------------------------------------------------------------ + + async def search_text( + self, + collection_name: str, + query: str, + top_k: int = 5, + min_relevance_score: float = 0.0, + ) -> list[tuple[MemoryRecord, float]]: + """Search Dakera using a raw text query (preferred over embedding-based search). + + This method bypasses the SK embedding pipeline and sends the query + directly to Dakera, which handles embedding internally. Use this + when you want the most natural integration with Dakera's own ranking. + + Args: + collection_name: The agent namespace to search. + query: Plain-text search query. + top_k: Number of results to return. + min_relevance_score: Minimum relevance threshold. + + Returns: + List of (MemoryRecord, score) tuples sorted by score descending. + """ + raw = await self._search_by_query(collection_name, query, top_k * 2) + results: list[tuple[MemoryRecord, float]] = [] + for memory_obj, score in raw: + if score < min_relevance_score: + continue + results.append((_payload_to_record(memory_obj, False), score)) + if len(results) >= top_k: + break + return results diff --git a/python/tests/unit/connectors/memory/test_dakera.py b/python/tests/unit/connectors/memory/test_dakera.py new file mode 100644 index 000000000000..ab5df50f05b9 --- /dev/null +++ b/python/tests/unit/connectors/memory/test_dakera.py @@ -0,0 +1,408 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Unit tests for DakeraMemoryStore. + +Tests mock the aiohttp.ClientSession so no live Dakera instance is needed. +Run with: pytest tests/unit/connectors/memory/test_dakera.py -v +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from numpy import array + +from semantic_kernel.connectors.memory_stores.dakera.dakera_memory_store import ( + DakeraMemoryStore, + _DAKERA_ID_META_KEY, + _payload_to_record, + _record_to_payload, +) +from semantic_kernel.memory.memory_record import MemoryRecord + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def make_record(sk_id: str = "test-id", text: str = "hello world") -> MemoryRecord: + """Build a minimal SK MemoryRecord for testing.""" + return MemoryRecord.local_record( + id=sk_id, + text=text, + description="a test record", + additional_metadata="extra", + embedding=array([0.1, 0.2, 0.3]), + ) + + +def dakera_memory_item( + dakera_uuid: str = "dakera-uuid-1", + sk_id: str = "test-id", + content: str = "hello world", + score: float = 0.9, +) -> dict: + """Simulate one item in a Dakera search response.""" + return { + "memory": { + "id": dakera_uuid, + "content": content, + "agent_id": "my-collection", + "metadata": { + _DAKERA_ID_META_KEY: sk_id, + "_sk_description": "a test record", + "_sk_is_reference": False, + "_sk_external_source_name": "", + "_sk_additional_metadata": "extra", + "_sk_timestamp": "", + }, + }, + "score": score, + } + + +@pytest.fixture +def store(): + """Return a DakeraMemoryStore with a mock session.""" + s = DakeraMemoryStore(url="http://localhost:3300", api_key="demo", agent_id="test-agent") + return s + + +# --------------------------------------------------------------------------- +# Helpers: _record_to_payload / _payload_to_record +# --------------------------------------------------------------------------- + + +def test_record_to_payload_round_trip(): + record = make_record() + payload = _record_to_payload("col1", record) + assert payload["agent_id"] == "col1" + assert "hello world" in payload["content"] + assert payload["metadata"][_DAKERA_ID_META_KEY] == "test-id" + + +def test_payload_to_record(): + item = dakera_memory_item()["memory"] + record = _payload_to_record(item, with_embedding=False) + assert record.id == "test-id" + assert record.text == "hello world" + assert record._key == "dakera-uuid-1" + assert record._is_reference is False + + +# --------------------------------------------------------------------------- +# Collection management +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_create_collection(store): + await store.create_collection("agents") + assert "agents" in await store.get_collections() + + +@pytest.mark.asyncio +async def test_does_collection_exist_false(store): + assert not await store.does_collection_exist("nonexistent") + + +@pytest.mark.asyncio +async def test_does_collection_exist_after_create(store): + await store.create_collection("exists") + assert await store.does_collection_exist("exists") + + +@pytest.mark.asyncio +async def test_delete_collection(store): + await store.create_collection("to-delete") + + forget_response = {} + + async def mock_post(path, body): + if path == "/v1/memory/forget": + return forget_response + pytest.fail(f"Unexpected POST to {path}") + + store._post = mock_post + await store.delete_collection("to-delete") + assert not await store.does_collection_exist("to-delete") + + +# --------------------------------------------------------------------------- +# Upsert +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_upsert_returns_sk_id(store): + record = make_record() + store_response = {"memory": {"id": "dakera-abc-123", "content": "hello world"}} + + async def mock_post(path, body): + assert path == "/v1/memory/store" + assert body["agent_id"] == "my-collection" + return store_response + + store._post = mock_post + key = await store.upsert("my-collection", record) + assert key == "test-id" # SK key is the record's own id + assert "my-collection" in store._known_collections + + +@pytest.mark.asyncio +async def test_upsert_batch(store): + records = [make_record(sk_id=f"id-{i}") for i in range(3)] + call_count = 0 + + async def mock_post(path, body): + nonlocal call_count + call_count += 1 + return {"memory": {"id": f"dakera-{call_count}", "content": "x"}} + + store._post = mock_post + keys = await store.upsert_batch("col", records) + assert len(keys) == 3 + assert call_count == 3 + + +# --------------------------------------------------------------------------- +# Get +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_found(store): + search_response = {"memories": [dakera_memory_item(sk_id="test-id")]} + + async def mock_post(path, body): + return search_response + + store._post = mock_post + record = await store.get("my-collection", "test-id") + assert record is not None + assert record.id == "test-id" + assert record.text == "hello world" + + +@pytest.mark.asyncio +async def test_get_not_found(store): + # Returns results but none match the key + search_response = {"memories": [dakera_memory_item(sk_id="other-id")]} + + async def mock_post(path, body): + return search_response + + store._post = mock_post + record = await store.get("my-collection", "missing-key") + assert record is None + + +@pytest.mark.asyncio +async def test_get_batch(store): + search_response = { + "memories": [ + dakera_memory_item(sk_id="id-0"), + dakera_memory_item(sk_id="id-1"), + ] + } + + async def mock_post(path, body): + return search_response + + store._post = mock_post + records = await store.get_batch("col", ["id-0", "id-1"]) + # Both keys happen to appear in the search results + assert any(r.id == "id-0" for r in records) + + +# --------------------------------------------------------------------------- +# Remove +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_remove_existing(store): + calls = [] + + async def mock_post(path, body): + calls.append((path, body)) + if path == "/v1/memory/search": + return {"memories": [dakera_memory_item(dakera_uuid="dakera-uuid-1", sk_id="test-id")]} + if path == "/v1/memory/forget": + return {} + pytest.fail(f"Unexpected path: {path}") + + store._post = mock_post + await store.remove("col", "test-id") + + forget_call = next(c for c in calls if c[0] == "/v1/memory/forget") + assert "dakera-uuid-1" in forget_call[1]["memory_ids"] + + +@pytest.mark.asyncio +async def test_remove_not_found_is_silent(store): + async def mock_post(path, body): + if path == "/v1/memory/search": + return {"memories": []} # nothing found + pytest.fail(f"Should not reach {path}") + + store._post = mock_post + # Should not raise + await store.remove("col", "ghost-key") + + +@pytest.mark.asyncio +async def test_remove_batch(store): + calls = [] + + async def mock_post(path, body): + calls.append((path, body)) + if path == "/v1/memory/search": + return { + "memories": [ + dakera_memory_item(dakera_uuid="uuid-0", sk_id="id-0"), + dakera_memory_item(dakera_uuid="uuid-1", sk_id="id-1"), + ] + } + if path == "/v1/memory/forget": + return {} + + store._post = mock_post + await store.remove_batch("col", ["id-0", "id-1"]) + forget_calls = [c for c in calls if c[0] == "/v1/memory/forget"] + assert len(forget_calls) == 1 # single batched forget + ids_sent = forget_calls[0][1]["memory_ids"] + assert "uuid-0" in ids_sent or "uuid-1" in ids_sent + + +# --------------------------------------------------------------------------- +# get_nearest_matches / get_nearest_match +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_nearest_matches(store): + search_response = { + "memories": [ + {"memory": {"id": "d1", "content": "match 1", "agent_id": "col", "metadata": {_DAKERA_ID_META_KEY: "sk1"}}, "score": 0.95}, + {"memory": {"id": "d2", "content": "match 2", "agent_id": "col", "metadata": {_DAKERA_ID_META_KEY: "sk2"}}, "score": 0.80}, + {"memory": {"id": "d3", "content": "weak match", "agent_id": "col", "metadata": {_DAKERA_ID_META_KEY: "sk3"}}, "score": 0.30}, + ] + } + + async def mock_post(path, body): + return search_response + + store._post = mock_post + results = await store.get_nearest_matches( + collection_name="col", + embedding=array([0.1, 0.2]), + limit=2, + min_relevance_score=0.5, + ) + assert len(results) == 2 + assert results[0][1] == 0.95 + assert results[1][1] == 0.80 + + +@pytest.mark.asyncio +async def test_get_nearest_match_returns_none_when_empty(store): + async def mock_post(path, body): + return {"memories": []} + + store._post = mock_post + result = await store.get_nearest_match( + collection_name="col", + embedding=array([0.1]), + min_relevance_score=0.9, + ) + assert result is None + + +@pytest.mark.asyncio +async def test_get_nearest_match_single(store): + search_response = { + "memories": [ + {"memory": {"id": "d1", "content": "best match", "agent_id": "col", "metadata": {_DAKERA_ID_META_KEY: "sk1"}}, "score": 0.99}, + ] + } + + async def mock_post(path, body): + return search_response + + store._post = mock_post + result = await store.get_nearest_match( + collection_name="col", + embedding=array([0.1]), + min_relevance_score=0.5, + ) + assert result is not None + record, score = result + assert score == 0.99 + assert record.text == "best match" + + +# --------------------------------------------------------------------------- +# search_text (convenience method) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_search_text(store): + search_response = { + "memories": [ + {"memory": {"id": "d1", "content": "result", "agent_id": "col", "metadata": {_DAKERA_ID_META_KEY: "sk1"}}, "score": 0.88}, + ] + } + + async def mock_post(path, body): + assert body["query"] == "my query" + return search_response + + store._post = mock_post + results = await store.search_text("col", "my query", top_k=5) + assert len(results) == 1 + assert results[0][1] == 0.88 + + +# --------------------------------------------------------------------------- +# HTTP error handling +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_post_raises_on_4xx(store): + from semantic_kernel.exceptions import ServiceResponseException + + mock_resp = AsyncMock() + mock_resp.status = 401 + mock_resp.text = AsyncMock(return_value="Unauthorized") + mock_resp.__aenter__ = AsyncMock(return_value=mock_resp) + mock_resp.__aexit__ = AsyncMock(return_value=False) + + mock_session = MagicMock() + mock_session.post = MagicMock(return_value=mock_resp) + store._session = mock_session + + with pytest.raises(ServiceResponseException, match="401"): + await store._post("/v1/memory/store", {"content": "x", "agent_id": "y"}) + + +# --------------------------------------------------------------------------- +# Context manager lifecycle +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_context_manager_creates_and_closes_session(): + store = DakeraMemoryStore(url="http://localhost:3300") + assert store._session is None + + with patch("aiohttp.ClientSession") as MockSession: + instance = AsyncMock() + MockSession.return_value = instance + async with store: + assert store._session is not None + + instance.close.assert_called_once() From a10ab6d2fbc0068e8b92e49730d15fa1f5771907 Mon Sep 17 00:00:00 2001 From: Dakera Ops Date: Fri, 3 Jul 2026 07:35:20 +0000 Subject: [PATCH 2/2] =?UTF-8?q?fix(memory):=20address=20review=20=E2=80=94?= =?UTF-8?q?=20official=20Dakera=20SDK,=20faithful=20vector=20search,=20exa?= =?UTF-8?q?ct=20round-trip?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reworks the DakeraMemoryStore connector to resolve all review feedback on the initial revision and to make the integration correct and idiomatic. ## Dependency (unsatisfiable extra → official SDK) - `pyproject.toml` extra was `dakera = ["aiohttp >= 3.9"]`, which both re-pinned a core dependency and read oddly against convention. Replaced with the official async SDK: `dakera = ["dakera[async] >= 0.12.9, < 1.0"]`, mirroring how every other connector declares its vendor library (`chroma → chromadb`, `milvus → pymilvus`). The connector now talks to Dakera through `dakera.AsyncDakeraClient` instead of hand-rolled aiohttp. ## Correctness - **Exact text round-trip.** Writes store the raw SK `text` as Dakera `content` (the description lives only in metadata), so reads reconstruct `MemoryRecord.text` without the previous `"text | description"` pollution. - **`get_nearest_matches` honours the query embedding.** The old code sent `query=collection_name` to a text-search endpoint, so results ignored the caller's query. `MemoryStoreBase` only provides an embedding (never the query text), so the connector now persists the SK-side embedding at write time and ranks candidates client-side with cosine similarity (wiring up the previously-unused `_cosine_similarity` helper) — ranking in the caller's own embedding space. - **`with_embeddings=True` now returns real vectors** hydrated from the persisted embedding instead of an empty array. - **`delete_collection` no longer fails the server safety guard.** It now issues a tag-filtered batch forget (`tags=[collection]`) rather than an unfiltered forget, which the server rejects. - **`upsert` docstring** now matches behaviour: it returns the SK record id (the addressing key); the Dakera UUID is kept on `record._key`. - Removed unused imports (`deepcopy`, `ServiceResourceNotFoundError`, and `json` in tests). ## Docs - Corrected the self-host instructions to the `dakera-deploy` compose stack on port 3000 (the server needs the object store the compose provisions; a bare `docker run` is not sufficient). - Added a Dakera-native `search_text()` path documented alongside the drop-in `get_nearest_matches` compatibility path. ## Tests - 24 unit tests mock the async client (no live server / no network) and cover round-trip fidelity, embedding-based ranking, min-relevance filtering, tag-filtered delete, both-ids batch removal, error wrapping, and client lifecycle ownership. Co-Authored-By: Paperclip --- python/pyproject.toml | 2 +- .../dakera/dakera_memory_store.py | 677 ++++++++---------- .../unit/connectors/memory/test_dakera.py | 480 ++++++------- 3 files changed, 530 insertions(+), 629 deletions(-) diff --git a/python/pyproject.toml b/python/pyproject.toml index 1e1012d5bb31..46eee6da0dfa 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -80,7 +80,7 @@ chroma = [ "chromadb >= 0.5,< 1.4" ] dakera = [ - "aiohttp >= 3.9" + "dakera[async] >= 0.12.9, < 1.0" ] copilotstudio = [ "microsoft-agents-copilotstudio-client >= 0.3.1", diff --git a/python/semantic_kernel/connectors/memory_stores/dakera/dakera_memory_store.py b/python/semantic_kernel/connectors/memory_stores/dakera/dakera_memory_store.py index edcb2ef40937..164f071c6153 100644 --- a/python/semantic_kernel/connectors/memory_stores/dakera/dakera_memory_store.py +++ b/python/semantic_kernel/connectors/memory_stores/dakera/dakera_memory_store.py @@ -2,146 +2,172 @@ """DakeraMemoryStore — Semantic Kernel memory connector for Dakera. -Dakera is a self-hosted REST memory server (https://dakera.ai). -Run it with: +`Dakera `_ is a self-hosted memory server that adds +persistent, decay-weighted vector recall across agent sessions: memories are +importance-scored and decay over time, so stale context stops competing with +fresh, relevant facts. - docker run -p 3300:3300 -e DAKERA_API_KEY=demo ghcr.io/dakera-ai/dakera:latest +Run it locally with the public ``dakera-deploy`` docker-compose stack (the +server needs the object store the compose file provisions, so a bare +``docker run`` is not sufficient):: -The SK ``collection_name`` concept maps to Dakera's ``agent_id`` namespace. -Each collection is an independent agent memory silo inside the same Dakera -instance, so you can have as many logical collections as you like without -needing separate deployments. + git clone https://github.com/dakera-ai/dakera-deploy + cd dakera-deploy + docker compose up -d # serves the REST API on http://localhost:3000 -Note: the legacy ``MemoryStoreBase`` API is deprecated upstream (it will be -removed in a future SK release). New code should prefer the VectorStore API. -This connector targets the legacy API for maximum compatibility with existing -pipelines that still use ``SemanticTextMemory``. +Install the connector extra, which pulls the async Dakera SDK:: + + pip install semantic-kernel[dakera] + +The SK ``collection_name`` concept maps 1-to-1 to Dakera's ``agent_id`` +namespace: each collection is a fully isolated memory silo inside the same +Dakera instance, so you can have as many logical collections as you like +without deploying more than once. + +**Two retrieval paths** + +``MemoryStoreBase.get_nearest_matches`` only receives a query *embedding* +(never the original query text), so this connector persists the SK-side +embedding alongside each memory and ranks candidates client-side with cosine +similarity. That keeps drop-in compatibility with ``SemanticTextMemory`` and +ranks in the caller's own embedding space. + +For the richer, Dakera-native experience use :meth:`DakeraMemoryStore.search_text`, +which sends the raw query text to Dakera's ``recall`` endpoint and benefits +from server-side decay- and importance-weighted ranking. + +Note: the legacy ``MemoryStoreBase`` API is deprecated upstream and will be +removed in a future SK release. New code should prefer the ``VectorStore`` API. +This connector targets the legacy API for compatibility with existing pipelines +that still use ``SemanticTextMemory``. """ -import json import logging import sys -from copy import deepcopy -from typing import Any +from datetime import datetime +from typing import Any, Final -import aiohttp +from dakera import AsyncDakeraClient +from dakera.exceptions import DakeraError, NotFoundError +from dakera.models import BatchForgetRequest, BatchMemoryFilter, BatchRecallRequest from numpy import array, linalg, ndarray -if sys.version_info >= (3, 12): - from typing import override -else: - from typing_extensions import override +from semantic_kernel.exceptions import ServiceResponseException +from semantic_kernel.memory.memory_record import MemoryRecord +from semantic_kernel.memory.memory_store_base import MemoryStoreBase if sys.version_info >= (3, 13): from warnings import deprecated else: from typing_extensions import deprecated -from semantic_kernel.exceptions import ServiceResourceNotFoundError, ServiceResponseException -from semantic_kernel.memory.memory_record import MemoryRecord -from semantic_kernel.memory.memory_store_base import MemoryStoreBase - logger: logging.Logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- -# Internal helpers +# Metadata keys — SK record fields are round-tripped through Dakera's metadata +# blob so a stored record can be reconstructed losslessly on read. # --------------------------------------------------------------------------- -_DAKERA_ID_META_KEY = "_sk_id" -_DAKERA_DESC_META_KEY = "_sk_description" -_DAKERA_IS_REF_META_KEY = "_sk_is_reference" -_DAKERA_EXT_SRC_META_KEY = "_sk_external_source_name" -_DAKERA_ADD_META_KEY = "_sk_additional_metadata" -_DAKERA_TIMESTAMP_META_KEY = "_sk_timestamp" +_SK_ID: Final[str] = "_sk_id" +_SK_DESCRIPTION: Final[str] = "_sk_description" +_SK_IS_REFERENCE: Final[str] = "_sk_is_reference" +_SK_EXTERNAL_SOURCE: Final[str] = "_sk_external_source_name" +_SK_ADDITIONAL_METADATA: Final[str] = "_sk_additional_metadata" +_SK_TIMESTAMP: Final[str] = "_sk_timestamp" +_SK_EMBEDDING: Final[str] = "_sk_embedding" +# Tag applied to every memory this connector writes. It doubles as the +# collection marker (we also tag with the collection name) so a whole +# collection can be listed or purged with a single tag filter. +_SK_TAG: Final[str] = "sk" -def _record_to_payload(collection_name: str, record: MemoryRecord) -> dict[str, Any]: - """Convert a MemoryRecord into a Dakera ``/v1/memory/store`` payload.""" - # Dakera content field receives the primary text. We embed all SK - # metadata fields as a JSON blob in the Dakera ``metadata`` dict so - # they survive a round-trip through the API. - metadata: dict[str, Any] = { - _DAKERA_ID_META_KEY: record._id, - _DAKERA_DESC_META_KEY: record._description or "", - _DAKERA_IS_REF_META_KEY: record._is_reference, - _DAKERA_EXT_SRC_META_KEY: record._external_source_name or "", - _DAKERA_ADD_META_KEY: record._additional_metadata or "", - _DAKERA_TIMESTAMP_META_KEY: str(record._timestamp) if record._timestamp else "", - } +# How many memories to page in when listing/scanning a collection. Dakera's +# batch-recall endpoint is filter-based (no query embedding required); this +# bounds the working set for client-side cosine ranking and id resolution. +_DEFAULT_FETCH_LIMIT: Final[int] = 1000 - # Build a human-readable content string that contains both the text and - # the description, because Dakera's semantic search operates on ``content``. - parts = [] - if record._text: - parts.append(record._text) - if record._description: - parts.append(record._description) - content = " | ".join(parts) if parts else record._id - - return { - "content": content, - "agent_id": collection_name, - # Use the SK record id as the Dakera session_id so that - # ``get()``/``remove()`` can address individual records. - # This is a lightweight mapping; a production connector might - # use a dedicated index field instead. - "metadata": metadata, - # importance range 0–1; default to middle if not set - "importance": 0.5, - "tags": [collection_name, "sk"], - } +# Backwards-compatible alias kept for callers/tests that imported the old name. +_DAKERA_ID_META_KEY = _SK_ID -def _payload_to_record(item: dict[str, Any], with_embedding: bool = False) -> MemoryRecord: - """Convert a Dakera memory item dict back into a MemoryRecord. +# --------------------------------------------------------------------------- +# Conversion helpers +# --------------------------------------------------------------------------- - ``item`` is the ``memory`` object returned by Dakera's API, shaped as: - ``{id, content, agent_id, metadata, created_at, ...}`` - """ - meta: dict[str, Any] = item.get("metadata") or {} - sk_id: str = meta.get(_DAKERA_ID_META_KEY) or item.get("id", "") - description: str | None = meta.get(_DAKERA_DESC_META_KEY) or None - is_reference: bool = bool(meta.get(_DAKERA_IS_REF_META_KEY, False)) - external_source_name: str | None = meta.get(_DAKERA_EXT_SRC_META_KEY) or None - additional_metadata: str | None = meta.get(_DAKERA_ADD_META_KEY) or None - timestamp_str: str = meta.get(_DAKERA_TIMESTAMP_META_KEY, "") +def _record_to_metadata(record: MemoryRecord) -> dict[str, Any]: + """Serialize the SK-specific fields of a MemoryRecord into a metadata dict. + + The embedding is stored too (as a plain float list) so that + ``get_nearest_matches`` can rank client-side and ``with_embedding=True`` + reads can hydrate the original vector. + """ + metadata: dict[str, Any] = { + _SK_ID: record._id, + _SK_DESCRIPTION: record._description or "", + _SK_IS_REFERENCE: bool(record._is_reference), + _SK_EXTERNAL_SOURCE: record._external_source_name or "", + _SK_ADDITIONAL_METADATA: record._additional_metadata or "", + _SK_TIMESTAMP: record._timestamp.isoformat() if record._timestamp else "", + } + embedding = record._embedding + if embedding is not None and getattr(embedding, "size", 0): + metadata[_SK_EMBEDDING] = array(embedding, dtype=float).tolist() + return metadata + + +def _to_record( + dakera_id: str, + content: str, + metadata: dict[str, Any] | None, + with_embedding: bool, +) -> MemoryRecord: + """Reconstruct a MemoryRecord from a Dakera memory (id + content + metadata). + + ``content`` maps straight back to ``MemoryRecord.text`` — the connector + never packs the description into the content, so the round-trip is exact. + """ + meta: dict[str, Any] = metadata or {} - from datetime import datetime + embedding: ndarray = array([]) + if with_embedding: + raw = meta.get(_SK_EMBEDDING) + if raw: + embedding = array(raw, dtype=float) - timestamp = None + timestamp: datetime | None = None + timestamp_str: str = meta.get(_SK_TIMESTAMP) or "" if timestamp_str: try: timestamp = datetime.fromisoformat(timestamp_str) except ValueError: - pass - - # Dakera does not persist raw float vectors — it handles embedding - # internally. We return a zero-length ndarray when embeddings are - # not available (which is normal for Dakera-backed stores). - embedding: ndarray = array([]) if not with_embedding else array([]) + logger.debug("DakeraMemoryStore: could not parse timestamp %r", timestamp_str) return MemoryRecord( - is_reference=is_reference, - external_source_name=external_source_name, - id=sk_id, - description=description, - text=item.get("content", ""), - additional_metadata=additional_metadata, + is_reference=bool(meta.get(_SK_IS_REFERENCE)), + external_source_name=meta.get(_SK_EXTERNAL_SOURCE) or None, + id=meta.get(_SK_ID) or dakera_id, + description=meta.get(_SK_DESCRIPTION) or None, + text=content, + additional_metadata=meta.get(_SK_ADDITIONAL_METADATA) or None, embedding=embedding, - key=item.get("id"), # Dakera UUID as the internal key + key=dakera_id, # Dakera's server-assigned UUID timestamp=timestamp, ) def _cosine_similarity(query: ndarray, vectors: ndarray) -> ndarray: - """Compute cosine similarity between a query vector and a matrix of vectors.""" + """Cosine similarity between a query vector and a matrix of row vectors. + + Rows whose norm is zero (or that mismatch the query dimensionality) score + ``-1.0`` so they sort last rather than raising. + """ query_norm = linalg.norm(query) col_norms = linalg.norm(vectors, axis=1) - valid = (query_norm != 0) & (col_norms != 0) scores = array([-1.0] * vectors.shape[0]) + valid = (query_norm != 0) & (col_norms != 0) if valid.any(): - scores[valid] = query.dot(vectors[valid].T) / (query_norm * col_norms[valid]) + scores[valid] = vectors[valid].dot(query) / (col_norms[valid] * query_norm) return scores @@ -152,66 +178,53 @@ def _cosine_similarity(query: ndarray, vectors: ndarray) -> ndarray: @deprecated( "DakeraMemoryStore uses the deprecated MemoryStoreBase API which will be " - "removed in a future SK version. Consider migrating to the VectorStore API." + "removed in a future SK version. Consider migrating to the VectorStore API." ) class DakeraMemoryStore(MemoryStoreBase): - """Semantic Kernel memory store backed by Dakera. + """Semantic Kernel memory store backed by a Dakera server. - Dakera is a self-hosted, decay-weighted vector memory server. This - connector bridges SK's legacy ``MemoryStoreBase`` contract to Dakera's - REST API so that any pipeline built on ``SemanticTextMemory`` can use - Dakera as its backing store with zero code changes outside the - constructor call. + This connector bridges SK's legacy ``MemoryStoreBase`` contract to Dakera's + REST API (via the official async ``dakera`` SDK) so that any pipeline built + on ``SemanticTextMemory`` can use Dakera as its backing store with no code + changes outside the constructor. **Collection → agent_id mapping** - SK's concept of a *collection* maps 1-to-1 to Dakera's ``agent_id``. - Each ``agent_id`` is a fully isolated memory namespace inside the - Dakera instance. ``create_collection`` is a no-op (Dakera lazily - creates namespaces on first write); ``delete_collection`` calls - ``POST /v1/memory/forget`` with no specific memory IDs, which wipes - the entire agent namespace. - - **Embedding / similarity** - - Dakera manages its own HNSW embedding index server-side. - ``get_nearest_matches`` therefore delegates directly to Dakera's - ``POST /v1/memory/search`` endpoint and uses the returned relevance - scores. The ``embedding`` ndarray parameter is accepted for API - compatibility but is **not** sent to Dakera (Dakera re-embeds the - query text internally). If the caller also needs to use the raw - embedding for downstream cosine ranking (e.g. to re-rank Dakera - results), this can be layered on top, but by default the Dakera - relevance scores are used directly. - - Args: - url: Base URL of the Dakera server, e.g. ``http://localhost:3300``. - api_key: Optional bearer token. If ``None`` or empty, the - ``Authorization`` header is omitted (useful when Dakera is - deployed without auth). - agent_id: Default agent namespace used when no collection name is - passed explicitly. In SK's usage pattern the collection name - always overrides this. - session: Optional existing ``aiohttp.ClientSession``. If ``None``, - a new session is created and owned by this instance. + An SK *collection* maps 1-to-1 to a Dakera ``agent_id`` — a fully isolated + memory namespace. ``create_collection`` is a no-op registration (Dakera + creates namespaces lazily on first write); ``delete_collection`` purges the + namespace with a single tag-filtered batch forget. """ def __init__( self, - url: str = "http://localhost:3300", + url: str = "http://localhost:3000", api_key: str | None = None, - agent_id: str = "sk-default", - session: aiohttp.ClientSession | None = None, + *, + client: AsyncDakeraClient | None = None, + fetch_limit: int = _DEFAULT_FETCH_LIMIT, ) -> None: - self._url = url.rstrip("/") - self._api_key = api_key or "" - self._default_agent_id = agent_id - self._owns_session = session is None - self._session = session # may be None until first use - - # Track which collections (agent_ids) we have seen so we can - # answer get_collections() and does_collection_exist() without a - # dedicated Dakera list-agents endpoint. + """Initialize a new instance of the DakeraMemoryStore class. + + Args: + url: Base URL of the Dakera server, e.g. ``http://localhost:3000``. + api_key: Optional API key (looks like ``dk-...``). When ``None`` or + empty, requests are sent without an ``Authorization`` header, + which is fine for a Dakera instance deployed without auth. + client: An existing :class:`~dakera.AsyncDakeraClient` to reuse. When + provided, the store does not own its lifecycle and will not close + it. When omitted, a client is created and owned by this instance. + fetch_limit: Maximum number of memories paged in when listing or + scanning a collection (client-side cosine ranking and id + resolution). Defaults to ``1000``. + """ + self._client: AsyncDakeraClient = client or AsyncDakeraClient(base_url=url, api_key=api_key or None) + self._owns_client: bool = client is None + self._fetch_limit: int = fetch_limit + + # Dakera has no list-all-agents endpoint, so we track collections that + # have been created or written to during this store's lifetime to + # answer get_collections()/does_collection_exist() without a round-trip. self._known_collections: set[str] = set() # ------------------------------------------------------------------ @@ -219,194 +232,144 @@ def __init__( # ------------------------------------------------------------------ async def __aenter__(self) -> "DakeraMemoryStore": - if self._session is None: - self._session = self._make_session() + """Enter the async context, returning this store.""" return self async def __aexit__(self, *args: Any) -> None: + """Exit the async context, closing an owned client.""" await self.close() async def close(self) -> None: - """Close the underlying aiohttp session (only if we created it).""" - if self._owns_session and self._session is not None: - await self._session.close() - self._session = None - - def _make_session(self) -> aiohttp.ClientSession: - headers: dict[str, str] = {"Content-Type": "application/json", "Accept": "application/json"} - if self._api_key: - headers["Authorization"] = f"Bearer {self._api_key}" - return aiohttp.ClientSession(headers=headers) - - def _get_session(self) -> aiohttp.ClientSession: - """Return (and lazily create) the HTTP session.""" - if self._session is None: - self._session = self._make_session() - return self._session - - async def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]: - """Execute a POST request and return the parsed JSON body. - - Raises: - ServiceResponseException: if the server returns a non-2xx status. - """ - session = self._get_session() - url = f"{self._url}{path}" - async with session.post(url, data=json.dumps(body)) as resp: - raw = await resp.text() - if resp.status >= 400: - raise ServiceResponseException( - f"Dakera API error {resp.status} on POST {path}: {raw[:500]}" - ) - try: - return json.loads(raw) - except json.JSONDecodeError as exc: - raise ServiceResponseException( - f"Dakera returned non-JSON response on POST {path}: {raw[:200]}" - ) from exc + """Close the underlying Dakera client (only if this store created it).""" + if self._owns_client: + await self._client.close() # ------------------------------------------------------------------ # Collection management # ------------------------------------------------------------------ - @override async def create_collection(self, collection_name: str) -> None: """Register a collection name locally. - Dakera has no explicit collection-creation endpoint; namespaces - are created lazily on the first write. We record the name so - ``get_collections`` and ``does_collection_exist`` can answer - without a round-trip. + Dakera creates namespaces lazily on the first write, so this only + records the name for ``get_collections``/``does_collection_exist``. """ self._known_collections.add(collection_name) logger.debug("DakeraMemoryStore: registered collection '%s'", collection_name) - @override async def get_collections(self) -> list[str]: - """Return all collection names that have been used in this session. + """Return the collection names seen during this store's lifetime. - Note: Dakera does not currently expose a list-all-agents endpoint. - This method returns the in-process set of collections that were - created or written to since the store was instantiated. For a - persistent view across restarts, callers should call + Dakera does not expose a list-all-agents endpoint, so this returns the + in-process set of collections created or written to since the store was + instantiated. For a persistent view across restarts, call ``create_collection`` for each known collection at startup. """ return list(self._known_collections) - @override async def delete_collection(self, collection_name: str) -> None: - """Delete all memories in the Dakera agent namespace for this collection. + """Delete every memory this connector stored in the collection. - Calls ``POST /v1/memory/forget`` with only the ``agent_id`` set - (no specific memory IDs), which instructs Dakera to wipe all - memories for that agent. + Issues a single tag-filtered batch forget (``tags=[collection_name]``), + which satisfies Dakera's safety guard that a bulk delete must carry at + least one filter predicate. """ + request = BatchForgetRequest( + agent_id=collection_name, + filter=BatchMemoryFilter(tags=[collection_name]), + ) try: - await self._post("/v1/memory/forget", {"agent_id": collection_name}) - except ServiceResponseException as exc: - # A 404 means there is nothing to delete — treat as success. - if "404" in str(exc): - pass - else: - raise + await self._client.batch_forget(request) + except NotFoundError: + # Nothing to delete — treat as success. + pass + except DakeraError as exc: + raise ServiceResponseException(f"Dakera batch_forget failed for '{collection_name}': {exc}") from exc self._known_collections.discard(collection_name) logger.debug("DakeraMemoryStore: deleted collection '%s'", collection_name) - @override async def does_collection_exist(self, collection_name: str) -> bool: - """Return True if we have previously seen this collection in the current session.""" + """Return True if this collection has been seen during the store's lifetime.""" return collection_name in self._known_collections # ------------------------------------------------------------------ # Write operations # ------------------------------------------------------------------ - @override async def upsert(self, collection_name: str, record: MemoryRecord) -> str: """Store a single MemoryRecord in Dakera. Returns: - str: The Dakera-assigned UUID for the stored memory. + str: The SK record id — the key that addresses this record in + subsequent ``get``/``remove`` calls. (Dakera's own server-assigned + UUID is retained on ``record._key`` for internal use.) """ - payload = _record_to_payload(collection_name, record) - response = await self._post("/v1/memory/store", payload) + content = record._text or record._description or record._id + try: + memory = await self._client.store_memory( + agent_id=collection_name, + content=content, + metadata=_record_to_metadata(record), + tags=[collection_name, _SK_TAG], + ) + except DakeraError as exc: + raise ServiceResponseException(f"Dakera store_memory failed: {exc}") from exc - memory_obj = response.get("memory") or {} - dakera_id: str = memory_obj.get("id", "") + dakera_id = memory.get("id") if isinstance(memory, dict) else None if not dakera_id: - raise ServiceResponseException( - f"Dakera did not return a memory id after store. Response: {response}" - ) + raise ServiceResponseException(f"Dakera did not return a memory id after store. Response: {memory!r}") - # Keep the SK-level key pointing to the original record id. - record._key = record._id + record._key = dakera_id self._known_collections.add(collection_name) - return record._key + return record._id - @override async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]: - """Store multiple MemoryRecords in Dakera. - - Dakera has no batch-store endpoint, so we issue one POST per record. - For large batches consider chunking with asyncio.gather. - """ + """Store multiple MemoryRecords, returning their SK ids in order.""" import asyncio - tasks = [self.upsert(collection_name, record) for record in records] - return list(await asyncio.gather(*tasks)) + return list(await asyncio.gather(*(self.upsert(collection_name, r) for r in records))) # ------------------------------------------------------------------ # Read operations # ------------------------------------------------------------------ - @override async def get(self, collection_name: str, key: str, with_embedding: bool = False) -> MemoryRecord | None: """Retrieve a single MemoryRecord by its SK id. - Dakera's ``/v1/memory/search`` is used with the ``key`` as the - query; we then match on the stored ``_sk_id`` metadata field. - This is an approximation: Dakera returns semantically similar - results, so we filter the list to find the exact match. + Resolves the SK id against the collection's stored ``_sk_id`` metadata + (Dakera addresses records by its own UUID, not the SK id). Returns: - MemoryRecord if found, None otherwise. + The MemoryRecord if found, otherwise ``None``. """ - results = await self._search_by_query( - collection_name=collection_name, - query=key, - top_k=20, - ) - for memory_obj, _score in results: - meta = memory_obj.get("metadata") or {} - if meta.get(_DAKERA_ID_META_KEY) == key: - return _payload_to_record(memory_obj, with_embedding) + for memory in await self._list_collection(collection_name): + meta = memory.metadata or {} + if meta.get(_SK_ID) == key: + return _to_record(memory.id, memory.content, meta, with_embedding) return None - @override async def get_batch( self, collection_name: str, keys: list[str], with_embeddings: bool = False, ) -> list[MemoryRecord]: - """Retrieve multiple MemoryRecords by their SK ids.""" - import asyncio - - tasks = [self.get(collection_name, key, with_embeddings) for key in keys] - results = await asyncio.gather(*tasks) - return [r for r in results if r is not None] + """Retrieve multiple MemoryRecords by their SK ids (single collection scan).""" + wanted = set(keys) + results: list[MemoryRecord] = [] + for memory in await self._list_collection(collection_name): + meta = memory.metadata or {} + sk_id = meta.get(_SK_ID) + if sk_id in wanted: + results.append(_to_record(memory.id, memory.content, meta, with_embeddings)) + return results # ------------------------------------------------------------------ # Delete operations # ------------------------------------------------------------------ - @override async def remove(self, collection_name: str, key: str) -> None: - """Remove a single MemoryRecord by its SK id. - - We first search for the Dakera UUID that corresponds to this SK id, - then call ``POST /v1/memory/forget`` with that specific UUID. - """ + """Remove a single MemoryRecord by its SK id (no-op if not found).""" dakera_id = await self._resolve_dakera_id(collection_name, key) if dakera_id is None: logger.warning( @@ -415,33 +378,26 @@ async def remove(self, collection_name: str, key: str) -> None: collection_name, ) return - await self._post( - "/v1/memory/forget", - {"agent_id": collection_name, "memory_ids": [dakera_id]}, - ) + await self._forget(collection_name, dakera_id) - @override async def remove_batch(self, collection_name: str, keys: list[str]) -> None: - """Remove multiple MemoryRecords by their SK ids.""" + """Remove multiple MemoryRecords by their SK ids (single collection scan).""" import asyncio - tasks = [self._resolve_dakera_id(collection_name, key) for key in keys] - dakera_ids_or_none = await asyncio.gather(*tasks) - dakera_ids = [d for d in dakera_ids_or_none if d is not None] - + wanted = set(keys) + dakera_ids = [ + memory.id + for memory in await self._list_collection(collection_name) + if (memory.metadata or {}).get(_SK_ID) in wanted + ] if not dakera_ids: return - - await self._post( - "/v1/memory/forget", - {"agent_id": collection_name, "memory_ids": dakera_ids}, - ) + await asyncio.gather(*(self._forget(collection_name, did) for did in dakera_ids)) # ------------------------------------------------------------------ # Similarity search # ------------------------------------------------------------------ - @override async def get_nearest_matches( self, collection_name: str, @@ -450,46 +406,54 @@ async def get_nearest_matches( min_relevance_score: float = 0.0, with_embeddings: bool = False, ) -> list[tuple[MemoryRecord, float]]: - """Find the nearest memories using Dakera's server-side vector search. + """Find the nearest memories to ``embedding`` by cosine similarity. - The ``embedding`` ndarray is accepted for interface compatibility but - is not forwarded to Dakera — Dakera re-embeds the query text on the - server side. The query text is reconstructed from the description - metadata if available; otherwise a generic search is performed. - - Note: because the embedding cannot be round-tripped through Dakera's - API, ``with_embeddings=True`` will return records with empty - embeddings (``array([])``). If you need raw embeddings, maintain a - separate embedding store alongside Dakera. + ``MemoryStoreBase`` passes only the query embedding (not the query + text), so results are ranked client-side against the embeddings this + connector persisted at write time — i.e. in the caller's own embedding + space. For Dakera's native, server-side decay-weighted ranking over a + text query, use :meth:`search_text` instead. Args: - collection_name: The Dakera agent_id namespace to search in. - embedding: The query embedding (used for interface compat only). + collection_name: The Dakera ``agent_id`` namespace to search. + embedding: The query embedding. limit: Maximum number of matches to return. - min_relevance_score: Minimum relevance score threshold. - with_embeddings: If True, attempt to include embeddings (will be empty arrays). + min_relevance_score: Minimum cosine score (inclusive) to keep. + with_embeddings: When True, hydrate each result's stored embedding. Returns: - List of (MemoryRecord, score) tuples sorted by score descending. + ``(MemoryRecord, score)`` tuples sorted by score descending. """ - raw_results = await self._search_by_query( - collection_name=collection_name, - query=collection_name, # fallback: search by collection name itself - top_k=limit * 2, # over-fetch to allow filtering + memories = await self._list_collection(collection_name) + + scored: list[tuple[Any, list[float]]] = [] + for memory in memories: + raw = (memory.metadata or {}).get(_SK_EMBEDDING) + if raw: + scored.append((memory, raw)) + if not scored: + return [] + + query = array(embedding, dtype=float) + matrix = array([vec for _memory, vec in scored], dtype=float) + similarities = _cosine_similarity(query, matrix) + + ranked = sorted( + zip((memory for memory, _vec in scored), similarities), + key=lambda pair: pair[1], + reverse=True, ) output: list[tuple[MemoryRecord, float]] = [] - for memory_obj, score in raw_results: + for memory, score in ranked: + score = float(score) if score < min_relevance_score: continue - record = _payload_to_record(memory_obj, with_embeddings) - output.append((record, score)) + output.append((_to_record(memory.id, memory.content, memory.metadata, with_embeddings), score)) if len(output) >= limit: break - return output - @override async def get_nearest_match( self, collection_name: str, @@ -497,11 +461,7 @@ async def get_nearest_match( min_relevance_score: float = 0.0, with_embedding: bool = False, ) -> tuple[MemoryRecord, float] | None: - """Find the single nearest memory. - - Returns: - A (MemoryRecord, score) tuple or None if no match meets the threshold. - """ + """Find the single nearest memory, or ``None`` if none meet the threshold.""" results = await self.get_nearest_matches( collection_name=collection_name, embedding=embedding, @@ -512,58 +472,7 @@ async def get_nearest_match( return results[0] if results else None # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ - - async def _search_by_query( - self, - collection_name: str, - query: str, - top_k: int = 10, - ) -> list[tuple[dict[str, Any], float]]: - """Execute ``POST /v1/memory/search`` and return (memory_obj, score) pairs.""" - body: dict[str, Any] = { - "agent_id": collection_name, - "query": query, - "top_k": top_k, - } - try: - response = await self._post("/v1/memory/search", body) - except ServiceResponseException as exc: - if "404" in str(exc): - return [] - raise - - items = response.get("memories") or [] - results: list[tuple[dict[str, Any], float]] = [] - for item in items: - memory_obj = item.get("memory") or item - score = float(item.get("score", 0.0)) - results.append((memory_obj, score)) - return results - - async def _resolve_dakera_id( - self, - collection_name: str, - sk_key: str, - ) -> str | None: - """Look up the Dakera internal UUID for a given SK record id. - - Returns the Dakera UUID string, or None if not found. - """ - candidates = await self._search_by_query( - collection_name=collection_name, - query=sk_key, - top_k=20, - ) - for memory_obj, _score in candidates: - meta = memory_obj.get("metadata") or {} - if meta.get(_DAKERA_ID_META_KEY) == sk_key: - return memory_obj.get("id") - return None - - # ------------------------------------------------------------------ - # Convenience search method (extends base contract) + # Dakera-native text search (extends the base contract) # ------------------------------------------------------------------ async def search_text( @@ -573,27 +482,73 @@ async def search_text( top_k: int = 5, min_relevance_score: float = 0.0, ) -> list[tuple[MemoryRecord, float]]: - """Search Dakera using a raw text query (preferred over embedding-based search). + """Search Dakera with a raw text query (recommended for Dakera-native apps). - This method bypasses the SK embedding pipeline and sends the query - directly to Dakera, which handles embedding internally. Use this - when you want the most natural integration with Dakera's own ranking. + Unlike :meth:`get_nearest_matches`, this sends the query text straight + to Dakera's ``recall`` endpoint, so ranking uses Dakera's server-side + decay- and importance-weighted scoring rather than a client-side + cosine over SK embeddings. Args: - collection_name: The agent namespace to search. + collection_name: The Dakera ``agent_id`` namespace to search. query: Plain-text search query. top_k: Number of results to return. - min_relevance_score: Minimum relevance threshold. + min_relevance_score: Minimum relevance score (inclusive) to keep. Returns: - List of (MemoryRecord, score) tuples sorted by score descending. + ``(MemoryRecord, score)`` tuples sorted by score descending. """ - raw = await self._search_by_query(collection_name, query, top_k * 2) + try: + response = await self._client.recall(agent_id=collection_name, query=query, top_k=top_k) + except NotFoundError: + return [] + except DakeraError as exc: + raise ServiceResponseException(f"Dakera recall failed for '{collection_name}': {exc}") from exc + results: list[tuple[MemoryRecord, float]] = [] - for memory_obj, score in raw: + for memory in response.memories: + score = float(memory.score) if score < min_relevance_score: continue - results.append((_payload_to_record(memory_obj, False), score)) + results.append((_to_record(memory.id, memory.content, memory.metadata, False), score)) if len(results) >= top_k: break return results + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + async def _list_collection(self, collection_name: str) -> list[Any]: + """List the memories this connector stored in a collection. + + Uses Dakera's filter-based batch recall (``tags=[collection_name]``), + which needs no query embedding and returns deterministic results. + """ + request = BatchRecallRequest( + agent_id=collection_name, + filter=BatchMemoryFilter(tags=[collection_name]), + limit=self._fetch_limit, + ) + try: + response = await self._client.batch_recall(request) + except NotFoundError: + return [] + except DakeraError as exc: + raise ServiceResponseException(f"Dakera batch_recall failed for '{collection_name}': {exc}") from exc + return list(response.memories) + + async def _resolve_dakera_id(self, collection_name: str, sk_key: str) -> str | None: + """Return the Dakera UUID for an SK record id, or ``None`` if absent.""" + for memory in await self._list_collection(collection_name): + if (memory.metadata or {}).get(_SK_ID) == sk_key: + return memory.id + return None + + async def _forget(self, collection_name: str, dakera_id: str) -> None: + try: + await self._client.forget(collection_name, dakera_id) + except NotFoundError: + pass + except DakeraError as exc: + raise ServiceResponseException(f"Dakera forget failed for '{dakera_id}': {exc}") from exc diff --git a/python/tests/unit/connectors/memory/test_dakera.py b/python/tests/unit/connectors/memory/test_dakera.py index ab5df50f05b9..0c49386989ea 100644 --- a/python/tests/unit/connectors/memory/test_dakera.py +++ b/python/tests/unit/connectors/memory/test_dakera.py @@ -2,27 +2,29 @@ """Unit tests for DakeraMemoryStore. -Tests mock the aiohttp.ClientSession so no live Dakera instance is needed. -Run with: pytest tests/unit/connectors/memory/test_dakera.py -v +The tests inject a mock ``AsyncDakeraClient`` so no live Dakera instance (and no +network) is required. Run with: + pytest tests/unit/connectors/memory/test_dakera.py -v """ -import json -from unittest.mock import AsyncMock, MagicMock, patch +from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest from numpy import array +from numpy.linalg import norm from semantic_kernel.connectors.memory_stores.dakera.dakera_memory_store import ( + _SK_EMBEDDING, + _SK_ID, DakeraMemoryStore, - _DAKERA_ID_META_KEY, - _payload_to_record, - _record_to_payload, + _record_to_metadata, + _to_record, ) from semantic_kernel.memory.memory_record import MemoryRecord - # --------------------------------------------------------------------------- -# Fixtures +# Fixtures / helpers # --------------------------------------------------------------------------- @@ -37,96 +39,111 @@ def make_record(sk_id: str = "test-id", text: str = "hello world") -> MemoryReco ) -def dakera_memory_item( +def dakera_memory( dakera_uuid: str = "dakera-uuid-1", sk_id: str = "test-id", content: str = "hello world", - score: float = 0.9, -) -> dict: - """Simulate one item in a Dakera search response.""" - return { - "memory": { - "id": dakera_uuid, - "content": content, - "agent_id": "my-collection", - "metadata": { - _DAKERA_ID_META_KEY: sk_id, - "_sk_description": "a test record", - "_sk_is_reference": False, - "_sk_external_source_name": "", - "_sk_additional_metadata": "extra", - "_sk_timestamp": "", - }, - }, - "score": score, + score: float | None = None, + embedding: list[float] | None = None, +): + """Simulate a memory object returned by the Dakera SDK (attribute access).""" + metadata = { + _SK_ID: sk_id, + "_sk_description": "a test record", + "_sk_is_reference": False, + "_sk_external_source_name": "", + "_sk_additional_metadata": "extra", + "_sk_timestamp": "", } + if embedding is not None: + metadata[_SK_EMBEDDING] = embedding + obj = SimpleNamespace(id=dakera_uuid, content=content, metadata=metadata, importance=0.5) + if score is not None: + obj.score = score + return obj @pytest.fixture -def store(): - """Return a DakeraMemoryStore with a mock session.""" - s = DakeraMemoryStore(url="http://localhost:3300", api_key="demo", agent_id="test-agent") - return s +def mock_client() -> AsyncMock: + """An AsyncDakeraClient test double.""" + client = AsyncMock() + client.store_memory = AsyncMock(return_value={"id": "dakera-uuid-1"}) + client.batch_recall = AsyncMock(return_value=SimpleNamespace(memories=[], total=0, filtered=0)) + client.batch_forget = AsyncMock(return_value=SimpleNamespace(deleted_count=0)) + client.forget = AsyncMock(return_value={}) + client.recall = AsyncMock(return_value=SimpleNamespace(memories=[])) + client.close = AsyncMock() + return client -# --------------------------------------------------------------------------- -# Helpers: _record_to_payload / _payload_to_record -# --------------------------------------------------------------------------- +@pytest.fixture +def store(mock_client: AsyncMock) -> DakeraMemoryStore: + return DakeraMemoryStore(url="http://localhost:3000", api_key="dk-demo", client=mock_client) -def test_record_to_payload_round_trip(): - record = make_record() - payload = _record_to_payload("col1", record) - assert payload["agent_id"] == "col1" - assert "hello world" in payload["content"] - assert payload["metadata"][_DAKERA_ID_META_KEY] == "test-id" +def _recall_response(memories): + return SimpleNamespace(memories=memories) -def test_payload_to_record(): - item = dakera_memory_item()["memory"] - record = _payload_to_record(item, with_embedding=False) - assert record.id == "test-id" - assert record.text == "hello world" - assert record._key == "dakera-uuid-1" - assert record._is_reference is False +def _batch_response(memories): + return SimpleNamespace(memories=memories, total=len(memories), filtered=len(memories)) # --------------------------------------------------------------------------- -# Collection management +# Conversion helpers — round-trip fidelity # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_create_collection(store): - await store.create_collection("agents") - assert "agents" in await store.get_collections() +def test_record_to_metadata_stores_embedding_and_fields(): + metadata = _record_to_metadata(make_record()) + assert metadata[_SK_ID] == "test-id" + assert metadata["_sk_description"] == "a test record" + assert metadata[_SK_EMBEDDING] == pytest.approx([0.1, 0.2, 0.3]) -@pytest.mark.asyncio -async def test_does_collection_exist_false(store): - assert not await store.does_collection_exist("nonexistent") +def test_text_round_trip_is_exact(): + """text must not be polluted with the description on read-back.""" + record = make_record(text="just the text") + metadata = _record_to_metadata(record) + # Dakera stores the raw text as `content`; description lives in metadata. + reconstructed = _to_record("dakera-uuid-1", "just the text", metadata, with_embedding=False) + assert reconstructed.text == "just the text" + assert reconstructed.description == "a test record" + assert reconstructed.id == "test-id" + assert reconstructed._key == "dakera-uuid-1" -@pytest.mark.asyncio -async def test_does_collection_exist_after_create(store): - await store.create_collection("exists") - assert await store.does_collection_exist("exists") +def test_to_record_hydrates_embedding_only_when_requested(): + metadata = _record_to_metadata(make_record()) + without = _to_record("d1", "hello world", metadata, with_embedding=False) + assert without._embedding.size == 0 + with_emb = _to_record("d1", "hello world", metadata, with_embedding=True) + assert with_emb._embedding.tolist() == pytest.approx([0.1, 0.2, 0.3]) + + +# --------------------------------------------------------------------------- +# Collection management +# --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_delete_collection(store): - await store.create_collection("to-delete") - forget_response = {} +async def test_create_and_exists(store: DakeraMemoryStore): + assert not await store.does_collection_exist("agents") + await store.create_collection("agents") + assert "agents" in await store.get_collections() + assert await store.does_collection_exist("agents") - async def mock_post(path, body): - if path == "/v1/memory/forget": - return forget_response - pytest.fail(f"Unexpected POST to {path}") - store._post = mock_post +async def test_delete_collection_uses_tag_filter(store: DakeraMemoryStore, mock_client: AsyncMock): + await store.create_collection("to-delete") await store.delete_collection("to-delete") + assert not await store.does_collection_exist("to-delete") + mock_client.batch_forget.assert_awaited_once() + request = mock_client.batch_forget.await_args.args[0] + # Safety guard: a bulk delete must carry at least one filter predicate. + assert request.agent_id == "to-delete" + assert request.filter.tags == ["to-delete"] # --------------------------------------------------------------------------- @@ -134,36 +151,26 @@ async def mock_post(path, body): # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_upsert_returns_sk_id(store): +async def test_upsert_returns_sk_id_and_keeps_dakera_key(store: DakeraMemoryStore, mock_client: AsyncMock): record = make_record() - store_response = {"memory": {"id": "dakera-abc-123", "content": "hello world"}} - - async def mock_post(path, body): - assert path == "/v1/memory/store" - assert body["agent_id"] == "my-collection" - return store_response - - store._post = mock_post key = await store.upsert("my-collection", record) - assert key == "test-id" # SK key is the record's own id - assert "my-collection" in store._known_collections + assert key == "test-id" # SK id addresses the record + assert record._key == "dakera-uuid-1" # Dakera UUID retained internally + assert "my-collection" in store._known_collections -@pytest.mark.asyncio -async def test_upsert_batch(store): - records = [make_record(sk_id=f"id-{i}") for i in range(3)] - call_count = 0 + call = mock_client.store_memory.await_args.kwargs + assert call["agent_id"] == "my-collection" + assert call["content"] == "hello world" # raw text, not "text | description" + assert call["tags"] == ["my-collection", "sk"] + assert call["metadata"][_SK_ID] == "test-id" - async def mock_post(path, body): - nonlocal call_count - call_count += 1 - return {"memory": {"id": f"dakera-{call_count}", "content": "x"}} - store._post = mock_post +async def test_upsert_batch(store: DakeraMemoryStore, mock_client: AsyncMock): + records = [make_record(sk_id=f"id-{i}") for i in range(3)] keys = await store.upsert_batch("col", records) - assert len(keys) == 3 - assert call_count == 3 + assert keys == ["id-0", "id-1", "id-2"] + assert mock_client.store_memory.await_count == 3 # --------------------------------------------------------------------------- @@ -171,49 +178,28 @@ async def mock_post(path, body): # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_get_found(store): - search_response = {"memories": [dakera_memory_item(sk_id="test-id")]} - - async def mock_post(path, body): - return search_response - - store._post = mock_post +async def test_get_found(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.batch_recall.return_value = _batch_response([dakera_memory(sk_id="test-id")]) record = await store.get("my-collection", "test-id") assert record is not None assert record.id == "test-id" assert record.text == "hello world" -@pytest.mark.asyncio -async def test_get_not_found(store): - # Returns results but none match the key - search_response = {"memories": [dakera_memory_item(sk_id="other-id")]} +async def test_get_not_found(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.batch_recall.return_value = _batch_response([dakera_memory(sk_id="other-id")]) + assert await store.get("my-collection", "missing-key") is None - async def mock_post(path, body): - return search_response - - store._post = mock_post - record = await store.get("my-collection", "missing-key") - assert record is None - - -@pytest.mark.asyncio -async def test_get_batch(store): - search_response = { - "memories": [ - dakera_memory_item(sk_id="id-0"), - dakera_memory_item(sk_id="id-1"), - ] - } - async def mock_post(path, body): - return search_response - - store._post = mock_post +async def test_get_batch(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.batch_recall.return_value = _batch_response([ + dakera_memory(dakera_uuid="d0", sk_id="id-0"), + dakera_memory(dakera_uuid="d1", sk_id="id-1"), + ]) records = await store.get_batch("col", ["id-0", "id-1"]) - # Both keys happen to appear in the search results - assert any(r.id == "id-0" for r in records) + assert {r.id for r in records} == {"id-0", "id-1"} + # A single collection scan resolves the whole batch. + mock_client.batch_recall.assert_awaited_once() # --------------------------------------------------------------------------- @@ -221,188 +207,148 @@ async def mock_post(path, body): # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_remove_existing(store): - calls = [] - - async def mock_post(path, body): - calls.append((path, body)) - if path == "/v1/memory/search": - return {"memories": [dakera_memory_item(dakera_uuid="dakera-uuid-1", sk_id="test-id")]} - if path == "/v1/memory/forget": - return {} - pytest.fail(f"Unexpected path: {path}") - - store._post = mock_post +async def test_remove_existing(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.batch_recall.return_value = _batch_response([ + dakera_memory(dakera_uuid="dakera-uuid-1", sk_id="test-id") + ]) await store.remove("col", "test-id") - - forget_call = next(c for c in calls if c[0] == "/v1/memory/forget") - assert "dakera-uuid-1" in forget_call[1]["memory_ids"] + mock_client.forget.assert_awaited_once_with("col", "dakera-uuid-1") -@pytest.mark.asyncio -async def test_remove_not_found_is_silent(store): - async def mock_post(path, body): - if path == "/v1/memory/search": - return {"memories": []} # nothing found - pytest.fail(f"Should not reach {path}") +async def test_remove_not_found_is_silent(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.batch_recall.return_value = _batch_response([]) + await store.remove("col", "ghost-key") # must not raise + mock_client.forget.assert_not_awaited() - store._post = mock_post - # Should not raise - await store.remove("col", "ghost-key") - -@pytest.mark.asyncio -async def test_remove_batch(store): - calls = [] - - async def mock_post(path, body): - calls.append((path, body)) - if path == "/v1/memory/search": - return { - "memories": [ - dakera_memory_item(dakera_uuid="uuid-0", sk_id="id-0"), - dakera_memory_item(dakera_uuid="uuid-1", sk_id="id-1"), - ] - } - if path == "/v1/memory/forget": - return {} - - store._post = mock_post +async def test_remove_batch_forgets_all_resolved_ids(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.batch_recall.return_value = _batch_response([ + dakera_memory(dakera_uuid="uuid-0", sk_id="id-0"), + dakera_memory(dakera_uuid="uuid-1", sk_id="id-1"), + ]) await store.remove_batch("col", ["id-0", "id-1"]) - forget_calls = [c for c in calls if c[0] == "/v1/memory/forget"] - assert len(forget_calls) == 1 # single batched forget - ids_sent = forget_calls[0][1]["memory_ids"] - assert "uuid-0" in ids_sent or "uuid-1" in ids_sent + + forgotten = {call.args[1] for call in mock_client.forget.await_args_list} + # Both resolved UUIDs must be forgotten (regression guard). + assert "uuid-0" in forgotten and "uuid-1" in forgotten # --------------------------------------------------------------------------- -# get_nearest_matches / get_nearest_match +# get_nearest_matches / get_nearest_match — client-side cosine ranking # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_get_nearest_matches(store): - search_response = { - "memories": [ - {"memory": {"id": "d1", "content": "match 1", "agent_id": "col", "metadata": {_DAKERA_ID_META_KEY: "sk1"}}, "score": 0.95}, - {"memory": {"id": "d2", "content": "match 2", "agent_id": "col", "metadata": {_DAKERA_ID_META_KEY: "sk2"}}, "score": 0.80}, - {"memory": {"id": "d3", "content": "weak match", "agent_id": "col", "metadata": {_DAKERA_ID_META_KEY: "sk3"}}, "score": 0.30}, - ] - } - - async def mock_post(path, body): - return search_response - - store._post = mock_post - results = await store.get_nearest_matches( - collection_name="col", - embedding=array([0.1, 0.2]), - limit=2, - min_relevance_score=0.5, - ) +async def test_get_nearest_matches_ranks_by_query_embedding(store: DakeraMemoryStore, mock_client: AsyncMock): + # Three stored vectors; the query is colinear with the middle one. + mock_client.batch_recall.return_value = _batch_response([ + dakera_memory(dakera_uuid="d1", sk_id="sk1", content="up", embedding=[1.0, 0.0]), + dakera_memory(dakera_uuid="d2", sk_id="sk2", content="diag", embedding=[1.0, 1.0]), + dakera_memory(dakera_uuid="d3", sk_id="sk3", content="right", embedding=[0.0, 1.0]), + ]) + results = await store.get_nearest_matches("col", embedding=array([1.0, 1.0]), limit=2) assert len(results) == 2 - assert results[0][1] == 0.95 - assert results[1][1] == 0.80 + # The colinear "diag" vector must rank first with score ~1.0. + assert results[0][0].text == "diag" + assert results[0][1] == pytest.approx(1.0) + # Results are sorted descending by score. + assert results[0][1] >= results[1][1] + + +async def test_get_nearest_matches_applies_min_relevance(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.batch_recall.return_value = _batch_response([ + dakera_memory(dakera_uuid="d1", sk_id="sk1", embedding=[1.0, 0.0]), + dakera_memory(dakera_uuid="d2", sk_id="sk2", embedding=[0.0, 1.0]), + ]) + # Query aligned with d1; d2 is orthogonal (score 0) and filtered out. + results = await store.get_nearest_matches("col", embedding=array([1.0, 0.0]), limit=5, min_relevance_score=0.5) + assert len(results) == 1 + assert results[0][0].id == "sk1" -@pytest.mark.asyncio -async def test_get_nearest_match_returns_none_when_empty(store): - async def mock_post(path, body): - return {"memories": []} +async def test_get_nearest_matches_ignores_memories_without_embedding(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.batch_recall.return_value = _batch_response([dakera_memory(sk_id="sk1")]) # no embedding stored + assert await store.get_nearest_matches("col", embedding=array([1.0, 0.0]), limit=5) == [] - store._post = mock_post - result = await store.get_nearest_match( - collection_name="col", - embedding=array([0.1]), - min_relevance_score=0.9, - ) - assert result is None +async def test_get_nearest_matches_hydrates_embeddings_when_requested(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.batch_recall.return_value = _batch_response([ + dakera_memory(dakera_uuid="d1", sk_id="sk1", embedding=[0.3, 0.4]) + ]) + results = await store.get_nearest_matches("col", embedding=array([0.3, 0.4]), limit=1, with_embeddings=True) + record, score = results[0] + assert record._embedding.tolist() == pytest.approx([0.3, 0.4]) + assert score == pytest.approx(float((array([0.3, 0.4]) @ array([0.3, 0.4])) / norm([0.3, 0.4]) ** 2)) -@pytest.mark.asyncio -async def test_get_nearest_match_single(store): - search_response = { - "memories": [ - {"memory": {"id": "d1", "content": "best match", "agent_id": "col", "metadata": {_DAKERA_ID_META_KEY: "sk1"}}, "score": 0.99}, - ] - } - async def mock_post(path, body): - return search_response +async def test_get_nearest_match_returns_none_when_empty(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.batch_recall.return_value = _batch_response([]) + assert await store.get_nearest_match("col", embedding=array([0.1, 0.2]), min_relevance_score=0.9) is None - store._post = mock_post - result = await store.get_nearest_match( - collection_name="col", - embedding=array([0.1]), - min_relevance_score=0.5, - ) + +async def test_get_nearest_match_single(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.batch_recall.return_value = _batch_response([ + dakera_memory(dakera_uuid="d1", sk_id="sk1", content="best match", embedding=[1.0, 0.0]) + ]) + result = await store.get_nearest_match("col", embedding=array([1.0, 0.0]), min_relevance_score=0.5) assert result is not None record, score = result - assert score == 0.99 assert record.text == "best match" + assert score == pytest.approx(1.0) # --------------------------------------------------------------------------- -# search_text (convenience method) +# search_text — Dakera-native recall path # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_search_text(store): - search_response = { - "memories": [ - {"memory": {"id": "d1", "content": "result", "agent_id": "col", "metadata": {_DAKERA_ID_META_KEY: "sk1"}}, "score": 0.88}, - ] - } - - async def mock_post(path, body): - assert body["query"] == "my query" - return search_response - - store._post = mock_post +async def test_search_text_uses_recall(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.recall.return_value = _recall_response([ + dakera_memory(dakera_uuid="d1", sk_id="sk1", content="result", score=0.88) + ]) results = await store.search_text("col", "my query", top_k=5) + + mock_client.recall.assert_awaited_once_with(agent_id="col", query="my query", top_k=5) assert len(results) == 1 - assert results[0][1] == 0.88 + assert results[0][0].text == "result" + assert results[0][1] == pytest.approx(0.88) -# --------------------------------------------------------------------------- -# HTTP error handling -# --------------------------------------------------------------------------- +async def test_search_text_applies_min_relevance(store: DakeraMemoryStore, mock_client: AsyncMock): + mock_client.recall.return_value = _recall_response([ + dakera_memory(dakera_uuid="d1", sk_id="sk1", content="strong", score=0.9), + dakera_memory(dakera_uuid="d2", sk_id="sk2", content="weak", score=0.1), + ]) + results = await store.search_text("col", "q", top_k=5, min_relevance_score=0.5) + assert [r[0].text for r in results] == ["strong"] -@pytest.mark.asyncio -async def test_post_raises_on_4xx(store): - from semantic_kernel.exceptions import ServiceResponseException +# --------------------------------------------------------------------------- +# Error handling / lifecycle +# --------------------------------------------------------------------------- - mock_resp = AsyncMock() - mock_resp.status = 401 - mock_resp.text = AsyncMock(return_value="Unauthorized") - mock_resp.__aenter__ = AsyncMock(return_value=mock_resp) - mock_resp.__aexit__ = AsyncMock(return_value=False) - mock_session = MagicMock() - mock_session.post = MagicMock(return_value=mock_resp) - store._session = mock_session +async def test_store_error_is_wrapped(store: DakeraMemoryStore, mock_client: AsyncMock): + from dakera.exceptions import DakeraError - with pytest.raises(ServiceResponseException, match="401"): - await store._post("/v1/memory/store", {"content": "x", "agent_id": "y"}) + from semantic_kernel.exceptions import ServiceResponseException - -# --------------------------------------------------------------------------- -# Context manager lifecycle -# --------------------------------------------------------------------------- + mock_client.store_memory.side_effect = DakeraError("boom") + with pytest.raises(ServiceResponseException): + await store.upsert("col", make_record()) -@pytest.mark.asyncio -async def test_context_manager_creates_and_closes_session(): - store = DakeraMemoryStore(url="http://localhost:3300") - assert store._session is None +async def test_owned_client_is_closed_on_exit(): + from unittest.mock import patch - with patch("aiohttp.ClientSession") as MockSession: + with patch("semantic_kernel.connectors.memory_stores.dakera.dakera_memory_store.AsyncDakeraClient") as mock_ctor: instance = AsyncMock() - MockSession.return_value = instance - async with store: - assert store._session is not None + mock_ctor.return_value = instance + async with DakeraMemoryStore(url="http://localhost:3000") as s: + assert s is not None + instance.close.assert_awaited_once() + - instance.close.assert_called_once() +async def test_injected_client_is_not_closed(mock_client: AsyncMock): + async with DakeraMemoryStore(client=mock_client): + pass + mock_client.close.assert_not_awaited()