mirror of
https://github.com/supermemoryai/supermemory.git
synced 2026-10-26 06:04:08 +00:00
fix(python): community fixes for the openai, cartesia and pipecat wrappers (#1799)
Bundles 10 community fixes for the Python wrappers, each commit kept with its original author. Ships supermemory-openai-sdk 1.0.10, supermemory-cartesia 0.1.5 and supermemory-pipecat 0.1.5. - openai-sdk-python: shared OpenAI clients no longer leak memories across tenants (#1503), memory search fails open instead of crashing the chat (#1288), the requests fallback runs off the event loop (#1660), get_profile works for users with no memories (#1759). - cartesia: the wrapped agent runs exactly once per event (#1280, #1670), plain-string and chunk-only search results are kept (#1683, #1693). - pipecat: memory text stays literal when replacing the injected block (#1647), `InputParams` is exported from the package root (#1575). - #1503 was hand-ported onto main's newer middleware. #1751 is left out: the v5 agent-framework migration (#1786) appends to the conversation document by design. Co-Authored-By: abhinav7x94 <204053250+abhinav7x94@users.noreply.github.com> Co-Authored-By: shamAnimates <145093437+shamAnimates@users.noreply.github.com> Co-Authored-By: aloktripathi1 <142301030+aloktripathi1@users.noreply.github.com> Co-Authored-By: Cintu07 <178455858+Cintu07@users.noreply.github.com> Co-Authored-By: Atharvsinh-codez <201381760+Atharvsinh-codez@users.noreply.github.com> Co-Authored-By: abhay-codes07 <182421137+abhay-codes07@users.noreply.github.com> Co-Authored-By: nkootstra <545768+nkootstra@users.noreply.github.com> Co-Authored-By: Agnik47 <140933190+Agnik47@users.noreply.github.com>
This commit is contained in:
parent
669d2697a9
commit
5323f427d0
23 changed files with 1915 additions and 127 deletions
|
|
@ -28,8 +28,7 @@ You can obtain an API key from [console.supermemory.ai](https://console.supermem
|
|||
Supermemory integration is provided through the `SupermemoryPipecatService` class in Pipecat:
|
||||
|
||||
```python
|
||||
from supermemory_pipecat import SupermemoryPipecatService
|
||||
from supermemory_pipecat.service import InputParams
|
||||
from supermemory_pipecat import InputParams, SupermemoryPipecatService
|
||||
|
||||
memory = SupermemoryPipecatService(
|
||||
api_key=os.getenv("SUPERMEMORY_API_KEY"),
|
||||
|
|
@ -157,8 +156,7 @@ from pipecat.transports.websocket.fastapi import (
|
|||
FastAPIWebsocketTransport,
|
||||
)
|
||||
|
||||
from supermemory_pipecat import SupermemoryPipecatService
|
||||
from supermemory_pipecat.service import InputParams
|
||||
from supermemory_pipecat import InputParams, SupermemoryPipecatService
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "supermemory-cartesia"
|
||||
version = "0.1.4"
|
||||
version = "0.1.5"
|
||||
description = "Supermemory integration for Cartesia Line - memory-enhanced voice agents"
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
|
@ -13,7 +13,7 @@ authors = [
|
|||
{ name = "Supermemory", email = "support@supermemory.ai" }
|
||||
]
|
||||
keywords = [
|
||||
"supermemory>=3.16.0,<5",
|
||||
"supermemory",
|
||||
"cartesia",
|
||||
"line",
|
||||
"memory",
|
||||
|
|
@ -32,7 +32,7 @@ classifiers = [
|
|||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
]
|
||||
dependencies = [
|
||||
"supermemory>=3.16.0",
|
||||
"supermemory>=3.16.0,<5",
|
||||
"cartesia-line>=0.2.0,<0.3.0",
|
||||
"pydantic>=2.10.0",
|
||||
"loguru>=0.7.3",
|
||||
|
|
|
|||
|
|
@ -56,7 +56,7 @@ try:
|
|||
__version__ = version("supermemory-cartesia")
|
||||
except PackageNotFoundError:
|
||||
# Source checkouts do not have installed distribution metadata.
|
||||
__version__ = "0.1.3"
|
||||
__version__ = "0.1.5"
|
||||
|
||||
__all__ = [
|
||||
# Main agent
|
||||
|
|
|
|||
|
|
@ -477,12 +477,12 @@ class SupermemoryCartesiaAgent:
|
|||
Yields:
|
||||
Output events from the wrapped agent.
|
||||
"""
|
||||
try:
|
||||
if type(event).__name__ == "UserTurnEnded":
|
||||
if type(event).__name__ == "UserTurnEnded":
|
||||
memory_context = None
|
||||
try:
|
||||
logger.info("[Supermemory] Processing UserTurnEnded event")
|
||||
event, memory_context = await self._enrich_event_with_memories(event)
|
||||
|
||||
# Store conversation in background
|
||||
if hasattr(event, 'history') and event.history:
|
||||
new_messages = self._new_history_messages(event.history)
|
||||
if new_messages:
|
||||
|
|
@ -505,15 +505,12 @@ class SupermemoryCartesiaAgent:
|
|||
task = asyncio.create_task(self._store_messages(new_messages))
|
||||
self._background_tasks.add(task)
|
||||
task.add_done_callback(self._background_tasks.discard)
|
||||
except Exception as e:
|
||||
logger.error(f"[Supermemory] Failed to enrich or queue the conversation turn: {e}")
|
||||
|
||||
async for output in self._process_agent(env, event, memory_context):
|
||||
yield output
|
||||
else:
|
||||
async for output in self.agent.process(env, event):
|
||||
yield output
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[Supermemory] Error in process: {e}")
|
||||
async for output in self._process_agent(env, event, memory_context):
|
||||
yield output
|
||||
else:
|
||||
async for output in self.agent.process(env, event):
|
||||
yield output
|
||||
|
||||
|
|
|
|||
|
|
@ -5,11 +5,12 @@ from datetime import datetime, timezone
|
|||
from typing import Any, Dict, List, Union
|
||||
|
||||
|
||||
def get_last_user_message(messages: List[Dict[str, str]]) -> str | None:
|
||||
def get_last_user_message(messages: List[Dict[str, Any]]) -> str | None:
|
||||
"""Extract the last user message content from a list of messages."""
|
||||
for msg in reversed(messages):
|
||||
if msg["role"] == "user":
|
||||
return msg["content"]
|
||||
content = msg.get("content")
|
||||
if msg.get("role") == "user" and isinstance(content, str):
|
||||
return content
|
||||
return None
|
||||
|
||||
|
||||
|
|
@ -70,6 +71,17 @@ def _field(item: Any, *names: str, default: Any = None) -> Any:
|
|||
return default
|
||||
|
||||
|
||||
def _memory_text(item: Any) -> str:
|
||||
"""First non-empty memory field; `_field` alone would stop at an empty `memory`."""
|
||||
if isinstance(item, str):
|
||||
return item.strip()
|
||||
for name in ("memory", "chunk", "content"):
|
||||
value = _field(item, name)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value.strip()
|
||||
return ""
|
||||
|
||||
|
||||
_MEMORY_DATE_PREFIX = re.compile(
|
||||
r"^\s*(?:\[recent\]\s*)?(?:\[\d{4}-\d{2}-\d{2}\]\s*)?",
|
||||
re.IGNORECASE,
|
||||
|
|
@ -124,10 +136,7 @@ def deduplicate_memories(
|
|||
out = []
|
||||
for r in results:
|
||||
# v4 search.memories/hybrid uses `memory` or `chunk`.
|
||||
memory = _field(r, "memory", "chunk", "content", default="")
|
||||
if not isinstance(memory, str):
|
||||
memory = ""
|
||||
memory = memory.strip()
|
||||
memory = _memory_text(r)
|
||||
key = _memory_key(memory)
|
||||
if key and key not in seen:
|
||||
seen.add(key)
|
||||
|
|
@ -174,7 +183,7 @@ def format_memories_to_text(
|
|||
lines.append(f"- {item}")
|
||||
continue
|
||||
|
||||
memory = _field(item, "memory", "chunk", "content", default="")
|
||||
memory = _memory_text(item)
|
||||
updated_at = _field(item, "updatedAt", "updated_at", default="")
|
||||
time_str = format_relative_time(updated_at) if updated_at else ""
|
||||
if time_str:
|
||||
|
|
|
|||
|
|
@ -92,6 +92,38 @@ class TestSupermemoryCartesiaNullProfile(unittest.IsolatedAsyncioTestCase):
|
|||
self.assertIsNotNone(context)
|
||||
self.assertIn(fact, context)
|
||||
|
||||
async def test_process_does_not_retry_agent_after_output(self) -> None:
|
||||
class UserTurnEnded:
|
||||
pass
|
||||
|
||||
class FailingAfterOutputAgent:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
async def process(self, _env, _event, **_kwargs):
|
||||
self.calls += 1
|
||||
yield "partial"
|
||||
raise RuntimeError("agent failed after output")
|
||||
|
||||
wrapped_agent = FailingAfterOutputAgent()
|
||||
agent = SupermemoryCartesiaAgent(
|
||||
agent=wrapped_agent,
|
||||
api_key="mock_key",
|
||||
container_tag="user-123",
|
||||
custom_id="conversation-456",
|
||||
add_memory="never",
|
||||
)
|
||||
event = UserTurnEnded()
|
||||
agent._enrich_event_with_memories = AsyncMock(return_value=(event, None))
|
||||
|
||||
outputs = []
|
||||
with self.assertRaisesRegex(RuntimeError, "agent failed after output"):
|
||||
async for output in agent.process(None, event):
|
||||
outputs.append(output)
|
||||
|
||||
self.assertEqual(outputs, ["partial"])
|
||||
self.assertEqual(wrapped_agent.calls, 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
|
|||
125
packages/cartesia-sdk-python/tests/test_single_agent_run.py
Normal file
125
packages/cartesia-sdk-python/tests/test_single_agent_run.py
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
"""Regression tests: the wrapped agent must run exactly once per event."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
|
||||
def _install_test_stubs() -> None:
|
||||
if "loguru" not in sys.modules:
|
||||
loguru_module = types.ModuleType("loguru")
|
||||
|
||||
class _Logger:
|
||||
def info(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
def debug(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
def warning(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
def error(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
loguru_module.logger = _Logger()
|
||||
sys.modules["loguru"] = loguru_module
|
||||
|
||||
if "pydantic" not in sys.modules:
|
||||
pydantic_module = types.ModuleType("pydantic")
|
||||
|
||||
class BaseModel:
|
||||
def __init__(self, **kwargs):
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
def Field(*, default=None, **_kwargs):
|
||||
return default
|
||||
|
||||
pydantic_module.BaseModel = BaseModel
|
||||
pydantic_module.Field = Field
|
||||
sys.modules["pydantic"] = pydantic_module
|
||||
|
||||
|
||||
_install_test_stubs()
|
||||
|
||||
from supermemory_cartesia.agent import SupermemoryCartesiaAgent
|
||||
|
||||
|
||||
class _StreamingAgent:
|
||||
"""Inner agent that yields chunks; optionally dies mid-stream."""
|
||||
|
||||
def __init__(self, fail_after_chunks: bool = False):
|
||||
self.fail_after_chunks = fail_after_chunks
|
||||
self.run_count = 0
|
||||
|
||||
async def process(self, env, event):
|
||||
self.run_count += 1
|
||||
yield "chunk-1"
|
||||
yield "chunk-2"
|
||||
if self.fail_after_chunks:
|
||||
raise RuntimeError("stream dropped")
|
||||
|
||||
|
||||
def _wrap(inner):
|
||||
return SupermemoryCartesiaAgent(
|
||||
agent=inner,
|
||||
api_key="mock_key",
|
||||
container_tag="user-123",
|
||||
custom_id="conversation-456",
|
||||
)
|
||||
|
||||
|
||||
def _user_turn_ended_event():
|
||||
return type("UserTurnEnded", (), {})()
|
||||
|
||||
|
||||
class TestSingleAgentRun(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_mid_stream_agent_error_is_not_retried(self) -> None:
|
||||
inner = _StreamingAgent(fail_after_chunks=True)
|
||||
wrapper = _wrap(inner)
|
||||
|
||||
outputs = []
|
||||
with self.assertRaises(RuntimeError):
|
||||
async for out in wrapper.process(None, SimpleNamespace()):
|
||||
outputs.append(out)
|
||||
|
||||
self.assertEqual(outputs, ["chunk-1", "chunk-2"])
|
||||
self.assertEqual(inner.run_count, 1)
|
||||
|
||||
async def test_mid_stream_error_on_user_turn_is_not_retried(self) -> None:
|
||||
inner = _StreamingAgent(fail_after_chunks=True)
|
||||
wrapper = _wrap(inner)
|
||||
wrapper._enrich_event_with_memories = AsyncMock(
|
||||
side_effect=lambda event: (event, None)
|
||||
)
|
||||
|
||||
outputs = []
|
||||
with self.assertRaises(RuntimeError):
|
||||
async for out in wrapper.process(None, _user_turn_ended_event()):
|
||||
outputs.append(out)
|
||||
|
||||
self.assertEqual(outputs, ["chunk-1", "chunk-2"])
|
||||
self.assertEqual(inner.run_count, 1)
|
||||
|
||||
async def test_enrichment_failure_still_runs_agent_once(self) -> None:
|
||||
inner = _StreamingAgent()
|
||||
wrapper = _wrap(inner)
|
||||
wrapper._enrich_event_with_memories = AsyncMock(
|
||||
side_effect=RuntimeError("enrichment exploded")
|
||||
)
|
||||
|
||||
outputs = [
|
||||
out async for out in wrapper.process(None, _user_turn_ended_event())
|
||||
]
|
||||
|
||||
self.assertEqual(outputs, ["chunk-1", "chunk-2"])
|
||||
self.assertEqual(inner.run_count, 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
189
packages/cartesia-sdk-python/tests/test_utils.py
Normal file
189
packages/cartesia-sdk-python/tests/test_utils.py
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
|
||||
def _install_test_stubs() -> None:
|
||||
if "loguru" not in sys.modules:
|
||||
loguru_module = types.ModuleType("loguru")
|
||||
|
||||
class _Logger:
|
||||
def info(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
def warning(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
def error(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
loguru_module.logger = _Logger()
|
||||
sys.modules["loguru"] = loguru_module
|
||||
|
||||
if "pydantic" not in sys.modules:
|
||||
pydantic_module = types.ModuleType("pydantic")
|
||||
|
||||
class BaseModel:
|
||||
def __init__(self, **kwargs):
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
def Field(*, default=None, **_kwargs):
|
||||
return default
|
||||
|
||||
pydantic_module.BaseModel = BaseModel
|
||||
pydantic_module.Field = Field
|
||||
sys.modules["pydantic"] = pydantic_module
|
||||
|
||||
|
||||
_install_test_stubs()
|
||||
|
||||
from supermemory_cartesia.utils import (
|
||||
deduplicate_memories,
|
||||
format_memories_to_text,
|
||||
get_last_user_message,
|
||||
)
|
||||
|
||||
|
||||
class TestGetLastUserMessage(unittest.TestCase):
|
||||
def test_returns_the_most_recent_user_message(self) -> None:
|
||||
messages = [
|
||||
{"role": "user", "content": "first"},
|
||||
{"role": "assistant", "content": "reply"},
|
||||
{"role": "user", "content": "second"},
|
||||
]
|
||||
|
||||
self.assertEqual(get_last_user_message(messages), "second")
|
||||
|
||||
def test_skips_entries_without_role_or_content(self) -> None:
|
||||
messages = [
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "assistant", "tool_calls": []},
|
||||
{"content": "provider event"},
|
||||
{"role": "user"},
|
||||
]
|
||||
|
||||
self.assertEqual(get_last_user_message(messages), "hello")
|
||||
|
||||
def test_skips_non_string_content(self) -> None:
|
||||
messages = [{"role": "user", "content": [{"type": "text", "text": "hello"}]}]
|
||||
|
||||
self.assertIsNone(get_last_user_message(messages))
|
||||
|
||||
|
||||
class TestDeduplicateMemories(unittest.TestCase):
|
||||
def test_keeps_plain_string_search_results(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=["User prefers tea"],
|
||||
)
|
||||
|
||||
self.assertEqual(result["search_results"], ["User prefers tea"])
|
||||
self.assertIn("User prefers tea", format_memories_to_text(result))
|
||||
|
||||
def test_still_reads_the_memory_field_of_models_and_dicts(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[
|
||||
SimpleNamespace(memory="From a model"),
|
||||
{"chunk": "From a dict"},
|
||||
],
|
||||
)
|
||||
|
||||
self.assertEqual(len(result["search_results"]), 2)
|
||||
rendered = format_memories_to_text(result)
|
||||
self.assertIn("From a model", rendered)
|
||||
self.assertIn("From a dict", rendered)
|
||||
|
||||
def test_deduplicates_a_string_result_already_in_the_profile(self) -> None:
|
||||
fact = "User is allergic to peanuts"
|
||||
result = deduplicate_memories(
|
||||
static=[fact],
|
||||
dynamic=[],
|
||||
search_results=[fact],
|
||||
)
|
||||
|
||||
self.assertEqual(result["static"], [fact])
|
||||
self.assertEqual(result["search_results"], [])
|
||||
|
||||
|
||||
class TestSearchResultExtraction(unittest.TestCase):
|
||||
"""v4 hybrid search returns chunk hits whose `memory` field is empty."""
|
||||
|
||||
def test_falls_through_to_chunk_when_memory_is_empty(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[{"memory": "", "chunk": "The user's dog is called Rex"}],
|
||||
)
|
||||
|
||||
self.assertEqual(len(result["search_results"]), 1)
|
||||
self.assertIn("The user's dog is called Rex", format_memories_to_text(result))
|
||||
|
||||
def test_falls_through_to_chunk_when_memory_is_whitespace(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[{"memory": " ", "chunk": "Chunk body"}],
|
||||
)
|
||||
|
||||
self.assertEqual(len(result["search_results"]), 1)
|
||||
self.assertIn("Chunk body", format_memories_to_text(result))
|
||||
|
||||
def test_falls_through_on_models_too(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[SimpleNamespace(memory="", chunk="Chunk body")],
|
||||
)
|
||||
|
||||
self.assertEqual(len(result["search_results"]), 1)
|
||||
self.assertIn("Chunk body", format_memories_to_text(result))
|
||||
|
||||
def test_never_renders_an_empty_bullet(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[
|
||||
{
|
||||
"memory": "",
|
||||
"chunk": "Chunk body",
|
||||
"updatedAt": "2020-01-01T00:00:00Z",
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
rendered = format_memories_to_text(result)
|
||||
self.assertIn("Chunk body", rendered)
|
||||
for line in rendered.splitlines():
|
||||
if line.startswith("- "):
|
||||
self.assertNotRegex(line, r"^- (\[[^\]]*\] )?$")
|
||||
|
||||
def test_memory_still_wins_over_chunk(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[{"memory": "Memory body", "chunk": "Chunk body"}],
|
||||
)
|
||||
|
||||
rendered = format_memories_to_text(result)
|
||||
self.assertIn("Memory body", rendered)
|
||||
self.assertNotIn("Chunk body", rendered)
|
||||
|
||||
def test_entry_with_no_usable_text_is_still_dropped(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[{"memory": "", "chunk": ""}, {}, None],
|
||||
)
|
||||
|
||||
self.assertEqual(result["search_results"], [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "supermemory-openai-sdk"
|
||||
version = "1.0.9"
|
||||
version = "1.0.10"
|
||||
description = "Memory tools for OpenAI function calling with supermemory"
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
import asyncio
|
||||
import inspect
|
||||
import os
|
||||
import types
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal, Optional, Union, cast
|
||||
|
|
@ -186,13 +187,75 @@ def _update_chat_memory_contexts(
|
|||
return enhanced
|
||||
|
||||
|
||||
class _ResourceFacade:
|
||||
"""Delegate an SDK resource while overriding selected attributes."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
resource: Any,
|
||||
lazy_overrides: Optional[dict[str, Any]] = None,
|
||||
**overrides: Any,
|
||||
) -> None:
|
||||
self._resource = resource
|
||||
self._lazy_overrides = lazy_overrides or {}
|
||||
for name, value in overrides.items():
|
||||
setattr(self, name, value)
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
factory = self._lazy_overrides.get(name)
|
||||
if factory is not None:
|
||||
value = factory()
|
||||
setattr(self, name, value)
|
||||
return value
|
||||
return getattr(self._resource, name)
|
||||
|
||||
|
||||
class _AsyncMemoryResponseContextManager:
|
||||
"""Apply async middleware before entering a streaming response context."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
wrapper: "SupermemoryOpenAIWrapper",
|
||||
original_create: Any,
|
||||
kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
self._wrapper = wrapper
|
||||
self._original_create = original_create
|
||||
self._kwargs = kwargs
|
||||
self._response_context: Any = None
|
||||
|
||||
async def __aenter__(self) -> Any:
|
||||
async def enter_original(**kwargs: Any) -> Any:
|
||||
self._response_context = self._original_create(**kwargs)
|
||||
return await self._response_context.__aenter__()
|
||||
|
||||
return await self._wrapper._create_with_memory_async(
|
||||
enter_original,
|
||||
**self._kwargs,
|
||||
)
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: Any,
|
||||
exc_value: Any,
|
||||
traceback: Any,
|
||||
) -> Any:
|
||||
if self._response_context is None:
|
||||
return None
|
||||
return await self._response_context.__aexit__(
|
||||
exc_type,
|
||||
exc_value,
|
||||
traceback,
|
||||
)
|
||||
|
||||
|
||||
async def supermemory_profile_search(
|
||||
container_tag: str,
|
||||
query_text: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
) -> SupermemoryProfileSearch:
|
||||
"""Search for memories using the SuperMemory profile API."""
|
||||
"""Search memories via the profile API; the requests fallback runs on a thread."""
|
||||
payload = {
|
||||
"containerTag": container_tag,
|
||||
"include": ["static", "dynamic"],
|
||||
|
|
@ -227,10 +290,10 @@ async def supermemory_profile_search(
|
|||
return SupermemoryProfileSearch(data)
|
||||
|
||||
except ImportError:
|
||||
# Fallback to requests if aiohttp not available
|
||||
import requests
|
||||
|
||||
response = requests.post(
|
||||
response = await asyncio.to_thread(
|
||||
requests.post,
|
||||
profile_url,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
|
|
@ -359,15 +422,12 @@ async def add_memory_tool(
|
|||
custom_id: Optional[str],
|
||||
logger: Logger,
|
||||
) -> None:
|
||||
"""Add a new memory to the SuperMemory system."""
|
||||
"""Add a memory without running the sync client call on the event loop."""
|
||||
try:
|
||||
kwargs = {"content": content, "container_tag": container_tag}
|
||||
if custom_id is not None:
|
||||
kwargs["custom_id"] = custom_id
|
||||
|
||||
# The wrapper currently constructs the synchronous Supermemory client for
|
||||
# both OpenAI variants. Never execute that network call on an async event
|
||||
# loop; mocks or future async clients can still return an awaitable.
|
||||
if inspect.iscoroutinefunction(client.add):
|
||||
result = client.add(**kwargs)
|
||||
else:
|
||||
|
|
@ -410,14 +470,22 @@ class SupermemoryOpenAIWrapper:
|
|||
openai_client: Union[OpenAI, AsyncOpenAI],
|
||||
options: OpenAIMiddlewareOptions,
|
||||
):
|
||||
self._client: Union[OpenAI, AsyncOpenAI] = openai_client
|
||||
self._client: Union[OpenAI, AsyncOpenAI] = getattr(
|
||||
openai_client,
|
||||
"__supermemory_openai_base_client__",
|
||||
openai_client,
|
||||
)
|
||||
self.__supermemory_openai_base_client__ = self._client
|
||||
self._container_tag: str = options.container_tag
|
||||
self._options: OpenAIMiddlewareOptions = options
|
||||
self._logger: Logger = create_logger(self._options.verbose)
|
||||
self._api_key = self._resolve_api_key(options.api_key)
|
||||
self._base_url = self._resolve_base_url(options.base_url)
|
||||
base_create = self._client.chat.completions.create
|
||||
self._is_async_client = isinstance(
|
||||
self._client, AsyncOpenAI
|
||||
) or inspect.iscoroutinefunction(inspect.unwrap(base_create))
|
||||
|
||||
# Track background tasks to ensure they complete
|
||||
self._background_tasks: set[asyncio.Task] = set()
|
||||
|
||||
if not hasattr(supermemory, "Supermemory"):
|
||||
|
|
@ -436,8 +504,7 @@ class SupermemoryOpenAIWrapper:
|
|||
f"Failed to initialize Supermemory client: {e}", e
|
||||
)
|
||||
|
||||
# Wrap the chat completions create method
|
||||
self._wrap_chat_completions()
|
||||
self.chat = self._create_chat_facade()
|
||||
|
||||
@staticmethod
|
||||
def _resolve_api_key(configured_api_key: Optional[str]) -> str:
|
||||
|
|
@ -462,35 +529,124 @@ class SupermemoryOpenAIWrapper:
|
|||
)
|
||||
return base_url.rstrip("/")
|
||||
|
||||
def _wrap_chat_completions(self) -> None:
|
||||
"""Wrap the chat completions create method with memory injection."""
|
||||
original_create = self._client.chat.completions.create
|
||||
def _create_chat_facade(self) -> _ResourceFacade:
|
||||
"""Create isolated chat facades; stream() is rebound so it calls the wrapped create."""
|
||||
completions_resource = self._client.chat.completions
|
||||
completions = _ResourceFacade(
|
||||
completions_resource,
|
||||
lazy_overrides={
|
||||
"with_raw_response": lambda: self._create_completion_variant_facade(
|
||||
"with_raw_response"
|
||||
),
|
||||
"with_streaming_response": lambda: self._create_completion_variant_facade(
|
||||
"with_streaming_response"
|
||||
),
|
||||
},
|
||||
create=self._create_completion_method(completions_resource.create),
|
||||
)
|
||||
stream = getattr(type(completions_resource), "stream", None)
|
||||
if stream is not None:
|
||||
completions.stream = types.MethodType(stream, completions)
|
||||
return _ResourceFacade(
|
||||
self._client.chat,
|
||||
lazy_overrides={
|
||||
"with_raw_response": lambda: self._create_chat_variant_facade(
|
||||
"with_raw_response"
|
||||
),
|
||||
"with_streaming_response": lambda: self._create_chat_variant_facade(
|
||||
"with_streaming_response"
|
||||
),
|
||||
},
|
||||
completions=completions,
|
||||
)
|
||||
|
||||
if asyncio.iscoroutinefunction(original_create):
|
||||
def _create_completion_variant_facade(self, name: str) -> _ResourceFacade:
|
||||
"""Preserve raw and streaming response behavior on isolated facades."""
|
||||
resource = getattr(self._client.chat.completions, name)
|
||||
return _ResourceFacade(
|
||||
resource,
|
||||
create=self._create_completion_method(
|
||||
resource.create,
|
||||
streaming_response=name == "with_streaming_response",
|
||||
),
|
||||
)
|
||||
|
||||
async def create_with_memory(
|
||||
def _create_chat_variant_facade(self, name: str) -> _ResourceFacade:
|
||||
"""Wrap completions reached through a chat response variant."""
|
||||
chat_resource = getattr(self._client.chat, name)
|
||||
completions_resource = chat_resource.completions
|
||||
return _ResourceFacade(
|
||||
chat_resource,
|
||||
completions=_ResourceFacade(
|
||||
completions_resource,
|
||||
create=self._create_completion_method(
|
||||
completions_resource.create,
|
||||
streaming_response=name == "with_streaming_response",
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
def _create_client_variant_facade(self, name: str) -> _ResourceFacade:
|
||||
"""Wrap completions reached through a client response variant."""
|
||||
client_resource = getattr(self._client, name)
|
||||
chat_resource = client_resource.chat
|
||||
completions_resource = chat_resource.completions
|
||||
return _ResourceFacade(
|
||||
client_resource,
|
||||
chat=_ResourceFacade(
|
||||
chat_resource,
|
||||
completions=_ResourceFacade(
|
||||
completions_resource,
|
||||
create=self._create_completion_method(
|
||||
completions_resource.create,
|
||||
streaming_response=name == "with_streaming_response",
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
def _create_completion_method(
|
||||
self,
|
||||
original_create: Any,
|
||||
*,
|
||||
streaming_response: bool = False,
|
||||
) -> Any:
|
||||
"""Wrap one completion create implementation with memory injection."""
|
||||
if self._is_async_client and streaming_response:
|
||||
|
||||
def create_streaming_with_memory(
|
||||
**kwargs: Any,
|
||||
) -> _AsyncMemoryResponseContextManager:
|
||||
return _AsyncMemoryResponseContextManager(
|
||||
self,
|
||||
original_create,
|
||||
kwargs,
|
||||
)
|
||||
|
||||
return create_streaming_with_memory
|
||||
|
||||
if self._is_async_client:
|
||||
|
||||
async def create_async_with_memory(
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
return await self._create_with_memory_async(original_create, **kwargs)
|
||||
|
||||
else:
|
||||
return create_async_with_memory
|
||||
|
||||
def create_with_memory(
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
return self._create_with_memory_sync(original_create, **kwargs)
|
||||
def create_sync_with_memory(
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
return self._create_with_memory_sync(original_create, **kwargs)
|
||||
|
||||
# Replace the create method with our wrapper
|
||||
setattr(self._client.chat.completions, "create", create_with_memory)
|
||||
return create_sync_with_memory
|
||||
|
||||
async def _create_with_memory_async(
|
||||
self,
|
||||
original_create: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Async version of create with memory injection."""
|
||||
# OpenAI accepts any Iterable here. Materialize it once because memory
|
||||
# extraction and injection both traverse the messages.
|
||||
"""Async create with memory injection; memory failures fail open."""
|
||||
messages = list(kwargs.get("messages", []))
|
||||
kwargs["messages"] = messages
|
||||
|
||||
|
|
@ -508,7 +664,6 @@ class SupermemoryOpenAIWrapper:
|
|||
else None
|
||||
)
|
||||
|
||||
# Create background task for memory storage
|
||||
task = asyncio.create_task(
|
||||
add_memory_tool(
|
||||
self._supermemory_client,
|
||||
|
|
@ -519,11 +674,9 @@ class SupermemoryOpenAIWrapper:
|
|||
)
|
||||
)
|
||||
|
||||
# Track the task and set up cleanup
|
||||
self._background_tasks.add(task)
|
||||
task.add_done_callback(self._background_tasks.discard)
|
||||
|
||||
# Log any exceptions but don't fail the main request
|
||||
def handle_task_exception(task_obj):
|
||||
try:
|
||||
if task_obj.exception() is not None:
|
||||
|
|
@ -568,14 +721,27 @@ class SupermemoryOpenAIWrapper:
|
|||
},
|
||||
)
|
||||
|
||||
enhanced_messages = await add_system_prompt(
|
||||
messages,
|
||||
self._container_tag,
|
||||
self._logger,
|
||||
self._options.mode,
|
||||
self._api_key,
|
||||
self._base_url,
|
||||
)
|
||||
try:
|
||||
enhanced_messages = await add_system_prompt(
|
||||
messages,
|
||||
self._container_tag,
|
||||
self._logger,
|
||||
self._options.mode,
|
||||
self._api_key,
|
||||
self._base_url,
|
||||
)
|
||||
except (SupermemoryNetworkError, SupermemoryAPIError) as e:
|
||||
self._logger.warn(
|
||||
"Memory search failed, proceeding without memory injection",
|
||||
{"error": str(e), "type": type(e).__name__},
|
||||
)
|
||||
return await original_create(**kwargs)
|
||||
except Exception as e:
|
||||
self._logger.error(
|
||||
"Unexpected error during memory search, proceeding without memory injection",
|
||||
{"error": str(e), "type": type(e).__name__},
|
||||
)
|
||||
return await original_create(**kwargs)
|
||||
|
||||
kwargs["messages"] = enhanced_messages
|
||||
return await original_create(**kwargs)
|
||||
|
|
@ -585,12 +751,10 @@ class SupermemoryOpenAIWrapper:
|
|||
original_create: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Sync version of create with memory injection."""
|
||||
# For sync clients, we implement a simplified version without background tasks
|
||||
"""Sync create with memory injection; memory failures fail open."""
|
||||
messages = list(kwargs.get("messages", []))
|
||||
kwargs["messages"] = messages
|
||||
|
||||
# Handle memory addition synchronously if needed
|
||||
if self._options.add_memory == "always":
|
||||
user_message = get_last_user_message(messages)
|
||||
if user_message and user_message.strip():
|
||||
|
|
@ -605,7 +769,6 @@ class SupermemoryOpenAIWrapper:
|
|||
else None
|
||||
)
|
||||
|
||||
# Use asyncio.run() for the memory addition
|
||||
try:
|
||||
asyncio.run(
|
||||
add_memory_tool(
|
||||
|
|
@ -618,7 +781,6 @@ class SupermemoryOpenAIWrapper:
|
|||
)
|
||||
except RuntimeError as e:
|
||||
if "cannot be called from a running event loop" in str(e):
|
||||
# We're in an async context, log warning and skip memory saving
|
||||
self._logger.warn(
|
||||
"Cannot save memory in sync client from async context",
|
||||
{"error": str(e)},
|
||||
|
|
@ -626,19 +788,15 @@ class SupermemoryOpenAIWrapper:
|
|||
else:
|
||||
raise
|
||||
except SupermemoryNetworkError as e:
|
||||
# Network errors are expected, log as warning
|
||||
self._logger.warn("Network error saving memory", {"error": str(e)})
|
||||
except (SupermemoryAPIError, SupermemoryMemoryOperationError) as e:
|
||||
# API/memory errors are concerning, log as error
|
||||
self._logger.error("Failed to save memory", {"error": str(e)})
|
||||
except Exception as e:
|
||||
# Unexpected errors should be investigated
|
||||
self._logger.error(
|
||||
"Unexpected error saving memory",
|
||||
{"error": str(e), "type": type(e).__name__},
|
||||
)
|
||||
|
||||
# Handle memory search and injection
|
||||
if self._options.mode != "profile":
|
||||
user_message = get_last_user_message(messages)
|
||||
if not user_message:
|
||||
|
|
@ -655,38 +813,51 @@ class SupermemoryOpenAIWrapper:
|
|||
},
|
||||
)
|
||||
|
||||
# Use asyncio.run() for memory search and injection
|
||||
try:
|
||||
enhanced_messages = asyncio.run(
|
||||
add_system_prompt(
|
||||
messages,
|
||||
self._container_tag,
|
||||
self._logger,
|
||||
self._options.mode,
|
||||
self._api_key,
|
||||
self._base_url,
|
||||
)
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "cannot be called from a running event loop" in str(e):
|
||||
# We're in an async context, run in a separate thread
|
||||
import concurrent.futures
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(
|
||||
asyncio.run,
|
||||
add_system_prompt(
|
||||
messages,
|
||||
self._container_tag,
|
||||
self._logger,
|
||||
self._options.mode,
|
||||
self._api_key,
|
||||
self._base_url,
|
||||
),
|
||||
try:
|
||||
enhanced_messages = asyncio.run(
|
||||
add_system_prompt(
|
||||
messages,
|
||||
self._container_tag,
|
||||
self._logger,
|
||||
self._options.mode,
|
||||
self._api_key,
|
||||
self._base_url,
|
||||
)
|
||||
enhanced_messages = future.result()
|
||||
else:
|
||||
raise
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "cannot be called from a running event loop" in str(e):
|
||||
import concurrent.futures
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=1
|
||||
) as executor:
|
||||
future = executor.submit(
|
||||
asyncio.run,
|
||||
add_system_prompt(
|
||||
messages,
|
||||
self._container_tag,
|
||||
self._logger,
|
||||
self._options.mode,
|
||||
self._api_key,
|
||||
self._base_url,
|
||||
),
|
||||
)
|
||||
enhanced_messages = future.result()
|
||||
else:
|
||||
raise
|
||||
except (SupermemoryNetworkError, SupermemoryAPIError) as e:
|
||||
self._logger.warn(
|
||||
"Memory search failed, proceeding without memory injection",
|
||||
{"error": str(e), "type": type(e).__name__},
|
||||
)
|
||||
return original_create(**kwargs)
|
||||
except Exception as e:
|
||||
self._logger.error(
|
||||
"Unexpected error during memory search, proceeding without memory injection",
|
||||
{"error": str(e), "type": type(e).__name__},
|
||||
)
|
||||
return original_create(**kwargs)
|
||||
|
||||
kwargs["messages"] = enhanced_messages
|
||||
return original_create(**kwargs)
|
||||
|
|
@ -722,7 +893,6 @@ class SupermemoryOpenAIWrapper:
|
|||
self._logger.warn(
|
||||
f"Background tasks did not complete within {timeout}s timeout"
|
||||
)
|
||||
# Cancel remaining tasks
|
||||
tasks_to_cancel = [
|
||||
task for task in self._background_tasks if not task.done()
|
||||
]
|
||||
|
|
@ -763,11 +933,9 @@ class SupermemoryOpenAIWrapper:
|
|||
"""Sync context manager exit - attempt to wait for background tasks."""
|
||||
if self._background_tasks:
|
||||
try:
|
||||
# Try to wait for background tasks in sync context
|
||||
asyncio.run(self.wait_for_background_tasks(timeout=5.0))
|
||||
except RuntimeError as e:
|
||||
if "cannot be called from a running event loop" in str(e):
|
||||
# In async context, just cancel the tasks
|
||||
self._logger.warn(
|
||||
"Cannot wait for background tasks in sync context from async environment. "
|
||||
"Use async context manager or call wait_for_background_tasks() manually."
|
||||
|
|
@ -783,6 +951,10 @@ class SupermemoryOpenAIWrapper:
|
|||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
"""Delegate all other attributes to the wrapped client."""
|
||||
if name in {"with_raw_response", "with_streaming_response"}:
|
||||
value = self._create_client_variant_facade(name)
|
||||
setattr(self, name, value)
|
||||
return value
|
||||
return getattr(self._client, name)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -528,7 +528,11 @@ class SupermemoryTools:
|
|||
|
||||
return ProfileResult(
|
||||
success=True,
|
||||
profile=_model_to_dict(response.profile),
|
||||
profile=(
|
||||
_model_to_dict(response.profile)
|
||||
if response.profile is not None
|
||||
else {"static": [], "dynamic": []}
|
||||
),
|
||||
search_results=(
|
||||
_model_to_dict(response.search_results)
|
||||
if response.search_results is not None
|
||||
|
|
|
|||
641
packages/openai-sdk-python/tests/test_client_isolation.py
Normal file
641
packages/openai-sdk-python/tests/test_client_isolation.py
Normal file
|
|
@ -0,0 +1,641 @@
|
|||
"""Regression tests for shared OpenAI client middleware isolation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import gc
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import warnings
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Generator, Literal, Optional
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from supermemory_openai import OpenAIMiddlewareOptions, with_supermemory
|
||||
|
||||
|
||||
def middleware_options(
|
||||
container_tag: str,
|
||||
add_memory: Literal["always", "never"] = "never",
|
||||
) -> OpenAIMiddlewareOptions:
|
||||
return OpenAIMiddlewareOptions(
|
||||
container_tag=container_tag,
|
||||
custom_id=f"thread-{container_tag}",
|
||||
add_memory=add_memory,
|
||||
)
|
||||
|
||||
|
||||
def create_sync_client() -> tuple[Any, Mock]:
|
||||
create = Mock(return_value={"id": "chat-response"})
|
||||
completions = SimpleNamespace(create=create, marker="completions-marker")
|
||||
chat = SimpleNamespace(completions=completions, marker="chat-marker")
|
||||
return SimpleNamespace(chat=chat), create
|
||||
|
||||
|
||||
def create_async_client() -> tuple[Any, AsyncMock]:
|
||||
create = AsyncMock(return_value={"id": "chat-response"})
|
||||
completions = SimpleNamespace(create=create)
|
||||
chat = SimpleNamespace(completions=completions)
|
||||
return SimpleNamespace(chat=chat), create
|
||||
|
||||
|
||||
class RawResponse:
|
||||
def __init__(self, label: str) -> None:
|
||||
self.label = label
|
||||
|
||||
def parse(self) -> str:
|
||||
return f"parsed-{self.label}"
|
||||
|
||||
|
||||
class SyncStreamContext:
|
||||
def __init__(self, label: str) -> None:
|
||||
self.label = label
|
||||
|
||||
def __enter__(self) -> "SyncStreamContext":
|
||||
return self
|
||||
|
||||
def __exit__(self, *args: Any) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class AsyncStreamContext:
|
||||
def __init__(self, label: str) -> None:
|
||||
self.label = label
|
||||
|
||||
async def __aenter__(self) -> "AsyncStreamContext":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def attach_response_variants(
|
||||
client: Any,
|
||||
raw_create_factory: Any,
|
||||
streaming_create_factory: Any,
|
||||
) -> dict[str, Any]:
|
||||
calls: dict[str, Any] = {}
|
||||
|
||||
def completion_resource(label: str, factory: Any) -> Any:
|
||||
create = factory(label)
|
||||
calls[label] = create
|
||||
return SimpleNamespace(create=create)
|
||||
|
||||
client.chat.completions.with_raw_response = completion_resource(
|
||||
"completions-raw", raw_create_factory
|
||||
)
|
||||
client.chat.with_raw_response = SimpleNamespace(
|
||||
completions=completion_resource("chat-raw", raw_create_factory)
|
||||
)
|
||||
client.with_raw_response = SimpleNamespace(
|
||||
chat=SimpleNamespace(
|
||||
completions=completion_resource("client-raw", raw_create_factory)
|
||||
)
|
||||
)
|
||||
client.chat.completions.with_streaming_response = completion_resource(
|
||||
"completions-stream", streaming_create_factory
|
||||
)
|
||||
client.chat.with_streaming_response = SimpleNamespace(
|
||||
completions=completion_resource("chat-stream", streaming_create_factory)
|
||||
)
|
||||
client.with_streaming_response = SimpleNamespace(
|
||||
chat=SimpleNamespace(
|
||||
completions=completion_resource("client-stream", streaming_create_factory)
|
||||
)
|
||||
)
|
||||
return calls
|
||||
|
||||
|
||||
def response_variant_creates(client: Any, name: str) -> list[Any]:
|
||||
return [
|
||||
getattr(client.chat.completions, name).create,
|
||||
getattr(client.chat, name).completions.create,
|
||||
getattr(client, name).chat.completions.create,
|
||||
]
|
||||
|
||||
|
||||
REAL_ASYNC_COMPLETION_PATHS = (
|
||||
"normal",
|
||||
"completions.raw",
|
||||
"chat.raw",
|
||||
"client.raw",
|
||||
"completions.streaming",
|
||||
"chat.streaming",
|
||||
"client.streaming",
|
||||
)
|
||||
|
||||
|
||||
def real_async_completion_create(client: Any, path: str) -> Any:
|
||||
if path == "normal":
|
||||
return client.chat.completions.create
|
||||
if path == "completions.raw":
|
||||
return client.chat.completions.with_raw_response.create
|
||||
if path == "chat.raw":
|
||||
return client.chat.with_raw_response.completions.create
|
||||
if path == "client.raw":
|
||||
return client.with_raw_response.chat.completions.create
|
||||
if path == "completions.streaming":
|
||||
return client.chat.completions.with_streaming_response.create
|
||||
if path == "chat.streaming":
|
||||
return client.chat.with_streaming_response.completions.create
|
||||
if path == "client.streaming":
|
||||
return client.with_streaming_response.chat.completions.create
|
||||
raise AssertionError(f"Unknown completion path: {path}")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True) # type: ignore[untyped-decorator]
|
||||
def supermemory_api_key() -> Generator[None, None, None]:
|
||||
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
|
||||
yield
|
||||
|
||||
|
||||
def test_shared_sync_client_does_not_stack_tenant_middleware() -> None:
|
||||
base_client, original_create = create_sync_client()
|
||||
lookups: list[str] = []
|
||||
|
||||
async def fake_prompt(
|
||||
messages: list[Any],
|
||||
container_tag: str,
|
||||
logger: Any,
|
||||
mode: Any,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
) -> list[Any]:
|
||||
lookups.append(container_tag)
|
||||
return [
|
||||
{"role": "system", "content": f"secret-{container_tag}"},
|
||||
*messages,
|
||||
]
|
||||
|
||||
with patch(
|
||||
"supermemory_openai.middleware.supermemory.Supermemory",
|
||||
return_value=Mock(),
|
||||
), patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
side_effect=fake_prompt,
|
||||
):
|
||||
tenant_a: Any = with_supermemory(
|
||||
base_client,
|
||||
middleware_options("tenant-a"),
|
||||
)
|
||||
tenant_b: Any = with_supermemory(
|
||||
base_client,
|
||||
middleware_options("tenant-b"),
|
||||
)
|
||||
|
||||
assert base_client.chat.completions.create is original_create
|
||||
assert tenant_a.chat is not base_client.chat
|
||||
assert tenant_b.chat.marker == "chat-marker"
|
||||
assert tenant_b.chat.completions.marker == "completions-marker"
|
||||
|
||||
tenant_b.chat.completions.create(
|
||||
model="gpt-test",
|
||||
messages=[{"role": "user", "content": "private"}],
|
||||
)
|
||||
|
||||
assert lookups == ["tenant-b"]
|
||||
messages = original_create.call_args.kwargs["messages"]
|
||||
assert messages[0]["content"] == "secret-tenant-b"
|
||||
assert all("tenant-a" not in str(message) for message in messages)
|
||||
|
||||
base_client.chat.completions.create(
|
||||
model="gpt-test",
|
||||
messages=[{"role": "user", "content": "unwrapped"}],
|
||||
)
|
||||
assert lookups == ["tenant-b"]
|
||||
|
||||
|
||||
def test_rewrapping_a_facade_recovers_the_pristine_client() -> None:
|
||||
base_client, original_create = create_sync_client()
|
||||
lookups: list[str] = []
|
||||
|
||||
async def fake_prompt(
|
||||
messages: list[Any],
|
||||
container_tag: str,
|
||||
logger: Any,
|
||||
mode: Any,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
) -> list[Any]:
|
||||
lookups.append(container_tag)
|
||||
return [
|
||||
{"role": "system", "content": f"secret-{container_tag}"},
|
||||
*messages,
|
||||
]
|
||||
|
||||
with patch(
|
||||
"supermemory_openai.middleware.supermemory.Supermemory",
|
||||
return_value=Mock(),
|
||||
), patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
side_effect=fake_prompt,
|
||||
):
|
||||
tenant_a: Any = with_supermemory(
|
||||
base_client,
|
||||
middleware_options("tenant-a"),
|
||||
)
|
||||
tenant_b: Any = with_supermemory(
|
||||
tenant_a,
|
||||
middleware_options("tenant-b"),
|
||||
)
|
||||
|
||||
tenant_b.chat.completions.create(
|
||||
model="gpt-test",
|
||||
messages=[{"role": "user", "content": "private"}],
|
||||
)
|
||||
|
||||
assert lookups == ["tenant-b"]
|
||||
assert original_create.call_count == 1
|
||||
assert base_client.chat.completions.create is original_create
|
||||
|
||||
|
||||
def test_raw_and_streaming_response_facades_remain_memory_aware() -> None:
|
||||
base_client, _ = create_sync_client()
|
||||
calls = attach_response_variants(
|
||||
base_client,
|
||||
lambda label: Mock(return_value=RawResponse(label)),
|
||||
lambda label: Mock(return_value=SyncStreamContext(label)),
|
||||
)
|
||||
lookups: list[str] = []
|
||||
|
||||
async def fake_prompt(
|
||||
messages: list[Any],
|
||||
container_tag: str,
|
||||
logger: Any,
|
||||
mode: Any,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
) -> list[Any]:
|
||||
lookups.append(container_tag)
|
||||
return [
|
||||
{"role": "system", "content": f"secret-{container_tag}"},
|
||||
*messages,
|
||||
]
|
||||
|
||||
with patch(
|
||||
"supermemory_openai.middleware.supermemory.Supermemory",
|
||||
return_value=Mock(),
|
||||
), patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
side_effect=fake_prompt,
|
||||
):
|
||||
tenant_b: Any = with_supermemory(
|
||||
base_client,
|
||||
middleware_options("tenant-b"),
|
||||
)
|
||||
|
||||
for create in response_variant_creates(tenant_b, "with_raw_response"):
|
||||
response = create(
|
||||
model="gpt-test",
|
||||
messages=[{"role": "user", "content": "raw"}],
|
||||
)
|
||||
assert response.parse().startswith("parsed-")
|
||||
|
||||
for create in response_variant_creates(tenant_b, "with_streaming_response"):
|
||||
with create(
|
||||
model="gpt-test",
|
||||
messages=[{"role": "user", "content": "streaming"}],
|
||||
) as stream:
|
||||
assert stream.label.endswith("stream")
|
||||
|
||||
assert lookups == ["tenant-b"] * 6
|
||||
for create in calls.values():
|
||||
assert "secret-tenant-b" in str(create.call_args.kwargs["messages"])
|
||||
|
||||
|
||||
def test_async_raw_and_streaming_response_prefixes_remain_memory_aware() -> None:
|
||||
base_client, _ = create_async_client()
|
||||
calls = attach_response_variants(
|
||||
base_client,
|
||||
lambda label: AsyncMock(return_value=RawResponse(label)),
|
||||
lambda label: Mock(return_value=AsyncStreamContext(label)),
|
||||
)
|
||||
lookups: list[str] = []
|
||||
|
||||
async def fake_prompt(
|
||||
messages: list[Any],
|
||||
container_tag: str,
|
||||
logger: Any,
|
||||
mode: Any,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
) -> list[Any]:
|
||||
lookups.append(container_tag)
|
||||
return [
|
||||
{"role": "system", "content": f"secret-{container_tag}"},
|
||||
*messages,
|
||||
]
|
||||
|
||||
async def call_raw(create: Any) -> None:
|
||||
response = await create(
|
||||
model="gpt-test",
|
||||
messages=[{"role": "user", "content": "raw"}],
|
||||
)
|
||||
assert response.parse().startswith("parsed-")
|
||||
|
||||
async def consume_stream(stream_context: Any) -> None:
|
||||
async with stream_context as stream:
|
||||
assert stream.label.endswith("stream")
|
||||
|
||||
with patch(
|
||||
"supermemory_openai.middleware.supermemory.Supermemory",
|
||||
return_value=Mock(),
|
||||
), patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
side_effect=fake_prompt,
|
||||
):
|
||||
tenant_b: Any = with_supermemory(
|
||||
base_client,
|
||||
middleware_options("tenant-b"),
|
||||
)
|
||||
|
||||
for create in response_variant_creates(tenant_b, "with_raw_response"):
|
||||
asyncio.run(call_raw(create))
|
||||
|
||||
for create in response_variant_creates(tenant_b, "with_streaming_response"):
|
||||
stream_context = create(
|
||||
model="gpt-test",
|
||||
messages=[{"role": "user", "content": "streaming"}],
|
||||
)
|
||||
asyncio.run(consume_stream(stream_context))
|
||||
|
||||
assert lookups == ["tenant-b"] * 6
|
||||
for create in calls.values():
|
||||
assert "secret-tenant-b" in str(create.call_args.kwargs["messages"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", REAL_ASYNC_COMPLETION_PATHS)
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_async_openai_paths_use_async_middleware(path: str) -> None:
|
||||
lookups: list[str] = []
|
||||
writes: list[tuple[str, Optional[str], str]] = []
|
||||
sent_messages: list[list[Any]] = []
|
||||
|
||||
async def fake_prompt(
|
||||
messages: list[Any],
|
||||
container_tag: str,
|
||||
logger: Any,
|
||||
mode: Any,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
) -> list[Any]:
|
||||
lookups.append(container_tag)
|
||||
return [
|
||||
{"role": "system", "content": f"secret-{container_tag}"},
|
||||
*messages,
|
||||
]
|
||||
|
||||
async def fake_add_memory(
|
||||
client: Any,
|
||||
container_tag: str,
|
||||
content: str,
|
||||
custom_id: Optional[str],
|
||||
logger: Any,
|
||||
) -> None:
|
||||
writes.append((container_tag, custom_id, content))
|
||||
|
||||
def handle_request(request: httpx.Request) -> httpx.Response:
|
||||
body = json.loads(request.content)
|
||||
sent_messages.append(body["messages"])
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=request,
|
||||
headers={"content-type": "application/json"},
|
||||
json={
|
||||
"id": "chatcmpl-test",
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": "gpt-test",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "ok"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
http_client = httpx.AsyncClient(transport=httpx.MockTransport(handle_request))
|
||||
base_client = AsyncOpenAI(api_key="openai-test", http_client=http_client)
|
||||
|
||||
try:
|
||||
with patch(
|
||||
"supermemory_openai.middleware.supermemory.Supermemory",
|
||||
return_value=Mock(),
|
||||
), patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
new=fake_prompt,
|
||||
), patch(
|
||||
"supermemory_openai.middleware.add_memory_tool",
|
||||
new=fake_add_memory,
|
||||
), warnings.catch_warnings(
|
||||
record=True
|
||||
) as caught:
|
||||
warnings.simplefilter("always", RuntimeWarning)
|
||||
wrapped: Any = with_supermemory(
|
||||
base_client,
|
||||
middleware_options("tenant-real", add_memory="always"),
|
||||
)
|
||||
create = real_async_completion_create(wrapped, path)
|
||||
kwargs = {
|
||||
"model": "gpt-test",
|
||||
"messages": [{"role": "user", "content": "private message"}],
|
||||
}
|
||||
|
||||
if path == "normal":
|
||||
response = await create(**kwargs)
|
||||
assert response.id == "chatcmpl-test"
|
||||
elif path.endswith(".raw"):
|
||||
raw_response = await create(**kwargs)
|
||||
assert raw_response.parse().id == "chatcmpl-test"
|
||||
else:
|
||||
response_context = create(**kwargs)
|
||||
assert not inspect.isawaitable(response_context)
|
||||
async with response_context as streaming_response:
|
||||
assert streaming_response.status_code == 200
|
||||
|
||||
await wrapped.wait_for_background_tasks()
|
||||
await asyncio.sleep(0)
|
||||
gc.collect()
|
||||
|
||||
runtime_warnings = [
|
||||
warning
|
||||
for warning in caught
|
||||
if issubclass(warning.category, RuntimeWarning)
|
||||
]
|
||||
|
||||
assert lookups == ["tenant-real"]
|
||||
assert writes == [
|
||||
(
|
||||
"tenant-real",
|
||||
"conversation:thread-tenant-real",
|
||||
"User: private message",
|
||||
)
|
||||
]
|
||||
assert len(sent_messages) == 1
|
||||
assert sent_messages[0][0]["content"] == "secret-tenant-real"
|
||||
assert runtime_warnings == []
|
||||
finally:
|
||||
await base_client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio # type: ignore[untyped-decorator]
|
||||
async def test_shared_async_client_saves_only_for_selected_tenant() -> None:
|
||||
base_client, original_create = create_async_client()
|
||||
lookups: list[str] = []
|
||||
writes: list[tuple[str, Optional[str], str]] = []
|
||||
|
||||
async def fake_prompt(
|
||||
messages: list[Any],
|
||||
container_tag: str,
|
||||
logger: Any,
|
||||
mode: Any,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
) -> list[Any]:
|
||||
lookups.append(container_tag)
|
||||
return [
|
||||
{"role": "system", "content": f"secret-{container_tag}"},
|
||||
*messages,
|
||||
]
|
||||
|
||||
async def fake_add_memory(
|
||||
client: Any,
|
||||
container_tag: str,
|
||||
content: str,
|
||||
custom_id: Optional[str],
|
||||
logger: Any,
|
||||
) -> None:
|
||||
writes.append((container_tag, custom_id, content))
|
||||
|
||||
with patch(
|
||||
"supermemory_openai.middleware.supermemory.Supermemory",
|
||||
return_value=Mock(),
|
||||
), patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
side_effect=fake_prompt,
|
||||
), patch(
|
||||
"supermemory_openai.middleware.add_memory_tool",
|
||||
side_effect=fake_add_memory,
|
||||
):
|
||||
tenant_a: Any = with_supermemory(
|
||||
base_client,
|
||||
middleware_options("tenant-a", add_memory="always"),
|
||||
)
|
||||
tenant_b: Any = with_supermemory(
|
||||
base_client,
|
||||
middleware_options("tenant-b", add_memory="always"),
|
||||
)
|
||||
|
||||
await tenant_b.chat.completions.create(
|
||||
model="gpt-test",
|
||||
messages=[{"role": "user", "content": "private tenant B message"}],
|
||||
)
|
||||
await tenant_a.wait_for_background_tasks()
|
||||
await tenant_b.wait_for_background_tasks()
|
||||
|
||||
assert base_client.chat.completions.create is original_create
|
||||
assert lookups == ["tenant-b"]
|
||||
assert writes == [
|
||||
(
|
||||
"tenant-b",
|
||||
"conversation:thread-tenant-b",
|
||||
"User: private tenant B message",
|
||||
)
|
||||
]
|
||||
assert original_create.call_count == 1
|
||||
|
||||
|
||||
def sse_chat_response(request: httpx.Request) -> httpx.Response:
|
||||
chunk = {
|
||||
"id": "chatcmpl-test",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 0,
|
||||
"model": "gpt-test",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"role": "assistant", "content": "ok"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
}
|
||||
body = f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n"
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=request,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
content=body.encode(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_async", [False, True])
|
||||
def test_real_completions_stream_helper_uses_middleware(is_async: bool) -> None:
|
||||
from openai import OpenAI
|
||||
|
||||
lookups: list[str] = []
|
||||
sent_messages: list[list[Any]] = []
|
||||
|
||||
async def fake_prompt(
|
||||
messages: list[Any],
|
||||
container_tag: str,
|
||||
logger: Any,
|
||||
mode: Any,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
) -> list[Any]:
|
||||
lookups.append(container_tag)
|
||||
return [{"role": "system", "content": f"secret-{container_tag}"}, *messages]
|
||||
|
||||
def handle_request(request: httpx.Request) -> httpx.Response:
|
||||
sent_messages.append(json.loads(request.content)["messages"])
|
||||
return sse_chat_response(request)
|
||||
|
||||
transport = httpx.MockTransport(handle_request)
|
||||
if is_async:
|
||||
base_client: Any = AsyncOpenAI(
|
||||
api_key="openai-test", http_client=httpx.AsyncClient(transport=transport)
|
||||
)
|
||||
else:
|
||||
base_client = OpenAI(
|
||||
api_key="openai-test", http_client=httpx.Client(transport=transport)
|
||||
)
|
||||
kwargs = {
|
||||
"model": "gpt-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
|
||||
with patch(
|
||||
"supermemory_openai.middleware.supermemory.Supermemory",
|
||||
return_value=Mock(),
|
||||
), patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
new=fake_prompt,
|
||||
):
|
||||
wrapped: Any = with_supermemory(base_client, middleware_options("tenant-a"))
|
||||
if is_async:
|
||||
|
||||
async def consume() -> None:
|
||||
for client in (wrapped, base_client):
|
||||
async with client.chat.completions.stream(**kwargs) as stream:
|
||||
async for _ in stream:
|
||||
pass
|
||||
|
||||
asyncio.run(consume())
|
||||
else:
|
||||
with wrapped.chat.completions.stream(**kwargs) as stream:
|
||||
for _ in stream:
|
||||
pass
|
||||
with base_client.chat.completions.stream(**kwargs) as stream:
|
||||
for _ in stream:
|
||||
pass
|
||||
|
||||
assert lookups == ["tenant-a"]
|
||||
assert sent_messages[0][0] == {"role": "system", "content": "secret-tenant-a"}
|
||||
assert sent_messages[1] == [{"role": "user", "content": "hi"}]
|
||||
|
|
@ -1,8 +1,11 @@
|
|||
"""Tests for middleware module."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import pytest
|
||||
import requests
|
||||
import asyncio
|
||||
import threading
|
||||
from unittest.mock import AsyncMock, Mock, patch, MagicMock
|
||||
from typing import Dict, Any
|
||||
|
||||
|
|
@ -26,10 +29,16 @@ except ImportError:
|
|||
SupermemoryOpenAIWrapper,
|
||||
)
|
||||
|
||||
from supermemory_openai.middleware import supermemory_profile_search
|
||||
from openai import OpenAI, AsyncOpenAI
|
||||
from openai.types.chat import ChatCompletion, ChatCompletionMessage
|
||||
from openai.types import CompletionUsage
|
||||
|
||||
from supermemory_openai.exceptions import (
|
||||
SupermemoryAPIError,
|
||||
SupermemoryNetworkError,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_openai_client():
|
||||
|
|
@ -555,10 +564,12 @@ class TestErrorHandling:
|
|||
async def test_supermemory_api_error_handling(
|
||||
self, mock_async_openai_client, mock_openai_response
|
||||
):
|
||||
"""Test handling of Supermemory API errors."""
|
||||
"""Supermemory API errors fail open with the original messages."""
|
||||
original_create = AsyncMock(return_value=mock_openai_response)
|
||||
mock_async_openai_client.chat.completions.create = original_create
|
||||
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
|
||||
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
|
||||
with patch("supermemory_openai.middleware.supermemory_profile_search") as mock_search:
|
||||
mock_search.side_effect = Exception("API Error")
|
||||
|
|
@ -569,11 +580,14 @@ class TestErrorHandling:
|
|||
)
|
||||
|
||||
# Should not raise exception, should fall back gracefully
|
||||
with pytest.raises(Exception):
|
||||
await wrapped_client.chat.completions.create(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "Hello"}]
|
||||
)
|
||||
result = await wrapped_client.chat.completions.create(
|
||||
model="gpt-4",
|
||||
messages=messages
|
||||
)
|
||||
|
||||
assert result == mock_openai_response
|
||||
original_create.assert_called_once()
|
||||
assert original_create.call_args[1]["messages"] == messages
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_user_message_handling(
|
||||
|
|
@ -812,3 +826,325 @@ class TestBackgroundTaskManagement:
|
|||
)
|
||||
|
||||
# Should complete without error
|
||||
|
||||
|
||||
class TestMemorySearchFailOpen:
|
||||
"""The memory-search path must fail open."""
|
||||
|
||||
ORIGINAL_MESSAGES = [{"role": "user", "content": "What do you know about me?"}]
|
||||
ENHANCED_MESSAGES = [
|
||||
{"role": "system", "content": "memories about the user"},
|
||||
{"role": "user", "content": "What do you know about me?"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"exc, expected_log",
|
||||
[
|
||||
(SupermemoryNetworkError("network down"), "warn"),
|
||||
(SupermemoryAPIError("api boom", status_code=500), "warn"),
|
||||
(ValueError("totally unexpected"), "error"),
|
||||
],
|
||||
)
|
||||
async def test_async_fail_open(
|
||||
self, mock_async_openai_client, mock_openai_response, exc, expected_log
|
||||
):
|
||||
original_create = AsyncMock(return_value=mock_openai_response)
|
||||
mock_async_openai_client.chat.completions.create = original_create
|
||||
|
||||
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
|
||||
with patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
new=AsyncMock(side_effect=exc),
|
||||
):
|
||||
wrapped_client = with_supermemory(
|
||||
mock_async_openai_client,
|
||||
OpenAIMiddlewareOptions(
|
||||
container_tag="user-123",
|
||||
custom_id="test-conv",
|
||||
add_memory="never",
|
||||
),
|
||||
)
|
||||
wrapped_client._logger.warn = Mock()
|
||||
wrapped_client._logger.error = Mock()
|
||||
|
||||
result = await wrapped_client.chat.completions.create(
|
||||
model="gpt-4", messages=self.ORIGINAL_MESSAGES
|
||||
)
|
||||
|
||||
assert result == mock_openai_response
|
||||
|
||||
original_create.assert_called_once()
|
||||
assert (
|
||||
original_create.call_args[1]["messages"] == self.ORIGINAL_MESSAGES
|
||||
)
|
||||
|
||||
if expected_log == "warn":
|
||||
wrapped_client._logger.warn.assert_called_once()
|
||||
wrapped_client._logger.error.assert_not_called()
|
||||
else:
|
||||
wrapped_client._logger.error.assert_called_once()
|
||||
wrapped_client._logger.warn.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_happy_path_unchanged(
|
||||
self, mock_async_openai_client, mock_openai_response
|
||||
):
|
||||
"""Regression: when search succeeds, enhanced messages are still used."""
|
||||
original_create = AsyncMock(return_value=mock_openai_response)
|
||||
mock_async_openai_client.chat.completions.create = original_create
|
||||
|
||||
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
|
||||
with patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
new=AsyncMock(return_value=self.ENHANCED_MESSAGES),
|
||||
):
|
||||
wrapped_client = with_supermemory(
|
||||
mock_async_openai_client,
|
||||
OpenAIMiddlewareOptions(
|
||||
container_tag="user-123",
|
||||
custom_id="test-conv",
|
||||
add_memory="never",
|
||||
),
|
||||
)
|
||||
wrapped_client._logger.warn = Mock()
|
||||
wrapped_client._logger.error = Mock()
|
||||
|
||||
result = await wrapped_client.chat.completions.create(
|
||||
model="gpt-4", messages=self.ORIGINAL_MESSAGES
|
||||
)
|
||||
|
||||
assert result == mock_openai_response
|
||||
original_create.assert_called_once()
|
||||
assert (
|
||||
original_create.call_args[1]["messages"] == self.ENHANCED_MESSAGES
|
||||
)
|
||||
wrapped_client._logger.warn.assert_not_called()
|
||||
wrapped_client._logger.error.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"exc, expected_log",
|
||||
[
|
||||
(SupermemoryNetworkError("network down"), "warn"),
|
||||
(SupermemoryAPIError("api boom", status_code=500), "warn"),
|
||||
(ValueError("totally unexpected"), "error"),
|
||||
],
|
||||
)
|
||||
def test_sync_fail_open(
|
||||
self, mock_openai_client, mock_openai_response, exc, expected_log
|
||||
):
|
||||
original_create = Mock(return_value=mock_openai_response)
|
||||
mock_openai_client.chat.completions.create = original_create
|
||||
|
||||
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
|
||||
with patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
new=AsyncMock(side_effect=exc),
|
||||
):
|
||||
wrapped_client = with_supermemory(
|
||||
mock_openai_client,
|
||||
OpenAIMiddlewareOptions(
|
||||
container_tag="user-123",
|
||||
custom_id="test-conv",
|
||||
add_memory="never",
|
||||
),
|
||||
)
|
||||
wrapped_client._logger.warn = Mock()
|
||||
wrapped_client._logger.error = Mock()
|
||||
|
||||
result = wrapped_client.chat.completions.create(
|
||||
model="gpt-4", messages=self.ORIGINAL_MESSAGES
|
||||
)
|
||||
|
||||
assert result == mock_openai_response
|
||||
original_create.assert_called_once()
|
||||
assert (
|
||||
original_create.call_args[1]["messages"] == self.ORIGINAL_MESSAGES
|
||||
)
|
||||
|
||||
if expected_log == "warn":
|
||||
wrapped_client._logger.warn.assert_called_once()
|
||||
wrapped_client._logger.error.assert_not_called()
|
||||
else:
|
||||
wrapped_client._logger.error.assert_called_once()
|
||||
wrapped_client._logger.warn.assert_not_called()
|
||||
|
||||
def test_sync_fail_open_in_async_context(
|
||||
self, mock_openai_client, mock_openai_response
|
||||
):
|
||||
"""Fail-open must also cover the ThreadPoolExecutor fallback branch."""
|
||||
|
||||
async def run_in_async():
|
||||
original_create = Mock(return_value=mock_openai_response)
|
||||
mock_openai_client.chat.completions.create = original_create
|
||||
|
||||
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
|
||||
with patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
new=AsyncMock(side_effect=SupermemoryNetworkError("network down")),
|
||||
):
|
||||
wrapped_client = with_supermemory(
|
||||
mock_openai_client,
|
||||
OpenAIMiddlewareOptions(
|
||||
container_tag="user-123",
|
||||
custom_id="test-conv",
|
||||
add_memory="never",
|
||||
),
|
||||
)
|
||||
wrapped_client._logger.warn = Mock()
|
||||
wrapped_client._logger.error = Mock()
|
||||
|
||||
result = wrapped_client.chat.completions.create(
|
||||
model="gpt-4", messages=self.ORIGINAL_MESSAGES
|
||||
)
|
||||
|
||||
assert result == mock_openai_response
|
||||
original_create.assert_called_once()
|
||||
assert (
|
||||
original_create.call_args[1]["messages"]
|
||||
== self.ORIGINAL_MESSAGES
|
||||
)
|
||||
wrapped_client._logger.warn.assert_called_once()
|
||||
wrapped_client._logger.error.assert_not_called()
|
||||
|
||||
asyncio.run(run_in_async())
|
||||
|
||||
def test_sync_happy_path_unchanged(self, mock_openai_client, mock_openai_response):
|
||||
"""Regression: when search succeeds, enhanced messages are still used."""
|
||||
original_create = Mock(return_value=mock_openai_response)
|
||||
mock_openai_client.chat.completions.create = original_create
|
||||
|
||||
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
|
||||
with patch(
|
||||
"supermemory_openai.middleware.add_system_prompt",
|
||||
new=AsyncMock(return_value=self.ENHANCED_MESSAGES),
|
||||
):
|
||||
wrapped_client = with_supermemory(
|
||||
mock_openai_client,
|
||||
OpenAIMiddlewareOptions(
|
||||
container_tag="user-123",
|
||||
custom_id="test-conv",
|
||||
add_memory="never",
|
||||
),
|
||||
)
|
||||
wrapped_client._logger.warn = Mock()
|
||||
wrapped_client._logger.error = Mock()
|
||||
|
||||
result = wrapped_client.chat.completions.create(
|
||||
model="gpt-4", messages=self.ORIGINAL_MESSAGES
|
||||
)
|
||||
|
||||
assert result == mock_openai_response
|
||||
original_create.assert_called_once()
|
||||
assert (
|
||||
original_create.call_args[1]["messages"] == self.ENHANCED_MESSAGES
|
||||
)
|
||||
wrapped_client._logger.warn.assert_not_called()
|
||||
wrapped_client._logger.error.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_black_box_async_connection_error(
|
||||
self, mock_async_openai_client, mock_openai_response
|
||||
):
|
||||
"""A real aiohttp connection error must not crash create()."""
|
||||
aiohttp = pytest.importorskip("aiohttp")
|
||||
|
||||
original_create = AsyncMock(return_value=mock_openai_response)
|
||||
mock_async_openai_client.chat.completions.create = original_create
|
||||
|
||||
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
|
||||
with patch(
|
||||
"supermemory_openai.middleware.add_memory_tool",
|
||||
new=AsyncMock(return_value=None),
|
||||
):
|
||||
with patch(
|
||||
"aiohttp.ClientSession.post",
|
||||
side_effect=aiohttp.ClientConnectionError("connection refused"),
|
||||
):
|
||||
wrapped_client = with_supermemory(
|
||||
mock_async_openai_client,
|
||||
OpenAIMiddlewareOptions(
|
||||
container_tag="user-123", custom_id="test-conv"
|
||||
),
|
||||
)
|
||||
|
||||
result = await wrapped_client.chat.completions.create(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
)
|
||||
|
||||
assert result == mock_openai_response
|
||||
original_create.assert_called_once()
|
||||
|
||||
def test_black_box_sync_connection_error(
|
||||
self, mock_openai_client, mock_openai_response
|
||||
):
|
||||
"""A real requests connection error must not crash create()."""
|
||||
import sys
|
||||
|
||||
import requests
|
||||
|
||||
original_create = Mock(return_value=mock_openai_response)
|
||||
mock_openai_client.chat.completions.create = original_create
|
||||
|
||||
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
|
||||
with patch(
|
||||
"supermemory_openai.middleware.add_memory_tool",
|
||||
new=AsyncMock(return_value=None),
|
||||
):
|
||||
with patch.dict(sys.modules, {"aiohttp": None}):
|
||||
with patch(
|
||||
"requests.post",
|
||||
side_effect=requests.exceptions.ConnectionError(
|
||||
"connection refused"
|
||||
),
|
||||
):
|
||||
wrapped_client = with_supermemory(
|
||||
mock_openai_client,
|
||||
OpenAIMiddlewareOptions(
|
||||
container_tag="user-123", custom_id="test-conv"
|
||||
),
|
||||
)
|
||||
|
||||
result = wrapped_client.chat.completions.create(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
)
|
||||
|
||||
assert result == mock_openai_response
|
||||
original_create.assert_called_once()
|
||||
|
||||
|
||||
class TestProfileSearchFallback:
|
||||
"""Test the requests fallback used when aiohttp is not installed."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_requests_fallback_does_not_block_event_loop(self):
|
||||
"""requests.post must run off the event loop so other tasks keep running."""
|
||||
released = threading.Event()
|
||||
|
||||
def slow_post(*args, **kwargs):
|
||||
if not released.wait(timeout=2):
|
||||
raise AssertionError("requests.post blocked the event loop")
|
||||
response = Mock(status_code=200)
|
||||
response.json.return_value = {
|
||||
"profile": {"static": ["Prefers vegetarian food"], "dynamic": []},
|
||||
"searchResults": {"results": []},
|
||||
}
|
||||
return response
|
||||
|
||||
with patch.dict(sys.modules, {"aiohttp": None}):
|
||||
with patch.object(requests, "post", side_effect=slow_post):
|
||||
search = asyncio.create_task(
|
||||
supermemory_profile_search(
|
||||
"user-123", "", "test-key", "https://api.supermemory.ai"
|
||||
)
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
released.set()
|
||||
result = await search
|
||||
|
||||
assert result.profile["static"] == ["Prefers vegetarian food"]
|
||||
|
|
|
|||
|
|
@ -255,6 +255,23 @@ class TestMemoryOperationsUnit:
|
|||
q="tea",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_profile_handles_user_with_no_memories(self):
|
||||
"""A null profile (new user, no memories yet) is an empty profile, not a failure."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
tools = SupermemoryTools("test-key", {"container_tags": ["unit-tag"]})
|
||||
tools.client.profile = AsyncMock(
|
||||
return_value=SimpleNamespace(profile=None, search_results=None)
|
||||
)
|
||||
|
||||
result = await tools.get_profile()
|
||||
|
||||
assert result["success"] is True
|
||||
assert result["profile"] == {"static": [], "dynamic": []}
|
||||
assert result["search_results"] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_list_uses_client_documents_list(self):
|
||||
"""document_list must call client.documents.list with container tag."""
|
||||
|
|
|
|||
2
packages/openai-sdk-python/uv.lock
generated
2
packages/openai-sdk-python/uv.lock
generated
|
|
@ -1287,7 +1287,7 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "supermemory-openai-sdk"
|
||||
version = "1.0.9"
|
||||
version = "1.0.10"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "openai" },
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "supermemory-pipecat"
|
||||
version = "0.1.4"
|
||||
version = "0.1.5"
|
||||
description = "Supermemory integration for Pipecat - memory-enhanced conversational AI pipelines"
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
|
@ -13,7 +13,7 @@ authors = [
|
|||
{ name = "Supermemory", email = "support@supermemory.ai" }
|
||||
]
|
||||
keywords = [
|
||||
"supermemory>=3.16.0,<5",
|
||||
"supermemory",
|
||||
"pipecat",
|
||||
"memory",
|
||||
"conversational-ai",
|
||||
|
|
@ -32,7 +32,7 @@ classifiers = [
|
|||
]
|
||||
dependencies = [
|
||||
"pipecat-ai>=0.0.98,<2.0.0",
|
||||
"supermemory>=3.16.0",
|
||||
"supermemory>=3.16.0,<5",
|
||||
"pydantic>=2.10.0",
|
||||
"loguru>=0.7.3",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -42,15 +42,18 @@ from .utils import (
|
|||
get_last_user_message,
|
||||
)
|
||||
|
||||
InputParams = SupermemoryPipecatService.InputParams
|
||||
|
||||
try:
|
||||
__version__ = version("supermemory-pipecat")
|
||||
except PackageNotFoundError:
|
||||
# Source-tree fallback; built wheels always use package metadata above.
|
||||
__version__ = "0.1.3"
|
||||
__version__ = "0.1.5"
|
||||
|
||||
__all__ = [
|
||||
# Main service
|
||||
"SupermemoryPipecatService",
|
||||
"InputParams",
|
||||
# Exceptions
|
||||
"SupermemoryPipecatError",
|
||||
"ConfigurationError",
|
||||
|
|
|
|||
|
|
@ -416,7 +416,7 @@ class SupermemoryPipecatService(FrameProcessor):
|
|||
existing_content = messages[system_idx].get("content", "")
|
||||
if MEMORY_TAG_PATTERN.search(existing_content):
|
||||
messages[system_idx]["content"] = MEMORY_TAG_PATTERN.sub(
|
||||
tagged_memory, existing_content
|
||||
lambda _match: tagged_memory, existing_content
|
||||
)
|
||||
else:
|
||||
messages[system_idx]["content"] = f"{existing_content}\n\n{tagged_memory}"
|
||||
|
|
|
|||
|
|
@ -90,6 +90,17 @@ def _field(item: Any, *names: str, default: Any = None) -> Any:
|
|||
return default
|
||||
|
||||
|
||||
def _memory_text(item: Any) -> str:
|
||||
"""First non-empty memory field; `_field` alone would stop at an empty `memory`."""
|
||||
if isinstance(item, str):
|
||||
return item.strip()
|
||||
for name in ("memory", "chunk", "content"):
|
||||
value = _field(item, name)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value.strip()
|
||||
return ""
|
||||
|
||||
|
||||
def deduplicate_memories(
|
||||
static: List[str],
|
||||
dynamic: List[str],
|
||||
|
|
@ -125,12 +136,7 @@ def deduplicate_memories(
|
|||
out: List[Any] = []
|
||||
for r in results:
|
||||
# v4 search.memories/hybrid uses `memory` or `chunk`.
|
||||
memory = (
|
||||
r if isinstance(r, str) else _field(r, "memory", "chunk", "content", default="")
|
||||
)
|
||||
if not isinstance(memory, str):
|
||||
memory = ""
|
||||
memory = memory.strip()
|
||||
memory = _memory_text(r)
|
||||
key = comparison_key(memory)
|
||||
if key and key not in seen:
|
||||
seen.add(key)
|
||||
|
|
@ -177,7 +183,7 @@ def format_memories_to_text(
|
|||
lines.append(f"- {item}")
|
||||
continue
|
||||
|
||||
memory = _field(item, "memory", "chunk", "content", default="")
|
||||
memory = _memory_text(item)
|
||||
updated_at = _field(item, "updatedAt", "updated_at", default="")
|
||||
time_str = format_relative_time(updated_at) if updated_at else ""
|
||||
if time_str:
|
||||
|
|
|
|||
97
packages/pipecat-sdk-python/tests/conftest.py
Normal file
97
packages/pipecat-sdk-python/tests/conftest.py
Normal file
|
|
@ -0,0 +1,97 @@
|
|||
"""Import stubs for heavy optional deps; no-ops in CI where the real ones are installed."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types
|
||||
|
||||
|
||||
def _install_test_stubs() -> None:
|
||||
if "loguru" not in sys.modules:
|
||||
loguru_module = types.ModuleType("loguru")
|
||||
|
||||
class _Logger:
|
||||
def warning(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
def error(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
loguru_module.logger = _Logger()
|
||||
sys.modules["loguru"] = loguru_module
|
||||
|
||||
if "pydantic" not in sys.modules:
|
||||
pydantic_module = types.ModuleType("pydantic")
|
||||
|
||||
class BaseModel:
|
||||
def __init__(self, **kwargs):
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
def Field(*, default=None, **_kwargs):
|
||||
return default
|
||||
|
||||
pydantic_module.BaseModel = BaseModel
|
||||
pydantic_module.Field = Field
|
||||
sys.modules["pydantic"] = pydantic_module
|
||||
|
||||
if "pipecat" not in sys.modules:
|
||||
pipecat_module = types.ModuleType("pipecat")
|
||||
sys.modules["pipecat"] = pipecat_module
|
||||
|
||||
frames_module = types.ModuleType("pipecat.frames.frames")
|
||||
|
||||
class Frame: # pragma: no cover - import stub
|
||||
pass
|
||||
|
||||
class InputAudioRawFrame: # pragma: no cover - import stub
|
||||
pass
|
||||
|
||||
class LLMContextFrame: # pragma: no cover - import stub
|
||||
pass
|
||||
|
||||
class LLMMessagesFrame: # pragma: no cover - import stub
|
||||
pass
|
||||
|
||||
frames_module.Frame = Frame
|
||||
frames_module.InputAudioRawFrame = InputAudioRawFrame
|
||||
frames_module.LLMContextFrame = LLMContextFrame
|
||||
frames_module.LLMMessagesFrame = LLMMessagesFrame
|
||||
|
||||
llm_context_module = types.ModuleType("pipecat.processors.aggregators.llm_context")
|
||||
|
||||
class LLMContext: # pragma: no cover - import stub
|
||||
pass
|
||||
|
||||
llm_context_module.LLMContext = LLMContext
|
||||
|
||||
openai_context_module = types.ModuleType(
|
||||
"pipecat.processors.aggregators.openai_llm_context"
|
||||
)
|
||||
|
||||
class OpenAILLMContextFrame: # pragma: no cover - import stub
|
||||
pass
|
||||
|
||||
openai_context_module.OpenAILLMContextFrame = OpenAILLMContextFrame
|
||||
|
||||
frame_processor_module = types.ModuleType("pipecat.processors.frame_processor")
|
||||
|
||||
class FrameDirection: # pragma: no cover - import stub
|
||||
pass
|
||||
|
||||
class FrameProcessor:
|
||||
def __init__(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
frame_processor_module.FrameDirection = FrameDirection
|
||||
frame_processor_module.FrameProcessor = FrameProcessor
|
||||
|
||||
sys.modules["pipecat.frames.frames"] = frames_module
|
||||
sys.modules["pipecat.processors.aggregators.llm_context"] = llm_context_module
|
||||
sys.modules[
|
||||
"pipecat.processors.aggregators.openai_llm_context"
|
||||
] = openai_context_module
|
||||
sys.modules["pipecat.processors.frame_processor"] = frame_processor_module
|
||||
|
||||
|
||||
_install_test_stubs()
|
||||
|
|
@ -154,3 +154,48 @@ class TestSupermemoryPipecatNullProfile(unittest.IsolatedAsyncioTestCase):
|
|||
self.assertTrue(
|
||||
any(fact in message.get("content", "") for message in context.messages)
|
||||
)
|
||||
|
||||
def test_system_injection_replaces_stale_block_with_backslash_memory(self) -> None:
|
||||
"""A memory containing backslashes must not be read as a regex template."""
|
||||
backslash = chr(92)
|
||||
fact = "User's repo is at C:" + backslash + "Users" + backslash + "alice"
|
||||
service = SupermemoryPipecatService(
|
||||
api_key="mock_key",
|
||||
user_id="user-123",
|
||||
session_id="conversation-456",
|
||||
params=SupermemoryPipecatService.InputParams(inject_mode="system"),
|
||||
)
|
||||
newline = chr(10)
|
||||
stale = newline.join(["<user_memories>", "Stale fact", "</user_memories>"])
|
||||
|
||||
class Context:
|
||||
def __init__(self):
|
||||
self.messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are helpful." + newline + newline + stale,
|
||||
},
|
||||
{"role": "user", "content": "Where is my repo?"},
|
||||
]
|
||||
|
||||
def get_messages(self):
|
||||
return self.messages
|
||||
|
||||
def add_message(self, message):
|
||||
self.messages.append(message)
|
||||
|
||||
context = Context()
|
||||
service._enhance_context_with_memories(
|
||||
context,
|
||||
"Where is my repo?",
|
||||
{
|
||||
"profile": {"static": [fact], "dynamic": []},
|
||||
"search_results": [],
|
||||
},
|
||||
)
|
||||
|
||||
system_content = context.messages[0]["content"]
|
||||
self.assertIn(fact, system_content)
|
||||
self.assertNotIn("Stale fact", system_content)
|
||||
self.assertIn("You are helpful.", system_content)
|
||||
self.assertEqual(system_content.count("<user_memories>"), 1)
|
||||
|
|
|
|||
31
packages/pipecat-sdk-python/tests/test_public_exports.py
Normal file
31
packages/pipecat-sdk-python/tests/test_public_exports.py
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
"""Documented public names, like `InputParams`, must import from the package root."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from .test_empty_profile import _install_test_stubs
|
||||
|
||||
_install_test_stubs()
|
||||
|
||||
import supermemory_pipecat
|
||||
from supermemory_pipecat import InputParams, SupermemoryPipecatService
|
||||
|
||||
|
||||
class TestPublicExports(unittest.TestCase):
|
||||
def test_input_params_is_exported_from_the_package_root(self) -> None:
|
||||
self.assertIn("InputParams", supermemory_pipecat.__all__)
|
||||
self.assertIs(InputParams, SupermemoryPipecatService.InputParams)
|
||||
|
||||
def test_documented_configuration_example_constructs(self) -> None:
|
||||
params = InputParams(
|
||||
mode="full",
|
||||
search_limit=10,
|
||||
search_threshold=0.1,
|
||||
system_prompt="Based on previous conversations:\n\n",
|
||||
)
|
||||
|
||||
self.assertEqual(params.mode, "full")
|
||||
self.assertEqual(params.search_limit, 10)
|
||||
self.assertEqual(params.search_threshold, 0.1)
|
||||
self.assertEqual(params.system_prompt, "Based on previous conversations:\n\n")
|
||||
86
packages/pipecat-sdk-python/tests/test_utils.py
Normal file
86
packages/pipecat-sdk-python/tests/test_utils.py
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from supermemory_pipecat.utils import (
|
||||
deduplicate_memories,
|
||||
format_memories_to_text,
|
||||
)
|
||||
|
||||
|
||||
class TestSearchResultExtraction(unittest.TestCase):
|
||||
"""v4 hybrid search returns chunk hits whose `memory` field is empty."""
|
||||
|
||||
def test_falls_through_to_chunk_when_memory_is_empty(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[{"memory": "", "chunk": "The user's dog is called Rex"}],
|
||||
)
|
||||
|
||||
self.assertEqual(len(result["search_results"]), 1)
|
||||
self.assertIn("The user's dog is called Rex", format_memories_to_text(result))
|
||||
|
||||
def test_falls_through_to_chunk_when_memory_is_whitespace(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[{"memory": " ", "chunk": "Chunk body"}],
|
||||
)
|
||||
|
||||
self.assertEqual(len(result["search_results"]), 1)
|
||||
self.assertIn("Chunk body", format_memories_to_text(result))
|
||||
|
||||
def test_falls_through_on_models_too(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[SimpleNamespace(memory="", chunk="Chunk body")],
|
||||
)
|
||||
|
||||
self.assertEqual(len(result["search_results"]), 1)
|
||||
self.assertIn("Chunk body", format_memories_to_text(result))
|
||||
|
||||
def test_never_renders_an_empty_bullet(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[
|
||||
{
|
||||
"memory": "",
|
||||
"chunk": "Chunk body",
|
||||
"updatedAt": "2020-01-01T00:00:00Z",
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
rendered = format_memories_to_text(result)
|
||||
self.assertIn("Chunk body", rendered)
|
||||
for line in rendered.splitlines():
|
||||
if line.startswith("- "):
|
||||
self.assertNotRegex(line, r"^- (\[[^\]]*\] )?$")
|
||||
|
||||
def test_memory_still_wins_over_chunk(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[{"memory": "Memory body", "chunk": "Chunk body"}],
|
||||
)
|
||||
|
||||
rendered = format_memories_to_text(result)
|
||||
self.assertIn("Memory body", rendered)
|
||||
self.assertNotIn("Chunk body", rendered)
|
||||
|
||||
def test_entry_with_no_usable_text_is_still_dropped(self) -> None:
|
||||
result = deduplicate_memories(
|
||||
static=[],
|
||||
dynamic=[],
|
||||
search_results=[{"memory": "", "chunk": ""}, {}, None],
|
||||
)
|
||||
|
||||
self.assertEqual(result["search_results"], [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Loading…
Add table
Reference in a new issue