scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 57

Merge 92bd3a85e4 into a3e74a5223

Browse source
This commit is contained in:
Mohamed Amine Ferhi 2026-10-06 21:10:15 +00:00 • committed by GitHub
commit ae721ddd25
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 447 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

238
litellm/integrations/dakera_memory.py Normal file
View file

@ -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}")

209
tests/test_litellm/integrations/test_dakera_memory.py Normal file
View file

@ -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