KaiserShultz/Ankelodon_AI_Multi_task_agentic_system
1
1import os2from src.state import AgentState3from src.tools.tools import preprocess_files4from typing import Optional5from langgraph.prebuilt import ToolNode6 7from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, ToolMessage8 9from src.prompts.prompts import (10 SYSTEM_PROMPT_PLANNER,11 SYSTEM_EXECUTOR_PROMPT,12 COMPLEXITY_ASSESSOR_PROMPT,13 CRITIC_PROMPT,14)15 16from src.config import llm_reasoning, TOOLS, planner_llm, llm_with_tools, llm_deterministic, llm_criticist, llm_simple_executor, llm_simple_with_tools17from src.schemas import PlannerPlan, ComplexityLevel, CritiqueFeedback, ExecutionReport, ToolExecution18 19from src.utils.utils import (20 format_final_answer,21 clean_message_history,22 log_stage,23 log_key_values,24 display_plan,25 format_plan_overview,26)27 28def _build_planner_prompt(state: AgentState, extra_context: Optional[str] = None) -> str:29 tool_catalogue = ", ".join(sorted(tool.name for tool in TOOLS))30 file_paths = state.get("files", [])31 file_list = ", ".join(os.path.basename(path) for path in file_paths) if file_paths else "none provided"32 extra = extra_context.strip() if extra_context else "None"33 return SYSTEM_PROMPT_PLANNER.format(34 tool_catalogue=tool_catalogue,35 file_list=file_list,36 extra_context=extra,37 ).strip()38 39def query_input(state : AgentState) -> AgentState:40 log_stage("USER QUERY", icon="💡")41 42 files = state.get("files", [])43 if files:44 log_stage("FILE PREPARATION", subtitle=f"Processing {len(files)} file(s)", icon="📁")45 file_info = preprocess_files(files)46 47 for file_path, info in file_info.items():48 #print(f" - {file_path}: {info['type']} ({info['size']} bytes) -> {info['suggested_tool']}")49 log_key_values(50 [51 ("path", file_path),52 ("type", info["type"]),53 ("size", f"{info['size']} bytes"),54 ("suggested_tool", info["suggested_tool"]),55 ]56 )57 state["file_contents"] = file_info58 file_context = "\n\n=== AVAILABLE FILES FOR ANALYSIS ===\n"59 for file_path, info in file_info.items():60 filename = os.path.basename(file_path)61 file_context += f"File: {filename}\n"62 file_context += f" - Type: {info['type']}\n" 63 file_context += f" - Size: {info['size']} bytes\n"64 file_context += f" - Suggested tool: {info['suggested_tool']}\n"65 if info.get("preview"):66 file_context += f" - Preview: {info['preview']}\n"67 file_context += "\n"68 69 # Добавляем инструкции по работе с файлами70 file_context += "IMPORTANT: Use the suggested tools to analyze these files before processing their data.\n"71 file_context += "File paths are available in the agent state and can be passed directly to analysis tools.\n"72 73 else:74 log_key_values([("files", "none provided")])75 file_context = ""76 original_query = state.get("query", "")77 state["query"] = original_query + file_context78 return state79 80 81def planner(state : AgentState) -> AgentState:82 83 log_stage("PLANNING", icon="🧭")84 planner_prompt = _build_planner_prompt(state)85 86 sys_stack = [87 SystemMessage(content=planner_prompt),88 HumanMessage(content=state["query"]),89 ]90 plan: PlannerPlan = planner_llm.invoke(sys_stack)91 92 #print("=== GENERATED PLAN ===")93 display_plan(plan)94 return {95 "messages": state["messages"] + sys_stack,96 "plan": plan,97 "current_step": 0,98 "reasoning_done": False,99 }100 101 102def agent(state: AgentState) -> AgentState:103 104 105 current_step = state.get("current_step", 0)106 reasoning_done = state.get("reasoning_done", False)107 plan: Optional[PlannerPlan] = state.get("plan")108 previous_tool_results = state.get("previous_tool_results", {})109 110 #steps = state["plan"].steps111 112 if not plan or not hasattr(plan, 'steps'):113 log_stage("PLAN VALIDATION", subtitle="Planner returned no actionable steps", icon="⚠️")114 warning = AIMessage(content="No valid plan available. <FINAL_ANSWER>")115 return {116 "messages": state["messages"] + [warning],117 "reasoning_done": False,118 }119 120 steps = plan.steps121 122 total_steps = len(steps)123 124 if total_steps == 0:125 log_stage("PLAN VALIDATION", subtitle="Plan indicates direct answer", icon="ℹ️")126 direct = AIMessage(content="Plan has no steps; respond directly. <FINAL_ANSWER>")127 return {128 "messages": state["messages"] + [direct],129 "reasoning_done": False,130 }131 132 if current_step >= total_steps:133 log_stage("PLAN COMPLETE", subtitle="All steps executed", icon="✅")134 completion = AIMessage(content="All plan steps completed. <FINAL_ANSWER>")135 return {136 "messages": state["messages"] + [completion],137 "reasoning_done": False,138 }139 140 current_step_info = steps[current_step]141 142 log_stage(143 "EXECUTION",144 subtitle=f"Step {current_step + 1}/{total_steps}: {current_step_info.goal}",145 icon="🤖",146 )147 log_key_values(148 [149 ("step_id", current_step_info.id),150 ("tool", current_step_info.tool or "none"),151 ("expected", current_step_info.expected_result),152 ]153 )154 155 plan_overview = format_plan_overview(plan)156 tool_catalogue = ", ".join(sorted(tool.name for tool in TOOLS))157 file_contents = state.get("file_contents", {})158 file_list = ", ".join(file_contents.keys()) if file_contents else "none provided"159 160 # Добавляем информацию о предыдущих результатах (UPDATE)161 previous_results_context = ""162 if previous_tool_results:163 previous_results_context = f"\n\nPREVIOUS CALCULATION RESULTS:\n"164 for tool_call_id, result in previous_tool_results.items():165 previous_results_context += f"- {tool_call_id}: {result}\n"166 previous_results_context += "You can reference these results in your calculations.\n"167 168 169 system_message = SystemMessage(170 content=SYSTEM_EXECUTOR_PROMPT.format(171 plan_summary=plan.summary,172 plan_overview=plan_overview,173 current_step_id=current_step_info.id,174 step_goal=current_step_info.goal,175 step_tool=current_step_info.tool or "no tool (respond directly)",176 tool_catalogue=tool_catalogue,177 file_list=file_list,178 ).strip()179 )180 181 182 if not reasoning_done:183 184 log_stage("REASONING", subtitle=f"{current_step_info.id}", icon="🧠")185 #print(reasoning_response.content)186 187 file_context = ""188 file_contents = state.get("file_contents", {})189 if file_contents:190 file_context = "\n\nAVAILABLE FILES IN CURRENT SESSION:\n"191 for filepath, info in file_contents.items():192 filename = os.path.basename(filepath)193 file_context += f"- {filename}: {info['type']} file, suggested tool: {info['suggested_tool']}\n"194 file_context += f" Path: {filepath}\n"195 196 reasoning_prompt = f"""197 {SYSTEM_EXECUTOR_PROMPT}198 199 CURRENT TASK: You must perform reasoning for step {current_step + 1}.200 201 STEP INFO: {current_step_info}\n\n202 203 FILE CONTEXT: {file_contents}204 205 CRITICAL: You MUST output your reasoning in <REASONING> tags, but DO NOT call any tools yet.206 Explain what you need to do and why, then end your response.207 208 REASONING IS IMPERATIVE BEFORE ANY TOOL CALLS.209 FOR MORE COMPLEX UNDERSTANDING -> USE RESULTS AND INSIGHTS FROM PREVIOUS STEPS.210 """211 212 sys_msg = SystemMessage(content = reasoning_prompt)213 stack = [sys_msg] + state["messages"]214 215 step = llm_reasoning.invoke(stack)216 #print("=== REASONING STEP ===")217 #print(step.content)218 219 return {220 "messages" : state["messages"] + [step],221 "reasoning_done" : True222 }223 224 else:225 tool_prompt = f"""226 Now execute the tool for step {current_step + 1}.227 228 You have already done the reasoning. Now call the appropriate tool with the correct parameters.229 Available file paths: {list(state.get("file_contents", {}).keys())}\n230 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!231 AVAILABLE TOOLS: {', '.join([tool.name for tool in TOOLS])}232 """ 233 234 sys_msg = SystemMessage(content=tool_prompt)235 stack = [sys_msg] + state["messages"] # Берем последние сообщения включая reasoning236 237 # Используем модель С инструментами для выполнения238 step = llm_with_tools.invoke(stack)239 print("=== TOOL EXECUTION ===")240 #print(step)241 print(f"Tool calls: {step.tool_calls}")242 243 return {244 "messages": state["messages"] + [step],245 "current_step": current_step + 1 if step.tool_calls else current_step,246 "reasoning_done": False # Сбрасываем для следующего шага247 }248 249def should_continue(state : AgentState) -> bool:250 251 last_message = state["messages"][-1]252 #print(f"=== LAST MESSAGE WAS: {last_message} ===")253 reasoning_done = state.get("reasoning_done", False)254 plan = state.get("plan", None)255 current_step = state.get("current_step", 0)256 257 print(f"=== SHOULD_CONTINUE DEBUG ===")258 print(f"Current step: {current_step}")259 print(f"Plan steps: {len(plan.steps) if plan else 0}")260 print(f"Reasoning done: {reasoning_done}")261 print(f"Last message type: {type(last_message).__name__}")262 263 #ПРИОРИТЕТ 1: Если есть tool_calls - выполняем их264 if hasattr(last_message, "tool_calls") and last_message.tool_calls:265 return "tools"266 267 # ПРИОРИТЕТ 2: Явный сигнал завершения268 if hasattr(last_message, "content") and "<FINAL_ANSWER>" in last_message.content:269 return "final_answer"270 271 # ПРИОРИТЕТ 3: Логика reasoning/execution272 if not reasoning_done and hasattr(last_message, 'content') and "<REASONING>" in last_message.content:273 # Reasoning выполнен, но инструменты еще не вызваны274 return "agent"275 elif reasoning_done:276 # Reasoning выполнен, теперь нужно вызвать инструменты277 return "agent"278 elif not reasoning_done:279 # Нужно сделать reasoning280 return "agent"281 282 # ПРИОРИТЕТ 4: Проверяем завершение плана (только если нет активных tool_calls)283 if plan and current_step >= len(plan.steps):284 return "final_answer"285 286 # По умолчанию продолжаем выполнение287 return "agent"288 289# 6. Добавить отладочную информацию в TOOL_NODE290class DebuggingToolNode(ToolNode):291 def __init__(self, tools):292 super().__init__(tools)293 294 def __call__(self, state):295 print("=== TOOL EXECUTION STARTED ===")296 result = super().__call__(state)297 print("=== TOOL EXECUTION COMPLETED ===")298 return result299 300 301def enhanced_finalizer(state: AgentState) -> AgentState:302 """Generate comprehensive execution report for critic evaluation."""303 print("=== GENERATING EXECUTION REPORT ===")304 305 # Extract tool execution information306 tools_executed = []307 data_sources = []308 309 for msg in state["messages"]:310 if hasattr(msg, 'tool_calls') and msg.tool_calls:311 for tool_call in msg.tool_calls:312 tools_executed.append(ToolExecution(313 tool_name=tool_call['name'],314 arguments=str(tool_call['args']),315 call_id=tool_call['id']316 ))317 318 # Extract data sources from tool results319 if hasattr(msg, 'content') and isinstance(msg.content, str):320 # Look for URLs, file names, or other sources321 import re322 urls = re.findall(r'https?://[^\s]+', msg.content)323 data_sources.extend(urls)324 325 # Get plan information if available326 plan = state.get("plan")327 approach_used = "Direct execution"328 assumptions_made = []329 330 if plan:331 approach_used = f"{plan.task_type} approach with {len(plan.steps)} steps"332 assumptions_made = plan.assumptions333 334 # Generate structured report (КОСТЫЛЬ ЗДЕСЬ!)335 report_generator_prompt = f"""336 Generate a comprehensive execution report for the following query processing:337 338 ORIGINAL QUERY: {state['query']}339 340 EXECUTION CONTEXT:341 - Complexity Level: {state.get('complexity_assessment', {}).level}342 - Plan Used: {plan if plan else {}}343 - Tools Executed: {tools_executed}344 - Available Files: {list(state.get('file_contents', {}).keys())}345 346 CONVERSATION HISTORY:347 {[msg.content[:200] + "..." if len(msg.content) > 200 else msg.content 348 for msg in state['messages'][-5:]]} # Last 5 messages for context349 350 Based on this information, create a structured execution report that includes:351 1. Query summary352 2. Approach used353 3. Key findings from the execution354 4. Data sources used355 5. Your confidence level in the results356 6. Any limitations or caveats357 7. The final answer358 359 Be thorough but concise. This report will be evaluated by a critic for quality assurance.360 """361 362 report_llm = llm_deterministic.with_structured_output(ExecutionReport)363 364 execution_report = report_llm.invoke([365 SystemMessage(content=report_generator_prompt),366 HumanMessage(content="Generate the execution report.")367 ])368 369 print(f"Report generated - Confidence: {execution_report.confidence_level}")370 print(f"Key findings: {len(execution_report.key_findings)}")371 print(f"Data sources: {len(execution_report.data_sources)}")372 373 # Format final answer for user374 formatted_answer = format_final_answer(execution_report, state.get('complexity_assessment', {}))375 #print(execution_report)376 print(f"FINAL ANSWER FOR EVALUATOR: {execution_report.final_answer}")377 378 return {379 "execution_report": execution_report,380 "final_answer": formatted_answer381 }382 383 384def simple_executor(state: AgentState) -> AgentState:385 """Handle simple queries directly without planning."""386 print("=== SIMPLE EXECUTION ===")387 388 # For simple queries, use the LLM with tools directly389 simple_prompt = f"""390 Answer this simple query directly and efficiently: {state['query']}391 392 You have access to tools if needed, but try to answer directly when possible.393 If you need files, they are available at: {list(state.get('file_contents', {}).keys())}394 395 Provide a clear, concise answer.396 """397 398 response = llm_simple_with_tools.invoke([399 SystemMessage(content=simple_prompt),400 HumanMessage(content=state['query'])401 ])402 403 print("Response generated for simple query.")404 405 return {406 "messages": state["messages"] + [response],407 "final_answer": response.content408 }409 410def should_use_tools_simple_executor(state: AgentState) -> str:411 """Decide whether to use tools or answer directly in simple executor."""412 last_message = state["messages"][-1]413 414 if hasattr(last_message, "tool_calls") and last_message.tool_calls:415 return "tools"416 417 if hasattr(last_message, "content") and "<FINAL_ANSWER>" in last_message.content:418 return "final_answer"419 420 return "final_answer"421 422 423def should_use_planning(state: AgentState) -> str:424 """Route based on complexity assessment."""425 complexity = state["complexity_assessment"]426 427 if complexity.level == "simple" and not complexity.needs_planning:428 return "simple_executor"429 else:430 return "planner"431 432 433def critic_evaluator(state: AgentState) -> AgentState:434 """Enhanced critic that evaluates execution reports."""435 print("=== ENHANCED ANSWER CRITIQUE ===")436 437 report = state.get("execution_report")438 critic_llm = llm_criticist.with_structured_output(CritiqueFeedback)439 440 critique_prompt = CRITIC_PROMPT.format(441 query=report.query_summary,442 approach=report.approach_used,443 tools=report.tools_executed,444 findings=report.key_findings,445 sources=report.data_sources,446 confidence=report.confidence_level,447 limitations=report.limitations,448 answer=report.final_answer449 )450 451 critique = critic_llm.invoke([452 SystemMessage(content=critique_prompt),453 HumanMessage(content="Evaluate this execution report thoroughly.")454 ])455 456 print(f"Quality Score: {critique.quality_score}/10")457 print(f"Complete: {critique.is_complete}")458 print(f"Accurate: {critique.is_accurate}")459 460 if critique.errors_found:461 print(f"Issues found: {critique.errors_found}")462 463 if critique.needs_replanning:464 print(f"Replanning needed: {critique.replan_instructions}")465 466 return {467 "critique_feedback": critique,468 "iteration_count": state.get("iteration_count", 0) + 1469 }470 471 472 473def should_replan(state: AgentState) -> str:474 """Decide whether to accept answer, replan, or stop."""475 critique = state.get("critique_feedback")476 iteration_count = state.get("iteration_count", 0)477 max_iterations = state.get("max_iterations", 3)478 activator = state.get("critic_replan", False)479 480 print(f"=== REPLAN DECISION ===")481 print(f"Iteration: {iteration_count}/{max_iterations}")482 print(f"Quality score: {critique.quality_score if critique else 'N/A'}")483 print(f"Needs replanning: {critique.needs_replanning if critique else 'N/A'}")484 485 if not activator:486 return "end"487 488 if not critique:489 return "end"490 491 # Stop if max iterations reached492 if iteration_count >= max_iterations:493 print(f"Max iterations ({max_iterations}) reached. Accepting current answer.")494 return "end"495 496 # Accept if quality is good enough497 if critique.quality_score >= 7 or not critique.needs_replanning:498 print("Quality acceptable, ending execution")499 return "end"500 501 # Replan if quality is poor and we haven't exceeded max iterations502 if critique.needs_replanning and iteration_count < max_iterations:503 print("Replanning due to critic feedback...")504 return "replan"505 506 return "end"507 508def replanner_old(state: AgentState) -> AgentState:509 """Create a revised plan based on critic feedback."""510 print("=== REPLANNING ===")511 512 critique = state["critique_feedback"]513 previous_plan = state.get("plan")514 515 replan_prompt = f"""516 {SYSTEM_PROMPT_PLANNER}517 518 REPLANNING CONTEXT:519 Original Query: {state['query']}520 Previous Plan: {previous_plan if previous_plan else {}}521 522 CRITIC FEEDBACK:523 - Quality Score: {critique.quality_score}/10524 - Issues Found: {critique.errors_found}525 - Missing Elements: {critique.missing_elements}526 - Improvement Suggestions: {critique.suggested_improvements}527 - Specific Instructions: {critique.replan_instructions}528 529 Create a REVISED plan that addresses these issues. Focus on fixing the identified problems.530 """531 532 revised_plan = planner_llm.invoke([533 SystemMessage(content=replan_prompt),534 HumanMessage(content="Create a revised plan based on the feedback.")535 ])536 537 print("Plan revised based on critic feedback")538 539 # Очищаем историю сообщений от неполных tool_calls540 current_messages = state.get("messages", [])541 cleaned_messages = clean_message_history(current_messages)542 543 # Оставляем только системные сообщения и начальный запрос544 essential_messages = []545 for msg in cleaned_messages:546 if isinstance(msg, (SystemMessage, HumanMessage)):547 # Сохраняем системные сообщения и пользовательские запросы548 if ("complexity" in msg.content.lower() or 549 "assess" in msg.content.lower() or550 isinstance(msg, HumanMessage)):551 essential_messages.append(msg)552 553 #print(f"Cleaned message history: {len(current_messages)} -> {len(essential_messages)} messages")554 #print("=== ESSENTIAL MESSAGES ===")555 #print(essential_messages)556 #print("=== AGENT STATE ===")557 #print(state["messages"])558 559 return {560 "plan": revised_plan,561 "current_step": 0,562 "reasoning_done": False,563 "messages": essential_messages,564 "execution_report": None565 }566 567def replanner(state: AgentState) -> AgentState:568 """Create a revised plan based on critic feedback."""569 print("=== REPLANNING ===")570 571 critique = state["critique_feedback"]572 previous_plan = state.get("plan")573 574 replan_prompt = f"""575 {SYSTEM_PROMPT_PLANNER}576 577 REPLANNING CONTEXT:578 Original Query: {state['query']}579 Previous Plan: {previous_plan if previous_plan else {}}580 581 CRITIC FEEDBACK:582 - Quality Score: {critique.quality_score}/10583 - Issues Found: {critique.errors_found}584 - Missing Elements: {critique.missing_elements}585 - Improvement Suggestions: {critique.suggested_improvements}586 - Specific Instructions: {critique.replan_instructions}587 588 Create a REVISED plan that addresses these issues. Focus on fixing the identified problems.589 """590 591 revised_plan = planner_llm.invoke([592 SystemMessage(content=replan_prompt),593 HumanMessage(content="Create a revised plan based on the feedback.")594 ])595 596 print("Plan revised based on critic feedback")597 598 # ИСПРАВЛЕНИЕ: Сохраняем важные результаты инструментов599 current_messages = state.get("messages", [])600 state["previous_final_answer"] = state.get("final_answer", "")601 # Находим полезные результаты инструментов602 preserved_messages = []603 tool_results = {}604 605 for i, msg in enumerate(current_messages):606 # Сохраняем системные сообщения и пользовательские запросы607 if isinstance(msg, (SystemMessage, HumanMessage)):608 # Фильтруем только исходные запросы, не промпты планировщика609 if (isinstance(msg, HumanMessage) or 610 ("complexity" in msg.content.lower() and "assessor" in msg.content.lower())):611 preserved_messages.append(msg)612 613 # Сохраняем успешные результаты инструментов614 elif isinstance(msg, ToolMessage) and msg.content and msg.content.strip():615 # Проверяем, что это полезный результат616 try:617 # Если результат можно преобразовать в число - это вычисление618 float(msg.content.strip())619 preserved_messages.append(msg)620 tool_results[msg.tool_call_id] = msg.content621 622 # Также нужно сохранить соответствующий AIMessage с tool_call623 for j in range(i-1, -1, -1):624 if (isinstance(current_messages[j], AIMessage) and 625 hasattr(current_messages[j], 'tool_calls') and626 current_messages[j].tool_calls):627 for tool_call in current_messages[j].tool_calls:628 if tool_call['id'] == msg.tool_call_id:629 if current_messages[j] not in preserved_messages:630 preserved_messages.insert(-1, current_messages[j])631 break632 break633 except (ValueError, AttributeError):634 # Если не число, но содержательный результат, тоже сохраняем635 if len(msg.content.strip()) > 1: # Минимальная длина для сохранения636 preserved_messages.append(msg)637 638 print(f"Preserved {len(tool_results)} tool results")639 #print(f"Cleaned message history: {len(current_messages)} -> {len(preserved_messages)} messages")640 641 # Добавляем контекст о доступных результатах642 if tool_results:643 context_msg = HumanMessage(644 content=f"Previous calculation results available: {tool_results}"645 )646 preserved_messages.append(context_msg)647 648 return {649 "plan": revised_plan,650 "current_step": 0,651 "reasoning_done": False,652 "messages": preserved_messages,653 "execution_report": None,654 # Сохраняем важную информацию о предыдущих вычислениях655 "previous_tool_results": tool_results656 }657 658def complexity_assessor(state: AgentState) -> AgentState:659 """Assess query complexity and determine if planning is needed."""660 print("=== COMPLEXITY ASSESSMENT ===")661 662 complexity_llm = llm_deterministic.with_structured_output(ComplexityLevel)663 664 assessment_message = [665 SystemMessage(content=COMPLEXITY_ASSESSOR_PROMPT.strip()),666 HumanMessage(content=f"Query: {state['query']}")667 ]668 669 assessment = complexity_llm.invoke(assessment_message)670 671 print(f"Complexity: {assessment.level}")672 print(f"Needs planning: {assessment.needs_planning}")673 print(f"Reasoning: {assessment.reasoning}")674 675 return {676 "complexity_assessment": assessment,677 "messages": state["messages"] + assessment_message678 }