diff --git a/open_notebook/graphs/ask.py b/open_notebook/graphs/ask.py index 2bdfc20..6c14e62 100644 --- a/open_notebook/graphs/ask.py +++ b/open_notebook/graphs/ask.py @@ -12,6 +12,7 @@ from typing_extensions import TypedDict from open_notebook.ai.provision import provision_langchain_model from open_notebook.domain.notebook import vector_search from open_notebook.utils import clean_thinking_content +from open_notebook.utils.text_utils import extract_text_content class SubGraphState(TypedDict): @@ -62,11 +63,7 @@ async def call_model_with_messages(state: ThreadState, config: RunnableConfig) - ai_message = await model.ainvoke(system_prompt) # Clean the thinking content from the response - message_content = ( - ai_message.content - if isinstance(ai_message.content, str) - else str(ai_message.content) - ) + message_content = extract_text_content(ai_message.content) cleaned_content = clean_thinking_content(message_content) # Parse the cleaned JSON content @@ -109,11 +106,7 @@ async def provide_answer(state: SubGraphState, config: RunnableConfig) -> dict: max_tokens=2000, ) ai_message = await model.ainvoke(system_prompt) - ai_content = ( - ai_message.content - if isinstance(ai_message.content, str) - else str(ai_message.content) - ) + ai_content = extract_text_content(ai_message.content) return {"answers": [clean_thinking_content(ai_content)]} @@ -126,11 +119,7 @@ async def write_final_answer(state: ThreadState, config: RunnableConfig) -> dict max_tokens=2000, ) ai_message = await model.ainvoke(system_prompt) - final_content = ( - ai_message.content - if isinstance(ai_message.content, str) - else str(ai_message.content) - ) + final_content = extract_text_content(ai_message.content) return {"final_answer": clean_thinking_content(final_content)} diff --git a/open_notebook/graphs/chat.py b/open_notebook/graphs/chat.py index 4c32570..070499f 100644 --- a/open_notebook/graphs/chat.py +++ b/open_notebook/graphs/chat.py @@ -14,6 +14,7 @@ from open_notebook.ai.provision import provision_langchain_model from open_notebook.config import LANGGRAPH_CHECKPOINT_FILE from open_notebook.domain.notebook import Notebook from open_notebook.utils import clean_thinking_content +from open_notebook.utils.text_utils import extract_text_content class ThreadState(TypedDict): @@ -69,11 +70,7 @@ def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict 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) - ) + content = extract_text_content(ai_message.content) cleaned_content = clean_thinking_content(content) cleaned_message = ai_message.model_copy(update={"content": cleaned_content}) diff --git a/open_notebook/graphs/prompt.py b/open_notebook/graphs/prompt.py index 00bc356..cd80d10 100644 --- a/open_notebook/graphs/prompt.py +++ b/open_notebook/graphs/prompt.py @@ -7,7 +7,7 @@ from langgraph.graph import END, START, StateGraph from typing_extensions import TypedDict from open_notebook.ai.provision import provision_langchain_model -from open_notebook.utils.text_utils import clean_thinking_content +from open_notebook.utils.text_utils import clean_thinking_content, extract_text_content class PatternChainState(TypedDict): @@ -33,7 +33,7 @@ async def call_model(state: dict, config: RunnableConfig) -> dict: response = await chain.ainvoke(payload) # Clean thinking tags from response (handles extended thinking models) - output = clean_thinking_content(str(response.content)) + output = clean_thinking_content(extract_text_content(response.content)) return {"output": output} diff --git a/open_notebook/utils/text_utils.py b/open_notebook/utils/text_utils.py index 3846924..ff7ea14 100644 --- a/open_notebook/utils/text_utils.py +++ b/open_notebook/utils/text_utils.py @@ -117,3 +117,29 @@ def clean_thinking_content(content: str) -> str: """ _, cleaned_content = parse_thinking_content(content) return cleaned_content + + +def extract_text_content(content) -> str: + """Extract text from LLM response content. + + Handles both plain string responses and structured content formats + (e.g. Gemini's envelope format): + [{'type': 'text', 'text': '...', 'extras': {...}}] + + Args: + content: The content from an AI message, either a string or a list of parts. + + Returns: + The extracted text content as a string. + """ + if isinstance(content, str): + return content + if isinstance(content, list): + text_parts = [] + for part in content: + if isinstance(part, dict) and "text" in part: + text_parts.append(part["text"]) + elif isinstance(part, str): + text_parts.append(part) + return "".join(text_parts) + return str(content)