From 855e730577cb1a59256b4130802af129c9eabdda Mon Sep 17 00:00:00 2001 From: Justin Florentine Date: Fri, 19 Dec 2025 20:08:12 -0500 Subject: [PATCH] fix: preserve AIMessage metadata when cleaning thinking content MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Use model_copy() instead of creating new AIMessage to preserve response_metadata, id, usage_metadata, etc. Also adds test coverage for malformed thinking tags pattern. Addresses PR #333 feedback from lfnovo and cubic-dev-ai. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- open_notebook/graphs/chat.py | 2 +- open_notebook/graphs/source_chat.py | 2 +- tests/test_utils.py | 8 ++++++++ 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/open_notebook/graphs/chat.py b/open_notebook/graphs/chat.py index 810d29f..442396a 100644 --- a/open_notebook/graphs/chat.py +++ b/open_notebook/graphs/chat.py @@ -72,7 +72,7 @@ def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict # Clean thinking content from AI response (e.g., ... tags) content = ai_message.content if isinstance(ai_message.content, str) else str(ai_message.content) cleaned_content = clean_thinking_content(content) - cleaned_message = AIMessage(content=cleaned_content) + cleaned_message = ai_message.model_copy(update={"content": cleaned_content}) return {"messages": cleaned_message} diff --git a/open_notebook/graphs/source_chat.py b/open_notebook/graphs/source_chat.py index a173a56..99217a8 100644 --- a/open_notebook/graphs/source_chat.py +++ b/open_notebook/graphs/source_chat.py @@ -159,7 +159,7 @@ def call_model_with_source_context( # Clean thinking content from AI response (e.g., ... tags) content = ai_message.content if isinstance(ai_message.content, str) else str(ai_message.content) cleaned_content = clean_thinking_content(content) - cleaned_message = AIMessage(content=cleaned_content) + cleaned_message = ai_message.model_copy(update={"content": cleaned_content}) # Update state with context information return { diff --git a/tests/test_utils.py b/tests/test_utils.py index 6052724..65e12e1 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -94,6 +94,14 @@ class TestTextUtilities: assert thinking == "" assert cleaned == "Just regular content" + def test_parse_thinking_content_malformed_no_open_tag(self): + """Test parsing malformed output where opening tag is missing.""" + content = "Some thinking contentHere is my answer" + thinking, cleaned = parse_thinking_content(content) + + assert thinking == "Some thinking content" + assert cleaned == "Here is my answer" + def test_parse_thinking_content_invalid_input(self): """Test parsing with invalid input types.""" # Non-string input