diff --git a/packages/cartesia-sdk-python/src/supermemory_cartesia/agent.py b/packages/cartesia-sdk-python/src/supermemory_cartesia/agent.py index 9019193e8..380a6915d 100644 --- a/packages/cartesia-sdk-python/src/supermemory_cartesia/agent.py +++ b/packages/cartesia-sdk-python/src/supermemory_cartesia/agent.py @@ -377,6 +377,55 @@ async def _enrich_event_with_memories(self, event: Any) -> tuple[Any, Optional[s logger.error(f"[Supermemory] Error in memory enrichment: {e}") return event, None + async def _prepare_user_turn(self, event: Event) -> Event: + """Prepare a user turn with memory context and background persistence.""" + logger.info("[Supermemory] Processing UserTurnEnded event") + event, memory_context = await self._enrich_event_with_memories(event) + + # Clean up old memory context and inject new one if available + if hasattr(self.agent, 'config'): + original_prompt = getattr(self.agent.config, 'system_prompt', '') + # Always remove old memory context if present to prevent stale data + if MEMORY_TAG_START in original_prompt: + original_prompt = re.sub( + rf'{re.escape(MEMORY_TAG_START)}.*?{re.escape(MEMORY_TAG_END)}\s*', + '', + original_prompt, + flags=re.DOTALL + ) + logger.debug("[Supermemory] Removed old memory context from system prompt") + + # Inject new memory context if available + if memory_context: + self.agent.config.system_prompt = f"{memory_context}\n\n{original_prompt}" + logger.info("[Supermemory] Injected new memory context into system prompt") + else: + # No new memories, but we cleaned up old ones + self.agent.config.system_prompt = original_prompt + logger.debug("[Supermemory] No new memories to inject, using clean prompt") + + # Store conversation in background + if hasattr(event, 'history') and event.history: + messages = self._extract_conversation_from_history(event.history) + unsent = messages[self._messages_sent_count:] + if unsent: + logger.info(f"[Supermemory] Queuing {len(unsent)} messages for storage") + task = asyncio.create_task(self._store_messages(unsent)) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + self._messages_sent_count = len(messages) + else: + # No history yet, store just the current user message + user_content = self._extract_user_message(event) + if user_content: + logger.info(f"[Supermemory] No history, storing current user message: {user_content[:50]}...") + task = asyncio.create_task(self._store_messages([{"role": "user", "content": user_content}])) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + self._messages_sent_count = 1 # CRITICAL: Increment counter to prevent duplicate storage + + return event + async def process(self, env: Any, event: Event) -> AsyncGenerator[Event, None]: """Process events with memory enrichment. @@ -387,63 +436,14 @@ async def process(self, env: Any, event: Event) -> AsyncGenerator[Event, None]: Yields: Output events from the wrapped agent. """ - try: - if type(event).__name__ == "UserTurnEnded": - logger.info("[Supermemory] Processing UserTurnEnded event") - event, memory_context = await self._enrich_event_with_memories(event) - - # Clean up old memory context and inject new one if available - if hasattr(self.agent, 'config'): - original_prompt = getattr(self.agent.config, 'system_prompt', '') - # Always remove old memory context if present to prevent stale data - if MEMORY_TAG_START in original_prompt: - original_prompt = re.sub( - rf'{re.escape(MEMORY_TAG_START)}.*?{re.escape(MEMORY_TAG_END)}\s*', - '', - original_prompt, - flags=re.DOTALL - ) - logger.debug("[Supermemory] Removed old memory context from system prompt") - - # Inject new memory context if available - if memory_context: - self.agent.config.system_prompt = f"{memory_context}\n\n{original_prompt}" - logger.info("[Supermemory] Injected new memory context into system prompt") - else: - # No new memories, but we cleaned up old ones - self.agent.config.system_prompt = original_prompt - logger.debug("[Supermemory] No new memories to inject, using clean prompt") - - # Store conversation in background - if hasattr(event, 'history') and event.history: - messages = self._extract_conversation_from_history(event.history) - unsent = messages[self._messages_sent_count:] - if unsent: - logger.info(f"[Supermemory] Queuing {len(unsent)} messages for storage") - task = asyncio.create_task(self._store_messages(unsent)) - self._background_tasks.add(task) - task.add_done_callback(self._background_tasks.discard) - self._messages_sent_count = len(messages) - else: - # No history yet, store just the current user message - user_content = self._extract_user_message(event) - if user_content: - logger.info(f"[Supermemory] No history, storing current user message: {user_content[:50]}...") - task = asyncio.create_task(self._store_messages([{"role": "user", "content": user_content}])) - self._background_tasks.add(task) - task.add_done_callback(self._background_tasks.discard) - self._messages_sent_count = 1 # CRITICAL: Increment counter to prevent duplicate storage - - async for output in self.agent.process(env, event): - yield output - else: - async for output in self.agent.process(env, event): - yield output + if type(event).__name__ == "UserTurnEnded": + try: + event = await self._prepare_user_turn(event) + except Exception as e: + logger.error(f"[Supermemory] Error preparing user turn: {e}") - except Exception as e: - logger.error(f"[Supermemory] Error in process: {e}") - async for output in self.agent.process(env, event): - yield output + async for output in self.agent.process(env, event): + yield output def reset_memory_tracking(self) -> None: """Reset memory tracking for a new conversation.""" diff --git a/packages/cartesia-sdk-python/tests/test_process.py b/packages/cartesia-sdk-python/tests/test_process.py new file mode 100644 index 000000000..bb76bc7fe --- /dev/null +++ b/packages/cartesia-sdk-python/tests/test_process.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import sys +import types +import unittest +from unittest.mock import AsyncMock + + +def _install_test_stubs() -> None: + if "loguru" not in sys.modules: + loguru_module = types.ModuleType("loguru") + + class _Logger: + def debug(self, *_args, **_kwargs): + return None + + 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.agent import SupermemoryCartesiaAgent + + +class UserTurnEnded: + def __init__(self, content: str = "Hello") -> None: + self.content = content + self.history = [] + + +class _FailingAgent: + def __init__(self, *, fail_after_output: bool) -> None: + self.calls = 0 + self.fail_after_output = fail_after_output + + async def process(self, _env, _event): + self.calls += 1 + if self.fail_after_output: + yield "first output" + raise RuntimeError("inner agent failed") + + +class _SuccessfulAgent: + def __init__(self) -> None: + self.calls = 0 + + async def process(self, _env, _event): + self.calls += 1 + yield "output" + + +def _wrapper(agent) -> SupermemoryCartesiaAgent: + return SupermemoryCartesiaAgent( + agent=agent, + api_key="mock_key", + container_tag="user-123", + custom_id="conversation-456", + ) + + +class TestSupermemoryCartesiaProcess(unittest.IsolatedAsyncioTestCase): + async def test_does_not_restart_agent_after_partial_output_failure(self) -> None: + inner_agent = _FailingAgent(fail_after_output=True) + wrapper = _wrapper(inner_agent) + event = UserTurnEnded() + wrapper._prepare_user_turn = AsyncMock(return_value=event) + + outputs = [] + with self.assertRaisesRegex(RuntimeError, "inner agent failed"): + async for output in wrapper.process(object(), event): + outputs.append(output) + + self.assertEqual(outputs, ["first output"]) + self.assertEqual(inner_agent.calls, 1) + + async def test_does_not_restart_agent_when_it_fails_before_output(self) -> None: + inner_agent = _FailingAgent(fail_after_output=False) + wrapper = _wrapper(inner_agent) + event = UserTurnEnded() + wrapper._prepare_user_turn = AsyncMock(return_value=event) + + with self.assertRaisesRegex(RuntimeError, "inner agent failed"): + async for _output in wrapper.process(object(), event): + pass + + self.assertEqual(inner_agent.calls, 1) + + async def test_preparation_failure_falls_back_to_one_agent_call(self) -> None: + inner_agent = _SuccessfulAgent() + wrapper = _wrapper(inner_agent) + event = UserTurnEnded() + wrapper._prepare_user_turn = AsyncMock( + side_effect=RuntimeError("memory preparation failed") + ) + + outputs = [output async for output in wrapper.process(object(), event)] + + self.assertEqual(outputs, ["output"]) + self.assertEqual(inner_agent.calls, 1) + + +if __name__ == "__main__": + unittest.main()