From a43160d983980fc1c4af553f198dd7438541cddd Mon Sep 17 00:00:00 2001 From: makeitdouble Date: Fri, 29 May 2026 15:36:40 +0300 Subject: [PATCH] chore: bugfixes - fix for handling incorrect runtime action rendering - fix for initial runtime state rendering - fix for summarizing system prompt - fix for runtime action tag rendering inside think block - add separate summarizing request timeout --- clients/brain_client.py | 13 +- clients/service_client.py | 18 +- config.example.py | 12 +- memory/message_memory.py | 80 ++++++-- runtime/runtime_client.py | 7 +- runtime/runtime_context.py | 8 +- templates/index.html | 2 +- tests/test_message_memory.py | 130 +++++++++++- tests/test_runtime_actions.py | 269 ++++++++++++++++++++++++ tests/test_search_flow.py | 19 ++ utils/runtime_actions.py | 370 +++++++++++++++++++++++++++++----- 11 files changed, 837 insertions(+), 91 deletions(-) diff --git a/clients/brain_client.py b/clients/brain_client.py index 931ca273..9d1688ed 100644 --- a/clients/brain_client.py +++ b/clients/brain_client.py @@ -839,7 +839,8 @@ async def ask_brain_stream( ) thinking_filter = RuntimeActionStreamFilter( - enabled_actions=enabled_actions + enabled_actions=enabled_actions, + preserve_action_text=True, ) content_filter = RuntimeActionStreamFilter( enabled_actions=enabled_actions @@ -902,6 +903,16 @@ async def filter_runtime_action_chunk( ) stop_for_runtime_action = True + + if ( + chunk_type == "thinking" + and result.text + ): + return { + **action_chunk, + "content": result.text, + } + return None if not result.text: diff --git a/clients/service_client.py b/clients/service_client.py index a917db1d..59896698 100644 --- a/clients/service_client.py +++ b/clients/service_client.py @@ -23,16 +23,24 @@ async def ask_service_model( system_prompt: str = "", temperature: float, max_tokens: int, + timeout: float | None = None, ): - return await client.ask( - system_prompt=( + request = { + "system_prompt": ( system_prompt or build_service_system_prompt() ), - user_prompt=user_prompt, - temperature=temperature, - max_tokens=max_tokens, + "user_prompt": user_prompt, + "temperature": temperature, + "max_tokens": max_tokens, + } + + if timeout is not None: + request["timeout"] = timeout + + return await client.ask( + **request ) diff --git a/config.example.py b/config.example.py index 3f305372..5fe4618e 100644 --- a/config.example.py +++ b/config.example.py @@ -15,7 +15,9 @@ BRAIN_REQUEST_TIMEOUT = 200.0 -BRAIN_CONTEXT_WINDOW = 32768 +BRAIN_CONTEXT_WINDOW = 8192 + +NIGHT_BRAIN_CONTEXT_WINDOW = 16384 BRAIN_TEMPERATURE = 0.7 @@ -31,9 +33,11 @@ SERVICE_REQUEST_TIMEOUT = 200.0 -SERVICE_CONTEXT_WINDOW = 8192 +RUNTIME_MEMORY_REQUEST_TIMEOUT = 120.0 + +SERVICE_CONTEXT_WINDOW = 4096 -SERVICE_TEMPERATURE = 0.15 +SERVICE_TEMPERATURE = 0.1 SERVICE_MAX_TOKENS = 1024 @@ -60,7 +64,7 @@ TRANSLATOR_REQUEST_TIMEOUT = 120 -TRANSLATOR_CONTEXT_WINDOW = 4096 +TRANSLATOR_CONTEXT_WINDOW = 2048 TRANSLATION_RETRIES = 1 diff --git a/memory/message_memory.py b/memory/message_memory.py index 426b020b..10547189 100644 --- a/memory/message_memory.py +++ b/memory/message_memory.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import traceback from clients.service_client import ( ask_service_model, @@ -11,6 +12,41 @@ ResponseExtractor, ) + +DEFAULT_RUNTIME_MEMORY = ( + "This session has just begun. " + "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, @@ -70,8 +106,8 @@ def build_runtime_memory_system_prompt() -> str: "Do not explain your reasoning or the summarization process.\n" "Write memory as atomic bullet lines, one semantic entity per line.\n" "Each line should start with a compact semantic label such as topic, " - "user intent, active topic, open reference, pending choice, offered options, " - "preference, decision, pattern, or interrupted response.\n" + "user intents, potential motifs, focus, priority, active topics, open references, pending choices, " + "offered options, preferences, expectations, current concern, decisions, patterns, failures or interruptions.\n" "Avoid writing about JIN's role unless the role itself changed or matters. " "Describe assistant actions neutrally instead.\n" "Keep memory actionable: write what helps the next answer, not a recap of " @@ -100,21 +136,20 @@ def build_runtime_memory_system_prompt() -> str: "Prefer compact continuity over transcript-like detail.\n" "Remove noise, implementation chatter, and one-off details unless they change " "what JIN should understand next.\n" - "Keep the memory under 1200 characters unless the latest turn is highly important.\n" - "The final memory should feel like live state, not chat history.\n" + "The final memory snapshot should feel like current live trusted state.\n" ) def build_runtime_memory_user_prompt( - *, - current_memory: str, - user_message: str, - assistant_message: str, + *, + current_memory: str, + user_message: str, + assistant_message: str, ) -> str: return ( "Current runtime memory:\n" - f"{current_memory.strip() or 'User and JIN just started interacting.'}\n\n" + f"{current_memory.strip() or DEFAULT_RUNTIME_MEMORY}\n\n" "Latest user message:\n" f"{user_message.strip()}\n\n" "Latest JIN answer:\n" @@ -250,6 +285,7 @@ async def ask_runtime_memory_model( max_tokens=( config.SERVICE_MAX_TOKENS ), + timeout=get_memory_request_timeout(), ) @@ -328,17 +364,15 @@ async def summarize_runtime_memory( ) return current_memory + updates_counter = getattr( + context, + "runtime_memory_updates", + 0, + ) - if updated_memory: + if updated_memory or updates_counter == 0: context.runtime_memory = updated_memory - context.runtime_memory_updates = ( - getattr( - context, - "runtime_memory_updates", - 0, - ) - + 1 - ) + context.runtime_memory_updates = updates_counter + 1 logger = getattr( context, @@ -369,7 +403,11 @@ async def summarize_runtime_memory( except asyncio.CancelledError: raise - except Exception as error: + except Exception: + formatted_traceback = ( + traceback.format_exc() + ) + logger = getattr( context, "logger", @@ -384,9 +422,7 @@ async def summarize_runtime_memory( await safe_call( log_error, "[MEMORY] runtime memory update failed", - details=str( - error - ), + details=formatted_traceback, ) return getattr( diff --git a/runtime/runtime_client.py b/runtime/runtime_client.py index a564891e..e1799833 100644 --- a/runtime/runtime_client.py +++ b/runtime/runtime_client.py @@ -83,6 +83,7 @@ async def ask( user_prompt: str, temperature: float, max_tokens: int, + timeout: float | None = None, ): payload = self.build_payload( @@ -99,7 +100,11 @@ async def ask( settings.CHAT_ENDPOINT, ), json=payload, - timeout=self.timeout, + timeout=( + self.timeout + if timeout is None + else timeout + ), ) response.raise_for_status() diff --git a/runtime/runtime_context.py b/runtime/runtime_context.py index 60c80bd6..ae4bcca0 100644 --- a/runtime/runtime_context.py +++ b/runtime/runtime_context.py @@ -6,6 +6,10 @@ RuntimeEmitter, ) +from memory.message_memory import ( + DEFAULT_RUNTIME_MEMORY, +) + @dataclass class RuntimeContext: @@ -44,9 +48,7 @@ class RuntimeContext: default_factory=list ) - runtime_memory: str = ( - "User and JIN just started interacting." - ) + runtime_memory: str = DEFAULT_RUNTIME_MEMORY runtime_memory_updates: int = 0 diff --git a/templates/index.html b/templates/index.html index 0ede7b1b..a06edd03 100644 --- a/templates/index.html +++ b/templates/index.html @@ -345,7 +345,7 @@ 

User and JIN just started interacting.
+ class="min-h-28 whitespace-pre-wrap break-words text-[11px] leading-relaxed text-zinc-100">This session has just begun. You have no history with the user yet. diff --git a/tests/test_message_memory.py b/tests/test_message_memory.py index d45571a9..3fb0e09b 100644 --- a/tests/test_message_memory.py +++ b/tests/test_message_memory.py @@ -1,16 +1,23 @@ import unittest from types import SimpleNamespace +from unittest.mock import patch from clients.brain_client import ( build_brain_system_prompt, ) from memory.message_memory import ( + DEFAULT_RUNTIME_MEMORY, 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, ) +from runtime.runtime_context import ( + RuntimeContext, +) from settings.config_loader import ( config, ) @@ -38,6 +45,7 @@ async def ask( user_prompt: str, temperature: float, max_tokens: int, + timeout: float | None = None, ): self.calls.append({ @@ -45,8 +53,15 @@ async def ask( "user_prompt": user_prompt, "temperature": temperature, "max_tokens": max_tokens, + "timeout": timeout, }) + if isinstance( + self.response_text, + Exception, + ): + raise self.response_text + if isinstance( self.response_text, list, @@ -111,6 +126,62 @@ class MessageMemoryTests( unittest.IsolatedAsyncioTestCase ): + def test_runtime_memory_user_prompt_uses_session_fallback(self): + + prompt = build_runtime_memory_user_prompt( + current_memory="", + user_message="hello", + assistant_message="hi", + ) + + self.assertIn( + DEFAULT_RUNTIME_MEMORY, + 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( + websocket=object(), + emitter=object(), + logger=object(), + clients={}, + ) + + prompt = build_brain_system_prompt( + context=context, + runtime_actions={ + "CAN_SEARCH": False, + "CAN_DEEP_THOUGHT": False, + }, + ) + + self.assertIn( + "", + prompt, + ) + self.assertIn( + DEFAULT_RUNTIME_MEMORY, + prompt, + ) + def test_runtime_memory_prompt_focuses_on_summary_depth(self): prompt = build_runtime_memory_system_prompt() @@ -144,7 +215,7 @@ def test_runtime_memory_prompt_focuses_on_summary_depth(self): prompt, ) self.assertIn( - "interrupted response", + "failures or interruptions", prompt, ) self.assertIn( @@ -235,7 +306,7 @@ async def test_summarizer_updates_runtime_memory(self): emit=None, ), logger=logger, - runtime_memory="User and JIN just started interacting.", + runtime_memory="", runtime_memory_updates=0, ) @@ -291,7 +362,7 @@ async def emit(event): ], ) - async def test_summarizer_uses_service_max_tokens(self): + async def test_summarizer_uses_service_max_tokens_and_memory_timeout(self): service_client = FakeServiceClient( "- Active topic: available functions\n" @@ -329,6 +400,10 @@ async def test_summarizer_uses_service_max_tokens(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): @@ -368,6 +443,55 @@ async def test_summarizer_skips_incomplete_memory(self): logger.errors ) + async def test_summarizer_failure_logs_traceback_details(self): + + class SilentError(Exception): + def __str__(self): + return "" + + logger = FakeLogger() + service_client = FakeServiceClient( + SilentError() + ) + context = SimpleNamespace( + clients={ + "service": service_client, + }, + logger=logger, + runtime_memory="Initial memory.", + runtime_memory_updates=0, + ) + + updated_memory = await summarize_runtime_memory( + context=context, + user_message="Remember this.", + assistant_message="I will remember it.", + ) + + self.assertEqual( + updated_memory, + "Initial memory.", + ) + self.assertEqual( + len(logger.errors), + 1, + ) + + message, details = logger.errors[0] + + self.assertEqual( + message, + "[MEMORY] runtime memory update failed", + ) + self.assertIn( + "Traceback (most recent call last):", + details, + ) + self.assertIn( + "SilentError", + details, + ) + async def test_scheduled_update_is_background_task(self): service_client = FakeServiceClient( diff --git a/tests/test_runtime_actions.py b/tests/test_runtime_actions.py index 5fb261e1..2f70cd97 100644 --- a/tests/test_runtime_actions.py +++ b/tests/test_runtime_actions.py @@ -140,6 +140,68 @@ def test_extracts_enabled_search_action(self): ), ) + def test_extracts_tool_call_prefixed_search_action(self): + + result = extract_runtime_actions( + ( + "before " + '<|tool_call>call:RUNTIME_ACTION:SEARCH>' + '{"query":"marijuana"}' + f"{SEARCH_ACTION_CLOSE} after" + ), + enabled_actions=[ + "CAN_SEARCH", + ], + ) + + self.assertEqual( + result.text, + "before after", + ) + + self.assertEqual( + result.count("SEARCH"), + 1, + ) + + self.assertEqual( + result.search_queries, + ( + "marijuana", + ), + ) + + def test_preserves_canonical_search_action_text_when_requested(self): + + result = extract_runtime_actions( + ( + "before " + '<|tool_call>call:RUNTIME_ACTION:SEARCH>' + '{"query":"marijuana"}' + f"{SEARCH_ACTION_CLOSE} after" + ), + enabled_actions=[ + "CAN_SEARCH", + ], + preserve_action_text=True, + ) + + self.assertEqual( + result.text, + ( + "before " + f'{SEARCH_ACTION_OPEN}{{"query":"marijuana"}}' + f"{SEARCH_ACTION_CLOSE} after" + ), + ) + + self.assertEqual( + result.search_queries, + ( + "marijuana", + ), + ) + def test_ignores_disabled_search_action(self): text = ( @@ -205,6 +267,213 @@ def test_stream_filter_handles_split_search_action(self): "", ) + def test_stream_filter_handles_split_tool_call_prefixed_search_action(self): + + stream_filter = RuntimeActionStreamFilter( + enabled_actions=[ + "CAN_SEARCH", + ], + ) + + first = stream_filter.filter( + ( + "before " + '<|tool_call>call:RUNTIME_ACTION:SEARCH>' + '{"query":"mari' + ) + ) + + second = stream_filter.filter( + f'juana"}}{SEARCH_ACTION_CLOSE} after' + ) + + self.assertEqual( + first.text, + "before ", + ) + + self.assertEqual( + first.count("SEARCH"), + 0, + ) + + self.assertEqual( + second.text, + " after", + ) + + self.assertEqual( + second.search_queries, + ( + "marijuana", + ), + ) + + self.assertEqual( + stream_filter.flush(), + "", + ) + + def test_stream_filter_preserves_complete_search_action_when_requested(self): + + stream_filter = RuntimeActionStreamFilter( + enabled_actions=[ + "CAN_SEARCH", + ], + preserve_action_text=True, + ) + + first = stream_filter.filter( + ( + "before " + '<|tool_call>call:RUNTIME_ACTION:SEARCH>' + '{"query":"mari' + ) + ) + + second = stream_filter.filter( + f'juana"}}{SEARCH_ACTION_CLOSE} after' + ) + + self.assertEqual( + first.text, + "before ", + ) + + self.assertEqual( + second.text, + ( + f'{SEARCH_ACTION_OPEN}{{"query":"marijuana"}}' + f"{SEARCH_ACTION_CLOSE} after" + ), + ) + + self.assertEqual( + second.search_queries, + ( + "marijuana", + ), + ) + + def test_stream_filter_preserves_mentioned_search_tag_without_stalling(self): + + stream_filter = RuntimeActionStreamFilter( + enabled_actions=[ + "CAN_SEARCH", + ], + preserve_action_text=True, + ) + + first = stream_filter.filter( + "like `" + ) + + second = stream_filter.filter( + "`) is just a tag name" + ) + + self.assertEqual( + first.text, + "like `", + ) + + self.assertEqual( + first.actions, + (), + ) + + self.assertEqual( + second.text, + "`) is just a tag name", + ) + + self.assertEqual( + second.actions, + (), + ) + + def test_stream_filter_parses_search_when_payload_follows_preserved_open_tag(self): + + stream_filter = RuntimeActionStreamFilter( + enabled_actions=[ + "CAN_SEARCH", + ], + preserve_action_text=True, + ) + + first = stream_filter.filter( + f"before {SEARCH_ACTION_OPEN}" + ) + + second = stream_filter.filter( + f'{{"query":"marijuana"}}{SEARCH_ACTION_CLOSE} after' + ) + + self.assertEqual( + first.text, + f"before {SEARCH_ACTION_OPEN}", + ) + + self.assertEqual( + first.actions, + (), + ) + + self.assertEqual( + second.text, + f'{{"query":"marijuana"}}{SEARCH_ACTION_CLOSE} after', + ) + + self.assertEqual( + second.search_queries, + ( + "marijuana", + ), + ) + + def test_removes_stray_tool_call_marker(self): + + result = extract_runtime_actions( + "before <|tool_call> after" + ) + + self.assertEqual( + result.text, + "before after", + ) + + self.assertEqual( + result.actions, + (), + ) + + def test_stream_filter_removes_split_tool_call_marker(self): + + stream_filter = RuntimeActionStreamFilter() + + first = stream_filter.filter( + "before <|tool" + ) + + second = stream_filter.filter( + "_call> after" + ) + + self.assertEqual( + first.text, + "before ", + ) + + self.assertEqual( + second.text, + " after", + ) + + self.assertEqual( + stream_filter.flush(), + "", + ) + def test_apply_runtime_action_calls_stores_search_queries(self): class Context: diff --git a/tests/test_search_flow.py b/tests/test_search_flow.py index 3ed9c490..ef73b7a0 100644 --- a/tests/test_search_flow.py +++ b/tests/test_search_flow.py @@ -556,6 +556,16 @@ async def test_search_is_model_driven_for_explicit_user_request(self): ).messages if message.get("type") == "runtime_action" ] + thinking_chunks = [ + message.get( + "chunk", + "", + ) + for message in get_fake_websocket( + context + ).messages + if message.get("type") == "thinking_chunk" + ] self.assertEqual( len(runtime_events), @@ -583,6 +593,15 @@ async def test_search_is_model_driven_for_explicit_user_request(self): "tesla car price", ], ) + self.assertIn( + ( + f'{SEARCH_ACTION_OPEN}{{"query":"tesla car price"}}' + f"{SEARCH_ACTION_CLOSE}" + ), + "".join( + thinking_chunks + ), + ) self.assertIn( "action: search", "\n".join( diff --git a/utils/runtime_actions.py b/utils/runtime_actions.py index 4678efb6..7630cd18 100644 --- a/utils/runtime_actions.py +++ b/utils/runtime_actions.py @@ -31,6 +31,8 @@ ) ) +TOOL_CALL_MARKER = "<|tool_call>" + def build_runtime_action_id( action_name: str, @@ -226,6 +228,7 @@ def extract_search_query( def extract_runtime_actions( text: str, enabled_actions=None, + preserve_action_text: bool = False, ) -> RuntimeActionResult: if not text: @@ -245,16 +248,33 @@ def extract_runtime_actions( if action_name not in enabled_action_names: continue - call_count = clean_text.count( - marker + pattern = re.compile( + _self_closing_action_open_pattern( + action_name + ), + re.DOTALL, + ) + + matches = tuple( + pattern.finditer( + clean_text + ) + ) + + call_count = len( + matches ) if not call_count: continue - clean_text = clean_text.replace( - marker, - "", + clean_text = pattern.sub( + ( + marker + if preserve_action_text + else "" + ), + clean_text, ) actions.extend( @@ -275,10 +295,10 @@ def extract_runtime_actions( pattern = re.compile( ( - re.escape( - open_marker + _paired_action_open_pattern( + action_name ) - + r"(.*?)" + + r"(\s*\{.*?\})\s*" + re.escape( close_marker ) @@ -288,16 +308,25 @@ def extract_runtime_actions( def replace_action(match): + payload = ( + match.group(1) + .strip() + ) + actions.append( RuntimeActionCall( name=action_name, - payload=( - match.group(1) - .strip() - ), + payload=payload, ) ) + if preserve_action_text: + return ( + open_marker + + payload + + close_marker + ) + return "" clean_text = pattern.sub( @@ -305,6 +334,10 @@ def replace_action(match): clean_text, ) + clean_text = _strip_tool_call_markers( + clean_text + ) + return RuntimeActionResult( text=clean_text, actions=tuple( @@ -313,6 +346,125 @@ def replace_action(match): ) +def _runtime_action_tag_name( + action_name: str, +) -> str: + + return ( + "RUNTIME_ACTION:" + + normalize_runtime_action_name( + action_name + ) + ) + + +def _paired_action_open_pattern( + action_name: str, +) -> str: + + return ( + r"<(?!/)[^<]*?" + + re.escape( + _runtime_action_tag_name( + action_name + ) + ) + + r"\s*>" + ) + + +def _self_closing_action_open_pattern( + action_name: str, +) -> str: + + return ( + r"<(?!/)[^<]*?" + + re.escape( + _runtime_action_tag_name( + action_name + ) + ) + + r"\s*/?>" + ) + + +def _tool_call_action_open_marker( + action_name: str, +) -> str: + + return ( + TOOL_CALL_MARKER + + "call:" + + _runtime_action_tag_name( + action_name + ) + + ">" + ) + + +def _runtime_action_open_marker( + action_name: str, +) -> str: + + return ( + "<" + + _runtime_action_tag_name( + action_name + ) + + ">" + ) + + +def _strip_tool_call_markers( + text: str, +) -> str: + + return text.replace( + TOOL_CALL_MARKER, + "", + ) + + +def _complete_paired_action_open_at_end( + text: str, + enabled_actions=None, +) -> tuple[int, str] | None: + + enabled_action_names = normalize_runtime_action_names( + enabled_actions + ) + + for action_name, markers in PAIRED_ACTION_MARKERS.items(): + + if action_name not in enabled_action_names: + continue + + open_marker, _ = markers + + pattern = re.compile( + _paired_action_open_pattern( + action_name + ), + re.DOTALL, + ) + + for match in pattern.finditer( + text + ): + + if match.end() != len( + text + ): + continue + + return ( + match.start(), + open_marker, + ) + + return None + + def _enabled_action_start_markers( enabled_actions=None, ) -> tuple[str, ...]: @@ -326,9 +478,19 @@ def _enabled_action_start_markers( for action_name, marker in SELF_CLOSING_ACTION_MARKERS.items(): if action_name in enabled_action_names: + markers.append( + _runtime_action_open_marker( + action_name + ) + ) markers.append( marker ) + markers.append( + _tool_call_action_open_marker( + action_name + ) + ) for action_name, paired_markers in PAIRED_ACTION_MARKERS.items(): @@ -336,6 +498,15 @@ def _enabled_action_start_markers( markers.append( paired_markers[0] ) + markers.append( + _tool_call_action_open_marker( + action_name + ) + ) + + markers.append( + TOOL_CALL_MARKER + ) return tuple( markers @@ -352,11 +523,14 @@ def _trailing_marker_prefix_length( ) max_length = min( - [len(text)] - + [ - len(marker) - 1 - for marker in markers - ] + len(text), + max( + [ + len(marker) - 1 + for marker in markers + ], + default=0, + ), ) for length in range( @@ -391,25 +565,45 @@ def _unclosed_paired_action_start( if action_name not in enabled_action_names: continue - open_marker, close_marker = markers + _, close_marker = markers - start = text.rfind( - open_marker + pattern = re.compile( + _paired_action_open_pattern( + action_name + ), + re.DOTALL, ) - if start < 0: - continue + for match in pattern.finditer( + text + ): - close = text.find( - close_marker, - start + len(open_marker), - ) + start = match.start() - if close >= 0: - continue + close = text.find( + close_marker, + match.end(), + ) - if start > latest_start: - latest_start = start + if close >= 0: + continue + + tail = text[ + match.end(): + ] + + stripped_tail = tail.lstrip() + + if ( + stripped_tail + and not stripped_tail.startswith( + "{" + ) + ): + continue + + if start > latest_start: + latest_start = start if latest_start < 0: return None @@ -422,12 +616,46 @@ class RuntimeActionStreamFilter: def __init__( self, enabled_actions=None, + preserve_action_text: bool = False, ): self.pending = "" + self.pending_emitted_text = "" + self.preserve_action_text = preserve_action_text self.enabled_actions = normalize_runtime_action_names( enabled_actions ) + def _finalize_result( + self, + result: RuntimeActionResult, + fallback_text: str, + ) -> RuntimeActionResult: + + emitted_text = self.pending_emitted_text + self.pending_emitted_text = "" + + if not emitted_text: + return result + + if result.text.startswith( + emitted_text + ): + return RuntimeActionResult( + text=result.text[ + len(emitted_text): + ], + actions=result.actions, + ) + + if result.actions: + return result + + return extract_runtime_actions( + fallback_text, + enabled_actions=self.enabled_actions, + preserve_action_text=self.preserve_action_text, + ) + def filter( self, chunk: str, @@ -445,53 +673,93 @@ def filter( self.pending = "" - result = extract_runtime_actions( - combined, - enabled_actions=self.enabled_actions, - ) + if self.preserve_action_text: + + open_at_end = _complete_paired_action_open_at_end( + combined, + enabled_actions=self.enabled_actions, + ) + + if open_at_end is not None: + + open_start, open_text = open_at_end + self.pending = combined[ + open_start: + ] + self.pending_emitted_text = open_text + + prefix_result = extract_runtime_actions( + combined[ + :open_start + ], + enabled_actions=self.enabled_actions, + preserve_action_text=self.preserve_action_text, + ) + + return RuntimeActionResult( + text=( + prefix_result.text + + open_text + ), + actions=prefix_result.actions, + ) unclosed_start = _unclosed_paired_action_start( - result.text, + combined, enabled_actions=self.enabled_actions, ) if unclosed_start is not None: - self.pending = result.text[ + self.pending = combined[ unclosed_start: ] - return RuntimeActionResult( - text=result.text[ + return extract_runtime_actions( + combined[ :unclosed_start ], - actions=result.actions, + enabled_actions=self.enabled_actions, + preserve_action_text=self.preserve_action_text, ) hold_length = ( _trailing_marker_prefix_length( - result.text, + combined, enabled_actions=self.enabled_actions, ) ) - if not hold_length: - return result + if hold_length: + + self.pending = combined[ + -hold_length: + ] - self.pending = result.text[ - -hold_length: - ] + return extract_runtime_actions( + combined[ + :-hold_length + ], + enabled_actions=self.enabled_actions, + preserve_action_text=self.preserve_action_text, + ) - return RuntimeActionResult( - text=result.text[ - :-hold_length - ], - actions=result.actions, + result = extract_runtime_actions( + combined, + enabled_actions=self.enabled_actions, + preserve_action_text=self.preserve_action_text, + ) + return self._finalize_result( + result, + chunk, ) def flush(self) -> str: pending = self.pending self.pending = "" + self.pending_emitted_text = "" - return pending + return _strip_tool_call_markers( + pending + )