Add reasoning to LangGraph AI assistant

app/agent/chat.py

@@ -1,19 +1,49 @@
 
 import os
 import logging
-from typing import Optional
+from typing import Optional, AsyncIterator
 from datetime import datetime, timezone, timedelta
 
-from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
 from langchain_core.runnables import RunnableConfig
 from langchain_tavily import TavilySearch
 from langchain_core.tools import tool
+from langchain_core.callbacks.base import AsyncCallbackHandler
 
 from app.agent.types import AgentState
 from app.agent.model import get_llm
 from app.mcp.manager import mcp
 import platform
 
+
+class ThinkingCallbackHandler(AsyncCallbackHandler):
+    """Callback handler to capture and emit thinking/reasoning blocks during streaming."""
+
+    def __init__(self):
+        self.thinking_content = []
+
+    async def on_llm_new_token(self, token: str, **kwargs) -> None:
+        """Called when a new token is generated."""
+        # Check if this token is part of a thinking block
+        # Claude models with extended_thinking will have thinking blocks
+        # in the format: <thinking>...</thinking>
+        pass
+
+    async def on_llm_start(self, serialized, prompts, **kwargs) -> None:
+        """Called when LLM starts."""
+        self.thinking_content = []
+
+    async def on_llm_end(self, response, **kwargs) -> None:
+        """Called when LLM ends."""
+        # Extract thinking blocks from response if present
+        if hasattr(response, 'generations') and response.generations:
+            for generation in response.generations:
+                for gen in generation:
+                    if hasattr(gen, 'message') and hasattr(gen.message, 'additional_kwargs'):
+                        thinking = gen.message.additional_kwargs.get('thinking')
+                        if thinking:
+                            self.thinking_content.append(thinking)
+
 @tool
 def get_current_datetime() -> str:
     """Get the current date and time."""
@@ -86,19 +116,39 @@ async def chat_node(state: AgentState, config: RunnableConfig):
 
         # Custom Assistant Instructions
         {assistant.get("instructions")}
