mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
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.
This commit is contained in:
ferhimedamine
parent
ab4d252cc6
commit
92bd3a85e4
2 changed files with 263 additions and 52 deletions
|
|
@ -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
|
||||
|
|
|
|||
209
tests/test_litellm/integrations/test_dakera_memory.py
Normal file
209
tests/test_litellm/integrations/test_dakera_memory.py
Normal 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
|
||||
Loading…
Add table
Reference in a new issue