Added math sub agent

config/prompts.yaml

@@ -2,8 +2,12 @@ prompts:
   base_system:
     content: |
       You are a general AI assistant tasked with answering complex questions.
-
       Make sure you think step by step in order to answer the given question.
+                
+      Here is the the question you received:
+      <question>
+      {{question}}
+      </question>
       
       {{attachment}}
 
@@ -12,8 +16,10 @@ prompts:
       {{summary}}
       </summary>
 
+       For mathematical questions or problems delegate them to the math_tool.   
+      
        Include citations for all the information you retrieve, ensuring you know exactly where the data comes from.
-       If you have the information inside your knowledge, still call a tool in order to confirm it.
+       If you have the information inside your knowledge, still call a tool in order to confirm it.   
 
        **Guidelines for Conducting Research:**
 
@@ -61,11 +67,10 @@ prompts:
       
        If the value of chunked_last_tool_call is true, this means that the last tool execution returns a result formed from the concatenation
        of multiple chunks.
-       Current value of the chunked_last_tool_call is {{chunked_last_tool_call}}
-      
+       Current value of the chunked_last_tool_call is {{chunked_last_tool_call}}      
        If you generate python code make sure you print the value of the variable you are interested in.
     type: base_system
-    variables: ["summary", "chunked_last_tool_call", "attachment"]
+    variables: ["summary", "chunked_last_tool_call", "attachment", "question"]
     version: 1.0
     description: "Core system prompt for all interactions"
   final_answer_processor:
@@ -100,4 +105,44 @@ prompts:
     type: tool
     variables: []
     version: 1.0
-    description: "Prompt for audio tool"
+    description: "Prompt for audio tool"
+  math_agent_base_system:
+    content: |
+      You are a research assistant conducting research on the user's input topic.
+  
+      <Task>
+      Your job is to use tools to gather information about the user's input topic.
+      You can use any of the tools provided to you to find resources that can help answer the research question. You can call these tools in series or in parallel, your research is conducted in a tool-calling loop.
+      </Task>
+  
+      <Available Tools>
+      You have access to two main tools:
+      1. **execute_python_code**: For executing python code
+      2. **think_tool**: For reflection and strategic planning during problem solving
+
+      **CRITICAL: Use think_tool after each search to reflect on results and plan next steps**
+      If you generate python code make sure you print the value of the variable you are interested in.
+      </Available Tools>
+  
+        
+      <Hard Limits>
+      **Tool Call Budgets** (Prevent excessive searching):
+             - **Simple queries**: Use 2-3 search tool calls maximum
+             - **Complex queries**: Use up to 5 search tool calls maximum
+             - **Always stop**: After 5 python tool calls if you cannot find the right sources
+    
+      **Stop Immediately When**:
+             - You can answer the user's question comprehensively
+      </Hard Limits>
+  
+      <Show Your Thinking>
+      After each search tool call, use think_tool to analyze the results:
+      - What key information did I find?
+      - What's missing?
+      - Do I have enough to answer the question comprehensively?
+      - Should I search more or provide my answer?
+      </Show Your Thinking>
+    type: sub_agent
+    variables: []
+    version: 1.0
+    description: "Core system prompt for the math agent"

core/state.py

@@ -1,4 +1,8 @@
 from langgraph.graph import MessagesState
+import operator
+from typing_extensions import TypedDict, Annotated, List, Sequence
+from langchain_core.messages import BaseMessage
+from langgraph.graph.message import add_messages
 
 
 class State(MessagesState):
@@ -6,3 +10,28 @@ class State(MessagesState):
     question: str
     chunked_last_tool_call: bool
     file_reference: str  # Attachment file reference: a path, URL, or unique ID
+
+
+class MathAgentState(MessagesState):
+    """
+    State for the math agent containing message history and research metadata.
+
+    This state tracks the agent's conversation, iteration count for limiting
+    tool calls, the research topic being investigated, compressed findings,
+    and raw research notes for detailed analysis.
+    """
+    tool_call_iterations: int
+    question: str
+    compressed_research: str
+    raw_notes: Annotated[List[str], operator.add]
+
+
+class MathAgentOutputState(MessagesState):
+    """
+    Output state for the math agent containing final results.
+
+    This represents the final output of the solving process with steps
+    and raw notes from the solving process.
+    """
+    compressed_research: str
+    raw_notes: Annotated[List[str], operator.add]

