Initial commit

.env.example

@@ -0,0 +1,4 @@
+OPENAI_API_KEY=sk-yourkeyhere
+TAVILY_API_KEY=tvly-yourkeyhere
+GENAI_MODEL=gemma-3-27b-it
+GOOGLE_API_KEY="your-google-api-key"

.gitattributes

@@ -0,0 +1 @@
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,70 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+*.egg-info/
+
+# Virtual environments
+venv/
+env/
+.venv/
+
+# IDEs
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# OS
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db
+
+# Logs
+*.log
+logs/
+
+# Environment variables
+.env
+.env.local
+.env.*.local
+
+# Dependencies
+.pipenv/
+Pipfile.lock
+
+# Data and large files (your datasets)
+data/*.xlsx
+data/*.pdf
+data/*.csv  # Если большие; иначе удали
+data/*.docx
+output/     # Generated outputs
+
+# Notebooks (if not committing)
+*.ipynb_checkpoints/
+test_*.ipynb  # Тестовые notebooks
+
+# Build/dist
+build/
+dist/
+*.egg
+
+# Testing
+.coverage
+htmlcov/
+.tox/
+.cache/
+
+# Mac
+._*
+
+# Windows
+Thumbs.db
+ehthumbs.db

src/__init__.py

@@ -0,0 +1,19 @@
+"""ANKELODON: Core AI Agent Package.
+
+Import key components for easy use:
+from src import workflow, llm
+"""
+
+from .config import llm, TOOLS, CONFIG, TOOL_NODE, planner_llm
+from .agent import workflow, build_workflow, should_continue
+from .nodes import agent, planner, query_input, critique
+from .schemas import AgentState, PlannerPlan, ComplexityLevel, CritiqueFeedback
+
+__version__ = "0.1.0"
+__all__ = [
+    "llm", "TOOLS", "CONFIG", "TOOL_NODE", "planner_llm",
+    "workflow", "build_workflow", "should_continue",
+    "agent", "planner", "query_input", "critique",
+    "AgentState", "PlannerPlan", "ComplexityLevel", "CritiqueFeedback",
+    "__version__"
+]

src/agent.py

@@ -0,0 +1,51 @@
+#GRAPH BUILDING
+from nodes import (query_input, complexity_assessor, planner, agent, simple_executor, critic_evaluator, replanner, enhanced_finalizer)
+from state import AgentState
+from langgraph.graph import StateGraph, END
+from nodes import should_continue, should_use_planning, should_replan
+from langgraph.checkpoint.memory import MemorySaver
+from config import DEBUGGING_TOOL_NODE
+
+def build_workflow(checkpointer=None) -> StateGraph[AgentState]:
+    builder = StateGraph(AgentState)
+    builder.add_node("INPUT", query_input)
+    builder.add_node("COMPLEXITY_ASSESSOR", complexity_assessor)
+    builder.add_node("PLANNING", planner)
+    builder.add_node("AGENT", agent)
+    builder.add_node("TOOLS", DEBUGGING_TOOL_NODE)
+    builder.add_node("FINALIZER", enhanced_finalizer)
+    builder.add_node("SIMPLE_EXECUTOR", simple_executor)
+    builder.add_node("CRITIC", critic_evaluator)
+    builder.add_node("REPLANNER", replanner)
+
+    builder.set_entry_point("INPUT")
+    builder.add_edge("INPUT", "COMPLEXITY_ASSESSOR")
+
+    builder.add_conditional_edges(
+            "COMPLEXITY_ASSESSOR",
+            should_use_planning,
+            {"simple_executor": "SIMPLE_EXECUTOR", "planner": "PLANNING"},
+        )
+    builder.add_edge("SIMPLE_EXECUTOR", "FINALIZER")
+
+
+    builder.add_edge("PLANNING", "AGENT")
+    builder.add_conditional_edges(
+            "AGENT",
+            should_continue,
+            {"tools": "TOOLS", "agent": "AGENT", "final_answer": "FINALIZER"},
+        )
+    builder.add_edge("TOOLS", "AGENT")
+    builder.add_edge("FINALIZER", "CRITIC")
+    builder.add_conditional_edges(
+            "CRITIC",
+            should_replan,
+            {"end": END, "replan": "REPLANNER"},
+        )
+    builder.add_edge("REPLANNER", "AGENT")
+
+    if checkpointer:
+        system = builder.compile(checkpointer=checkpointer)
+    else:
+        system = builder.compile()
+    return system

src/config.py

@@ -0,0 +1,35 @@
+from langchain_openai import ChatOpenAI
+from tools.tools import *
+from tools.code_interpreter import safe_code_run
+from langgraph.prebuilt import ToolNode
+from schemas import PlannerPlan
+
+config = {"configurable": {"thread_id": "1"}, "recursion_limit" : 50}
+
+TOOLS = [download_file_from_url, web_search, 
+         arxiv_search, wiki_search, add, subtract, multiply, divide, 
+         power, analyze_excel_file, analyze_csv_file, analyze_docx_file, 
+         analyze_pdf_file, analyze_txt_file, analyze_image_file, 
+         vision_qa_gemma, safe_code_run]
+
+class DebuggingToolNode(ToolNode):
+    def __init__(self, tools):
+        super().__init__(tools)
+    
+    def __call__(self, state):
+        print("=== TOOL EXECUTION STARTED ===")
+        result = super().__call__(state)
+        print("=== TOOL EXECUTION COMPLETED ===")
+        return result
+
+
+TOOL_NODE = ToolNode(TOOLS)
+DEBUGGING_TOOL_NODE = DebuggingToolNode(TOOLS)
+
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.25)
+llm_with_tools = llm.bind_tools(TOOLS)
+planner_llm = llm.with_structured_output(PlannerPlan)
+
+
+
+

src/nodes.py

@@ -0,0 +1,461 @@
+import os
+from state import AgentState
+from tools.tools import preprocess_files
+from langgraph.prebuilt import ToolNode
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+from prompts.prompts import SYSTEM_PROMPT_PLANNER, SYSTEM_EXECUTOR_PROMPT, COMPLEXITY_ASSESSOR_PROMPT, CRITIC_PROMPT
+from config import llm, TOOLS, planner_llm, llm_with_tools
+from schemas import PlannerPlan, ComplexityLevel, CritiqueFeedback, ExecutionReport, ToolExecution
+from utils.utils import format_final_answer, clean_message_history
+
+def query_input(state : AgentState) -> AgentState:
+    print("=== USER QUERY TRANSFERED TO AGENT ===")
+
+    files = state.get("files", [])
+    if files:
+        print(f"Processing {len(files)} files:")
+        file_info = preprocess_files(files)
+    
+        for file_path, info in file_info.items():
+            print(f"  - {file_path}: {info['type']} ({info['size']} bytes) -> {info['suggested_tool']}")
+
+        state["file_contents"] = file_info
+        file_context = "\n\n=== AVAILABLE FILES FOR ANALYSIS ===\n"
+        for file_path, info in file_info.items():
+            filename = os.path.basename(file_path)
+            file_context += f"File: {filename}\n"
+            file_context += f"  - Type: {info['type']}\n"  
+            file_context += f"  - Size: {info['size']} bytes\n"
+            file_context += f"  - Suggested tool: {info['suggested_tool']}\n"
+            if info.get("preview"):
+                file_context += f"  - Preview: {info['preview']}\n"
+            file_context += "\n"
+        
+        # Добавляем инструкции по работе с файлами
+        file_context += "IMPORTANT: Use the suggested tools to analyze these files before processing their data.\n"
+        file_context += "File paths are available in the agent state and can be passed directly to analysis tools.\n"
+        
+        original_query = state.get("query", "")
+        state["query"] = original_query + file_context
+    return state
+
+
+def planner(state : AgentState) -> AgentState:
+    sys_stack = [
+            SystemMessage(content=SYSTEM_PROMPT_PLANNER.strip()),
+            HumanMessage(content=state["query"]),
+        ]
+    plan: PlannerPlan = planner_llm.invoke(sys_stack)
+    
+    print("=== GENERATED PLAN ===")
+    return {"messages" : sys_stack + state["messages"],
+            "plan": plan,
+            "current_step ": 0,
+            "reasoning_done": False}
+
+
+def agent(state: AgentState) -> AgentState:
+    
+    """
+    sys_msg = SystemMessage(
+        content=SYSTEM_EXECUTOR_PROMPT.strip().format(
+            plan=json.dumps(state["plan"], indent=2)
+        )
+    )
+    """
+    current_step = state.get("current_step", 0)
+    reasoning_done = state.get("reasoning_done", False)
+    plan = state.get("plan", {})
+    steps = state["plan"].steps
+
+    print(f"=== AGENT DEBUG ===")
+    print(f"Current step: {current_step}")
+    print(f"Reasoning done: {reasoning_done}")
+    print(f"Plan exists: {plan is not None}")
+    print(f"Total steps in plan: {len(plan.steps) if plan else 'No plan'}")
+
+    if not plan or not hasattr(plan, 'steps') or not plan.steps:
+        print("ERROR: No valid plan found!")
+        return {
+            "messages": state["messages"] + [AIMessage(content="No valid plan available. <FINAL_ANSWER>")],
+            "reasoning_done": False
+        }
+    
+    steps = plan.steps
+    
+    if current_step >= len(steps):
+        print("All plan steps completed, moving to finalization")
+        return {
+            "messages": state["messages"] + [AIMessage(content="All steps completed. <FINAL_ANSWER>")],
+            "reasoning_done": False
+        }
+
+    current_step_info = steps[current_step]
+    print(f"Executing step {current_step + 1}: {current_step_info.description}")
+
+    if not reasoning_done:
+
+        # ✅ ДОБАВЛЕНО: Специальный контекст для файлов
+        file_context = ""
+        file_contents = state.get("file_contents", {})
+        if file_contents:
+            file_context = "\n\nAVAILABLE FILES IN CURRENT SESSION:\n"
+            for filepath, info in file_contents.items():
+                filename = os.path.basename(filepath)
+                file_context += f"- {filename}: {info['type']} file, suggested tool: {info['suggested_tool']}\n"
+                file_context += f"  Path: {filepath}\n"
+
+        reasoning_prompt = f"""
+        {SYSTEM_EXECUTOR_PROMPT}
+        
+        CURRENT TASK: You must perform reasoning for step {current_step + 1}.
+        
+        STEP INFO: {current_step_info}\n\n
+
+        FILE CONTEXT: {file_contents}
+        
+        CRITICAL: You MUST output your reasoning in <REASONING> tags, but DO NOT call any tools yet.
+        Explain what you need to do and why, then end your response.
+
+        REASONING IS IMPERATIVE BEFORE ANY TOOL CALLS.
+        """
+
+        sys_msg = SystemMessage(content = reasoning_prompt)
+        stack = [sys_msg] + state["messages"]
+
+        step = llm.invoke(stack)
+        print("=== REASONING STEP ===")
+        print(step.content)
+
+        return {
+            "messages" : state["messages"] + [step],
+            "reasoning_done" : True
+        }
+    
+    else:
+        tool_prompt = f"""
+        Now execute the tool for step {current_step + 1}.
+        
+        You have already done the reasoning. Now call the appropriate tool with the correct parameters.
+        Available file paths: {list(state.get("file_contents", {}).keys())}\n
+        IMPORTANT NOTE: IF YOU DECIDED TO USE safe_code_run, MAKE SURE TO FINISH CALCULATIONS WITH print() or saving to a variable NAMED 'result' so that the output can be captured!
+        AVAILABLE TOOLS: {', '.join([tool.name for tool in TOOLS])}
+        """ 
+
+        sys_msg = SystemMessage(content=tool_prompt)
+        stack = [sys_msg] + state["messages"]  # Берем последние сообщения включая reasoning
+        
+        # Используем модель С инструментами для выполнения
+        step = llm_with_tools.invoke(stack)
+        print("=== TOOL EXECUTION ===")
+        print(f"Tool calls: {step.tool_calls}")
+        
+        return {
+            "messages": state["messages"] + [step],
+            "current_step": current_step + 1 if step.tool_calls else current_step,
+            "reasoning_done": False  # Сбрасываем для следующего шага
+        }
+    
+def should_continue(state : AgentState) -> bool:
+    
+    last_message = state["messages"][-1]
+    reasoning_done = state.get("reasoning_done", False)
+    plan = state.get("plan", None)
+    current_step = state.get("current_step", 0)
+
+    if plan and current_step >= len(plan.steps):
+        return "final_answer"
+
+
+    if hasattr(last_message, "content") and "<FINAL_ANSWER>" in last_message.content:
+        return "final_answer"
+    elif hasattr(last_message, "tool_calls") and last_message.tool_calls:
+        return "tools" 
+    elif not reasoning_done and hasattr(last_message, 'content') and "<REASONING>" in last_message.content:
+        # Reasoning выполнен, но инструменты еще не вызваны
+        return "agent"
+    elif reasoning_done:
+        # Reasoning выполнен, теперь нужно вызвать инструменты
+        return "agent"
+    else:
+        # Нужно сделать reasoning
+        return "agent"
+
+# 6. Добавить отладочную информацию в TOOL_NODE
+class DebuggingToolNode(ToolNode):
+    def __init__(self, tools):
+        super().__init__(tools)
+    
+    def __call__(self, state):
+        print("=== TOOL EXECUTION STARTED ===")
+        result = super().__call__(state)
+        print("=== TOOL EXECUTION COMPLETED ===")
+        return result
+    
+
+
+def enhanced_finalizer(state: AgentState) -> AgentState:
+    """Generate comprehensive execution report for critic evaluation."""
+    print("=== GENERATING EXECUTION REPORT ===")
+    
+    # Extract tool execution information
+    tools_executed = []
+    data_sources = []
+    
+    for msg in state["messages"]:
+        if hasattr(msg, 'tool_calls') and msg.tool_calls:
+            for tool_call in msg.tool_calls:
+                tools_executed.append(ToolExecution(
+                tool_name=tool_call['name'],
+                arguments=str(tool_call['args']),
+                call_id=tool_call['id']
+            ))
+        
+        # Extract data sources from tool results
+        if hasattr(msg, 'content') and isinstance(msg.content, str):
+            # Look for URLs, file names, or other sources
+            import re
+            urls = re.findall(r'https?://[^\s]+', msg.content)
+            data_sources.extend(urls)
+    
+    # Get plan information if available
+    plan = state.get("plan")
+    approach_used = "Direct execution"
+    assumptions_made = []
+    
+    if plan:
+        approach_used = f"{plan.task_type} approach with {len(plan.steps)} steps"
+        assumptions_made = plan.assumptions
+    
+    # Generate structured report (КОСТЫЛЬ ЗДЕСЬ!)
+    report_generator_prompt = f"""
+    Generate a comprehensive execution report for the following query processing:
+
+    ORIGINAL QUERY: {state['query']}
+    
+    EXECUTION CONTEXT:
+    - Complexity Level: {state.get('complexity_assessment', {}).level}
+    - Plan Used: {plan if plan else {}}
+    - Tools Executed: {tools_executed}
+    - Available Files: {list(state.get('file_contents', {}).keys())}
+    
+    CONVERSATION HISTORY:
+    {[msg.content[:200] + "..." if len(msg.content) > 200 else msg.content 
+      for msg in state['messages'][-5:]]}  # Last 5 messages for context
+    
+    Based on this information, create a structured execution report that includes:
+    1. Query summary
+    2. Approach used
+    3. Key findings from the execution
+    4. Data sources used
+    5. Your confidence level in the results
+    6. Any limitations or caveats
+    7. The final answer
+    
+    Be thorough but concise. This report will be evaluated by a critic for quality assurance.
+    """
+    
+    report_llm = llm.with_structured_output(ExecutionReport)
+    
+    execution_report = report_llm.invoke([
+        SystemMessage(content=report_generator_prompt),
+        HumanMessage(content="Generate the execution report.")
+    ])
+    
+    print(f"Report generated - Confidence: {execution_report.confidence_level}")
+    print(f"Key findings: {len(execution_report.key_findings)}")
+    print(f"Data sources: {len(execution_report.data_sources)}")
+    
+    # Format final answer for user
+    formatted_answer = format_final_answer(execution_report, state.get('complexity_assessment', {}))
+    print(execution_report)
+    return {
+        "execution_report": execution_report,
+        "final_answer": formatted_answer
+    }
+
+
+def simple_executor(state: AgentState) -> AgentState:
+    """Handle simple queries directly without planning."""
+    print("=== SIMPLE EXECUTION ===")
+    
+    # For simple queries, use the LLM with tools directly
+    simple_prompt = f"""
+    Answer this simple query directly and efficiently: {state['query']}
+    
+    You have access to tools if needed, but try to answer directly when possible.
+    If you need files, they are available at: {list(state.get('file_contents', {}).keys())}
+    
+    Provide a clear, concise answer.
+    """
+    
+    response = llm_with_tools.invoke([
+        SystemMessage(content=simple_prompt),
+        HumanMessage(content=state['query'])
+    ])
+    
+    return {
+        "messages": state["messages"] + [response],
+        "final_answer": response.content
+    }
+
+
+def should_use_planning(state: AgentState) -> str:
+    """Route based on complexity assessment."""
+    complexity = state["complexity_assessment"]
+    
+    if complexity.level == "simple" and not complexity.needs_planning:
+        return "simple_executor"
+    else:
+        return "planner"
+    
+
+def critic_evaluator(state: AgentState) -> AgentState:
+    """Enhanced critic that evaluates execution reports."""
+    print("=== ENHANCED ANSWER CRITIQUE ===")
+    
+    report = state.get("execution_report")
+    critic_llm = llm.with_structured_output(CritiqueFeedback)
+    
+    critique_prompt = CRITIC_PROMPT.format(
+        query=report.query_summary,
+        approach=report.approach_used,
+        tools=report.tools_executed,
+        findings=report.key_findings,
+        sources=report.data_sources,
+        confidence=report.confidence_level,
+        limitations=report.limitations,
+        answer=report.final_answer
+    )
+    
+    critique = critic_llm.invoke([
+        SystemMessage(content=critique_prompt),
+        HumanMessage(content="Evaluate this execution report thoroughly.")
+    ])
+    
+    print(f"Quality Score: {critique.quality_score}/10")
+    print(f"Complete: {critique.is_complete}")
+    print(f"Accurate: {critique.is_accurate}")
+    
+    if critique.errors_found:
+        print(f"Issues found: {critique.errors_found}")
+    
+    if critique.needs_replanning:
+        print(f"Replanning needed: {critique.replan_instructions}")
+    
+    return {
+        "critique_feedback": critique,
+        "iteration_count": state.get("iteration_count", 0) + 1
+    }
+
+
+
+def should_replan(state: AgentState) -> str:
+    """Decide whether to accept answer, replan, or stop."""
+    critique = state.get("critique_feedback")
+    iteration_count = state.get("iteration_count", 0)
+    max_iterations = state.get("max_iterations", 3)
+    
+
+    print(f"=== REPLAN DECISION ===")
+    print(f"Iteration: {iteration_count}/{max_iterations}")
+    print(f"Quality score: {critique.quality_score if critique else 'N/A'}")
+    print(f"Needs replanning: {critique.needs_replanning if critique else 'N/A'}")
+
+    if not critique:
+        return "end"
+    
+    # Stop if max iterations reached
+    if iteration_count >= max_iterations:
+        print(f"Max iterations ({max_iterations}) reached. Accepting current answer.")
+        return "end"
+    
+    # Accept if quality is good enough
+    if critique.quality_score >= 7 or not critique.needs_replanning:
+        print("Quality acceptable, ending execution")
+        return "end"
+    
+    # Replan if quality is poor and we haven't exceeded max iterations
+    if critique.needs_replanning and iteration_count < max_iterations:
+        print("Replanning due to critic feedback...")
+        return "replan"
+    
+    return "end"
+
+def replanner(state: AgentState) -> AgentState:
+    """Create a revised plan based on critic feedback."""
+    print("=== REPLANNING ===")
+    
+    critique = state["critique_feedback"]
+    previous_plan = state.get("plan")
+    
+    replan_prompt = f"""
+    {SYSTEM_PROMPT_PLANNER}
+    
+    REPLANNING CONTEXT:
+    Original Query: {state['query']}
+    Previous Plan: {previous_plan if previous_plan else {}}
+    
+    CRITIC FEEDBACK:
+    - Quality Score: {critique.quality_score}/10
+    - Issues Found: {critique.errors_found}
+    - Missing Elements: {critique.missing_elements}
+    - Improvement Suggestions: {critique.suggested_improvements}
+    - Specific Instructions: {critique.replan_instructions}
+    
+    Create a REVISED plan that addresses these issues. Focus on fixing the identified problems.
+    """
+    
+    revised_plan = planner_llm.invoke([
+        SystemMessage(content=replan_prompt),
+        HumanMessage(content="Create a revised plan based on the feedback.")
+    ])
+    
+    print("Plan revised based on critic feedback")
+    
+    # Очищаем историю сообщений от неполных tool_calls
+    current_messages = state.get("messages", [])
+    cleaned_messages = clean_message_history(current_messages)
+    
+    # Оставляем только системные сообщения и начальный запрос
+    essential_messages = []
+    for msg in cleaned_messages:
+        if isinstance(msg, (SystemMessage, HumanMessage)):
+            # Сохраняем системные сообщения и пользовательские запросы
+            if ("complexity" in msg.content.lower() or 
+                "assess" in msg.content.lower() or
+                isinstance(msg, HumanMessage)):
+                essential_messages.append(msg)
+    
+    print(f"Cleaned message history: {len(current_messages)} -> {len(essential_messages)} messages")
+    
+    return {
+        "plan": revised_plan,
+        "current_step": 0,
+        "reasoning_done": False,
+        "messages": essential_messages,
+        "execution_report": None
+    }
+
+
+def complexity_assessor(state: AgentState) -> AgentState:
+    """Assess query complexity and determine if planning is needed."""
+    print("=== COMPLEXITY ASSESSMENT ===")
+    
+    complexity_llm = llm.with_structured_output(ComplexityLevel)
+    
+    assessment_message = [
+        SystemMessage(content=COMPLEXITY_ASSESSOR_PROMPT.strip()),
+        HumanMessage(content=f"Query: {state['query']}")
+    ]
+    
+    assessment = complexity_llm.invoke(assessment_message)
+    
+    print(f"Complexity: {assessment.level}")
+    print(f"Needs planning: {assessment.needs_planning}")
+    print(f"Reasoning: {assessment.reasoning}")
+    
+    return {
+        "complexity_assessment": assessment,
+        "messages": state["messages"] + assessment_message
+    }

src/notebook_port/test_unstable_port.py

@@ -0,0 +1,772 @@
+# === AUTO-GENERATED FROM test_.ipynb (do not edit logic) ===
+# Only additive imports below to resolve package paths.
+import sys, os
+from pathlib import Path as _Path
+
+# Ensure project root is importable when running as a module
+_CUR = _Path(__file__).resolve()
+_SRC = _CUR.parent.parent
+_ROOT = _SRC.parent
+if str(_ROOT) not in sys.path:
+    sys.path.insert(0, str(_ROOT))
+
+# Prefer package-qualified imports; leave original notebook imports untouched below.
+try:
+    from src.prompts import *          # noqa: F401,F403
+    from src.schemas import *          # noqa: F401,F403
+    from src.tools import *            # noqa: F401,F403
+    from src.tools.code_interpreter import safe_code_run  # noqa: F401
+except Exception:
+    # Fallbacks if executed inside src as working directory
+    pass
+
+
+# === CELL 0 FROM NOTEBOOK ===
+import math
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
+import uuid
+from prompts import *
+from schemas import *
+from typing import Sequence
+from langchain_core.messages import BaseMessage
+from langgraph.graph.message import add_messages
+
+import os, io, json, base64
+from typing import Optional, Dict, Any, List
+from langchain_core.tools import tool
+
+# pip install google-generativeai pillow
+import google.generativeai as genai
+from PIL import Image
+from langgraph.prebuilt import ToolNode
+
+from dotenv import load_dotenv
+import pandas as pd
+from IPython.display import display, Image
+from langchain_community.document_loaders import DataFrameLoader, TextLoader
+from langchain_community.vectorstores import Chroma
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.schema import Document
+from langchain.schema.output_parser import StrOutputParser
+import pickle 
+
+
+from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
+from pydantic import BaseModel, Field
+
+from typing import List, TypedDict, Annotated, Literal, Optional, Union
+
+from langgraph.graph import StateGraph, END
+
+load_dotenv()
+import os
+import json
+import re
+import operator
+
+from langgraph.store.memory import InMemoryStore
+in_memory_store = InMemoryStore() #сохраняем состояние между запусками
+
+from IPython.display import Image, display
+
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import StateGraph, MessagesState, START, END
+from langgraph.store.base import BaseStore
+
+from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.runnables.config import RunnableConfig
+from PIL import Image, ImageStat, ExifTags
+import pandas as pd
+
+
+#TOOLS
+
+from tools import (web_search, arxiv_search, wiki_search, add, subtract, multiply, divide, power, 
+analyze_csv_file, analyze_docx_file, analyze_pdf_file, analyze_txt_file, analyze_image_file, vision_qa_gemma, analyze_excel_file, preprocess_files, save_and_read_file, download_file_from_url)
+
+from code_interpreter import safe_code_run
+
+
+# === CELL 1 FROM NOTEBOOK ===
+
+def clean_message_history(messages):
+    """
+    Очищает историю сообщений от неполных циклов tool_calls/responses.
+    Удаляет AIMessage с tool_calls, если нет соответствующих ToolMessage.
+    """
+    cleaned_messages = []
+    i = 0
+    
+    while i < len(messages):
+        msg = messages[i]
+        
+        # Если это AIMessage с tool_calls
+        if hasattr(msg, 'tool_calls') and msg.tool_calls:
+            # Ищем соответствующие ToolMessage
+            tool_call_ids = {tc['id'] for tc in msg.tool_calls}
+            found_responses = set()
+            
+            # Проверяем следующие сообщения на наличие ответов
+            j = i + 1
+            while j < len(messages) and isinstance(messages[j], ToolMessage):
+                if messages[j].tool_call_id in tool_call_ids:
+                    found_responses.add(messages[j].tool_call_id)
+                j += 1
+            
+            # Если все tool_calls имеют ответы, добавляем весь блок
+            if found_responses == tool_call_ids:
+                # Добавляем AIMessage и все соответствующие ToolMessage
+                cleaned_messages.append(msg)
+                for k in range(i + 1, j):
+                    cleaned_messages.append(messages[k])
+                i = j
+            else:
+                # Пропускаем неполный блок
+                print(f"Removing incomplete tool call block: {tool_call_ids - found_responses}")
+                i = j
+        else:
+            # Обычное сообщение - добавляем
+            cleaned_messages.append(msg)
+            i += 1
+    
+    return cleaned_messages
+
+# === CELL 2 FROM NOTEBOOK ===
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.25)
+TOOLS = [download_file_from_url, web_search, arxiv_search, wiki_search, add, subtract, multiply, divide, power, analyze_excel_file, analyze_csv_file, analyze_docx_file, analyze_pdf_file, analyze_txt_file, analyze_image_file, vision_qa_gemma, safe_code_run]
+
+# === CELL 3 FROM NOTEBOOK ===
+llm_with_tools = llm.bind_tools(TOOLS)
+config = {"configurable": {"thread_id": "1"}, "recursion_limit" : 50}
+TOOL_NODE = ToolNode(TOOLS)
+planner_llm = llm.with_structured_output(PlannerPlan)
+
+class AgentState(MessagesState):
+    query: str
+    final_answer: str
+    plan: Optional[PlannerPlan]
+    complexity_assessment: ComplexityLevel
+    current_step: int
+    reasoning_done: bool
+    messages : Annotated[Sequence[BaseMessage], add_messages]
+    files: List[str]
+    file_contents: Dict[str, Any]
+    critique_feedback: Optional[CritiqueFeedback]
+    iteration_count :int
+    max_iterations: int
+    execution_report : ExecutionReport
+
+
+
+def query_input(state : AgentState) -> AgentState:
+    print("=== USER QUERY TRANSFERED TO AGENT ===")
+
+    files = state.get("files", [])
+    if files:
+        print(f"Processing {len(files)} files:")
+        file_info = preprocess_files(files)
+    
+        for file_path, info in file_info.items():
+            print(f"  - {file_path}: {info['type']} ({info['size']} bytes) -> {info['suggested_tool']}")
+
+        state["file_contents"] = file_info
+        file_context = "\n\n=== AVAILABLE FILES FOR ANALYSIS ===\n"
+        for file_path, info in file_info.items():
+            filename = os.path.basename(file_path)
+            file_context += f"File: {filename}\n"
+            file_context += f"  - Type: {info['type']}\n"  
+            file_context += f"  - Size: {info['size']} bytes\n"
+            file_context += f"  - Suggested tool: {info['suggested_tool']}\n"
+            if info.get("preview"):
+                file_context += f"  - Preview: {info['preview']}\n"
+            file_context += "\n"
+        
+        # Добавляем инструкции по работе с файлами
+        file_context += "IMPORTANT: Use the suggested tools to analyze these files before processing their data.\n"
+        file_context += "File paths are available in the agent state and can be passed directly to analysis tools.\n"
+        
+        original_query = state.get("query", "")
+        state["query"] = original_query + file_context
+    return state
+
+
+def planner(state : AgentState) -> AgentState:
+    sys_stack = [
+            SystemMessage(content=SYSTEM_PROMPT_PLANNER.strip()),
+            HumanMessage(content=state["query"]),
+        ]
+    plan: PlannerPlan = planner_llm.invoke(sys_stack)
+    
+    print("=== GENERATED PLAN ===")
+    return {"messages" : sys_stack + state["messages"],
+            "plan": plan,
+            "current_step ": 0,
+            "reasoning_done": False}
+
+def agent(state: AgentState) -> AgentState:
+    
+    """
+    sys_msg = SystemMessage(
+        content=SYSTEM_EXECUTOR_PROMPT.strip().format(
+            plan=json.dumps(state["plan"], indent=2)
+        )
+    )
+    """
+    current_step = state.get("current_step", 0)
+    reasoning_done = state.get("reasoning_done", False)
+    plan = state.get("plan", {})
+    steps = state["plan"].steps
+
+    print(f"=== AGENT DEBUG ===")
+    print(f"Current step: {current_step}")
+    print(f"Reasoning done: {reasoning_done}")
+    print(f"Plan exists: {plan is not None}")
+    print(f"Total steps in plan: {len(plan.steps) if plan else 'No plan'}")
+
+    if not plan or not hasattr(plan, 'steps') or not plan.steps:
+        print("ERROR: No valid plan found!")
+        return {
+            "messages": state["messages"] + [AIMessage(content="No valid plan available. <FINAL_ANSWER>")],
+            "reasoning_done": False
+        }
+    
+    steps = plan.steps
+    
+    if current_step >= len(steps):
+        print("All plan steps completed, moving to finalization")
+        return {
+            "messages": state["messages"] + [AIMessage(content="All steps completed. <FINAL_ANSWER>")],
+            "reasoning_done": False
+        }
+
+    current_step_info = steps[current_step]
+    print(f"Executing step {current_step + 1}: {current_step_info.description}")
+
+    if not reasoning_done:
+
+        # ✅ ДОБАВЛЕНО: Специальный контекст для файлов
+        file_context = ""
+        file_contents = state.get("file_contents", {})
+        if file_contents:
+            file_context = "\n\nAVAILABLE FILES IN CURRENT SESSION:\n"
+            for filepath, info in file_contents.items():
+                filename = os.path.basename(filepath)
+                file_context += f"- {filename}: {info['type']} file, suggested tool: {info['suggested_tool']}\n"
+                file_context += f"  Path: {filepath}\n"
+
+        reasoning_prompt = f"""
+        {SYSTEM_EXECUTOR_PROMPT}
+        
+        CURRENT TASK: You must perform reasoning for step {current_step + 1}.
+        
+        STEP INFO: {current_step_info}\n\n
+
+        FILE CONTEXT: {file_contents}
+        
+        CRITICAL: You MUST output your reasoning in <REASONING> tags, but DO NOT call any tools yet.
+        Explain what you need to do and why, then end your response.
+
+        REASONING IS IMPERATIVE BEFORE ANY TOOL CALLS.
+        """
+
+        sys_msg = SystemMessage(content = reasoning_prompt)
+        stack = [sys_msg] + state["messages"]
+
+        step = llm.invoke(stack)
+        print("=== REASONING STEP ===")
+        print(step.content)
+
+        return {
+            "messages" : state["messages"] + [step],
+            "reasoning_done" : True
+        }
+    
+    else:
+        tool_prompt = f"""
+        Now execute the tool for step {current_step + 1}.
+        
+        You have already done the reasoning. Now call the appropriate tool with the correct parameters.
+        Available file paths: {list(state.get("file_contents", {}).keys())}\n
+        IMPORTANT NOTE: IF YOU DECIDED TO USE safe_code_run, MAKE SURE TO FINISH CALCULATIONS WITH print() or saving to a variable NAMED 'result' so that the output can be captured!
+        AVAILABLE TOOLS: {', '.join([tool.name for tool in TOOLS])}
+        """ 
+
+        sys_msg = SystemMessage(content=tool_prompt)
+        stack = [sys_msg] + state["messages"]  # Берем последние сообщения включая reasoning
+        
+        # Используем модель С инструментами для выполнения
+        step = llm_with_tools.invoke(stack)
+        print("=== TOOL EXECUTION ===")
+        print(f"Tool calls: {step.tool_calls}")
+        
+        return {
+            "messages": state["messages"] + [step],
+            "current_step": current_step + 1 if step.tool_calls else current_step,
+            "reasoning_done": False  # Сбрасываем для следующего шага
+        }
+
+
+def should_continue(state : AgentState) -> bool:
+    
+    last_message = state["messages"][-1]
+    reasoning_done = state.get("reasoning_done", False)
+    plan = state.get("plan", None)
+    current_step = state.get("current_step", 0)
+
+    if plan and current_step >= len(plan.steps):
+        return "final_answer"
+
+
+    if hasattr(last_message, "content") and "<FINAL_ANSWER>" in last_message.content:
+        return "final_answer"
+    elif hasattr(last_message, "tool_calls") and last_message.tool_calls:
+        return "tools" 
+    elif not reasoning_done and hasattr(last_message, 'content') and "<REASONING>" in last_message.content:
+        # Reasoning выполнен, но инструменты еще не вызваны
+        return "agent"
+    elif reasoning_done:
+        # Reasoning выполнен, теперь нужно вызвать инструменты
+        return "agent"
+    else:
+        # Нужно сделать reasoning
+        return "agent"
+
+# 6. Добавить отладочную информацию в TOOL_NODE
+class DebuggingToolNode(ToolNode):
+    def __init__(self, tools):
+        super().__init__(tools)
+    
+    def __call__(self, state):
+        print("=== TOOL EXECUTION STARTED ===")
+        result = super().__call__(state)
+        print("=== TOOL EXECUTION COMPLETED ===")
+        return result
+
+DEBUGGING_TOOL_NODE = DebuggingToolNode(TOOLS)
+
+
+
+"""
+def summary(state : AgentState) -> AgentState:
+    print("=== FINAL ANSWER ===")
+    summarizer_prompt = 
+    Now you have to provide final answer for the user query : {query}
+    In messages below you have all the context you need.
+
+    YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, Apply the rules above for each element (number or string), ensure there is exactly one space after each comma.
+    Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
+
+    Here is the context:
+    {messages}
+
+    REMEMBER AND STRICTLY FOLLOW THE FORMATTING RULES ABOVE. ALWAYS USE THIS FORMAT:
+    FINAL ANSWER: ...
+    
+
+    state["final_answer"] = llm.invoke([SystemMessage(content=summarizer_prompt.strip().format(query=state["query"], messages = state["messages"]))])
+    return state
+"""
+
+def enhanced_finalizer(state: AgentState) -> AgentState:
+    """Generate comprehensive execution report for critic evaluation."""
+    print("=== GENERATING EXECUTION REPORT ===")
+    
+    # Extract tool execution information
+    tools_executed = []
+    data_sources = []
+    
+    for msg in state["messages"]:
+        if hasattr(msg, 'tool_calls') and msg.tool_calls:
+            for tool_call in msg.tool_calls:
+                tools_executed.append(ToolExecution(
+                tool_name=tool_call['name'],
+                arguments=str(tool_call['args']),
+                call_id=tool_call['id']
+            ))
+        
+        # Extract data sources from tool results
+        if hasattr(msg, 'content') and isinstance(msg.content, str):
+            # Look for URLs, file names, or other sources
+            import re
+            urls = re.findall(r'https?://[^\s]+', msg.content)
+            data_sources.extend(urls)
+    
+    # Get plan information if available
+    plan = state.get("plan")
+    approach_used = "Direct execution"
+    assumptions_made = []
+    
+    if plan:
+        approach_used = f"{plan.task_type} approach with {len(plan.steps)} steps"
+        assumptions_made = plan.assumptions
+    
+    # Generate structured report (КОСТЫЛЬ ЗДЕСЬ!)
+    report_generator_prompt = f"""
+    Generate a comprehensive execution report for the following query processing:
+
+    ORIGINAL QUERY: {state['query']}
+    
+    EXECUTION CONTEXT:
+    - Complexity Level: {state.get('complexity_assessment', {}).level}
+    - Plan Used: {plan if plan else {}}
+    - Tools Executed: {tools_executed}
+    - Available Files: {list(state.get('file_contents', {}).keys())}
+    
+    CONVERSATION HISTORY:
+    {[msg.content[:200] + "..." if len(msg.content) > 200 else msg.content 
+      for msg in state['messages'][-5:]]}  # Last 5 messages for context
+    
+    Based on this information, create a structured execution report that includes:
+    1. Query summary
+    2. Approach used
+    3. Key findings from the execution
+    4. Data sources used
+    5. Your confidence level in the results
+    6. Any limitations or caveats
+    7. The final answer
+    
+    Be thorough but concise. This report will be evaluated by a critic for quality assurance.
+    """
+    
+    report_llm = llm.with_structured_output(ExecutionReport)
+    
+    execution_report = report_llm.invoke([
+        SystemMessage(content=report_generator_prompt),
+        HumanMessage(content="Generate the execution report.")
+    ])
+    
+    print(f"Report generated - Confidence: {execution_report.confidence_level}")
+    print(f"Key findings: {len(execution_report.key_findings)}")
+    print(f"Data sources: {len(execution_report.data_sources)}")
+    
+    # Format final answer for user
+    formatted_answer = format_final_answer(execution_report, state.get('complexity_assessment', {}))
+    print(execution_report)
+    return {
+        "execution_report": execution_report,
+        "final_answer": formatted_answer
+    }
+
+def format_final_answer(report: ExecutionReport, complexity: dict) -> str:
+    """Format the final answer based on complexity and report content."""
+    
+    if complexity.level == 'simple':
+        # For simple queries, just return the answer
+        return f"FINAL ANSWER: {report.final_answer}"
+    
+    # For complex queries, provide more detailed response
+    formatted = f"""FINAL ANSWER: {report.final_answer}
+
+SUMMARY:
+{report.query_summary}
+
+KEY FINDINGS:
+{chr(10).join(f"• {finding}" for finding in report.key_findings)}"""
+    
+    if report.data_sources:
+        formatted += f"""
+
+SOURCES:
+{chr(10).join(f"• {source}" for source in report.data_sources[:5])}"""  # Limit to 5 sources
+    
+    if report.limitations:
+        formatted += f"""
+
+LIMITATIONS:
+{chr(10).join(f"• {limitation}" for limitation in report.limitations)}"""
+    
+    return formatted
+
+
+def complexity_assessor(state: AgentState) -> AgentState:
+    """Assess query complexity and determine if planning is needed."""
+    print("=== COMPLEXITY ASSESSMENT ===")
+    
+    complexity_llm = llm.with_structured_output(ComplexityLevel)
+    
+    assessment_message = [
+        SystemMessage(content=COMPLEXITY_ASSESSOR_PROMPT.strip()),
+        HumanMessage(content=f"Query: {state['query']}")
+    ]
+    
+    assessment = complexity_llm.invoke(assessment_message)
+    
+    print(f"Complexity: {assessment.level}")
+    print(f"Needs planning: {assessment.needs_planning}")
+    print(f"Reasoning: {assessment.reasoning}")
+    
+    return {
+        "complexity_assessment": assessment,
+        "messages": state["messages"] + assessment_message
+    }
+
+
+def simple_executor(state: AgentState) -> AgentState:
+    """Handle simple queries directly without planning."""
+    print("=== SIMPLE EXECUTION ===")
+    
+    # For simple queries, use the LLM with tools directly
+    simple_prompt = f"""
+    Answer this simple query directly and efficiently: {state['query']}
+    
+    You have access to tools if needed, but try to answer directly when possible.
+    If you need files, they are available at: {list(state.get('file_contents', {}).keys())}
+    
+    Provide a clear, concise answer.
+    """
+    
+    response = llm_with_tools.invoke([
+        SystemMessage(content=simple_prompt),
+        HumanMessage(content=state['query'])
+    ])
+    
+    return {
+        "messages": state["messages"] + [response],
+        "final_answer": response.content
+    }
+
+
+def should_use_planning(state: AgentState) -> str:
+    """Route based on complexity assessment."""
+    complexity = state["complexity_assessment"]
+    
+    if complexity.level == "simple" and not complexity.needs_planning:
+        return "simple_executor"
+    else:
+        return "planner"
+    
+"""    
+def critic_evaluator(state: AgentState) -> AgentState:
+    
+    print("=== ANSWER CRITIQUE ===")
+    
+    critic_llm = llm.with_structured_output(CritiqueFeedback)
+    
+    # Gather tool execution results for context
+    tool_results = []
+    for msg in state["messages"]:
+        if hasattr(msg, 'tool_calls') and msg.tool_calls:
+            tool_results.extend([f"Tool: {tc['name']}, Args: {tc['args']}" for tc in msg.tool_calls])
+    
+    if state.get("plan"):
+        terra = state.get("plan")
+    else:
+        terra = "No plan used"
+    critique_prompt = CRITIC_PROMPT.format(
+        query=state["query"],
+        plan=terra,
+        answer=state["final_answer"],
+        tool_results=tool_results[:5]   #Limit context
+    )
+    
+    critique = critic_llm.invoke([
+        SystemMessage(content=critique_prompt),
+        HumanMessage(content="Please evaluate this answer.")
+    ])
+    
+    print(f"Quality Score: {critique.quality_score}/10")
+    print(f"Complete: {critique.is_complete}")
+    print(f"Accurate: {critique.is_accurate}")
+    if critique.errors_found:
+        print(f"Errors: {critique.errors_found}")
+    if critique.needs_replanning:
+        print(f"Needs replanning: {critique.replan_instructions}")
+    
+    return {
+        "critique_feedback": critique,
+        "iteration_count": state.get("iteration_count", 0) + 1
+    }
+"""
+
+def critic_evaluator(state: AgentState) -> AgentState:
+    """Enhanced critic that evaluates execution reports."""
+    print("=== ENHANCED ANSWER CRITIQUE ===")
+    
+    report = state.get("execution_report")
+    critic_llm = llm.with_structured_output(CritiqueFeedback)
+    
+    critique_prompt = CRITIC_PROMPT.format(
+        query=report.query_summary,
+        approach=report.approach_used,
+        tools=report.tools_executed,
+        findings=report.key_findings,
+        sources=report.data_sources,
+        confidence=report.confidence_level,
+        limitations=report.limitations,
+        answer=report.final_answer
+    )
+    
+    critique = critic_llm.invoke([
+        SystemMessage(content=critique_prompt),
+        HumanMessage(content="Evaluate this execution report thoroughly.")
+    ])
+    
+    print(f"Quality Score: {critique.quality_score}/10")
+    print(f"Complete: {critique.is_complete}")
+    print(f"Accurate: {critique.is_accurate}")
+    
+    if critique.errors_found:
+        print(f"Issues found: {critique.errors_found}")
+    
+    if critique.needs_replanning:
+        print(f"Replanning needed: {critique.replan_instructions}")
+    
+    return {
+        "critique_feedback": critique,
+        "iteration_count": state.get("iteration_count", 0) + 1
+    }
+
+
+
+def should_replan(state: AgentState) -> str:
+    """Decide whether to accept answer, replan, or stop."""
+    critique = state.get("critique_feedback")
+    iteration_count = state.get("iteration_count", 0)
+    max_iterations = state.get("max_iterations", 3)
+    
+
+    print(f"=== REPLAN DECISION ===")
+    print(f"Iteration: {iteration_count}/{max_iterations}")
+    print(f"Quality score: {critique.quality_score if critique else 'N/A'}")
+    print(f"Needs replanning: {critique.needs_replanning if critique else 'N/A'}")
+
+    if not critique:
+        return "end"
+    
+    # Stop if max iterations reached
+    if iteration_count >= max_iterations:
+        print(f"Max iterations ({max_iterations}) reached. Accepting current answer.")
+        return "end"
+    
+    # Accept if quality is good enough
+    if critique.quality_score >= 7 or not critique.needs_replanning:
+        print("Quality acceptable, ending execution")
+        return "end"
+    
+    # Replan if quality is poor and we haven't exceeded max iterations
+    if critique.needs_replanning and iteration_count < max_iterations:
+        print("Replanning due to critic feedback...")
+        return "replan"
+    
+    return "end"
+
+def replanner(state: AgentState) -> AgentState:
+    """Create a revised plan based on critic feedback."""
+    print("=== REPLANNING ===")
+    
+    critique = state["critique_feedback"]
+    previous_plan = state.get("plan")
+    
+    replan_prompt = f"""
+    {SYSTEM_PROMPT_PLANNER}
+    
+    REPLANNING CONTEXT:
+    Original Query: {state['query']}
+    Previous Plan: {previous_plan if previous_plan else {}}
+    
+    CRITIC FEEDBACK:
+    - Quality Score: {critique.quality_score}/10
+    - Issues Found: {critique.errors_found}
+    - Missing Elements: {critique.missing_elements}
+    - Improvement Suggestions: {critique.suggested_improvements}
+    - Specific Instructions: {critique.replan_instructions}
+    
+    Create a REVISED plan that addresses these issues. Focus on fixing the identified problems.
+    """
+    
+    revised_plan = planner_llm.invoke([
+        SystemMessage(content=replan_prompt),
+        HumanMessage(content="Create a revised plan based on the feedback.")
+    ])
+    
+    print("Plan revised based on critic feedback")
+    
+    # Очищаем историю сообщений от неполных tool_calls
+    current_messages = state.get("messages", [])
+    cleaned_messages = clean_message_history(current_messages)
+    
+    # Оставляем только системные сообщения и начальный запрос
+    essential_messages = []
+    for msg in cleaned_messages:
+        if isinstance(msg, (SystemMessage, HumanMessage)):
+            # Сохраняем системные сообщения и пользовательские запросы
+            if ("complexity" in msg.content.lower() or 
+                "assess" in msg.content.lower() or
+                isinstance(msg, HumanMessage)):
+                essential_messages.append(msg)
+    
+    print(f"Cleaned message history: {len(current_messages)} -> {len(essential_messages)} messages")
+    
+    return {
+        "plan": revised_plan,
+        "current_step": 0,
+        "reasoning_done": False,
+        "messages": essential_messages,
+        "execution_report": None
+    }
+
+# === CELL 4 FROM NOTEBOOK ===
+#GRAPH BUILDING
+
+builder = StateGraph(AgentState)
+builder.add_node("INPUT", query_input)
+builder.add_node("COMPLEXITY_ASSESSOR", complexity_assessor)
+builder.add_node("PLANNING", planner)
+builder.add_node("AGENT", agent)
+builder.add_node("TOOLS", DEBUGGING_TOOL_NODE)
+builder.add_node("FINALIZER", enhanced_finalizer)
+builder.add_node("SIMPLE_EXECUTOR", simple_executor)
+builder.add_node("CRITIC", critic_evaluator)
+builder.add_node("REPLANNER", replanner)
+
+builder.set_entry_point("INPUT")
+builder.add_edge("INPUT", "COMPLEXITY_ASSESSOR")
+
+builder.add_conditional_edges(
+        "COMPLEXITY_ASSESSOR",
+        should_use_planning,
+        {"simple_executor": "SIMPLE_EXECUTOR", "planner": "PLANNING"},
+    )
+builder.add_edge("SIMPLE_EXECUTOR", "FINALIZER")
+
+
+builder.add_edge("PLANNING", "AGENT")
+builder.add_conditional_edges(
+        "AGENT",
+        should_continue,
+        {"tools": "TOOLS", "agent": "AGENT", "final_answer": "FINALIZER"},
+    )
+builder.add_edge("TOOLS", "AGENT")
+builder.add_edge("FINALIZER", "CRITIC")
+builder.add_conditional_edges(
+        "CRITIC",
+        should_replan,
+        {"end": END, "replan": "REPLANNER"},
+    )
+builder.add_edge("REPLANNER", "AGENT")
+
+
+system = builder.compile(checkpointer=MemorySaver())
+
+# === CELL 5 FROM NOTEBOOK ===
+workflow = system.invoke({"query" : "How many cumulative milliliters of fluid is in all the opaque-capped vials without stickers in the 114 version of the kit that was used for the PromethION long-read sequencing in the paper De Novo-Whole Genome Assembly of the Roborovski Dwarf Hamster (Phodopus roborovskii) Genome?", "current_step": 0, "reasoning_done": False, "files" : [], "files_contents" : {}, "iteration_count" : 0, "max_iterations" : 10, "plan" : None} , config = config)
+
+# === CELL 6 FROM NOTEBOOK ===
+for message in workflow["messages"]:
+    message.pretty_print()
+
+print("\n=== FINAL ANSWER ===")
+
+# === CELL 7 FROM NOTEBOOK ===
+workflow["final_answer"]
+
+# === CELL 8 FROM NOTEBOOK ===
+workflow
+
+# === CELL 9 FROM NOTEBOOK ===
+#TO-DO:
+# - imrove image generation and plots/tables creation
+# - add more tools (e.g. calendar, email, pdf editing, file system)
+# - UI creation

src/prompts/prompts.py

@@ -0,0 +1,150 @@
+SYSTEM_PROMPT_PLANNER = """
+You are the PLANNER of a multi-tool agent (GAIA I–II level). Produce a minimal, reliable plan to solve the user's request using available tools. You DO NOT call tools; output ONLY a JSON plan. Tools are bound via .bind_tools()—use EXACT names.
+
+CORE RULES:
+- MINIMALITY: 1-3 steps max; chain only essentials (e.g., search → download → analyze).
+- ROUTING: Classify as info (web facts), calc (math on known data), table (CSV/Excel agg), doc_qa (PDF/DOCX/TXT extract), image_qa (IMG OCR/vision), multi_hop (anything cross-modality or research—default for unknowns).
+- PREREQUISITES: For external docs/images (e.g., "paper X", URLs): ALWAYS start with web_search/arxiv_search → download_file_from_url (local path like "paper.pdf") → analyze_*. NEVER assume local files—validate existence implicitly via chain.
+- COST-AWARE: Cheap first: search snippets > full download > compute. No raw files to safe_code_run—extract first.
+- EVIDENCE: Mandate citations/pages for facts; units/rounding explicit in guidelines.
+- FALLBACKS: Every step needs success_criteria; on_fail="replan" (default) or "sN" (jump). Add 1 fallback step if high-risk (e.g., no-results → alt query).
+
+ROUTING PATTERNS (MANDATORY CHAINS):
+- info: web_search/wiki_search/arxiv_search → cite snippets.
+- calc: If data missing, insert extract step → safe_code_run (e.g., "sum volumes from text").
+- table: analyze_csv_file/analyze_excel_file (preview) → safe_code_run (agg/query).
+- doc_qa: web_search("paper title PDF") → download_file_from_url → analyze_pdf_file/analyze_docx_file (query="vials fluid ml") → safe_code_run if sum needed.
+- image_qa: web_search → download_file_from_url → analyze_image_file/vision_qa_gemma → safe_code_run for chart-to-table.
+- multi_hop: Decompose (e.g., sub-query1: search; sub-query2: extract) → synthesize.
+
+Output ONLY valid JSON:
+{
+  "task_type": "info|calc|table|doc_qa|image_qa|multi_hop",
+  "assumptions": ["..."],  // 0-2 max; e.g., "Paper details vials explicitly"
+  "plan_rationale": "Brief: why route + key tools/chain",  // 1 sentence
+  "steps": [  // 1-3 only
+    {
+      "id": "s1",
+      "description": "Precise action + why (e.g., 'web_search for paper PDF to locate source')",
+      "evidence_needed": ["citations","page_numbers","stats_check"],  // 1-3
+      "success_criteria": "e.g., 'Top result has PDF URL; or data extracted'",
+      "on_fail": "replan|sN",  // Default: replan
+      "outputs_to_state": ["e.g., 'pdf_url', 'extracted_text'"]  // For chaining
+    }
+  ],
+  "answer_guidelines": {
+    "final_answer_template": "e.g., 'Cumulative volume: X mL (from [cite])'",
+    "citations_required": true,
+    "min_citations": 1,
+    "units_policy": "e.g., 'mL; convert if cm³'",
+    "rounding_policy": "e.g., 'Nearest integer'",
+    "include_artifacts": ["snippets","tables"]  // 0-2
+  }
+}
+
+CONSTRAINTS:
+- Valid JSON only—no extras. If query trivial (no tools), task_type="info" with 0 steps.
+- Exact tool names: web_search, download_file_from_url, analyze_pdf_file, safe_code_run, etc.
+- For research: If no chain, replan triggers auto-fix.
+"""
+
+SYSTEM_EXECUTOR_PROMPT = """
+ROLE: EXECUTOR of multi-tool agent (GAIA level). You follow the FIXED {plan} EXACTLY—no changes, no new steps. Current step: {current_step_id} ("{step_desc}"). Advance ONE step per response.
+
+EXECUTION RULES:
+- BEFORE EVERY TOOL: <REASONING> (2-3 sentences: What step? Why this tool? Exact inputs? Expected output?) </REASONING>
+- THEN: Tool call ONLY for this step (exact name/args from plan). NO OTHER OUTPUT.
+- NO TOOLS? Direct output (e.g., "Calc: 5 mL") + set reasoning_done=True.
+- Check state for priors (e.g., if s2 needs pdf_url from s1, wait/replan if missing).
+- On fail (bad output): <REASONING>Assess + on_fail action</REASONING> then tool or stop.
+- END STEP: If success, output "STEP COMPLETE: {outputs_to_state}" to advance.
+
+RESOURCE CHAIN (MANDATORY IF NEEDED):
+- External doc? Use plan's search→download before analyze.
+- NEVER guess paths—use state["files"] or replan.
+
+OUTPUT FORMAT: <REASONING>...</REASONING> [tool call or direct] [STEP COMPLETE if done]. NO JSON/PLANS/MARKDOWN.
+
+FAILSAFE: If unclear, <REASONING>Replan needed</REASONING> and stop.
+DO NOT FORGET TO ADD <FINAL_ANSWER> IF YOU THINK IT'S TIME TO ANSWER THE USER AND YOU HAVE ALL THE DATA FOR EXACT ANSWER.
+"""
+
+
+COMPLEXITY_ASSESSOR_PROMPT = """
+You are a COMPLEXITY ASSESSOR for a multi-tool agent system.
+Your job is to analyze user queries and determine their complexity level and processing requirements.
+
+COMPLEXITY LEVELS:
+1. SIMPLE: Direct questions that can be answered immediately without tools or with single tool use
+   - Examples: "What is 2+2?", "Define photosynthesis", "What's the capital of France?"
+   
+2. MODERATE: Questions requiring 1-3 tool calls or basic analysis
+   - Examples: "Search for recent news about AI", "Analyze this CSV file", "What's the weather tomorrow?"
+   
+3. COMPLEX: Multi-step problems requiring planning, multiple tools, or sophisticated reasoning
+   - Examples: Research tasks, multi-file analysis, calculations with dependencies, creative projects
+
+ASSESSMENT CRITERIA:
+- Number of steps likely needed
+- Tool complexity and dependencies
+- Data processing requirements
+- Need for intermediate reasoning
+- Risk of failure without proper planning
+
+RULES:
+- SIMPLE queries bypass planning entirely
+- MODERATE queries may use lightweight planning
+- COMPLEX queries require full planning with fallbacks
+- When in doubt, err toward higher complexity
+
+Analyze the query and respond with your assessment.
+"""
+
+CRITIC_PROMPT = """
+You are the CRITIC of a multi-tool agent system.
+Your job is to evaluate execution reports and provide detailed feedback.
+
+EVALUATION FRAMEWORK:
+
+1. COMPLETENESS (0-3 points):
+   - 3: Fully addresses all aspects of the query
+   - 2: Addresses main aspects, minor gaps
+   - 1: Partial answer, significant gaps
+   - 0: Incomplete or off-topic
+
+2. ACCURACY (0-3 points):
+   - 3: All information appears accurate and well-sourced
+   - 2: Mostly accurate, minor issues
+   - 1: Some accuracy concerns
+   - 0: Significant accuracy problems
+
+3. METHODOLOGY (0-2 points):
+   - 2: Appropriate tools and approach used
+   - 1: Acceptable approach, could be better
+   - 0: Poor methodology or tool selection
+
+4. EVIDENCE (0-2 points):
+   - 2: Strong evidence and sources provided
+   - 1: Some evidence provided
+   - 0: Insufficient evidence
+
+TOTAL SCORE: /10 points
+
+DECISION THRESHOLDS:
+- 8-10: Accept (excellent quality)
+- 6-7: Accept with minor notes
+- 4-5: Marginal, consider replanning
+- 0-3: Reject, requires replanning
+
+EXECUTION REPORT TO EVALUATE:
+Query: {query}
+Approach: {approach}
+Tools Used: {tools}
+Key Findings: {findings}
+Sources: {sources}
+Confidence: {confidence}
+Limitations: {limitations}
+Final Answer: {answer}
+
+Provide detailed critique focusing on what works well and what could be improved.
+"""

src/schemas.py

@@ -0,0 +1,73 @@
+from typing import Any, Dict, List, Optional, Literal, Iterable
+from pydantic import BaseModel, Field, ValidationError
+
+
+class ComplexityLevel(BaseModel):
+    level: Literal["simple", "moderate", "complex"] = Field(description="Complexity level of the query")
+    reasoning: str = Field(description="Explanation for the complexity assessment")
+    needs_planning: bool = Field(description="Whether this query requires detailed planning")
+    suggested_approach: str = Field(description="Recommended approach for handling this query")
+
+class CritiqueFeedback(BaseModel):
+    quality_score: int = Field(ge=1, le=10, description="Quality score from 1-10")
+    is_complete: bool = Field(description="Whether the answer is complete")
+    is_accurate: bool = Field(description="Whether the answer appears accurate")
+    missing_elements: List[str] = Field(default_factory=list, description="What's missing from the answer")
+    errors_found: List[str] = Field(default_factory=list, description="Potential errors identified")
+    suggested_improvements: List[str] = Field(default_factory=list, description="Suggestions for improvement")
+    needs_replanning: bool = Field(description="Whether the plan should be revised")
+    replan_instructions: Optional[str] = Field(default=None, description="Instructions for replanning")
+
+
+
+TaskType = Literal["info", "calc", "table", "doc_qa", "image_qa", "multi_hop"]
+EvidenceTag = Literal["citations", "page_numbers", "figure_captions", "stats_check", "unit_check"]
+
+class PlanStep(BaseModel):
+    id: str
+    description: str
+    #tool: Optional[str] = Field(default=None, description="Exact tool name or null for reasoning step")
+    #args_hint: Dict[str, Any] = Field(default_factory=dict)
+    evidence_needed: List[EvidenceTag] = Field(default_factory=list)
+    success_criteria: str
+    on_fail: str = Field(default="replan", description="One of: 'replan' | 'stop' | step-id")
+    outputs_to_state: List[str] = Field(default_factory=list)
+
+class AnswerGuidelines(BaseModel):
+    final_answer_template: str
+    citations_required: bool = False
+    min_citations: int = 0
+    units_policy: Optional[str] = None
+    rounding_policy: Optional[str] = None
+    include_artifacts: List[str] = Field(default_factory=list)
+
+class PlannerPlan(BaseModel):
+    task_type: TaskType
+    assumptions: List[str] = Field(default_factory=list)
+    plan_rationale: str
+    steps: List[PlanStep]
+    answer_guidelines: AnswerGuidelines
+
+
+class ToolExecution(BaseModel):
+    tool_name: str
+    arguments: str
+    call_id: str
+    
+    class Config:
+        extra = "forbid"
+
+class ExecutionReport(BaseModel):
+    """Structured report for critic evaluation."""
+    query_summary: str = Field(description="Brief summary of the user's query")
+    approach_used: str = Field(description="What approach/strategy was used")
+    tools_executed: List[ToolExecution] = Field(default_factory=list, description="List of tools used with results")
+    key_findings: List[str] = Field(default_factory=list, description="Main findings or results")
+    data_sources: List[str] = Field(default_factory=list, description="Sources of information used")
+    assumptions_made: List[str] = Field(default_factory=list, description="Any assumptions made during execution")
+    confidence_level: Literal["low", "medium", "high"] = Field(description="Confidence in the answer")
+    limitations: List[str] = Field(default_factory=list, description="Known limitations or caveats")
+    final_answer: str = Field(description="The actual answer to the user's query")
+
+    class Config:
+        extra = "forbid"

src/state.py

@@ -0,0 +1,22 @@
+from langgraph.graph import MessagesState
+from typing import List, Annotated, Optional, Dict, Any
+from schemas import PlannerPlan, ComplexityLevel, CritiqueFeedback, ExecutionReport
+from typing import Sequence
+from langchain_core.messages import BaseMessage
+from langgraph.graph.message import add_messages
+
+class AgentState(MessagesState):
+    query: str
+    final_answer: str
+    plan: Optional[PlannerPlan]
+    complexity_assessment: ComplexityLevel
+    current_step: int
+    reasoning_done: bool
+    messages : Annotated[Sequence[BaseMessage], add_messages]
+    files: List[str]
+    file_contents: Dict[str, Any]
+    critique_feedback: Optional[CritiqueFeedback]
+    iteration_count :int
+    max_iterations: int
+    execution_report : ExecutionReport
+

src/tools/code_interpreter.py

@@ -0,0 +1,325 @@
+# src/gaia_agent/tools/safe_code_run.py
+from __future__ import annotations
+import io, os, sys, uuid, base64, traceback, contextlib, tempfile, shutil
+import multiprocessing as mp
+from typing import Optional, Dict, Any, List
+from pydantic import ValidationError
+from langchain_core.tools import tool
+from utils.code_run import (
+    CodeRunRequest, CodeRunResult, EnvInfo,
+    PlotArtifact, DataFrameArtifact,
+)
+
+# ====================== HELPERS ======================
+
+def _b64_png(fig, dpi: int) -> str:
+    import matplotlib.pyplot as plt
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", dpi=dpi, bbox_inches="tight")
+    buf.seek(0)
+    data = base64.b64encode(buf.read()).decode("utf-8")
+    buf.close()
+    return data
+
+def _clip_df(df, max_rows: int, max_cols: int):
+    sub = df.iloc[:max_rows, :max_cols]
+    head = sub.to_dict(orient="records")
+    dtypes = {str(k): str(v) for k, v in sub.dtypes.to_dict().items()}
+    return head, list(df.shape), dtypes
+
+def _env_info() -> EnvInfo:
+    try:
+        import numpy as _np; nv = _np.__version__
+    except Exception:
+        nv = None
+    try:
+        import pandas as _pd; pv = _pd.__version__
+    except Exception:
+        pv = None
+    return EnvInfo(numpy=nv, pandas=pv)
+
+# ====================== CHILD PROCESS ======================
+
+def _child_exec(payload: Dict[str, Any], queue: mp.Queue):
+    """
+    Изолированное выполнение user-кода:
+      - урезанные builtins
+      - безопасный open (read-only в sandbox)
+      - белый список импортов
+      - запрет сети
+      - temp cwd + очистка
+      - RLIMIT CPU/AS (Unix)
+      - захват stdout/stderr
+      - сбор matplotlib и pandas.DataFrame (по флагам)
+    """
+    import builtins, importlib
+
+    code: str = payload["code"]
+    limits: Dict[str, Any] = payload["limits"]
+    allowed: List[str] = payload["allowed"]
+    return_plots: bool = payload["return_plots"]
+    return_dfs: bool = payload["return_dfs"]
+
+    # ---------- OS limits (Unix) ----------
+    try:
+        import resource
+        cpu = max(1, int(limits["timeout_seconds"]))
+        resource.setrlimit(resource.RLIMIT_CPU, (cpu, cpu + 1))
+        # мягкий лимит RAM ~1.5GB (подстрой при необходимости)
+        one_gb = 1024 * 1024 * 1024
+        resource.setrlimit(resource.RLIMIT_AS, (int(1.5 * one_gb), int(1.5 * one_gb)))
+        # ограничим размеры файлов
+        resource.setrlimit(resource.RLIMIT_FSIZE, (50 * 1024 * 1024, 50 * 1024 * 1024))
+    except Exception:
+        pass
+
+    # ---------- Sandbox FS ----------
+    workdir = tempfile.mkdtemp(prefix="ci_")
+    os.chdir(workdir)
+
+    # ---------- Network ban ----------
+    try:
+        import socket
+        class _NoNet(socket.socket):
+            def __init__(self, *a, **kw):
+                raise OSError("Network disabled in sandbox")
+        socket.socket = _NoNet  # type: ignore
+    except Exception:
+        pass
+
+    # ---------- Builtins ----------
+    safe_names = [
+        "abs","all","any","bool","dict","float","int","len","list","max","min",
+        "range","str","sum","print","enumerate","zip","map","filter","sorted",
+        "reversed","complex","pow","divmod"
+    ]
+    safe_builtins = {n: getattr(builtins, n) for n in safe_names}
+
+    # сохранём реальный open, потом подменим на безопасный
+    real_open = open
+
+    def _safe_open(path, mode="r", *a, **kw):
+        # Разрешаем ТОЛЬКО чтение, ТОЛЬКО внутри workdir
+        if any(m in mode for m in ("w", "a", "+", "x")):
+            raise PermissionError("Write access forbidden in sandbox")
+        abspath = os.path.abspath(path)
+        # запрещаем выход из песочницы и следование symlink наружу
+        if not abspath.startswith(workdir + os.sep) and abspath != workdir:
+            raise PermissionError("Access outside sandbox forbidden")
+        # запретим двоичный write по flags
+        return real_open(abspath, mode, *a, **kw)
+
+    # удалим опасные builtins и поставим наш open
+    for banned in ["exec","eval","__import__","compile","input","globals","locals","vars","dir","help","__build_class__"]:
+        safe_builtins.pop(banned, None)
+    safe_builtins["open"] = _safe_open
+
+    # ---------- Import whitelist ----------
+    real_import = builtins.__import__
+    ALLOWED = set(allowed)
+    def _safe_import(name, globals=None, locals=None, fromlist=(), level=0):
+        base = name.split(".")[0]
+        if (name not in ALLOWED) and (base not in ALLOWED):
+            raise ImportError(f"Module '{name}' is not allowed")
+        return real_import(name, globals, locals, fromlist, level)
+
+    glb: Dict[str, Any] = {"__builtins__": safe_builtins}
+    lcl: Dict[str, Any] = {}
+
+    # ---------- Matplotlib headless ----------
+    plt = None
+    if return_plots:
+        try:
+            import matplotlib
+            matplotlib.use("Agg")
+            import matplotlib.pyplot as _plt
+            plt = _plt
+        except Exception:
+            plt = None
+
+    # ---------- Preload whitelisted mods ----------
+    preloads = [
+        "math","random","statistics","datetime","re","json","fractions","decimal",
+        "numpy","pandas","cmath",
+        "matplotlib","matplotlib.pyplot"
+    ]
+    for mod in preloads:
+        try:
+            if (mod in ALLOWED) or (mod.split(".")[0] in ALLOWED):
+                glb[mod.split(".")[-1]] = importlib.import_module(mod)
+        except Exception:
+            pass
+
+    # включаем безопасный импорт
+    safe_builtins["__import__"] = _safe_import
+
+    # ---------- Execute ----------
+    out_buf, err_buf = io.StringIO(), io.StringIO()
+    status = "error"
+    result_repr: Optional[str] = None
+    plots: List[Dict[str, Any]] = []
+    dataframes: List[Dict[str, Any]] = []
+
+    try:
+        with contextlib.redirect_stdout(out_buf), contextlib.redirect_stderr(err_buf):
+            exec(code, glb, lcl)
+            status = "success"
+
+            # вернём repr результата, если есть _ или result
+            if "_" in lcl:
+                result_repr = repr(lcl["_"])
+            elif "result" in lcl:
+                result_repr = repr(lcl["result"])
+
+            # графики
+            if plt is not None and return_plots:
+                fig_nums = plt.get_fignums()[: int(limits["max_plots"])]
+                for num in fig_nums:
+                    fig = plt.figure(num)
+                    b64 = _b64_png(fig, dpi=int(limits["plot_dpi"]))
+                    plots.append({"data_base64": b64, "format": "png"})
+                plt.close("all")
+
+            # DataFrame’ы
+            if return_dfs:
+                try:
+                    import pandas as _pd
+                    for name, val in list(lcl.items()):
+                        if isinstance(val, _pd.DataFrame):
+                            if len(dataframes) >= int(limits["max_dataframes"]):
+                                break
+                            head, shape, dtypes = _clip_df(
+                                val,
+                                max_rows=int(limits["max_df_rows"]),
+                                max_cols=int(limits["max_df_cols"]),
+                            )
+                            dataframes.append({
+                                "name": str(name),
+                                "head": head,
+                                "shape": shape,
+                                "dtypes": dtypes,
+                            })
+                except Exception:
+                    pass
+
+    except Exception:
+        status = "error"
+        print(traceback.format_exc(), file=err_buf)
+    finally:
+        try:
+            shutil.rmtree(workdir, ignore_errors=True)
+        except Exception:
+            pass
+
+    queue.put({
+        "status": status,
+        "stdout": out_buf.getvalue(),
+        "stderr": err_buf.getvalue(),
+        "result_repr": result_repr,
+        "plots": plots,
+        "dataframes": dataframes,
+    })
+
+# ====================== HOST PROCESS ======================
+
+def run_python_in_subprocess(req: CodeRunRequest) -> CodeRunResult:
+    exec_id = str(uuid.uuid4())
+    ctx = mp.get_context("spawn")
+    q: mp.Queue = ctx.Queue()
+
+    payload = {
+        "code": req.code,
+        "limits": req.limits.model_dump(),
+        "allowed": list(req.allowed_modules),
+        "return_plots": bool(req.return_plots),
+        "return_dfs": bool(req.return_dataframes),
+    }
+
+    p = ctx.Process(target=_child_exec, args=(payload, q), daemon=True)
+    p.start()
+    p.join(req.limits.timeout_seconds)
+
+    status = "timeout"
+    stdout = ""
+    stderr = "Timed out."
+    result_repr = None
+    plots: List[PlotArtifact] = []
+    dataframes: List[DataFrameArtifact] = []
+
+    if p.is_alive():
+        p.terminate()
+        p.join(1)
+    else:
+        try:
+            msg = q.get_nowait()
+            status = msg.get("status", "error")
+            stdout = (msg.get("stdout") or "")[: req.limits.max_stdout_chars]
+            stderr = (msg.get("stderr") or "")[: req.limits.max_stderr_chars]
+            result_repr = msg.get("result_repr")
+            plots = [PlotArtifact(**p_) for p_ in msg.get("plots", [])]
+            dataframes = [DataFrameArtifact(**d_) for d_ in msg.get("dataframes", [])]
+        except Exception as e:
+            status = "error"
+            stderr = f"Worker crashed: {e}"
+
+    return CodeRunResult(
+        execution_id=exec_id,
+        status=status,
+        stdout=stdout,
+        stderr=stderr,
+        result_repr=result_repr,
+        plots=plots,
+        dataframes=dataframes,
+        env=_env_info(),
+    )
+
+# ====================== LangChain TOOL ======================
+
+@tool
+def safe_code_run(code:str) -> str:
+    """
+    Safely execute Python code in an isolated subprocess with security restrictions.
+    
+    IMPORTANT - To see output, you MUST:
+    - Use print() statements for output
+    - Assign final result to variable 'result' or '_' 
+    - Save data to variables for DataFrame/plot capture
+    
+    Examples:
+    ✅ Good:
+    result = 2 + 2
+    print(f"Answer: {result}")
+    
+    ✅ Good:
+    import numpy as np
+    arr = np.array([1, 2, 3])
+    print(arr.mean())
+    
+    ✅ Good:
+    import pandas as pd
+    df = pd.DataFrame({'x': [1, 2], 'y': [3, 4]})
+    print(df)
+    result = df.sum()
+    
+    ❌ Bad (no output):
+    2 + 2  # This won't show anything
+    
+    Security features:
+    - Whitelisted imports only (numpy, pandas, matplotlib, etc.)
+    - Read-only file access within sandbox
+    - Network disabled
+    - Memory/CPU limits
+    - Timeout protection
+    
+    Returns JSON with: status, stdout, stderr, result_repr, plots, dataframes, env info
+    """
+
+    # упаковываем запрос в JSON
+    req = CodeRunRequest(
+        code=code,
+        # для первого запуска дайте запас
+        limits=dict(timeout_seconds=35)  # или 45
+    ).model_dump_json()
+
+    res = run_python_in_subprocess(CodeRunRequest.model_validate_json(req))
+    return res.model_dump_json()

src/tools/tools.py

@@ -0,0 +1,883 @@
+from __future__ import annotations
+
+import os
+import json
+import base64
+import tldextract
+import tempfile
+from urllib.parse import urlparse
+import io
+import pandas as pd
+from typing import List, Optional, Dict, Any
+from datetime import datetime
+from PIL import Image, ImageStat, ExifTags
+import google.generativeai as genai
+from langchain_core.tools import tool
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import ArxivLoader
+from langchain_community.document_loaders import WikipediaLoader
+from PIL import ImageDraw, ImageFont, ImageEnhance, ImageFilter
+from utils.image_processing import *
+
+def _exif_dict(img: Image.Image) -> dict:
+    try:
+        exif = img._getexif() or {}
+        out = {}
+        for k, v in exif.items():
+            tag = ExifTags.TAGS.get(k, str(k))
+            out[tag] = v if isinstance(v, (int, float, str)) else str(v)
+        return out
+    except Exception:
+        return {}
+
+def _clip(text: str | None, n: int) -> str:
+    """Утилита: безопасно обрезаем длинные сниппеты."""
+    if not text:
+        return ""
+    text = text.strip()
+    return (text[: n - 1] + "…") if len(text) > n else text
+
+
+def _parse_dt(v) -> Optional[str]:
+    """[ИЗМЕНЕНИЕ] Приводим даты к ISO-строке, если возможно."""
+    try:
+        if isinstance(v, datetime):
+            return v.isoformat()
+        if isinstance(v, str) and v:
+            
+            return v
+    except Exception:
+        pass
+    return None
+
+def _read_text_best_effort(path: str, max_chars: int) -> tuple[str, str]:
+    # пробуем utf-8 → fallback latin-1 (без chardet)
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            s = f.read()
+            return s[:max_chars], "utf-8"
+    except Exception:
+        with open(path, "r", encoding="latin-1", errors="replace") as f:
+            s = f.read()
+            return s[:max_chars], "latin-1"
+
+# ИСПРАВЛЕНИЕ 3: Улучшить preprocess_files с более точным определением типов
+def preprocess_files(files: List[str]) -> Dict[str, Dict[str, Any]]:
+    """Анализирует файлы и возвращает их метаданные"""
+    file_info = {}
+    
+    for file_path in files:
+        if not os.path.exists(file_path):
+            print(f"Warning: File {file_path} not found")
+            continue
+            
+        file_ext = os.path.splitext(file_path)[1].lower()
+        file_size = os.path.getsize(file_path)
+        
+        info = {
+            "path": file_path,
+            "extension": file_ext,
+            "size": file_size,
+            "type": None,
+            "suggested_tool": None,  # ✅ ДОБАВЛЕНО: подсказка для reasoning
+            "preview": None
+        }
+        
+        # ✅ УЛУЧШЕНО: Более точное определение типов и инструментов
+        if file_ext in ['.csv']:
+            info["type"] = "table"
+            info["suggested_tool"] = "analyze_csv_file"
+        elif file_ext in ['.xlsx', '.xls']:
+            info["type"] = "excel"
+            info["suggested_tool"] = "analyze_excel_file"
+        elif file_ext in ['.pdf']:
+            info["type"] = "document"
+            info["suggested_tool"] = "analyze_pdf_file"
+        elif file_ext in ['.docx', '.doc']:
+            info["type"] = "document"
+            info["suggested_tool"] = "analyze_docx_file"
+        elif file_ext in ['.txt', '.md']:
+            info["type"] = "text"
+            info["suggested_tool"] = "analyze_txt_file"
+        elif file_ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
+            info["type"] = "image"
+            info["suggested_tool"] = "if its about image itself : analyze_image_file, if its aboutrt content or qa : vision_qa_gemma ONLY"
+        else:
+            info["type"] = "unknown"
+            info["suggested_tool"] = "analyze_txt_file (fallback)"
+            
+        # Безопасное превью для небольших текстовых файлов
+        if file_ext == '.txt' and file_size < 1000:
+            try:
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+                    info["preview"] = content[:200] + "..." if len(content) > 200 else content
+            except Exception as e:
+                info["preview"] = f"Error reading file: {e}"
+        
+        file_info[file_path] = info
+    
+    return file_info
+
+#----------------------------------------------WEB BROWSING TOOLS------------------------------------------------#
+
+
+#WIKIPEDIA SEARCH TOOL
+
+@tool
+def wiki_search(
+    query: str,
+    max_results: int = 3,
+    language: str = "en",
+    content_chars_max: int = 5000,
+    snippet_chars: int = 400,
+) -> str:
+    """
+    Search Wikipedia using LangChain's WikipediaLoader.
+    Returns a JSON string:
+    {
+      "query": "...",
+      "language": "en",
+      "items": [
+        {
+          "url": "https://en.wikipedia.org/wiki/...",
+          "title": "Title",
+          "snippet": "First N chars of page content",
+          "page_content": "...(clipped to content_chars_max)..."
+        }
+      ]
+    }
+    """
+    try:
+        docs = WikipediaLoader(
+            query=query,
+            load_max_docs=max_results,
+            lang=language,
+            doc_content_chars_max=content_chars_max,
+        ).load()
+
+        items: List[dict] = []
+        seen_urls = set()
+
+        for d in docs:
+            url = d.metadata.get("source") or ""
+            if not url or url in seen_urls:
+                continue
+            seen_urls.add(url)
+
+            title = d.metadata.get("title") or ""
+            page_content = d.page_content or ""
+            snippet = _clip(page_content, snippet_chars)
+
+            items.append(
+                {
+                    "url": url,
+                    "title": title,
+                    "snippet": snippet,
+                    "page_content": page_content,  # уже ограничен doc_content_chars_max
+                }
+            )
+
+        payload = {
+            "query": query,
+            "language": language,
+            "items": items,
+        }
+        return json.dumps(payload)
+
+    except Exception as e:
+        return json.dumps({"error": str(e), "query": query, "language": language})
+
+#TAVILY WEB SEARCH TOOL
+
+def _domain(url: str) -> str:
+    """Утилита: вытаскиваем домен в виде 'site.tld' (без поддоменов)."""
+    ext = tldextract.extract(url)
+    return ".".join([p for p in (ext.domain, ext.suffix) if p])
+
+@tool
+def web_search(
+    query: str,
+    max_results: int = 5,                          # [ИЗМЕНЕНИЕ] параметризуем число результатов (было зашито 3)
+    unique_domains: int = 5,                        # [ИЗМЕНЕНИЕ] хотим максимум N разных доменов (борьба с дубликатами)
+    snippet_chars: int = 400,                       # [ИЗМЕНЕНИЕ] ограничиваем длину сниппета
+    include_domains: Optional[List[str]] = None,    # [ИЗМЕНЕНИЕ] вайтлист доменов
+    exclude_domains: Optional[List[str]] = None,    # [ИЗМЕНЕНИЕ] блэклист доменов
+) -> str:
+    """
+    Structured web search via Tavily.
+
+    Возвращает JSON-строку такого вида:
+    {
+      "query": "...",
+      "provider": "tavily",
+      "items": [
+        {
+          "url": "...",
+          "title": "...",
+          "snippet": "...",
+          "published": "2024-05-01T10:00:00Z",   # если Tavily отдал
+          "source": "example.com"                # домен
+        }
+      ]
+    }
+    """
+    # [ИЗМЕНЕНИЕ] раньше возвращалась сыровая строка с разметкой; теперь — строгий JSON для удобства парсинга
+    try:
+        # [ИЗМЕНЕНИЕ] используем официальный LangChain-тул, но берём больше (max_results), чтобы потом отфильтровать домены
+        raw_results = TavilySearchResults(max_results=max_results).invoke(query)
+
+        items: List[dict] = []
+        seen_urls: set[str] = set()
+        seen_domains: set[str] = set()
+
+        inc = set(include_domains or [])          # [ИЗМЕНЕНИЕ] поддержка фильтров доменов (whitelist)
+        exc = set(exclude_domains or [])          # [ИЗМЕНЕНИЕ] поддержка фильтров доменов (blacklist)
+
+        for r in raw_results:
+            url = (r.get("url") or "").strip()
+            if not url:
+                continue
+
+            dom = _domain(url)
+
+            # [ИЗМЕНЕНИЕ] применяем include/exclude-фильтры доменов
+            if inc and dom not in inc:
+                continue
+            if dom in exc:
+                continue
+
+            # [ИЗМЕНЕНИЕ] дедупликация ссылок
+            if url in seen_urls:
+                continue
+
+            # [ИЗМЕНЕНИЕ] ограничиваем разнообразие доменов (часто Tavily даёт много результатов с одного сайта)
+            if unique_domains > 0 and dom in seen_domains:
+                # если домен уже встречался и лимит по доменам строгий, пропускаем
+                pass
+            else:
+                # засчитываем домен как использованный
+                seen_domains.add(dom)
+
+            title = (r.get("title") or "").strip()
+            content = r.get("content") or r.get("snippet") or ""
+            snippet = _clip(content, snippet_chars)  # [ИЗМЕНЕНИЕ] делаем аккуратный сниппет
+            published = r.get("published_date") or r.get("created_at")  # [ИЗМЕНЕНИЕ] пытаемся достать дату
+
+            items.append(
+                {
+                    "url": url,
+                    "title": title,
+                    "snippet": snippet,
+                    "published": published,
+                    "source": dom,  # [ИЗМЕНЕНИЕ] явный домен — удобно для форматтера/критика
+                }
+            )
+            seen_urls.add(url)
+
+            # [ИЗМЕНЕНИЕ] если мы уже собрали нужное число результатов после фильтрации — выходим
+            if len(items) >= max_results:
+                break
+
+        payload = {
+            "query": query,
+            "provider": "tavily",
+            "items": items,
+        }
+        return json.dumps(payload)
+
+    except Exception as e:
+        # [ИЗМЕНЕНИЕ] единый формат ошибок в JSON — проще логировать и обрабатывать в агенте
+        return json.dumps({"error": str(e), "query": query, "provider": "tavily"})
+
+
+#ARXIV SEARCH TOOL
+
+@tool
+def arxiv_search(
+    query: str,
+    max_results: int = 5,                
+) -> str:
+    """
+    Поиск по arXiv через LangChain ArxivLoader.
+
+    [ИЗМЕНЕНИЕ] Возвращает **строгий JSON** вида:
+    {
+      "query": "...",
+      "provider": "arxiv",
+      "items": [
+        {
+          "title": "...",
+          "authors": ["A. Author","B. Author"],
+          "published": "YYYY-MM-DDTHH:MM:SS",
+          "journal_ref": "…",                 # если есть
+          "comment": "…",                     # если есть
+          "snippet": "first N chars of summary",
+          "summary": "… (может быть клипнут ArxivLoader'ом по умолчанию)"
+        }
+      ]
+    }
+    """
+    try:
+        docs = ArxivLoader(
+            query=query,
+            load_max_docs=max_results,
+        ).load()
+
+        items: List[dict] = []
+
+        for d in docs:
+            md = d.metadata or {}
+
+            title = md.get("Title") or md.get("title") or ""
+            authors = md.get("Authors") or md.get("authors") or []
+            if isinstance(authors, str):
+                authors = [a.strip() for a in authors.split(",") if a.strip()]
+
+            published = _parse_dt(md.get("Published") or md.get("published"))
+            summary = d.page_content or ""
+
+            items.append(
+                {
+                    "title": title,
+                    "authors": authors,
+                    "published": published,
+                    "summary": summary,
+                }
+            )
+
+            if len(items) >= max_results:
+                break
+
+        payload = {
+            "query": query,
+            "provider": "arxiv",
+            "items": items,
+        }
+        return json.dumps(payload)
+
+    except Exception as e:
+        return json.dumps({"error": str(e), "query": query, "provider": "arxiv"})
+    
+
+
+#----------------------------------------------MATH TOOLS------------------------------------------------#
+    
+@tool
+def add(a: float, b: float) -> float:
+    """Returns the sum of two numbers.
+        Example: add(2, 3) -> 5
+    """
+    return a + b
+
+@tool
+def subtract(a: float, b: float) -> float:
+    """Returns the difference of two numbers.
+        Example: subtract(5, 3) -> 2
+    """
+    return a - b
+
+@tool
+def multiply(a: float, b: float) -> float:
+    """Returns the product of two numbers.
+        Example: multiply(2, 3) -> 6
+    """
+    return a * b
+
+@tool
+def divide(a: float, b: float) -> float:
+    """Returns the quotient of two numbers.
+        Example: divide(6, 3) -> 2
+    """
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def power(a: float, b: float) -> float:
+    """Returns a raised to the power of b.
+        Example: power(2, 3) -> 8
+    """
+    return a ** b
+
+
+#----------------------------------------------FILE PROCESSING TOOLS------------------------------------------------#
+
+@tool
+def analyze_csv_file(file_path: str, preview_rows: int = 20) -> str:
+    """
+    Analyze a CSV file: returns JSON with {kind, path, shape, columns, head, numeric_summary}.
+    - preview_rows: number of rows for preview (head)
+    """
+    if not os.path.exists(file_path):
+        return json.dumps({"error": "file not found", "path": file_path})
+    try:
+        df = pd.read_csv(file_path)
+        head = df.head(preview_rows).to_dict(orient="records")
+        numeric = df.select_dtypes("number").describe().to_dict()
+        payload = {
+            "kind": "csv",
+            "path": file_path,
+            "shape": list(df.shape),
+            "columns": list(map(str, df.columns)),
+            "head": head,
+            "numeric_summary": numeric,  # {col: {count, mean, std, ...}}
+        }
+        return json.dumps(payload)
+    except Exception as e:
+        return json.dumps({"error": str(e), "path": file_path})
+    
+@tool
+def analyze_excel_file(file_path: str, sheet: int | str | None = None, preview_rows: int = 20, list_sheets: bool = True) -> str:
+    """
+    Analyze an Excel file: {kind, path, sheets?, active_sheet, shape, columns, head}.
+    - sheet: sheet index or name (None -> first sheet)
+    - list_sheets: include all sheet names
+    """
+    if not os.path.exists(file_path):
+        return json.dumps({"error": "file not found", "path": file_path})
+    try:
+        xls = pd.ExcelFile(file_path)
+        target = sheet if sheet is not None else 0
+        df = pd.read_excel(xls, sheet_name=target)
+        head = df.head(preview_rows).to_dict(orient="records")
+        payload = {
+            "kind": "excel",
+            "path": file_path,
+            "active_sheet": target if isinstance(target, int) else str(target),
+            "shape": list(df.shape),
+            "columns": list(map(str, df.columns)),
+            "head": head,
+        }
+        if list_sheets:
+            payload["sheets"] = list(map(str, xls.sheet_names))
+        return json.dumps(payload)
+    except Exception as e:
+        return json.dumps({"error": str(e), "path": file_path})
+    
+@tool
+def analyze_docx_file(file_path: str, max_chars: int = 20000, join_with: str = "\n") -> str:
+    """
+    Extract text from DOCX: {kind, path, paragraphs, text[:max_chars]}.
+    """
+    if not os.path.exists(file_path):
+        return json.dumps({"error": "file not found", "path": file_path})
+    try:
+        from docx import Document  # pip install python-docx
+    except Exception as e:
+        return json.dumps({"error": f"python-docx not installed: {e}"})
+    try:
+        doc = Document(file_path)
+        paras = [p.text for p in doc.paragraphs if p.text is not None]
+        text = join_with.join(paras)
+        payload = {
+            "kind": "docx",
+            "path": file_path,
+            "paragraphs": len(paras),
+            "text": text[:max_chars],
+            "length": len(text),
+        }
+        return json.dumps(payload)
+    except Exception as e:
+        return json.dumps({"error": str(e), "path": file_path})
+    
+
+@tool
+def analyze_txt_file(file_path: str, max_chars: int = 20000) -> str:
+    """
+    Read plain text: {kind, path, encoding_guess, text[:max_chars], length}.
+    """
+    if not os.path.exists(file_path):
+        return json.dumps({"error": "file not found", "path": file_path})
+    try:
+        text, enc = _read_text_best_effort(file_path, max_chars=max_chars)
+        payload = {
+            "kind": "txt",
+            "path": file_path,
+            "encoding_guess": enc,
+            "text": text,
+            "length": os.path.getsize(file_path),
+        }
+        return json.dumps(payload)
+    except Exception as e:
+        return json.dumps({"error": str(e), "path": file_path})
+    
+@tool
+def analyze_pdf_file(file_path: str, max_chars: int = 20000) -> str:
+    """
+    Extract text & page count from PDF: {kind, path, pages, text[:max_chars]}.
+    Uses pdfminer.six for text and page counting.
+    """
+    if not os.path.exists(file_path):
+        return json.dumps({"error": "file not found", "path": file_path})
+    try:
+        # text
+        from pdfminer.high_level import extract_text
+        text = extract_text(file_path) or ""
+        # pages
+        from pdfminer.pdfpage import PDFPage
+        with open(file_path, "rb") as f:
+            pages = sum(1 for _ in PDFPage.get_pages(f))
+        payload = {
+            "kind": "pdf",
+            "path": file_path,
+            "pages": pages,
+            "text": text[:max_chars],
+            "length": len(text),
+        }
+        return json.dumps(payload)
+    except Exception as e:
+        return json.dumps({"error": str(e), "path": file_path})
+    
+
+#----------------------------------------------IMAGE PROCESSING TOOLS------------------------------------------------#
+
+@tool
+def analyze_image_file(file_path: str, ocr: bool = False, lang: Optional[str] = None, max_ocr_chars: int = 10000) -> str:
+    """
+    Analyze image: {kind, path, format, mode, size, mean_brightness, exif?, ocr_text?}.
+    - ocr: optional Tesseract OCR (pip install pytesseract + tesseract)
+    """
+    if not os.path.exists(file_path):
+        return json.dumps({"error": "file not found", "path": file_path})
+    try:
+        img = Image.open(file_path)
+        stat = ImageStat.Stat(img.convert("L"))
+        mean_brightness = float(stat.mean[0])  # 0..255
+        payload = {
+            "kind": "image",
+            "path": file_path,
+            "format": img.format,
+            "mode": img.mode,
+            "size": list(img.size),  # [width, height]
+            "mean_brightness": mean_brightness,
+        }
+        exif = _exif_dict(img)
+        if exif:
+            payload["exif"] = exif
+
+        if ocr:
+            try:
+                import pytesseract
+                conf = {}
+                if lang:
+                    conf["lang"] = lang
+                text = pytesseract.image_to_string(img, **conf) or ""
+                payload["ocr_text"] = text[:max_ocr_chars]
+                payload["ocr_length"] = len(text)
+            except Exception as e:
+                payload["ocr_error"] = str(e)
+
+        return json.dumps(payload)
+    except Exception as e:
+        return json.dumps({"error": str(e), "path": file_path})
+    
+
+
+
+# ------------------------- helpers for QA image TOOL -------------------------
+
+def _configure():
+    api_key = os.getenv("GOOGLE_API_KEY") or os.getenv("GENAI_API_KEY")
+    if not api_key:
+        raise RuntimeError("Missing GOOGLE_API_KEY (or GENAI_API_KEY) in environment")
+    genai.configure(api_key=api_key)
+
+def _image_bytes_to_part(img_bytes: bytes, mime: str = "image/png") -> Dict[str, Any]:
+    # формат, который понимает genai.generate_content
+    return {"mime_type": mime, "data": base64.b64encode(img_bytes).decode("utf-8")}
+
+def _ensure_png_bytes(img: Image.Image, max_pixels: int = 25_000_000) -> bytes:
+    # мягко даунскейлим огромные изображения (защита от "image bomb")
+    w, h = img.size
+    if w * h > max_pixels:
+        scale = (max_pixels / (w * h)) ** 0.5
+        nw, nh = max(1, int(w * scale)), max(1, int(h * scale))
+        img = img.resize((nw, nh), Image.LANCZOS)
+
+    # приводим к PNG (надёжно для SDK)
+    buf = io.BytesIO()
+    img.save(buf, format="PNG", optimize=True)
+    return buf.getvalue()
+
+def _load_image_as_png_bytes_from_path(path: str) -> bytes:
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"Image not found: {path}")
+    img = Image.open(path)
+    return _ensure_png_bytes(img)
+
+def _load_image_as_png_bytes_from_b64(b64: str) -> bytes:
+    raw = base64.b64decode(b64, validate=True)
+    img = Image.open(io.BytesIO(raw))
+    return _ensure_png_bytes(img)
+
+def _clean_json_text(s: str) -> str:
+    # вычищаем обёртки ```json ... ``` и забираем объект { ... }
+    s = s.strip()
+    if s.startswith("```"):
+        s = s.strip("`").replace("json", "", 1).strip()
+    # вырезать по внешним фигурным скобкам
+    start = s.find("{")
+    end = s.rfind("}")
+    if start != -1 and end != -1 and end > start:
+        return s[start:end+1]
+    return s
+
+_SINGLE_IMAGE_QA_PROMPT = (
+    "You will be given ONE image and a user question about it.\n"
+    "Answer STRICTLY and CONCISELY based only on the image content.\n"
+    "If the image does not contain enough information to answer, reply 'not enough information'.\n"
+    "If the answer is numeric, include units if visible.\n"
+    "Return ONLY valid JSON with the schema:\n"
+    "{\"answer\": string}\n"
+)
+
+def _call_model(parts: List[Any], temperature: float) -> Dict[str, Any]:
+    MODEL_NAME = "gemma-3-27b-it"
+    model = genai.GenerativeModel(MODEL_NAME)
+    resp = model.generate_content(parts, generation_config={"temperature": temperature})
+    text = (resp.text or "").strip()
+
+    # пробуем сразу распарсить
+    try:
+        return json.loads(_clean_json_text(text))
+    except Exception:
+        # второй шанс: попросим модель вернуть строгий JSON
+        fixer = genai.GenerativeModel(MODEL_NAME)
+        fix_prompt = (
+            "Convert the following text into STRICT valid JSON matching schema {\"answer\": string}. "
+            "Return ONLY JSON, no extra text:\n" + text
+        )
+        fix_resp = fixer.generate_content([{"text": fix_prompt}])
+        return json.loads(_clean_json_text((fix_resp.text or "").strip()))
+
+# --------------------------- TOOL ---------------------------
+
+@tool
+def vision_qa_gemma(
+    question: str,
+    image_path: Optional[str] = None,
+    image_base64: Optional[str] = None,
+    temperature: float = 0.2,
+) -> str:
+    """
+    Vision QA with Google GenAI (Gemma/Gemini). Returns JSON: {"answer": "..."}.
+
+    Args:
+      question: user question about the image.
+      image_path: local file path to the image (PNG/JPG/...).
+      image_base64: base64-encoded image (if no path).
+      temperature: decoding temperature (default 0.2).
+
+    Exactly ONE of (image_path, image_base64) must be provided.
+    """
+    import json as _json
+    try:
+        _configure()
+        if bool(image_path) == bool(image_base64):
+            return _json.dumps({"error": "Provide exactly ONE of image_path or image_base64"})
+
+        if image_path:
+            img_bytes = _load_image_as_png_bytes_from_path(image_path)
+        else:
+            img_bytes = _load_image_as_png_bytes_from_b64(image_base64)
+
+        parts = [
+            {"text": _SINGLE_IMAGE_QA_PROMPT + "\nQuestion: " + question.strip()},
+            _image_bytes_to_part(img_bytes, "image/png"),
+        ]
+
+        data = _call_model(parts, temperature)
+        # финальная защита: оставляем только "answer"
+        answer = data["answer"] if isinstance(data, dict) and "answer" in data else None
+        if not isinstance(answer, str):
+            answer = str(answer) if answer is not None else "not enough information"
+
+        return _json.dumps({
+            "answer": answer,
+        })
+
+    except Exception as e:
+        return _json.dumps({"error": str(e)})
+
+
+#-------------------------------------------------------------- ADDITIONAL TOOLS -------------------------------------------------------------#
+@tool
+def draw_on_image(
+    image_base64: str, drawing_type: str, params: Dict[str, Any]
+) -> Dict[str, Any]:
+    """
+    Draw shapes (rectangle, circle, line) or text onto an image.
+    Args:
+        image_base64 (str): Base64 encoded input image
+        drawing_type (str): Drawing type
+        params (Dict[str, Any]): Drawing parameters
+    Returns:
+        Dictionary with result image (base64)
+    """
+    try:
+        img = decode_image(image_base64)
+        draw = ImageDraw.Draw(img)
+        color = params.get("color", "red")
+
+        if drawing_type == "rectangle":
+            draw.rectangle(
+                [params["left"], params["top"], params["right"], params["bottom"]],
+                outline=color,
+                width=params.get("width", 2),
+            )
+        elif drawing_type == "circle":
+            x, y, r = params["x"], params["y"], params["radius"]
+            draw.ellipse(
+                (x - r, y - r, x + r, y + r),
+                outline=color,
+                width=params.get("width", 2),
+            )
+        elif drawing_type == "line":
+            draw.line(
+                (
+                    params["start_x"],
+                    params["start_y"],
+                    params["end_x"],
+                    params["end_y"],
+                ),
+                fill=color,
+                width=params.get("width", 2),
+            )
+        elif drawing_type == "text":
+            font_size = params.get("font_size", 20)
+            try:
+                font = ImageFont.truetype("arial.ttf", font_size)
+            except IOError:
+                font = ImageFont.load_default()
+            draw.text(
+                (params["x"], params["y"]),
+                params.get("text", "Text"),
+                fill=color,
+                font=font,
+            )
+        else:
+            return {"error": f"Unknown drawing type: {drawing_type}"}
+
+        result_path = save_image(img)
+        result_base64 = encode_image(result_path)
+        return {"result_image": result_base64}
+
+    except Exception as e:
+        return {"error": str(e)}
+    
+@tool
+def transform_image(
+    image_base64: str, operation: str, params: Optional[Dict[str, Any]] = None
+) -> Dict[str, Any]:
+    """
+    Apply transformations: resize, rotate, crop, flip, brightness, contrast, blur, sharpen, grayscale.
+    Args:
+        image_base64 (str): Base64 encoded input image
+        operation (str): Transformation operation
+        params (Dict[str, Any], optional): Parameters for the operation
+    Returns:
+        Dictionary with transformed image (base64)
+    """
+    try:
+        img = decode_image(image_base64)
+        params = params or {}
+
+        if operation == "resize":
+            img = img.resize(
+                (
+                    params.get("width", img.width // 2),
+                    params.get("height", img.height // 2),
+                )
+            )
+        elif operation == "rotate":
+            img = img.rotate(params.get("angle", 90), expand=True)
+        elif operation == "crop":
+            img = img.crop(
+                (
+                    params.get("left", 0),
+                    params.get("top", 0),
+                    params.get("right", img.width),
+                    params.get("bottom", img.height),
+                )
+            )
+        elif operation == "flip":
+            if params.get("direction", "horizontal") == "horizontal":
+                img = img.transpose(Image.FLIP_LEFT_RIGHT)
+            else:
+                img = img.transpose(Image.FLIP_TOP_BOTTOM)
+        elif operation == "adjust_brightness":
+            img = ImageEnhance.Brightness(img).enhance(params.get("factor", 1.5))
+        elif operation == "adjust_contrast":
+            img = ImageEnhance.Contrast(img).enhance(params.get("factor", 1.5))
+        elif operation == "blur":
+            img = img.filter(ImageFilter.GaussianBlur(params.get("radius", 2)))
+        elif operation == "sharpen":
+            img = img.filter(ImageFilter.SHARPEN)
+        elif operation == "grayscale":
+            img = img.convert("L")
+        else:
+            return {"error": f"Unknown operation: {operation}"}
+
+        result_path = save_image(img)
+        result_base64 = encode_image(result_path)
+        return {"transformed_image": result_base64}
+
+    except Exception as e:
+        return {"error": str(e)}
+
+@tool
+def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
+    """
+    Save content to a file and return the path.
+    Args:
+        content (str): the content to save to the file
+        filename (str, optional): the name of the file. If not provided, a random name file will be created.
+    """
+    temp_dir = tempfile.gettempdir()
+    if filename is None:
+        temp_file = tempfile.NamedTemporaryFile(delete=False, dir=temp_dir)
+        filepath = temp_file.name
+    else:
+        filepath = os.path.join(temp_dir, filename)
+
+    with open(filepath, "w") as f:
+        f.write(content)
+
+    return f"File saved to {filepath}. You can read this file to process its contents."
+
+
+import requests
+
+@tool
+def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
+    """
+    Download a file from a URL and save it to a temporary location.
+    Args:
+        url (str): the URL of the file to download.
+        filename (str, optional): the name of the file. If not provided, a random name file will be created.
+    """
+    try:
+        # Parse URL to get filename if not provided
+        if not filename:
+            path = urlparse(url).path
+            filename = os.path.basename(path)
+            if not filename:
+                filename = f"downloaded_{uuid.uuid4().hex[:8]}"
+
+        # Create temporary file
+        temp_dir = tempfile.gettempdir()
+        filepath = os.path.join(temp_dir, filename)
+
+        # Download the file
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+
+        # Save the file
+        with open(filepath, "wb") as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+
+        return f"File downloaded to {filepath}. You can read this file to process its contents."
+    except Exception as e:
+        return f"Error downloading file: {str(e)}"

src/utils/code_run.py

@@ -0,0 +1,56 @@
+from typing import List, Optional, Literal, Dict, Any
+from pydantic import BaseModel, Field
+import platform, sys
+
+class Limits(BaseModel):
+    timeout_seconds: int = Field(12, ge=1, le=120)
+    max_stdout_chars: int = Field(10000, ge=256, le=200_000)
+    max_stderr_chars: int = Field(10000, ge=256, le=200_000)
+    max_plots: int = Field(4, ge=0, le=10)
+    max_dataframes: int = Field(3, ge=0, le=10)
+    max_df_rows: int = Field(20, ge=1, le=200)
+    max_df_cols: int = Field(20, ge=1, le=200)
+    plot_dpi: int = Field(120, ge=72, le=300)
+    max_pixels: int = Field(25_000_000, ge=1)  # если вдруг юзер генерит большие изображения
+
+class CodeRunRequest(BaseModel):
+    language: Literal["python"] = "python"
+    code: str
+    # Явный allowlist модулей (верхнеуровневые имена)
+    allowed_modules: List[str] = Field(
+        default_factory=lambda: [
+            "math","random","statistics","datetime","re","json","fractions","decimal",
+            "numpy","pandas","cmath","matplotlib","matplotlib.pyplot", "seaborn","sklearn","sklearn.datasets","sklearn.model_selection", "sympy"
+        ]
+    )
+    # Флаги, что возвращать
+    return_plots: bool = True
+    return_dataframes: bool = True
+    # Ограничения
+    limits: Limits = Field(default_factory=Limits)
+
+class PlotArtifact(BaseModel):
+    data_base64: str
+    format: Literal["png"] = "png"
+
+class DataFrameArtifact(BaseModel):
+    name: str
+    head: List[Dict[str, Any]]
+    shape: List[int]
+    dtypes: Dict[str, str]
+
+class EnvInfo(BaseModel):
+    python: str = Field(default_factory=lambda: sys.version.split()[0])
+    numpy: Optional[str] = None
+    pandas: Optional[str] = None
+    platform: str = Field(default_factory=platform.platform)
+
+class CodeRunResult(BaseModel):
+    execution_id: str
+    status: Literal["success","error","timeout"]
+    stdout: str = ""
+    stderr: str = ""
+    result_repr: Optional[str] = None
+    plots: List[PlotArtifact] = Field(default_factory=list)
+    dataframes: List[DataFrameArtifact] = Field(default_factory=list)
+    env: EnvInfo

src/utils/image_processing.py

@@ -0,0 +1,26 @@
+import os
+import io
+import base64
+import uuid
+from PIL import Image
+
+# Helper functions for image processing
+def encode_image(image_path: str) -> str:
+    """Convert an image file to base64 string."""
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
+
+
+def decode_image(base64_string: str) -> Image.Image:
+    """Convert a base64 string to a PIL Image."""
+    image_data = base64.b64decode(base64_string)
+    return Image.open(io.BytesIO(image_data))
+
+
+def save_image(image: Image.Image, directory: str = "image_outputs") -> str:
+    """Save a PIL Image to disk and return the path."""
+    os.makedirs(directory, exist_ok=True)
+    image_id = str(uuid.uuid4())
+    image_path = os.path.join(directory, f"{image_id}.png")
+    image.save(image_path)
+    return image_path

src/utils/utils.py

@@ -0,0 +1,100 @@
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
+from schemas import ComplexityLevel, ExecutionReport
+from prompts.prompts import COMPLEXITY_ASSESSOR_PROMPT
+from config import llm
+from state import AgentState
+
+def clean_message_history(messages):
+    """
+    Очищает историю сообщений от неполных циклов tool_calls/responses.
+    Удаляет AIMessage с tool_calls, если нет соответствующих ToolMessage.
+    """
+    cleaned_messages = []
+    i = 0
+    
+    while i < len(messages):
+        msg = messages[i]
+        
+        # Если это AIMessage с tool_calls
+        if hasattr(msg, 'tool_calls') and msg.tool_calls:
+            # Ищем соответствующие ToolMessage
+            tool_call_ids = {tc['id'] for tc in msg.tool_calls}
+            found_responses = set()
+            
+            # Проверяем следующие сообщения на наличие ответов
+            j = i + 1
+            while j < len(messages) and isinstance(messages[j], ToolMessage):
+                if messages[j].tool_call_id in tool_call_ids:
+                    found_responses.add(messages[j].tool_call_id)
+                j += 1
+            
+            # Если все tool_calls имеют ответы, добавляем весь блок
+            if found_responses == tool_call_ids:
+                # Добавляем AIMessage и все соответствующие ToolMessage
+                cleaned_messages.append(msg)
+                for k in range(i + 1, j):
+                    cleaned_messages.append(messages[k])
+                i = j
+            else:
+                # Пропускаем неполный блок
+                print(f"Removing incomplete tool call block: {tool_call_ids - found_responses}")
+                i = j
+        else:
+            # Обычное сообщение - добавляем
+            cleaned_messages.append(msg)
+            i += 1
+    
+    return cleaned_messages
+
+def format_final_answer(report: ExecutionReport, complexity: dict) -> str:
+    """Format the final answer based on complexity and report content."""
+    
+    if complexity.level == 'simple':
+        # For simple queries, just return the answer
+        return f"FINAL ANSWER: {report.final_answer}"
+    
+    # For complex queries, provide more detailed response
+    formatted = f"""FINAL ANSWER: {report.final_answer}
+
+SUMMARY:
+{report.query_summary}
+
+KEY FINDINGS:
+{chr(10).join(f"• {finding}" for finding in report.key_findings)}"""
+    
+    if report.data_sources:
+        formatted += f"""
+
+SOURCES:
+{chr(10).join(f"• {source}" for source in report.data_sources[:5])}"""  # Limit to 5 sources
+    
+    if report.limitations:
+        formatted += f"""
+
+LIMITATIONS:
+{chr(10).join(f"• {limitation}" for limitation in report.limitations)}"""
+    
+    return formatted
+
+
+def complexity_assessor(state: AgentState) -> AgentState:
+    """Assess query complexity and determine if planning is needed."""
+    print("=== COMPLEXITY ASSESSMENT ===")
+    
+    complexity_llm = llm.with_structured_output(ComplexityLevel)
+    
+    assessment_message = [
+        SystemMessage(content=COMPLEXITY_ASSESSOR_PROMPT.strip()),
+        HumanMessage(content=f"Query: {state['query']}")
+    ]
+    
+    assessment = complexity_llm.invoke(assessment_message)
+    
+    print(f"Complexity: {assessment.level}")
+    print(f"Needs planning: {assessment.needs_planning}")
+    print(f"Reasoning: {assessment.reasoning}")
+    
+    return {
+        "complexity_assessment": assessment,
+        "messages": state["messages"] + assessment_message
+    }

src/workflow_test.ipynb

@@ -0,0 +1,96 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "d:\\REGNUM_SPECTRARUM_updated\\.venv\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "from agent import build_workflow\n",
+    "from config import config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "graph = build_workflow()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "=== USER QUERY TRANSFERED TO AGENT ===\n",
+      "=== COMPLEXITY ASSESSMENT ===\n",
+      "Complexity: simple\n",
+      "Needs planning: False\n",
+      "Reasoning: This query is a straightforward arithmetic calculation that can be answered immediately without any tools or complex reasoning.\n",
+      "=== SIMPLE EXECUTION ===\n",
+      "=== GENERATING EXECUTION REPORT ===\n",
+      "Report generated - Confidence: high\n",
+      "Key findings: 2\n",
+      "Data sources: 0\n",
+      "query_summary='The user requested the result of the arithmetic expression 2 + 2 - 2 + 2.' approach_used='The query was evaluated using basic arithmetic operations, following the standard order of operations.' tools_executed=[] key_findings=['The expression was simplified step-by-step: 2 + 2 = 4, then 4 - 2 = 2, and finally 2 + 2 = 4.', 'The final result of the expression is 4.'] data_sources=[] assumptions_made=['The user intended to use standard arithmetic rules without any additional context or modifications.'] confidence_level='high' limitations=['The query was straightforward, and no complex tools or external data sources were required.', 'The execution context did not involve any ambiguity or alternative interpretations.'] final_answer='4'\n",
+      "=== ENHANCED ANSWER CRITIQUE ===\n",
+      "Quality Score: 8/10\n",
+      "Complete: True\n",
+      "Accurate: True\n",
+      "=== REPLAN DECISION ===\n",
+      "Iteration: 1/10\n",
+      "Quality score: 8\n",
+      "Needs replanning: False\n",
+      "Quality acceptable, ending execution\n"
+     ]
+    }
+   ],
+   "source": [
+    "result = graph.invoke({\"query\" : \"2+2-2+2?\", \"current_step\": 0, \"reasoning_done\": False, \"files\" : [], \"files_contents\" : {}, \"iteration_count\" : 0, \"max_iterations\" : 10, \"plan\" : None} , config = config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

test_folder/test.py

@@ -0,0 +1,11 @@
+import math
+
+def calculate_area(radius):
+    """Calculate the area of a circle given its radius."""
+    if radius < 0:
+        raise ValueError("Radius cannot be negative")
+    return math.pi * (radius ** 2)
+
+radius = 5
+area = calculate_area(radius)
+print(area)

test_folder/test_run.py

@@ -0,0 +1,14 @@
+
+
+# === From notebook execution cells ===
+workflow = system.invoke({"query" : "How many cumulative milliliters of fluid is in all the opaque-capped vials without stickers in the 114 version of the kit that was used for the PromethION long-read sequencing in the paper De Novo-Whole Genome Assembly of the Roborovski Dwarf Hamster (Phodopus roborovskii) Genome?", "current_step": 0, "reasoning_done": False, "files" : [], "files_contents" : {}, "iteration_count" : 0, "max_iterations" : 10, "plan" : None} , config = config)
+
+for message in workflow["messages"]:
+    message.pretty_print()
+
+print("\n=== FINAL ANSWER ===")
+
+workflow["final_answer"]
+
+workflow
+