From 20631ad1f613e812cea8ea9e4e5e90a32526c757 Mon Sep 17 00:00:00 2001 From: ferhimedamine Date: Wed, 1 Jul 2026 06:26:51 +0000 Subject: [PATCH 1/8] =?UTF-8?q?feat(integrations):=20add=20DakeraMemoryLog?= =?UTF-8?q?ger=20=E2=80=94=20persistent=20cross-session=20memory=20for=20a?= =?UTF-8?q?ll=20LLM=20calls?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds litellm/integrations/dakera_memory.py implementing CustomLogger with: - async_pre_call_hook: recalls relevant Dakera memories and injects them as a system message before each LLM call - async_log_success_event: persists completed exchanges to Dakera after success Self-host: docker run -p 3000:3000 dakera/dakera:latest Usage: litellm.callbacks = [DakeraMemoryLogger(base_url=..., api_key=...)] --- litellm/integrations/dakera_memory.py | 184 ++++++++++++++++++++++++++ 1 file changed, 184 insertions(+) create mode 100644 litellm/integrations/dakera_memory.py diff --git a/litellm/integrations/dakera_memory.py b/litellm/integrations/dakera_memory.py new file mode 100644 index 00000000000..da4c572f83b --- /dev/null +++ b/litellm/integrations/dakera_memory.py @@ -0,0 +1,184 @@ +# What is this? +## Persistent memory for LLM calls via Dakera +## Recalls relevant prior context before each LLM call, stores exchanges after. +## +## Usage: +## import litellm +## from litellm.integrations.dakera_memory import DakeraMemoryLogger +## litellm.callbacks = [DakeraMemoryLogger(base_url="http://localhost:3000", api_key="dk_...")] +## +## Or as a named callback (add "dakera_memory" to litellm.success_callback): +## litellm.success_callback = ["dakera_memory"] + +import os +from datetime import datetime +from typing import Any, Dict, List, Optional + +import httpx + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm._logging import verbose_logger + + +class DakeraMemoryLogger(CustomLogger): + """Persistent LLM memory via self-hosted Dakera. + + Hooks into litellm's pre/post call lifecycle: + - async_pre_call_hook: recalls relevant memories and injects them as a + system message before the LLM call + - async_log_success_event: persists the completed exchange to Dakera + + Self-host Dakera: docker run -p 3000:3000 dakera/dakera:latest + + 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. + top_k: Number of memories to recall per call. Defaults to 5. + session_id_key: Metadata key used to group memories by session. + Reads from litellm call metadata dict. Defaults to "session_id". + """ + + def __init__( + self, + base_url: Optional[str] = None, + api_key: Optional[str] = None, + top_k: int = 5, + session_id_key: str = "session_id", + ): + self.base_url = (base_url or os.getenv("DAKERA_API_URL", "http://localhost:3000")).rstrip("/") + self.api_key = api_key or os.getenv("DAKERA_API_KEY", "") + self.top_k = top_k + self.session_id_key = session_id_key + + 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 _session_id(self, metadata: Optional[Dict]) -> str: + if not metadata: + return "default" + return str(metadata.get(self.session_id_key, "default")) + + # ------------------------------------------------------------------ + # Pre-call hook: inject recalled memories before the LLM 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 memories and prepend them as a system message.""" + try: + messages: List[Dict] = data.get("messages", []) + if not messages: + return data + + # Extract the last user message as the recall query + last_user = next( + (m["content"] for m in reversed(messages) if m.get("role") == "user"), + None, + ) + if not last_user: + return data + + session_id = self._session_id(data.get("metadata")) + + async with httpx.AsyncClient(timeout=5.0) as client: + resp = await client.post( + f"{self.base_url}/v1/memories/search", + headers=self._headers(), + json={ + "query": last_user if isinstance(last_user, str) else str(last_user), + "session_id": session_id, + "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 string + memory_lines = "\n".join(f"- {r['content']}" for r in results) + memory_msg = { + "role": "system", + "content": ( + f"Relevant context from prior sessions (retrieved from persistent memory):\n" + f"{memory_lines}" + ), + } + + # Inject before the first non-system message + 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(results)} memories for session={session_id}" + ) + except Exception as e: + verbose_logger.warning(f"DakeraMemoryLogger.async_pre_call_hook failed: {e}") + + return data + + # ------------------------------------------------------------------ + # Post-call hook: store the completed exchange to Dakera + # ------------------------------------------------------------------ + + async def async_log_success_event( + self, + kwargs: Dict, + response_obj: Any, + start_time: datetime, + end_time: datetime, + ) -> None: + """Persist the completed LLM exchange to Dakera after a successful call.""" + try: + messages: List[Dict] = kwargs.get("messages", []) + session_id = self._session_id(kwargs.get("metadata")) + + # Find last user message content to store + last_user = next( + (m.get("content", "") for m in reversed(messages) if m.get("role") == "user"), + None, + ) + if not last_user: + return + + # Also grab the assistant response if available + 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}" + + async with httpx.AsyncClient(timeout=5.0) as client: + await client.post( + 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", ""), + }, + }, + ) + verbose_logger.debug( + f"DakeraMemoryLogger: stored exchange for session={session_id}" + ) + except Exception as e: + verbose_logger.warning(f"DakeraMemoryLogger.async_log_success_event failed: {e}") From 2504a65143f857156d9f50ea70a1d3642b4ce92b Mon Sep 17 00:00:00 2001 From: ferhimedamine Date: Wed, 1 Jul 2026 07:03:23 +0000 Subject: [PATCH 2/8] fix(dakera): multimodal content + tenant isolation + _extract_text helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add _extract_text() helper to safely extract text from multimodal message content (list of parts) — prevents garbage queries from vision inputs - Apply _extract_text in both async_pre_call_hook and async_log_success_event - Improve tenant isolation: _session_id() now falls back to SHA256 hash of the caller's API key when no explicit session_id is supplied, preventing cross-tenant memory leakage in multi-user proxy deployments - Use get_async_httpx_client (litellm shared HTTP handler) instead of raw httpx - Document tenant isolation semantics in class docstring --- litellm/integrations/dakera_memory.py | 231 +++++++++++++++++--------- 1 file changed, 152 insertions(+), 79 deletions(-) diff --git a/litellm/integrations/dakera_memory.py b/litellm/integrations/dakera_memory.py index da4c572f83b..2431794ee60 100644 --- a/litellm/integrations/dakera_memory.py +++ b/litellm/integrations/dakera_memory.py @@ -1,42 +1,77 @@ -# What is this? -## Persistent memory for LLM calls via Dakera -## Recalls relevant prior context before each LLM call, stores exchanges after. -## -## Usage: -## import litellm -## from litellm.integrations.dakera_memory import DakeraMemoryLogger -## litellm.callbacks = [DakeraMemoryLogger(base_url="http://localhost:3000", api_key="dk_...")] -## -## Or as a named callback (add "dakera_memory" to litellm.success_callback): -## litellm.success_callback = ["dakera_memory"] +"""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 + +Self-host Dakera: + docker run -p 3000:3000 -e DAKERA_API_KEY=demo dakera/dakera:latest + +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 typing import Any, Dict, List, Optional, Union -import httpx - -import litellm -from litellm.integrations.custom_logger import CustomLogger 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, +) + + +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. - Hooks into litellm's pre/post call lifecycle: - - async_pre_call_hook: recalls relevant memories and injects them as a - system message before the LLM call - - async_log_success_event: persists the completed exchange to Dakera - - Self-host Dakera: docker run -p 3000:3000 dakera/dakera:latest + 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. 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. - top_k: Number of memories to recall per call. Defaults to 5. - session_id_key: Metadata key used to group memories by session. - Reads from litellm call metadata dict. Defaults to "session_id". + base_url: Dakera server URL. Defaults to ``DAKERA_API_URL`` env var. + api_key: Dakera API key. Defaults to ``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__( @@ -45,11 +80,17 @@ class DakeraMemoryLogger(CustomLogger): api_key: Optional[str] = None, top_k: int = 5, session_id_key: str = "session_id", - ): - self.base_url = (base_url or os.getenv("DAKERA_API_URL", "http://localhost:3000")).rstrip("/") + ) -> None: + super().__init__() + self.base_url = ( + base_url or os.getenv("DAKERA_API_URL", "http://localhost:3000") + ).rstrip("/") self.api_key = api_key or os.getenv("DAKERA_API_KEY", "") self.top_k = top_k self.session_id_key = session_id_key + self._http_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) def _headers(self) -> Dict[str, str]: h: Dict[str, str] = {"Content-Type": "application/json"} @@ -57,13 +98,26 @@ class DakeraMemoryLogger(CustomLogger): h["Authorization"] = f"Bearer {self.api_key}" return h - def _session_id(self, metadata: Optional[Dict]) -> str: - if not metadata: - return "default" - return str(metadata.get(self.session_id_key, "default")) + 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. + # This isolates memories per litellm API key when no explicit session is supplied. + caller_key: str = "" + if user_api_key_dict is not None: + caller_key = getattr(user_api_key_dict, "api_key", "") or "" + if caller_key: + return "key:" + hashlib.sha256(caller_key.encode()).hexdigest()[:16] + return "default" # ------------------------------------------------------------------ - # Pre-call hook: inject recalled memories before the LLM sees the prompt + # Pre-call hook: inject recalled memories before the model sees the prompt # ------------------------------------------------------------------ async def async_pre_call_hook( @@ -73,50 +127,62 @@ class DakeraMemoryLogger(CustomLogger): data: Dict, call_type: str, ) -> Optional[Dict]: - """Recall relevant memories and prepend them as a system message.""" + """Recall relevant prior memories and prepend them as a system message.""" try: messages: List[Dict] = data.get("messages", []) if not messages: return data - # Extract the last user message as the recall query - last_user = next( - (m["content"] for m in reversed(messages) if m.get("role") == "user"), + # 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, ) + 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")) + session_id = self._session_id(data.get("metadata"), user_api_key_dict) - async with httpx.AsyncClient(timeout=5.0) as client: - resp = await client.post( - f"{self.base_url}/v1/memories/search", - headers=self._headers(), - json={ - "query": last_user if isinstance(last_user, str) else str(last_user), - "session_id": session_id, - "top_k": self.top_k, - }, - ) - if resp.status_code != 200: - return data - results = resp.json().get("results", []) + 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, + ) + if resp.status_code != 200: + return data + + results = resp.json().get("results", []) if not results: return data - # Build memory context string - memory_lines = "\n".join(f"- {r['content']}" for r in results) + # 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") + ) + if not memory_lines: + return data + memory_msg = { "role": "system", "content": ( - f"Relevant context from prior sessions (retrieved from persistent memory):\n" + "Relevant context from prior sessions (retrieved from persistent memory):\n" f"{memory_lines}" ), } - - # Inject before the first non-system message 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 @@ -124,13 +190,13 @@ class DakeraMemoryLogger(CustomLogger): verbose_logger.debug( f"DakeraMemoryLogger: injected {len(results)} memories for session={session_id}" ) - except Exception as e: - verbose_logger.warning(f"DakeraMemoryLogger.async_pre_call_hook failed: {e}") + except Exception as exc: + verbose_logger.warning(f"DakeraMemoryLogger.async_pre_call_hook failed: {exc}") return data # ------------------------------------------------------------------ - # Post-call hook: store the completed exchange to Dakera + # Post-call hook: persist the completed exchange after a successful call # ------------------------------------------------------------------ async def async_log_success_event( @@ -140,20 +206,25 @@ class DakeraMemoryLogger(CustomLogger): start_time: datetime, end_time: datetime, ) -> None: - """Persist the completed LLM exchange to Dakera after a successful call.""" + """Store the completed LLM exchange in Dakera.""" try: messages: List[Dict] = kwargs.get("messages", []) session_id = self._session_id(kwargs.get("metadata")) - # Find last user message content to store - last_user = next( - (m.get("content", "") for m in reversed(messages) if m.get("role") == "user"), + 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 - # Also grab the assistant response if available assistant_content: Optional[str] = None if hasattr(response_obj, "choices") and response_obj.choices: choice = response_obj.choices[0] @@ -164,21 +235,23 @@ class DakeraMemoryLogger(CustomLogger): if assistant_content: content += f"\nAssistant: {assistant_content}" - async with httpx.AsyncClient(timeout=5.0) as client: - await client.post( - 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._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", ""), }, - ) + }, + timeout=5.0, + ) verbose_logger.debug( f"DakeraMemoryLogger: stored exchange for session={session_id}" ) - except Exception as e: - verbose_logger.warning(f"DakeraMemoryLogger.async_log_success_event failed: {e}") + except Exception as exc: + verbose_logger.warning( + f"DakeraMemoryLogger.async_log_success_event failed: {exc}" + ) From e202a8aae4fc50090bedf43db1ad839c8a679ccc Mon Sep 17 00:00:00 2001 From: ferhimedamine Date: Wed, 1 Jul 2026 07:12:56 +0000 Subject: [PATCH 3/8] style: apply ruff format --- litellm/integrations/dakera_memory.py | 45 ++++++--------------------- 1 file changed, 10 insertions(+), 35 deletions(-) diff --git a/litellm/integrations/dakera_memory.py b/litellm/integrations/dakera_memory.py index 2431794ee60..82a865bb33d 100644 --- a/litellm/integrations/dakera_memory.py +++ b/litellm/integrations/dakera_memory.py @@ -45,11 +45,7 @@ def _extract_text(content: Any) -> str: 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 " ".join(p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text").strip() return str(content) if content else "" @@ -82,15 +78,11 @@ 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:3000") - ).rstrip("/") + self.base_url = (base_url or os.getenv("DAKERA_API_URL", "http://localhost:3000")).rstrip("/") self.api_key = api_key or os.getenv("DAKERA_API_KEY", "") self.top_k = top_k self.session_id_key = session_id_key - self._http_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self._http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) def _headers(self) -> Dict[str, str]: h: Dict[str, str] = {"Content-Type": "application/json"} @@ -136,11 +128,7 @@ class DakeraMemoryLogger(CustomLogger): # 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" - ), + (m.get("content", "") for m in reversed(messages) if m.get("role") == "user"), None, ) if not raw_content: @@ -170,26 +158,21 @@ class DakeraMemoryLogger(CustomLogger): 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") - ) + memory_lines = "\n".join(f"- {r.get('content', '')}" for r in results if r.get("content")) if not memory_lines: return data memory_msg = { "role": "system", "content": ( - "Relevant context from prior sessions (retrieved from persistent memory):\n" - f"{memory_lines}" + 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(results)} memories for session={session_id}" - ) + verbose_logger.debug(f"DakeraMemoryLogger: injected {len(results)} memories for session={session_id}") except Exception as exc: verbose_logger.warning(f"DakeraMemoryLogger.async_pre_call_hook failed: {exc}") @@ -212,11 +195,7 @@ class DakeraMemoryLogger(CustomLogger): session_id = self._session_id(kwargs.get("metadata")) raw_content = next( - ( - m.get("content", "") - for m in reversed(messages) - if m.get("role") == "user" - ), + (m.get("content", "") for m in reversed(messages) if m.get("role") == "user"), None, ) if not raw_content: @@ -248,10 +227,6 @@ class DakeraMemoryLogger(CustomLogger): }, timeout=5.0, ) - verbose_logger.debug( - f"DakeraMemoryLogger: stored exchange for session={session_id}" - ) + verbose_logger.debug(f"DakeraMemoryLogger: stored exchange for session={session_id}") except Exception as exc: - verbose_logger.warning( - f"DakeraMemoryLogger.async_log_success_event failed: {exc}" - ) + verbose_logger.warning(f"DakeraMemoryLogger.async_log_success_event failed: {exc}") From 72ef1df38596f183bf723d8451a905dec2bbdb1f Mon Sep 17 00:00:00 2001 From: ferhimedamine Date: Wed, 1 Jul 2026 07:17:58 +0000 Subject: [PATCH 4/8] fix: use ghcr.io/dakera-ai/dakera:latest and port 3300 --- litellm/integrations/dakera_memory.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/dakera_memory.py b/litellm/integrations/dakera_memory.py index 82a865bb33d..eaa1b74296a 100644 --- a/litellm/integrations/dakera_memory.py +++ b/litellm/integrations/dakera_memory.py @@ -6,7 +6,7 @@ Hooks into litellm's pre/post call lifecycle: - async_log_success_event: persists the completed exchange to Dakera after success Self-host Dakera: - docker run -p 3000:3000 -e DAKERA_API_KEY=demo dakera/dakera:latest + docker run -p 3300:3300 -e DAKERA_API_KEY=demo ghcr.io/dakera-ai/dakera:latest Usage: import litellm @@ -14,7 +14,7 @@ Usage: litellm.callbacks = [ DakeraMemoryLogger( - base_url="http://localhost:3000", + base_url="http://localhost:3300", api_key="dk_your_key", top_k=5, ) @@ -78,7 +78,7 @@ 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:3000")).rstrip("/") + 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.top_k = top_k self.session_id_key = session_id_key From cf8ed3200e581ee86ad07327f3b341427eb6cd33 Mon Sep 17 00:00:00 2001 From: ferhimedamine Date: Wed, 1 Jul 2026 07:23:37 +0000 Subject: [PATCH 5/8] fix: align session_id between pre-call and post-call hooks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit async_log_success_event was calling _session_id without user_api_key_dict, so when no explicit session_id was in metadata the key-hash fallback silently degraded to 'default' — memories were stored under a different namespace than they were recalled from. Pass kwargs['user_api_key_dict'] to _session_id and update the method to handle both UserAPIKeyAuth objects and serialized dicts. Co-Authored-By: Paperclip --- litellm/integrations/dakera_memory.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/dakera_memory.py b/litellm/integrations/dakera_memory.py index eaa1b74296a..f241c5d0f2f 100644 --- a/litellm/integrations/dakera_memory.py +++ b/litellm/integrations/dakera_memory.py @@ -100,10 +100,14 @@ class DakeraMemoryLogger(CustomLogger): if sid: return str(sid) # Fall back to a stable per-API-key namespace to prevent cross-tenant leakage. - # This isolates memories per litellm API key when no explicit session is supplied. + # 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", "") or "" + 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" @@ -192,7 +196,7 @@ class DakeraMemoryLogger(CustomLogger): """Store the completed LLM exchange in Dakera.""" try: messages: List[Dict] = kwargs.get("messages", []) - session_id = self._session_id(kwargs.get("metadata")) + 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"), From 4beb1e0965084d723d262a683b7bd8b69077cb3f Mon Sep 17 00:00:00 2001 From: ferhimedamine Date: Wed, 1 Jul 2026 07:32:27 +0000 Subject: [PATCH 6/8] fix: remove unused Union import (ruff F401) Co-Authored-By: Paperclip --- litellm/integrations/dakera_memory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/integrations/dakera_memory.py b/litellm/integrations/dakera_memory.py index f241c5d0f2f..531acad753c 100644 --- a/litellm/integrations/dakera_memory.py +++ b/litellm/integrations/dakera_memory.py @@ -30,7 +30,7 @@ Usage: import hashlib import os from datetime import datetime -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger From ab4d252cc6d536a25b19864f93d99a88f2587c43 Mon Sep 17 00:00:00 2001 From: ferhimedamine Date: Wed, 1 Jul 2026 08:26:26 +0000 Subject: [PATCH 7/8] =?UTF-8?q?fix(lint):=20suppress=20BLE001=20in=20Daker?= =?UTF-8?q?aMemoryLogger=20hooks=20=E2=80=94=20broad=20catch=20intentional?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Logging hooks must never crash the main LLM request, so catching broadly is correct. Suppress BLE001 via noqa to stay within the strict-rule budget. Co-Authored-By: Paperclip --- litellm/integrations/dakera_memory.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/dakera_memory.py b/litellm/integrations/dakera_memory.py index 531acad753c..d0ccaae2d75 100644 --- a/litellm/integrations/dakera_memory.py +++ b/litellm/integrations/dakera_memory.py @@ -177,7 +177,7 @@ class DakeraMemoryLogger(CustomLogger): data["messages"] = existing_system + [memory_msg] + non_system verbose_logger.debug(f"DakeraMemoryLogger: injected {len(results)} memories for session={session_id}") - except Exception as exc: + except Exception as exc: # noqa: BLE001 verbose_logger.warning(f"DakeraMemoryLogger.async_pre_call_hook failed: {exc}") return data @@ -232,5 +232,5 @@ class DakeraMemoryLogger(CustomLogger): timeout=5.0, ) verbose_logger.debug(f"DakeraMemoryLogger: stored exchange for session={session_id}") - except Exception as exc: + except Exception as exc: # noqa: BLE001 verbose_logger.warning(f"DakeraMemoryLogger.async_log_success_event failed: {exc}") From 92bd3a85e482f9a6346d245864a83c04f7da3fa4 Mon Sep 17 00:00:00 2001 From: ferhimedamine Date: Thu, 2 Jul 2026 19:46:32 +0000 Subject: [PATCH 8/8] 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