# sub_agents/researcher.py from langchain.agents import create_openai_tools_agent, AgentExecutor from langchain_core.prompts import ChatPromptTemplate from state import AgentState from db_utils import initialize_researcher, update_researcher_report, extract_and_clean_json, get_document_object, update_document_object from tools import process_url from utils import filter_content_for_summarization from terminal_utils import print_colorful_break async def researcher_agent_node(state: AgentState): """The main node for the researcher workflow.""" print_colorful_break("RESEARCHER") logger = state['logger'] report_id = state.get("researcher_report_id") gaps_todo = state.get("researcher_gaps_todo", []) gaps_complete = state.get("researcher_gaps_complete", []) final_status = None if not report_id: try: init_result = initialize_researcher(state['timestamp']) report_id = init_result.get("researcher_report_id") gaps_todo = init_result.get("researcher_gaps_todo") gaps_complete = [] status = f"--- Initialized researcher state: {init_result} ---" logger.info(status) except Exception as e: status = f"Failed to initialize researcher state: {e}" logger.error(status) return {"status": status} # Define Planner Agent with open("prompts/planner_prompt.txt", "r") as f: planner_prompt_template = f.read() planner_prompt = ChatPromptTemplate.from_template(planner_prompt_template) try: planner_agent_runnable = create_openai_tools_agent(state['model'], [], planner_prompt) planner_executor = AgentExecutor(agent=planner_agent_runnable, tools=[], verbose=True) except Exception as e: status = f"Failed to create planner agent executor: {e}" logger.error(status) return {"status": status} # Define Refiner Agent with open("prompts/refiner_prompt.txt", "r") as f: refiner_prompt_template = f.read() refiner_prompt = ChatPromptTemplate.from_template(refiner_prompt_template) try: refiner_agent_runnable = create_openai_tools_agent(state['model'], [], refiner_prompt) refiner_executor = AgentExecutor(agent=refiner_agent_runnable, tools=[], verbose=True) except Exception as e: status = f"Failed to create refiner agent executor: {e}" logger.error(status) return {"status": status} # Define Summarizer Agent with open("prompts/summarizer_prompt.txt", "r") as f: summarizer_prompt_template = f.read() summarizer_prompt = ChatPromptTemplate.from_template(summarizer_prompt_template) try: summarizer_agent_runnable = create_openai_tools_agent(state['model'], [], summarizer_prompt) summarizer_executor = AgentExecutor(agent=summarizer_agent_runnable, tools=[], verbose=True) except Exception as e: status = f"Failed to create summarizer agent executor: {e}" logger.error(status) return {"status": status} # Get the tools google_search_tool = next((tool for tool in state['mcp_tools'] if tool.name == 'google_search'), None) if not google_search_tool: status = "google_search tool not found." logger.error(status) return {"status": status} # Main control loop try: if isinstance(gaps_todo, list) and gaps_todo: for current_gap in list(gaps_todo): gap_id = current_gap['gap_id'] research_topic = current_gap['research_topic'] all_searches_for_gap = [] research_topic_title = research_topic.get('title', 'No Title') status = f"Starting research for gap: {gap_id}, research topic: {research_topic_title}" logger.info(status) try: # 1. Planning Step status = f"Invoking planner for gap {gap_id}." logger.info(status) planner_result = await planner_executor.ainvoke({"input": research_topic}) planner_output = extract_and_clean_json(planner_result.get("output", "")) planned_searches = planner_output.get("searches", []) status = f"Planner for gap {gap_id} returned {len(planned_searches)} searches." logger.info(status) # 2. Initial Execution Step for planned_search in planned_searches: query = planned_search.get("query") rationale = planned_search.get("rationale") parameters = planned_search.get("parameters", {}) search_id = planned_search.get("search_id") if not query: continue if 'query' not in parameters: parameters['query'] = query try: status = f"Executing search for gap {gap_id}, with parameters: {parameters}" logger.info(status) raw_search_results = await google_search_tool.arun(parameters) search_results = extract_and_clean_json(raw_search_results) for i, result in enumerate(search_results): url = result.get('url') if url: logger.info(f"Attempting document store initialzation for URL: {url}") try: url_id, url_status = await process_url(url, logger) except Exception as e: status = f"Error processing URL {url}: {e}" logger.error(status, exc_info=True) continue search_results[i]['url_id'] = url_id if url_status == "new": status=f"New URL {url} added to the database with ID: {url_id}." else: status = f"URL {url} already exists in the database with ID: {url_id}." logger.info(status) search_object = { "search_id": search_id, "rationale": rationale, "parameters": parameters, "results": search_results } all_searches_for_gap.append(search_object) status = f"Search for gap {gap_id} finished for query: '{query}'" logger.info(status) except Exception as e: status = f"Search for gap {gap_id}, query '{query}' failed: {e}" logger.error(status) continue # 3. Refinement Step status = f"Invoking refiner for gap {gap_id}." logger.info(status) refiner_input = {"research_topic": research_topic, "search_results": all_searches_for_gap} refiner_result = await refiner_executor.ainvoke({"input": refiner_input}) refiner_output = extract_and_clean_json(refiner_result.get("output", "")) status_check = "" if isinstance(refiner_output, dict): status_check = refiner_output.get("status", "").lower() elif isinstance(refiner_output, str): if "insufficient" in refiner_output.lower(): status_check = "insufficient" if status_check == "insufficient": refined_searches = [] if isinstance(refiner_output, dict): refined_searches = refiner_output.get("searches", []) status = f"Refiner for gap {gap_id} returned {len(refined_searches)} new searches." logger.info(status) # 4. Refined Execution Step for refined_search in refined_searches: query = refined_search.get("query") rationale = refined_search.get("rationale") parameters = refined_search.get("parameters", {}) search_id = refined_search.get("search_id") if not query: continue if 'query' not in parameters: parameters['query'] = query try: status = f"Executing refined search for gap {gap_id}, with parameters: {parameters}" logger.info(status) raw_search_results = await google_search_tool.arun(parameters) search_results = extract_and_clean_json(raw_search_results) for i, result in enumerate(search_results): url = result.get('url') if url: logger.info(f"Attempting document store initialzation for URL: {url}") try: url_id, url_status = await process_url(url, logger) except Exception as e: status = f"Error processing URL {url}: {e}" logger.error(status, exc_info=True) continue search_results[i]['url_id'] = url_id if url_status == "new": status=f"New URL {url} added to the database with ID: {url_id}." else: status = f"URL {url} already exists in the database with ID: {url_id}." logger.info(status) search_object = { "search_id": search_id, "rationale": rationale, "parameters": parameters, "results": search_results } all_searches_for_gap.append(search_object) status = f"Refined search for gap {gap_id} finished for query: '{query}'" logger.info(status) except Exception as e: status = f"Refined search for gap {gap_id}, query '{query}' failed: {e}" logger.error(status) continue else: status = f"Refiner for gap {gap_id} deemed results sufficient." logger.info(status) # 5. Summarization Step status = f"Starting summarization for gap {gap_id}." logger.info(status) for search in all_searches_for_gap: for i, result in enumerate(search.get('results', [])): url = result.get('url') url_id = result.get('url_id') url_summary = get_document_object(url_id, type="summary") if url_summary: logger.info(f"Skipping summarization for url_id: {url_id}, summary already exists.") continue else: markdown_content = get_document_object(url_id, type="markdown_content") if not markdown_content or markdown_content.startswith("[MARKDOWN_GENERATION_FAILED"): logger.info(f"Skipping summarization for url_id: {url_id}, no valid markdown content available.") continue logger.info(f"Attempting summary for url_id: {url_id}") try: filtered_content = filter_content_for_summarization(markdown_content) summarizer_result = await summarizer_executor.ainvoke({"input": filtered_content}) summary_output = extract_and_clean_json(summarizer_result.get("output", "")) if isinstance(summary_output, dict): summary = summary_output.get('summary') else: summary = str(summary_output) update_document_object(url_id, type="summary", object=summary) status = f"Successfully summarized and updated document for url_id: {url_id}" logger.info(status) except Exception as e: status = f"Error summarizing url_id {url_id}: {e}" logger.error(status, exc_info=True) continue # 6. Update Step status = f"Preparing to update report for gap {gap_id} with {len(all_searches_for_gap)} searches." logger.info(status) try: update_researcher_report(report_id, gap_id, all_searches_for_gap) gaps_complete.append(gap_id) gaps_todo = [g for g in gaps_todo if g.get("gap_id") != gap_id] status = f"Updated researcher report for gap: {gap_id}" logger.info(status) except Exception as e: status = f"Error updating report for gap {gap_id}: {e}" logger.error(status, exc_info=True) continue status = f"--- Successfully completed research and report writing for gap: {gap_id} ---" logger.info(status) except Exception as e: status = f"An unexpected error occurred while processing gap {gap_id}: {e}" logger.error(status, exc_info=True) continue except Exception as e: final_status = f"Main loop failed: {e}" logger.error(final_status, exc_info=True) if not final_status: final_status = f"Successfully and incrementally completed researcher report with ID {report_id} and wrote report to DB." logger.info(final_status) return { "status": final_status, "researcher_report_id": report_id, "researcher_gaps_todo": gaps_todo, "researcher_gaps_complete": gaps_complete }