nodes/nodes.py

@@ -10,6 +10,7 @@ from core.state import State
 from nodes.chunking_node import OversizedContentHandler
 from tools.audio_tool import query_audio
 from tools.excel_tool import query_excel_file
+from tools.math_agent import math_tool
 from tools.python_executor import execute_python_code
 from tools.tavily_tools import llm_tools
 from utils.prompt_manager import prompt_mgmt
@@ -19,6 +20,7 @@ response_processing_model = ChatOpenAI(model="gpt-4.1-mini")
 llm_tools.append(query_audio)
 llm_tools.append(query_excel_file)
 llm_tools.append(execute_python_code)
+llm_tools.append(math_tool)
 model = model.bind_tools(llm_tools, parallel_tool_calls=False)
 
 
@@ -56,7 +58,7 @@ def assistant(state: State):
     if file_reference:
         attachment = f" you have access to the file with the following reference {file_reference}"
     prompt_params = {"summary": summary, "chunked_last_tool_call": state.get("chunked_last_tool_call", False),
-                     "attachment": attachment}
+                     "attachment": attachment, "question": question}
     sys_msg = SystemMessage(content=prompt_mgmt.render_template("base_system", prompt_params))
     try:
         response = model.invoke([sys_msg] + state["messages"])

tools/math_agent.py

