diff --git a/litellm/integrations/dakera_memory.py b/litellm/integrations/dakera_memory.py new file mode 100644 index 00000000000..caecf2a0758 --- /dev/null +++ b/litellm/integrations/dakera_memory.py @@ -0,0 +1,238 @@ +"""Persistent cross-session memory for LiteLLM via self-hosted Dakera. + +Hooks into litellm's pre/post call lifecycle: +- async_pre_call_hook: recalls relevant prior memories and prepends them as a + system message before the LLM sees the prompt +- async_log_success_event: persists the completed exchange to Dakera after success + +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:: + + import litellm + from litellm.integrations.dakera_memory import DakeraMemoryLogger + + litellm.callbacks = [ + DakeraMemoryLogger( + base_url="http://localhost:3000", + api_key="dk-your-key", + top_k=5, + ) + ] + + response = litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "What did we discuss last time?"}], + metadata={"session_id": "user-alice"}, + ) +""" + +import hashlib +import os +from datetime import datetime +from typing import Any, Dict, List, Optional + +from litellm._logging import verbose_logger +from litellm.integrations.custom_logger import CustomLogger + +_DEFAULT_BASE_URL = "http://localhost:3000" + + +def _extract_text(content: Any) -> str: + """Extract plain text from message content that may be str or list (multimodal).""" + if isinstance(content, str): + return content + if isinstance(content, list): + return " ".join(p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text").strip() + return str(content) if content else "" + + +class DakeraMemoryLogger(CustomLogger): + """Persistent LLM memory via self-hosted Dakera. + + 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, using the official ``dakera`` SDK. + + Args: + 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"``. + + **Tenant isolation:** Always pass a unique ``session_id`` in call metadata to + prevent cross-user memory leakage. When no ``session_id`` is provided the logger + derives a namespace from the caller's API key hash, which isolates memory per + litellm API key but still mixes sessions for the same key. For strict per-user + isolation, supply ``metadata={"session_id": user_id}`` on every call. + """ + + def __init__( + self, + base_url: Optional[str] = None, + api_key: Optional[str] = None, + top_k: int = 5, + session_id_key: str = "session_id", + ) -> None: + super().__init__() + 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._client: Any = None + + 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, + metadata: Optional[Dict], + user_api_key_dict: Any = None, + ) -> str: + if metadata: + sid = metadata.get(self.session_id_key) + if sid: + return str(sid) + # Fall back to a stable per-API-key namespace to prevent cross-tenant leakage. + # Handles both UserAPIKeyAuth objects (has .api_key attribute) and serialized + # dicts (kwargs["user_api_key_dict"] in async_log_success_event). + caller_key: str = "" + if user_api_key_dict is not None: + caller_key = getattr(user_api_key_dict, "api_key", None) + if caller_key is None and isinstance(user_api_key_dict, dict): + caller_key = user_api_key_dict.get("api_key", "") + caller_key = caller_key or "" + if caller_key: + return "key:" + hashlib.sha256(caller_key.encode()).hexdigest()[:16] + return "default" + + # ------------------------------------------------------------------ + # Pre-call hook: inject recalled memories before the model sees the prompt + # ------------------------------------------------------------------ + + async def async_pre_call_hook( + self, + user_api_key_dict: Any, + cache: Any, + data: Dict, + call_type: str, + ) -> Optional[Dict]: + """Recall relevant prior memories and prepend them as a system message.""" + try: + messages: List[Dict] = data.get("messages", []) + if not messages: + return data + + raw_content = next( + (m.get("content", "") for m in reversed(messages) if m.get("role") == "user"), + None, + ) + if not raw_content: + return data + last_user = _extract_text(raw_content) + if not last_user: + return data + + session_id = self._session_id(data.get("metadata"), user_api_key_dict) + + resp = await self._get_client().recall( + agent_id=session_id, + query=last_user, + top_k=self.top_k, + ) + + 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 + + memory_msg = { + "role": "system", + "content": ( + f"Relevant context from prior sessions (retrieved from persistent memory):\n{memory_lines}" + ), + } + existing_system = [m for m in messages if m.get("role") == "system"] + 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(memories)} memories for session={session_id}") + except Exception as exc: # noqa: BLE001 + verbose_logger.warning(f"DakeraMemoryLogger.async_pre_call_hook failed: {exc}") + + return data + + # ------------------------------------------------------------------ + # Post-call hook: persist the completed exchange after a successful call + # ------------------------------------------------------------------ + + async def async_log_success_event( + self, + kwargs: Dict, + response_obj: Any, + start_time: datetime, + end_time: datetime, + ) -> None: + """Store the completed LLM exchange in Dakera.""" + try: + messages: List[Dict] = kwargs.get("messages", []) + session_id = self._session_id(kwargs.get("metadata"), kwargs.get("user_api_key_dict")) + + raw_content = next( + (m.get("content", "") for m in reversed(messages) if m.get("role") == "user"), + None, + ) + if not raw_content: + return + last_user = _extract_text(raw_content) + if not last_user: + return + + assistant_content: Optional[str] = None + if hasattr(response_obj, "choices") and response_obj.choices: + choice = response_obj.choices[0] + if hasattr(choice, "message") and hasattr(choice.message, "content"): + assistant_content = choice.message.content + + content = f"User: {last_user}" + if assistant_content: + content += f"\nAssistant: {assistant_content}" + + 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", ""), + }, + ) + verbose_logger.debug(f"DakeraMemoryLogger: stored exchange for session={session_id}") + except Exception as exc: # noqa: BLE001 + verbose_logger.warning(f"DakeraMemoryLogger.async_log_success_event failed: {exc}") 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