-        
+
         Follow the custom instructions above while helping the user.
         """
     else:
         system_message = base_system_message
 
+    # Create callback handler for thinking
+    thinking_handler = ThinkingCallbackHandler()
+
     response = await llm_with_tools.ainvoke(
         [
             SystemMessage(content=system_message),
             *state["messages"]
         ],
         config=config,
     )
+
+    # Extract thinking blocks from Claude's extended thinking
+    thinking_blocks = []
+    if hasattr(response, 'content') and isinstance(response.content, list):
+        for content_block in response.content:
+            if isinstance(content_block, dict):
+                # Check for thinking block in Claude's response
+                if content_block.get('type') == 'thinking':
+                    thinking_blocks.append(content_block.get('thinking', ''))
+
+    # Store thinking blocks in the response for streaming
+    if thinking_blocks:
+        print(f"💭 Captured {len(thinking_blocks)} thinking blocks")
+        if not hasattr(response, 'additional_kwargs'):
+            response.additional_kwargs = {}
+        response.additional_kwargs['thinking_blocks'] = thinking_blocks
+
     print(response, "response in chat_node")
     return {
         **state,

app/agent/model.py

@@ -6,6 +6,7 @@
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_deepseek import ChatDeepSeek
 from langchain_openai import ChatOpenAI
+from langchain_anthropic import ChatAnthropic
 
 from app.agent.types import AgentState
 
@@ -26,6 +27,33 @@ def get_llm(state: AgentState) -> BaseChatModel:
 
     print(f"Model: {model_name}, Temperature: {temperature}, Max Tokens: {max_tokens}")
 
+    # Handle Claude/Anthropic models with extended thinking
+    if model_name.startswith("claude"):
+        api_key = os.environ.get("ANTHROPIC_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "ANTHROPIC_API_KEY environment variable is not set. "
+                "Please set it in your .env file or environment variables."
+            )
+
+        # Build model kwargs with extended thinking enabled
+        model_kwargs = {
+            "model": model_name,
+            "api_key": api_key,
+            "temperature": temperature,
+            "streaming": True,
+        }
+        if max_tokens is not None:
+            model_kwargs["max_tokens"] = max_tokens
+
+        # Enable extended thinking for supported models
+        # This shows the model's reasoning process before generating the final answer
+        if "sonnet" in model_name or "opus" in model_name:
+            model_kwargs["extended_thinking"] = True
+            print(f"✨ Extended thinking enabled for {model_name}")
+
+        return ChatAnthropic(**model_kwargs)
+
     # Handle OpenRouter models first (detected by :free suffix)
     if ":free" in model_name:
         api_key = os.environ.get("OPENROUTER_API_KEY")

app/agent/openai_reasoning_graph.py

@@ -0,0 +1,220 @@
+"""
+Ready-to-use LangGraph with reasoning for OpenAI models.
+This graph adds visible reasoning/thinking to your OpenAI-powered agent.
+
+Usage:
+    from app.agent.openai_reasoning_graph import openai_reasoning_graph
+
+    agent = LangGraphAGUIAgent(
+        name="mcpAssistant",
+        description="OpenAI Assistant with reasoning",
+        graph=openai_reasoning_graph
+    )
+"""
+
+from langgraph.graph import StateGraph, START, END
+from langgraph.checkpoint.memory import MemorySaver
+
+from app.agent.types import AgentState
+from app.agent.reasoning import reasoning_chat_node
+from app.agent.agent import async_tool_node, should_continue
+
+
+def create_openai_reasoning_graph():
+    """
+    Create a graph that adds reasoning capabilities to OpenAI models.
+
+    Flow:
+    START → chat (with reasoning) → [tools or END]
+                                ↓
+                             tools → back to chat
+
+    The chat node automatically:
+    1. Generates reasoning about the user's question
+    2. Uses that reasoning to create a better response
+    3. Returns both reasoning and response for frontend display
+    """
+    graph_builder = StateGraph(AgentState)
+
+    # Add nodes
+    # reasoning_chat_node combines reasoning + chat in one node
+    graph_builder.add_node("chat", reasoning_chat_node)
+    graph_builder.add_node("tools", async_tool_node)
+
+    # Add edges
+    graph_builder.add_edge(START, "chat")
+
+    # Conditional routing after chat
+    graph_builder.add_conditional_edges(
+        "chat",
+        should_continue,
+        {
+            "tools": "tools",
+            "end": END
+        }
+    )
+
+    # After tools, go back to chat
+    graph_builder.add_edge("tools", "chat")
+
+    # Compile with checkpointer for conversation memory
+    checkpointer = MemorySaver()
+    return graph_builder.compile(checkpointer=checkpointer)
+
+
+# Export the compiled graph
+openai_reasoning_graph = create_openai_reasoning_graph()
+
+
+# Alternative: Graph with explicit reasoning node (more control)
+def create_openai_reasoning_graph_explicit():
+    """
+    Create a graph with explicit reasoning node for more control.
+
+    Flow:
+    START → reasoning → chat → [tools or END]
+                               ↓
+                            tools → back to chat
+
+    This gives you more visibility and control over the reasoning step.
+    """
+    from app.agent.reasoning import reasoning_node
+
+    graph_builder = StateGraph(AgentState)
+
+    # Add nodes
+    graph_builder.add_node("reasoning", reasoning_node)  # Explicit reasoning
+    graph_builder.add_node("chat", reasoning_chat_node)  # Chat with reasoning context
+    graph_builder.add_node("tools", async_tool_node)
+
+    # Add edges
+    graph_builder.add_edge(START, "reasoning")  # Start with reasoning
+    graph_builder.add_edge("reasoning", "chat")  # Then chat
+
+    # Conditional routing after chat
+    graph_builder.add_conditional_edges(
+        "chat",
+        should_continue,
+        {
+            "tools": "tools",
+            "end": END
+        }
+    )
+
+    # After tools, skip reasoning and go directly to chat
+    # (reasoning only needed for initial user question)
+    graph_builder.add_edge("tools", "chat")
+
+    # Compile
+    checkpointer = MemorySaver()
+    return graph_builder.compile(checkpointer=checkpointer)
+
+
+# Export alternative graph
+openai_reasoning_graph_explicit = create_openai_reasoning_graph_explicit()
+
+
+# Conditional reasoning graph (toggle on/off)
+def create_openai_conditional_reasoning_graph():
+    """
+    Create a graph where reasoning can be toggled on/off via config.
+
+    Flow with reasoning enabled:
+    START → reasoning → chat → [tools or END]
+
+    Flow with reasoning disabled:
+    START → chat → [tools or END]
+
+    Toggle via assistant config:
+    {
+        "assistant": {
+            "config": {
+                "enable_reasoning": true  # or false
+            }
+        }
+    }
+    """
+    from app.agent.reasoning import reasoning_node
+    from app.agent.chat import chat_node
+
+    def should_use_reasoning(state: AgentState) -> str:
+        """Decide whether to use reasoning based on config."""
+        assistant_config = state.get("assistant", {}).get("config", {})
+        use_reasoning = assistant_config.get("enable_reasoning", False)
+
+        if use_reasoning:
+            return "reasoning"
+        else:
+            return "chat"
+
+    graph_builder = StateGraph(AgentState)
+
+    # Add nodes
+    graph_builder.add_node("reasoning", reasoning_node)
+    graph_builder.add_node("chat", reasoning_chat_node)
+    graph_builder.add_node("chat_no_reasoning", chat_node)  # Regular chat without reasoning
+    graph_builder.add_node("tools", async_tool_node)
+
+    # Conditional start - use reasoning or not
+    graph_builder.add_conditional_edges(
+        START,
+        should_use_reasoning,
+        {
+            "reasoning": "reasoning",
+            "chat": "chat_no_reasoning"
+        }
+    )
+
+    # After reasoning, go to chat
+    graph_builder.add_edge("reasoning", "chat")
+
+    # Conditional routing after both chat nodes
+    graph_builder.add_conditional_edges(
+        "chat",
+        should_continue,
+        {
+            "tools": "tools",
+            "end": END
+        }
+    )
+
+    graph_builder.add_conditional_edges(
+        "chat_no_reasoning",
+        should_continue,
+        {
+            "tools": "tools",
+            "end": END
+        }
+    )
+
+    # After tools, go back to appropriate chat node
+    def route_after_tools(state: AgentState) -> str:
+        assistant_config = state.get("assistant", {}).get("config", {})
+        if assistant_config.get("enable_reasoning", False):
+            return "chat"
+        return "chat_no_reasoning"
+
+    graph_builder.add_conditional_edges(
+        "tools",
+        route_after_tools,
+        {
+            "chat": "chat",
+            "chat_no_reasoning": "chat_no_reasoning"
+        }
+    )
+
+    # Compile
+    checkpointer = MemorySaver()
+    return graph_builder.compile(checkpointer=checkpointer)
+
+
+# Export conditional graph
+openai_conditional_reasoning_graph = create_openai_conditional_reasoning_graph()
+
+
+# Export all variants
+__all__ = [
+    "openai_reasoning_graph",  # Default - always use reasoning
+    "openai_reasoning_graph_explicit",  # Explicit reasoning node
+    "openai_conditional_reasoning_graph",  # Toggle reasoning on/off
+]