admin-valentin/open-notebook
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: preserve AIMessage metadata when cleaning thinking content

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
Justin Florentine 2025-12-19 20:08:12 -05:00
parent 869664a10b
commit 855e730577
No known key found for this signature in database
3 changed files with 10 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
open_notebook/graphs/chat.py
View file

@ -72,7 +72,7 @@ def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict
# Clean thinking content from AI response (e.g., <think>...</think> tags) # Clean thinking content from AI response (e.g., <think>...</think> tags)
content = ai_message.content if isinstance(ai_message.content, str) else str(ai_message.content) content = ai_message.content if isinstance(ai_message.content, str) else str(ai_message.content)
cleaned_content = clean_thinking_content(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} return {"messages": cleaned_message}

2
open_notebook/graphs/source_chat.py
View file

@ -159,7 +159,7 @@ def call_model_with_source_context(
# Clean thinking content from AI response (e.g., <think>...</think> tags) # Clean thinking content from AI response (e.g., <think>...</think> tags)
content = ai_message.content if isinstance(ai_message.content, str) else str(ai_message.content) content = ai_message.content if isinstance(ai_message.content, str) else str(ai_message.content)
cleaned_content = clean_thinking_content(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 # Update state with context information
return { return {

8
tests/test_utils.py
View file

@ -94,6 +94,14 @@ class TestTextUtilities:
assert thinking == "" assert thinking == ""
assert cleaned == "Just regular content" assert cleaned == "Just regular content"
def test_parse_thinking_content_malformed_no_open_tag(self):
"""Test parsing malformed output where opening <think> tag is missing."""
content = "Some thinking content</think>Here 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): def test_parse_thinking_content_invalid_input(self):
"""Test parsing with invalid input types.""" """Test parsing with invalid input types."""
# Non-string input # Non-string input