From 4bd0cc3161bf876e78e56ee75458f7e4c833c17d Mon Sep 17 00:00:00 2001 From: makeitdouble Date: Fri, 29 May 2026 17:59:43 +0300 Subject: [PATCH] feature: append message - add append last message before runtime memory update - remove separate runtime memory timeout using service timeout instead - fix tests --- config.example.py | 2 - memory/message_memory.py | 409 ++++++++++++++++++++++++++++------- runtime/runtime_context.py | 8 + tests/test_message_memory.py | 31 +-- utils/runtime_actions.py | 2 +- websocket.py | 5 + 6 files changed, 354 insertions(+), 103 deletions(-) diff --git a/config.example.py b/config.example.py index 5fe4618e..e6777886 100644 --- a/config.example.py +++ b/config.example.py @@ -33,8 +33,6 @@ SERVICE_REQUEST_TIMEOUT = 200.0 -RUNTIME_MEMORY_REQUEST_TIMEOUT = 120.0 - SERVICE_CONTEXT_WINDOW = 4096 SERVICE_TEMPERATURE = 0.1 diff --git a/memory/message_memory.py b/memory/message_memory.py index 10547189..f02419b1 100644 --- a/memory/message_memory.py +++ b/memory/message_memory.py @@ -18,39 +18,10 @@ "You have no history with the user yet." ) - -DEFAULT_MEMORY_REQUEST_TIMEOUT = 120.0 - - -def get_memory_request_timeout() -> float: - - configured_timeout = getattr( - config, - "RUNTIME_MEMORY_REQUEST_TIMEOUT", - None, - ) - - if configured_timeout is None: - configured_timeout = max( - float( - getattr( - config, - "SERVICE_REQUEST_TIMEOUT", - 0.0, - ) - ), - DEFAULT_MEMORY_REQUEST_TIMEOUT, - ) - - return float( - configured_timeout - ) - - async def safe_call( - call, - *args, - **kwargs, + call, + *args, + **kwargs, ): if call is None: @@ -64,7 +35,7 @@ async def safe_call( async def emit_runtime_memory_update( - context, + context, ) -> None: emitter = getattr( @@ -158,10 +129,61 @@ def build_runtime_memory_user_prompt( ) +def build_runtime_memory_batch_user_prompt( + *, + current_memory: str, + turns: list[dict], +) -> str: + + lines = [ + "Current runtime memory:", + current_memory.strip() or DEFAULT_RUNTIME_MEMORY, + "", + "New completed turns since that memory snapshot:", + ] + + for index, turn in enumerate( + turns, + start=1, + ): + lines.extend([ + "", + f"Turn {index}", + "Latest user message:", + ( + turn.get( + "user_message", + "", + ) + .strip() + ), + "", + "Latest JIN answer:", + ( + turn.get( + "assistant_message", + "", + ) + .strip() + ), + ]) + + lines.extend([ + "", + "Rewrite the runtime memory now as atomic bullet lines.", + "Use the current memory as the last stable snapshot.", + "Integrate all new completed turns in order.", + ]) + + return "\n".join( + lines + ) + + def build_interrupted_assistant_message( - *, - user_message: str, - assistant_message: str, + *, + user_message: str, + assistant_message: str, ) -> str: partial_text = assistant_message.strip() @@ -182,23 +204,23 @@ def build_interrupted_assistant_message( def extract_runtime_memory_text( - response: dict, + response: dict, ) -> str: text = ( - ResponseExtractor.extract_content_text( - response - ) - or ResponseExtractor.extract_reasoning_text( - response - ) + ResponseExtractor.extract_content_text( + response + ) + or ResponseExtractor.extract_reasoning_text( + response + ) ) return text.strip() def is_runtime_memory_response_truncated( - response: dict, + response: dict, ) -> bool: finish_reason = ( @@ -216,23 +238,23 @@ def is_runtime_memory_response_truncated( def looks_like_incomplete_runtime_memory( - text: str, + text: str, ) -> bool: stripped = ( - text - or "" + text + or "" ).strip() if not stripped: return True if stripped[-1] in ( - ",", - ":", - "(", - "[", - "{", + ",", + ":", + "(", + "[", + "{", ): return True @@ -260,11 +282,11 @@ def looks_like_incomplete_runtime_memory( async def ask_runtime_memory_model( - *, - service_client, - current_memory: str, - user_message: str, - assistant_message: str, + *, + service_client, + current_memory: str, + user_message: str, + assistant_message: str, ) -> dict: return await ask_service_model( @@ -285,15 +307,41 @@ async def ask_runtime_memory_model( max_tokens=( config.SERVICE_MAX_TOKENS ), - timeout=get_memory_request_timeout(), + ) + + +async def ask_runtime_memory_batch_model( + *, + service_client, + current_memory: str, + turns: list[dict], +) -> dict: + + return await ask_service_model( + client=service_client, + system_prompt=( + build_runtime_memory_system_prompt() + ), + user_prompt=( + build_runtime_memory_batch_user_prompt( + current_memory=current_memory, + turns=turns, + ) + ), + temperature=( + config.SERVICE_TEMPERATURE + ), + max_tokens=( + config.SERVICE_MAX_TOKENS + ), ) async def summarize_runtime_memory( - *, - context, - user_message: str, - assistant_message: str, + *, + context, + user_message: str, + assistant_message: str, ) -> str: if not assistant_message.strip(): @@ -340,12 +388,158 @@ async def summarize_runtime_memory( ) if ( - is_runtime_memory_response_truncated( - response + is_runtime_memory_response_truncated( + response + ) + or looks_like_incomplete_runtime_memory( + updated_memory + ) + ): + await safe_call( + getattr( + getattr( + context, + "logger", + None, + ), + "log_error", + None, + ), + "[MEMORY] runtime memory update skipped", + details=( + "Summarizer returned an incomplete memory update." + ), ) - or looks_like_incomplete_runtime_memory( - updated_memory + + return current_memory + + updates_counter = getattr( + context, + "runtime_memory_updates", + 0, + ) + + if updated_memory or updates_counter == 0: + context.runtime_memory = updated_memory + context.runtime_memory_stable = updated_memory + context.runtime_memory_updates = updates_counter + 1 + + logger = getattr( + context, + "logger", + None, + ) + log_service = getattr( + logger, + "log_service", + None, ) + + await safe_call( + log_service, + "[MEMORY] runtime memory updated", + ) + + await emit_runtime_memory_update( + context + ) + + return getattr( + context, + "runtime_memory", + "", + ) + + except asyncio.CancelledError: + raise + + except Exception: + formatted_traceback = ( + traceback.format_exc() + ) + + logger = getattr( + context, + "logger", + None, + ) + log_error = getattr( + logger, + "log_error", + None, + ) + + await safe_call( + log_error, + "[MEMORY] runtime memory update failed", + details=formatted_traceback, + ) + + return getattr( + context, + "runtime_memory", + "", + ) + + +async def summarize_runtime_memory_pending_turns( + *, + context, +) -> str: + + turns = list( + context.runtime_memory_pending_turns + ) + + if not turns: + return getattr( + context, + "runtime_memory", + "", + ) + + service_client = ( + getattr( + context, + "clients", + {}, + ) + .get( + "service" + ) + ) + + if service_client is None: + return getattr( + context, + "runtime_memory", + "", + ) + + initial_memory = getattr( + context, + "runtime_memory_stable", + "", + ) + + try: + response = await ask_runtime_memory_batch_model( + service_client=service_client, + current_memory=initial_memory, + turns=turns, + ) + + updated_memory = extract_runtime_memory_text( + response + ) + + if ( + is_runtime_memory_response_truncated( + response + ) + or looks_like_incomplete_runtime_memory( + updated_memory + ) ): await safe_call( getattr( @@ -363,7 +557,8 @@ async def summarize_runtime_memory( ), ) - return current_memory + return initial_memory + updates_counter = getattr( context, "runtime_memory_updates", @@ -372,8 +567,15 @@ async def summarize_runtime_memory( if updated_memory or updates_counter == 0: context.runtime_memory = updated_memory + context.runtime_memory_stable = updated_memory context.runtime_memory_updates = updates_counter + 1 + context.runtime_memory_pending_turns = [ + turn + for turn in context.runtime_memory_pending_turns + if turn not in turns + ] + logger = getattr( context, "logger", @@ -431,25 +633,53 @@ async def summarize_runtime_memory( "", ) + finally: + if ( + getattr( + context, + "runtime_memory_update_task", + None, + ) + is asyncio.current_task() + ): + context.runtime_memory_update_task = None + def schedule_runtime_memory_update( - *, - context, - user_message: str, - assistant_message: str, + *, + context, + user_message: str, + assistant_message: str, ) -> asyncio.Task | None: if not assistant_message.strip(): return None + context.runtime_memory_pending_turns.append({ + "user_message": user_message, + "assistant_message": assistant_message, + }) + + previous_task = getattr( + context, + "runtime_memory_update_task", + None, + ) + + if ( + previous_task is not None + and not previous_task.done() + ): + previous_task.cancel() + task = asyncio.create_task( - summarize_runtime_memory( + summarize_runtime_memory_pending_turns( context=context, - user_message=user_message, - assistant_message=assistant_message, ) ) + context.runtime_memory_update_task = task + background_tasks = getattr( context, "background_tasks", @@ -471,8 +701,8 @@ def schedule_runtime_memory_update( def schedule_interrupted_runtime_memory_update( - *, - context, + *, + context, ) -> asyncio.Task | None: user_message = getattr( @@ -500,3 +730,30 @@ def schedule_interrupted_runtime_memory_update( user_message=user_message, assistant_message=assistant_message, ) + + +async def cancel_runtime_memory_update( + context, +) -> None: + + task = getattr( + context, + "runtime_memory_update_task", + None, + ) + + if ( + task is None + or task.done() + ): + return + + task.cancel() + + with contextlib.suppress( + asyncio.CancelledError, + Exception, + ): + await task + + context.runtime_memory_update_task = None \ No newline at end of file diff --git a/runtime/runtime_context.py b/runtime/runtime_context.py index ae4bcca0..23900694 100644 --- a/runtime/runtime_context.py +++ b/runtime/runtime_context.py @@ -50,8 +50,16 @@ class RuntimeContext: runtime_memory: str = DEFAULT_RUNTIME_MEMORY + runtime_memory_stable: str = DEFAULT_RUNTIME_MEMORY + runtime_memory_updates: int = 0 + runtime_memory_pending_turns: list[dict] = field( + default_factory=list + ) + + runtime_memory_update_task: object | None = None + background_tasks: set = field( default_factory=set ) diff --git a/tests/test_message_memory.py b/tests/test_message_memory.py index 3fb0e09b..0e73546e 100644 --- a/tests/test_message_memory.py +++ b/tests/test_message_memory.py @@ -1,6 +1,5 @@ import unittest from types import SimpleNamespace -from unittest.mock import patch from clients.brain_client import ( build_brain_system_prompt, @@ -10,7 +9,6 @@ build_interrupted_assistant_message, build_runtime_memory_system_prompt, build_runtime_memory_user_prompt, - get_memory_request_timeout, schedule_interrupted_runtime_memory_update, schedule_runtime_memory_update, summarize_runtime_memory, @@ -139,23 +137,6 @@ def test_runtime_memory_user_prompt_uses_session_fallback(self): prompt, ) - def test_memory_request_timeout_has_background_default(self): - - with patch( - "memory.message_memory.config.SERVICE_REQUEST_TIMEOUT", - 30.0, - create=True, - ), patch( - "memory.message_memory.config.RUNTIME_MEMORY_REQUEST_TIMEOUT", - None, - create=True, - ): - - self.assertEqual( - get_memory_request_timeout(), - 120.0, - ) - def test_first_brain_prompt_includes_default_runtime_memory(self): context = RuntimeContext( @@ -362,7 +343,7 @@ async def emit(event): ], ) - async def test_summarizer_uses_service_max_tokens_and_memory_timeout(self): + async def test_summarizer_uses_service_max_tokens(self): service_client = FakeServiceClient( "- Active topic: available functions\n" @@ -400,10 +381,6 @@ async def test_summarizer_uses_service_max_tokens_and_memory_timeout(self): service_client.calls[0]["max_tokens"], config.SERVICE_MAX_TOKENS, ) - self.assertEqual( - service_client.calls[0]["timeout"], - get_memory_request_timeout(), - ) async def test_summarizer_skips_incomplete_memory(self): @@ -503,7 +480,10 @@ async def test_scheduled_update_is_background_task(self): }, logger=FakeLogger(), runtime_memory="Initial memory.", + runtime_memory_stable="Initial memory.", runtime_memory_updates=0, + runtime_memory_pending_turns=[], + runtime_memory_update_task=None, ) task = schedule_runtime_memory_update( @@ -547,7 +527,10 @@ async def test_interrupted_update_uses_partial_response(self): }, logger=FakeLogger(), runtime_memory="Initial memory.", + runtime_memory_stable="Initial memory.", runtime_memory_updates=0, + runtime_memory_pending_turns=[], + runtime_memory_update_task=None, runtime_turn_user_message="Tell me a story.", runtime_turn_assistant_response="Once upon a", ) diff --git a/utils/runtime_actions.py b/utils/runtime_actions.py index 7630cd18..84a6ab49 100644 --- a/utils/runtime_actions.py +++ b/utils/runtime_actions.py @@ -298,7 +298,7 @@ def extract_runtime_actions( _paired_action_open_pattern( action_name ) - + r"(\s*\{.*?\})\s*" + + r"(\s*\{.*?})\s*" + re.escape( close_marker ) diff --git a/websocket.py b/websocket.py index 6e48baba..92d04fc9 100644 --- a/websocket.py +++ b/websocket.py @@ -31,6 +31,7 @@ from memory.message_memory import ( schedule_interrupted_runtime_memory_update, schedule_runtime_memory_update, + cancel_runtime_memory_update, ) from utils.brain import ( @@ -453,6 +454,10 @@ async def websocket_endpoint( # START BACKGROUND TASK # ------------------------------------------------- + await cancel_runtime_memory_update( + context + ) + await refresh_pending_brain_usage( context, user_text,