From 869664a10b434e2a7e1d84adee24fe13b27996b4 Mon Sep 17 00:00:00 2001 From: Justin Florentine Date: Thu, 18 Dec 2025 16:31:23 -0500 Subject: [PATCH 1/2] fix: strip tags from chat responses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add thinking content cleaning to notebook and source chat graphs. Previously, models that output ... tags (like DeepSeek) or malformed variants without opening tags (like Nemotron) would leak reasoning content into user-visible responses. Changes: - chat.py: Clean AI response content before returning messages - source_chat.py: Same fix for source-specific chat - text_utils.py: Handle malformed output where opening tag is missing but is present 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- open_notebook/graphs/chat.py | 12 ++++++++-- open_notebook/graphs/source_chat.py | 11 +++++++-- open_notebook/utils/text_utils.py | 37 ++++++++++++++++++++--------- 3 files changed, 45 insertions(+), 15 deletions(-) diff --git a/open_notebook/graphs/chat.py b/open_notebook/graphs/chat.py index 52041ea..810d29f 100644 --- a/open_notebook/graphs/chat.py +++ b/open_notebook/graphs/chat.py @@ -3,8 +3,10 @@ import sqlite3 from typing import Annotated, Optional from ai_prompter import Prompter -from langchain_core.messages import SystemMessage +from langchain_core.messages import AIMessage, SystemMessage from langchain_core.runnables import RunnableConfig + +from open_notebook.utils import clean_thinking_content from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.graph import END, START, StateGraph from langgraph.graph.message import add_messages @@ -66,7 +68,13 @@ def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict ) ai_message = model.invoke(payload) - return {"messages": ai_message} + + # 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) + + return {"messages": cleaned_message} conn = sqlite3.connect( diff --git a/open_notebook/graphs/source_chat.py b/open_notebook/graphs/source_chat.py index 868164d..a173a56 100644 --- a/open_notebook/graphs/source_chat.py +++ b/open_notebook/graphs/source_chat.py @@ -3,8 +3,10 @@ import sqlite3 from typing import Annotated, Dict, List, Optional from ai_prompter import Prompter -from langchain_core.messages import SystemMessage +from langchain_core.messages import AIMessage, SystemMessage from langchain_core.runnables import RunnableConfig + +from open_notebook.utils import clean_thinking_content from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.graph import END, START, StateGraph from langgraph.graph.message import add_messages @@ -154,9 +156,14 @@ def call_model_with_source_context( ai_message = model.invoke(payload) + # 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) + # Update state with context information return { - "messages": ai_message, + "messages": cleaned_message, "source": source, "insights": insights, "context": formatted_context, diff --git a/open_notebook/utils/text_utils.py b/open_notebook/utils/text_utils.py index 0024ca3..b2a7720 100644 --- a/open_notebook/utils/text_utils.py +++ b/open_notebook/utils/text_utils.py @@ -11,8 +11,11 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter from .token_utils import token_count -# Pattern for matching thinking content in AI responses +# Patterns for matching thinking content in AI responses +# Standard pattern: ... THINK_PATTERN = re.compile(r"(.*?)", re.DOTALL) +# Pattern for malformed output: content (missing opening tag) +THINK_PATTERN_NO_OPEN = re.compile(r"^(.*?)", re.DOTALL) def split_text(txt: str, chunk_size=500): @@ -77,6 +80,9 @@ def parse_thinking_content(content: str) -> Tuple[str, str]: """ Parse message content to extract thinking content from tags. + Handles both well-formed tags and malformed output where the opening + tag is missing but is present. + Args: content (str): The original message content @@ -101,22 +107,31 @@ def parse_thinking_content(content: str) -> Tuple[str, str]: if len(content) > 100000: return "", content - # Find all thinking blocks + # Find all well-formed thinking blocks thinking_matches = THINK_PATTERN.findall(content) - if not thinking_matches: - return "", content + if thinking_matches: + # Join all thinking content with double newlines + thinking_content = "\n\n".join(match.strip() for match in thinking_matches) - # Join all thinking content with double newlines - thinking_content = "\n\n".join(match.strip() for match in thinking_matches) + # Remove all ... blocks from the original content + cleaned_content = THINK_PATTERN.sub("", content) - # Remove all ... blocks from the original content - cleaned_content = THINK_PATTERN.sub("", content) + # Clean up extra whitespace + cleaned_content = re.sub(r"\n\s*\n\s*\n", "\n\n", cleaned_content).strip() - # Clean up extra whitespace - cleaned_content = re.sub(r"\n\s*\n\s*\n", "\n\n", cleaned_content).strip() + return thinking_content, cleaned_content - return thinking_content, cleaned_content + # Handle malformed output: content (missing opening tag) + # Some models like Nemotron output thinking without the opening tag + malformed_match = THINK_PATTERN_NO_OPEN.match(content) + if malformed_match: + thinking_content = malformed_match.group(1).strip() + # Remove the thinking content and tag + cleaned_content = content[malformed_match.end():].strip() + return thinking_content, cleaned_content + + return "", content def clean_thinking_content(content: str) -> str: From 855e730577cb1a59256b4130802af129c9eabdda Mon Sep 17 00:00:00 2001 From: Justin Florentine Date: Fri, 19 Dec 2025 20:08:12 -0500 Subject: [PATCH 2/2] 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