@@ -0,0 +1,163 @@
+"""Research Agent Implementation.
+
+This module implements a research agent that can perform iterative web searches
+and synthesis to answer complex research questions.
+"""
+
+from typing import Literal
+
+from langchain.chat_models import init_chat_model
+from langchain.tools import tool
+from langchain_core.messages import SystemMessage, ToolMessage, filter_messages, HumanMessage
+from langgraph.graph import StateGraph, START, END
+
+from core.state import MathAgentState, MathAgentOutputState
+from tools.python_executor import execute_python_code
+from tools.think_tool import think_tool
+from utils.prompt_manager import prompt_mgmt
+
+# ===== CONFIGURATION =====
+
+# Set up tools and model binding
+tools = [execute_python_code, think_tool]
+tools_by_name = {tool.name: tool for tool in tools}
+
+# Initialize models
+math_model = init_chat_model(model="openai:gpt-5")
+model_with_tools = math_model.bind_tools(tools)
+summarization_model = init_chat_model(model="openai:gpt-4.1-mini")
+compress_model = init_chat_model(model="openai:gpt-4.1",
+                                 max_tokens=32000)
+
+
+# ===== AGENT NODES =====
+
+def llm_call(state: MathAgentState):
+    """Analyze current state and decide on next actions.
+
+    The model analyzes the current state and decides whether to:
+    1. Call search tools to gather more information
+    2. Provide a final answer based on gathered information
+
+    Returns updated state with the model's response.
+    """
+    return {
+        "messages": [
+            model_with_tools.invoke(
+                [SystemMessage(content=prompt_mgmt.render_template("math_agent_base_system", {}))] + state[
+                    "messages"]
+            )
+        ]
+    }
+
+
+def tool_node(state: MathAgentState):
+    """Execute all tool calls from the previous LLM response.
+
+    Executes all tool calls from the previous LLM responses.
+    Returns updated state with tool execution results.
+    """
+    tool_calls = state["messages"][-1].tool_calls
+
+    # Execute all tool calls
+    observations = []
+    for tool_call in tool_calls:
+        tool = tools_by_name[tool_call["name"]]
+        observations.append(tool.invoke(tool_call["args"]))
+
+    # Create tool message outputs
+    tool_outputs = [
+        ToolMessage(
+            content=observation,
+            name=tool_call["name"],
+            tool_call_id=tool_call["id"]
+        ) for observation, tool_call in zip(observations, tool_calls)
+    ]
+
+    return {"messages": tool_outputs}
+
+
+def compress_research(state: MathAgentState) -> dict:
+    """Compress research findings into a concise summary.
+
+    Takes all the research messages and tool outputs and creates
+    a compressed summary suitable for the supervisor's decision-making.
+    """
+
+    last_message = state.get("messages", [])[-1]
+
+    # Extract raw notes from tool and AI messages
+    raw_notes = [
+        str(m.content) for m in filter_messages(
+            state["messages"],
+            include_types=["tool", "ai"]
+        )
+    ]
+
+    return {
+        "compressed_research": str(last_message.content),
+        "raw_notes": ["\n".join(raw_notes)]
+    }
+
+
+# ===== ROUTING LOGIC =====
+
+def should_continue(state: MathAgentState) -> Literal["tool_node", "compress_research"]:
+    """Determine whether to continue research or provide final answer.
+
+    Determines whether the agent should continue the research loop or provide
+    a final answer based on whether the LLM made tool calls.
+
+    Returns:
+        "tool_node": Continue to tool execution
+        "compress_research": Stop and compress research
+    """
+    messages = state["messages"]
+    last_message = messages[-1]
+
+    # If the LLM makes a tool call, continue to tool execution
+    if last_message.tool_calls:
+        return "tool_node"
+    # Otherwise, we have a final answer
+    return "compress_research"
+
+
+# ===== GRAPH CONSTRUCTION =====
+
+# Build the agent workflow
+agent_builder = StateGraph(MathAgentState, output_schema=MathAgentOutputState)
+
+# Add nodes to the graph
+agent_builder.add_node("llm_call", llm_call)
+agent_builder.add_node("tool_node", tool_node)
+agent_builder.add_node("compress_research", compress_research)
+
+# Add edges to connect nodes
+agent_builder.add_edge(START, "llm_call")
+agent_builder.add_conditional_edges(
+    "llm_call",
+    should_continue,
+    {
+        "tool_node": "tool_node",  # Continue research loop
+        "compress_research": "compress_research",  # Provide final answer
+    },
+)
+agent_builder.add_edge("tool_node", "llm_call")  # Loop back for more research
+agent_builder.add_edge("compress_research", END)
+
+# Compile the agent
+math_agent = agent_builder.compile()
+
+
+@tool
+def math_tool(problem: str):
+    """
+    Tool for solving a mathematical problem
+    :param problem: The problem to be solved
+    :return: the solution to the given problem
+    """
+    response = math_agent.invoke({"messages": [HumanMessage(content=problem)], "question": problem})
+    # return the LLM response as a string (expected tool response format)
+    # this will be automatically turned to ToolMessage
+    # by the prebuilt create_react_agent (supervisor)
+    return response['compressed_research']

tools/think_tool.py

@@ -0,0 +1,29 @@
+from langchain.tools import tool
+
+
+@tool(parse_docstring=True)
+def think_tool(reflection: str) -> str:
+    """Tool for strategic reflection on research progress and decision-making.
+
+    Use this tool after each search to analyze results and plan next steps systematically.
+    This creates a deliberate pause in the research workflow for quality decision-making.
+
+    When to use:
+    - After receiving search results: What key information did I find?
+    - Before deciding next steps: Do I have enough to answer comprehensively?
+    - When assessing research gaps: What specific information am I still missing?
+    - Before concluding research: Can I provide a complete answer now?
+
+    Reflection should address:
+    1. Analysis of current findings - What concrete information have I gathered?
+    2. Gap assessment - What crucial information is still missing?
+    3. Quality evaluation - Do I have sufficient evidence/examples for a good answer?
+    4. Strategic decision - Should I continue searching or provide my answer?
+
+    Args:
+        reflection: Your detailed reflection on research progress, findings, gaps, and next steps
+
+    Returns:
+        Confirmation that reflection was recorded for decision-making
+    """
+    return f"Reflection recorded: {reflection}"

utils/prompt_manager.py

@@ -14,6 +14,7 @@ class PromptType(Enum):
     ANSWER_REFINEMENT = "answer_refinement"
     MEMORY_OPTIMIZATION = "memory_optimization"
     TOOL = "tool"
+    SUB_AGENT = "sub_agent"
 
 
 @dataclass