From 92bd3a85e482f9a6346d245864a83c04f7da3fa4 Mon Sep 17 00:00:00 2001 From: ferhimedamine Date: Thu, 2 Jul 2026 19:46:32 +0000 Subject: [PATCH] fix(integrations): use official dakera SDK in DakeraMemoryLogger Recall and storage previously hit hand-rolled REST paths on port 3300, which do not match the Dakera server API. Route them through the official dakera Python SDK instead (AsyncDakeraClient.recall / store_memory), default to the correct port 3000, and lazily import the optional dependency with a clear install hint. Add unit tests covering memory injection position, recall query selection, multimodal extraction, error swallowing, session-namespace isolation, and the missing-SDK error path. --- litellm/integrations/dakera_memory.py | 106 ++++----- .../integrations/test_dakera_memory.py | 209 ++++++++++++++++++ 2 files changed, 263 insertions(+), 52 deletions(-) create mode 100644 tests/test_litellm/integrations/test_dakera_memory.py diff --git a/litellm/integrations/dakera_memory.py b/litellm/integrations/dakera_memory.py index d0ccaae2d75..caecf2a0758 100644 --- a/litellm/integrations/dakera_memory.py +++ b/litellm/integrations/dakera_memory.py @@ -5,17 +5,28 @@ Hooks into litellm's pre/post call lifecycle: system message before the LLM sees the prompt - async_log_success_event: persists the completed exchange to Dakera after success -Self-host Dakera: - docker run -p 3300:3300 -e DAKERA_API_KEY=demo ghcr.io/dakera-ai/dakera:latest +Dakera is a self-hosted, decay-weighted vector memory server. Recall and storage +go through the official ``dakera`` Python SDK, so this logger always speaks the +same verified API as the rest of the Dakera ecosystem. + +Install the SDK (optional dependency, only needed when this callback is used):: + + pip install dakera + +Self-host Dakera with the public docker-compose (server + object store):: + + git clone https://github.com/dakera-ai/dakera-deploy && cd dakera-deploy + docker compose up -d # serves the API on http://localhost:3000 + +Usage:: -Usage: import litellm from litellm.integrations.dakera_memory import DakeraMemoryLogger litellm.callbacks = [ DakeraMemoryLogger( - base_url="http://localhost:3300", - api_key="dk_your_key", + base_url="http://localhost:3000", + api_key="dk-your-key", top_k=5, ) ] @@ -34,10 +45,8 @@ from typing import Any, Dict, List, Optional from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger -from litellm.llms.custom_httpx.http_handler import ( - get_async_httpx_client, - httpxSpecialProvider, -) + +_DEFAULT_BASE_URL = "http://localhost:3000" def _extract_text(content: Any) -> str: @@ -54,11 +63,12 @@ class DakeraMemoryLogger(CustomLogger): Dakera is a decay-weighted vector memory server you run on your own infrastructure. This logger gives every litellm call persistent cross-session - memory without any cloud dependency. + memory without any cloud dependency, using the official ``dakera`` SDK. Args: - base_url: Dakera server URL. Defaults to ``DAKERA_API_URL`` env var. - api_key: Dakera API key. Defaults to ``DAKERA_API_KEY`` env var. + base_url: Dakera server URL. Defaults to the ``DAKERA_API_URL`` env var, + then ``http://localhost:3000``. + api_key: Dakera API key. Defaults to the ``DAKERA_API_KEY`` env var. top_k: Number of memories to recall per LLM call. session_id_key: Key in litellm call metadata used to group memories by session. Defaults to ``"session_id"``. @@ -78,17 +88,27 @@ class DakeraMemoryLogger(CustomLogger): session_id_key: str = "session_id", ) -> None: super().__init__() - self.base_url = (base_url or os.getenv("DAKERA_API_URL", "http://localhost:3300")).rstrip("/") - self.api_key = api_key or os.getenv("DAKERA_API_KEY", "") + self.base_url = (base_url or os.getenv("DAKERA_API_URL") or _DEFAULT_BASE_URL).rstrip("/") + self.api_key = api_key or os.getenv("DAKERA_API_KEY") or None self.top_k = top_k self.session_id_key = session_id_key - self._http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) + self._client: Any = None - def _headers(self) -> Dict[str, str]: - h: Dict[str, str] = {"Content-Type": "application/json"} - if self.api_key: - h["Authorization"] = f"Bearer {self.api_key}" - return h + def _get_client(self) -> Any: + """Lazily construct and cache the async Dakera SDK client. + + Raises a clear, actionable error if the optional ``dakera`` package is + not installed, rather than failing silently at call time. + """ + if self._client is None: + try: + from dakera import AsyncDakeraClient + except ImportError as exc: + raise ImportError( + "DakeraMemoryLogger requires the 'dakera' package. Install it with: pip install dakera" + ) from exc + self._client = AsyncDakeraClient(base_url=self.base_url, api_key=self.api_key) + return self._client def _session_id( self, @@ -129,8 +149,6 @@ class DakeraMemoryLogger(CustomLogger): if not messages: return data - # Use the last user message as the semantic recall query. - # _extract_text handles multimodal content (list of parts) gracefully. raw_content = next( (m.get("content", "") for m in reversed(messages) if m.get("role") == "user"), None, @@ -143,26 +161,14 @@ class DakeraMemoryLogger(CustomLogger): session_id = self._session_id(data.get("metadata"), user_api_key_dict) - resp = await self._http_handler.post( - url=f"{self.base_url}/v1/memories/search", - headers=self._headers(), - json={ - "query": last_user, - "session_id": session_id, - "top_k": self.top_k, - }, - timeout=5.0, + resp = await self._get_client().recall( + agent_id=session_id, + query=last_user, + top_k=self.top_k, ) - if resp.status_code != 200: - return data - - results = resp.json().get("results", []) - if not results: - return data - - # Build memory context and inject before the first non-system message - memory_lines = "\n".join(f"- {r.get('content', '')}" for r in results if r.get("content")) + memories = getattr(resp, "memories", None) or [] + memory_lines = "\n".join(f"- {m.content}" for m in memories if getattr(m, "content", None)) if not memory_lines: return data @@ -176,7 +182,7 @@ class DakeraMemoryLogger(CustomLogger): non_system = [m for m in messages if m.get("role") != "system"] data["messages"] = existing_system + [memory_msg] + non_system - verbose_logger.debug(f"DakeraMemoryLogger: injected {len(results)} memories for session={session_id}") + verbose_logger.debug(f"DakeraMemoryLogger: injected {len(memories)} memories for session={session_id}") except Exception as exc: # noqa: BLE001 verbose_logger.warning(f"DakeraMemoryLogger.async_pre_call_hook failed: {exc}") @@ -218,18 +224,14 @@ class DakeraMemoryLogger(CustomLogger): if assistant_content: content += f"\nAssistant: {assistant_content}" - await self._http_handler.post( - url=f"{self.base_url}/v1/memories", - headers=self._headers(), - json={ - "content": content, - "session_id": session_id, - "metadata": { - "model": kwargs.get("model", ""), - "call_type": kwargs.get("call_type", ""), - }, + await self._get_client().store_memory( + agent_id=session_id, + content=content, + session_id=session_id, + metadata={ + "model": kwargs.get("model", ""), + "call_type": kwargs.get("call_type", ""), }, - timeout=5.0, ) verbose_logger.debug(f"DakeraMemoryLogger: stored exchange for session={session_id}") except Exception as exc: # noqa: BLE001 diff --git a/tests/test_litellm/integrations/test_dakera_memory.py b/tests/test_litellm/integrations/test_dakera_memory.py new file mode 100644 index 00000000000..84abc20015d --- /dev/null +++ b/tests/test_litellm/integrations/test_dakera_memory.py @@ -0,0 +1,209 @@ +"""Tests for the Dakera persistent-memory logger. + +Each test injects a mocked async Dakera client into DakeraMemoryLogger (via the +``_client`` attribute) so no real SDK or network is required. The assertions pin +the observable contract: the pre-call hook recalls with the right namespace/query +and injects a memory system message in the right position, and the success hook +persists the full exchange under the session namespace. +""" + +import os +import sys +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.integrations.dakera_memory import DakeraMemoryLogger, _extract_text + + +def _client_recalling(*contents: str) -> MagicMock: + client = MagicMock() + memories = [SimpleNamespace(content=c) for c in contents] + client.recall = AsyncMock(return_value=SimpleNamespace(memories=memories)) + client.store_memory = AsyncMock(return_value={}) + return client + + +def _logger_with(client: MagicMock, **kwargs) -> DakeraMemoryLogger: + logger = DakeraMemoryLogger(**kwargs) + logger._client = client + return logger + + +@pytest.mark.asyncio +async def test_pre_call_injects_recalled_memories_after_system_before_user(): + client = _client_recalling("alice prefers dark mode", "alice is in tokyo") + logger = _logger_with(client, top_k=3) + data = { + "messages": [ + {"role": "system", "content": "you are helpful"}, + {"role": "user", "content": "what are my prefs?"}, + ], + "metadata": {"session_id": "user-alice"}, + } + + out = await logger.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + client.recall.assert_awaited_once_with(agent_id="user-alice", query="what are my prefs?", top_k=3) + roles = [m["role"] for m in out["messages"]] + assert roles == ["system", "system", "user"] + injected = out["messages"][1]["content"] + assert "alice prefers dark mode" in injected + assert "alice is in tokyo" in injected + # original system prompt must survive and stay first + assert out["messages"][0]["content"] == "you are helpful" + + +@pytest.mark.asyncio +async def test_pre_call_no_memories_leaves_messages_untouched(): + client = _client_recalling() # empty recall + logger = _logger_with(client) + original = [{"role": "user", "content": "hi"}] + data = {"messages": list(original), "metadata": {"session_id": "s1"}} + + out = await logger.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert out["messages"] == original + + +@pytest.mark.asyncio +async def test_pre_call_uses_last_user_message_as_query(): + client = _client_recalling("m") + logger = _logger_with(client) + data = { + "messages": [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "reply"}, + {"role": "user", "content": "second and latest"}, + ], + "metadata": {"session_id": "s1"}, + } + + await logger.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert client.recall.await_args.kwargs["query"] == "second and latest" + + +@pytest.mark.asyncio +async def test_pre_call_extracts_multimodal_text_for_query(): + client = _client_recalling("m") + logger = _logger_with(client) + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this"}, + {"type": "image_url", "image_url": {"url": "http://x/y.png"}}, + ], + } + ], + "metadata": {"session_id": "s1"}, + } + + await logger.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert client.recall.await_args.kwargs["query"] == "describe this" + + +@pytest.mark.asyncio +async def test_pre_call_swallows_client_errors_and_returns_data(): + client = MagicMock() + client.recall = AsyncMock(side_effect=RuntimeError("dakera down")) + logger = _logger_with(client) + data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + + out = await logger.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert out is data # never raises into the request path + + +@pytest.mark.asyncio +async def test_success_event_persists_user_and_assistant_exchange(): + client = _client_recalling() + logger = _logger_with(client) + response_obj = SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="the answer"))]) + + await logger.async_log_success_event( + kwargs={ + "messages": [{"role": "user", "content": "the question"}], + "metadata": {"session_id": "user-bob"}, + "model": "gpt-4o", + }, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + client.store_memory.assert_awaited_once() + call = client.store_memory.await_args.kwargs + assert call["agent_id"] == "user-bob" + assert "User: the question" in call["content"] + assert "Assistant: the answer" in call["content"] + assert call["metadata"]["model"] == "gpt-4o" + + +@pytest.mark.asyncio +async def test_success_event_skips_store_when_no_user_message(): + client = _client_recalling() + logger = _logger_with(client) + + await logger.async_log_success_event( + kwargs={"messages": [{"role": "system", "content": "sys only"}]}, + response_obj=SimpleNamespace(choices=[]), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + client.store_memory.assert_not_awaited() + + +def test_session_id_prefers_metadata_key(): + logger = DakeraMemoryLogger() + assert logger._session_id({"session_id": "explicit"}) == "explicit" + + +def test_session_id_falls_back_to_stable_api_key_namespace(): + logger = DakeraMemoryLogger() + auth = SimpleNamespace(api_key="sk-secret") + + first = logger._session_id(None, auth) + second = logger._session_id({}, auth) + + assert first == second # deterministic per key + assert first.startswith("key:") + assert "sk-secret" not in first # raw key never leaked into the namespace + + +def test_session_id_defaults_when_no_identity(): + logger = DakeraMemoryLogger() + assert logger._session_id(None, None) == "default" + + +def test_get_client_raises_helpful_error_when_sdk_missing(monkeypatch): + # Simulate the optional 'dakera' package not being installed by masking it + # in sys.modules; importing it then raises ImportError, which _get_client + # must translate into an actionable install hint. + monkeypatch.setitem(sys.modules, "dakera", None) + logger = DakeraMemoryLogger() + + with pytest.raises(ImportError, match="pip install dakera"): + logger._get_client() + + +def test_extract_text_handles_str_list_and_empty(): + assert _extract_text("plain") == "plain" + assert _extract_text([{"type": "text", "text": "a"}, {"type": "image_url"}]) == "a" + assert _extract_text(None) == "" + + +def test_base_url_defaults_to_local_dakera_port_3000(monkeypatch): + monkeypatch.delenv("DAKERA_API_URL", raising=False) + monkeypatch.delenv("DAKERA_API_KEY", raising=False) + logger = DakeraMemoryLogger(base_url=None, api_key=None) + assert logger.base_url == "http://localhost:3000" + assert logger.api_key is None