import asyncio import traceback from typing import Any, Callable, List, Optional from agents import function_tool from openai.types.responses import ResponseTextDeltaEvent from pydantic import ValidationError from iterative_detail_plan import IterativeDetailPlan from iterative_research import IterativeResearcher from utils.baseclass import ResearchAgent, ResearchRunner from tools.long_writer_agent import LongWriterOutput, write_report_from_section_drafts from utils.schemas import ReportDraft, ReportDraftSection from tools.detail_plan_agent import CoreOutline, CoreSection from tools.writer_agent import ( checkout_section_agent, section_summary_agent, abstract_agent, translate_title_chinese_agent, ) from config_logger import logger # logger = logging.getLogger(__name__) class LiteratureReviewTool: """ 文献研究工具类，用于自动查询文章并生成研究报告。 """ def __init__( self, verbose: bool = True, thoughts_callback: Optional[Callable[[str], Any]] = None, results_callback: Optional[Callable[[str], Any]] = None, stop_event: Optional[asyncio.Event] = None, hooks=None, u_id: Optional[str] = None, is_web: Optional[bool] = False, is_pkb: Optional[bool] = False, language: str = "EN", # EN or CH ): self.verbose = verbose self.thoughts_callback = thoughts_callback self.results_callback = results_callback self.stop_event = stop_event self.hooks = hooks self.u_id = u_id self.is_web = is_web self.is_pkb = is_pkb self.language = language self.enrichquery = "" if thoughts_callback is None: async def noop(x): pass self.thoughts_callback = noop if results_callback is None: self.results_callback = thoughts_callback async def run( self, query: str, ) -> str: """ 运行文献研究工具，生成研究报告。 Args: query: 研究主题或问题 thoughts_callback: 用于报告进度和思考的异步回调函数 results_callback: 用于流式返回结果的异步回调函数 stop_event: 用于检查取消操作的异步事件 Returns: 生成的文献研究报告（Markdown格式） """ try: # 1. 构建报告计划 report_plan = await self._build_detail_report_plan(query) # await self.results_callback(f"########## report_plan{report_plan}") # 2. 为每个章节执行文献研究 research_results, found_references = await self._run_research_loops( report_plan, ) await self._log_message("Research_results loop down") # 3. 创建最终报告 logger.info(f"Creating final report... \n") final_report = await self._create_final_report( query, report_plan, research_results, found_references, self.thoughts_callback, self.language, ) logger.info(f"Final report created... \n") await self.results_callback("Final_report\n") # await self.results_callback(final_report) await self.stream_text(final_report) return final_report except Exception as e: error_msg = f"Research error: {str(e)}\n{traceback.format_exc()}" if self.thoughts_callback: await self.thoughts_callback(error_msg) return f"Research error: {str(e)}" async def stream_text(self, res: str, chunk_size: int = 100): for i in range(0, len(res), chunk_size): chunk = res[i : i + chunk_size] await asyncio.sleep(0.05) await self.results_callback(chunk) async def _build_detail_report_plan( self, query: str, ) -> CoreOutline: """构建详细报告计划，使用planner_agent_test生成报告计划""" await self._log_message("\n=== Building Detail Report Plan ===\n") # user_raw_query = query # 构建多个agent的循环输入 generator = IterativeDetailPlan( max_iterations=3, max_time_minutes=10, thoughts_callback=self.results_callback, ) logger.info(f"Building detail report plan... \n") detail_outline, enrichquery = await generator.run(query=query) self.enrichquery = enrichquery await self._log_message("\n=== Report Plan Built ===\n") return detail_outline async def _run_research_loops( self, report_plan: CoreOutline, ) -> tuple[Any, List[Any]]: """为每个章节执行文献研究并收集结果""" research_results = [] found_ref = [] await self._log_message("\n **Reasoning about Sections** \n") async def run_research_for_section(section: CoreSection): if self.stop_event and self.stop_event.is_set(): await self._log_message( f"\n **Study section {section.title} canceled** \n" ) return "Study canceled", [] await self._log_message( f"\n===Initializing Section: {section.title} Research Loops Study===\n" ) # 创建IterativeResearcher实例 iterative_researcher = IterativeResearcher( max_iterations=1, #2, # 可以根据需要调整 max_time_minutes=12, # 可以根据需要调整 verbose=True, thoughts_callback=self.thoughts_callback, hooks=self.hooks, u_id=self.u_id, ) # 准备IterativeResearcher的参数 args = { "query": self.enrichquery, "output_length": " 800", "output_instructions": section, "background_context": report_plan.background, } try: section_result, section_references = await iterative_researcher.run( **args ) await self._log_message( f"\nSection: {section.title} Research Loops Study completed\n" ) except Exception as e: error_msg = f"Section {section.title} error: {str(e)}" logger.error(error_msg) section_result = None section_references = None # return f"Error: {str(e)}", [] return section_result, section_references # await self._log_message("=== Initializing Research Loops ===") # 并发执行所有章节的研究 is_loop_iter = False if is_loop_iter: # 单次跑 # for section in report_plan.report_outline: # result = await run_research_for_section(section) # research_results.append(result) # max_tasks = 2 for i in range(0, len(report_plan.sections), max_tasks): bach_sections = report_plan.sections[i : i + max_tasks] batch_tasks = [ run_research_for_section(section) for section in bach_sections ] batch_results = await asyncio.gather(*batch_tasks) for section_result, section_references in batch_results: research_results.append(section_result) found_ref.extend(section_references) else: # 使用asyncio.gather并发执行所有章节的研究 batch_results = await asyncio.gather( *(run_research_for_section(section) for section in report_plan.sections) ) research_results = [] found_ref = [] for section_result, section_references in batch_results: # print(f"########## section_references {section_references},length {len(section_references)}.\n ########## section_result {section_result}") # print(f"########################################################") research_results.append(section_result) if section_references: found_ref.extend(section_references) return research_results, found_ref async def _create_final_report( self, query: str, report_plan: CoreOutline, section_drafts: List[LongWriterOutput], ref: List[Any], thoughts_callback: Optional[Callable[[str], Any]] = None, language: str = "EN", # EN or CH ) -> str: """从报告计划和章节草稿创建最终报告""" # 构建ReportDraft对象 logger.info( f"########## found_references length {len(ref)},\n research_results length {len(section_drafts)}" ) report_draft = ReportDraft(sections=[]) async def check_section(section_draft: LongWriterOutput, ins_query: str, section_title: str): logger.info(f"Checking section {section_title}... \n") await self.results_callback(f"Checking section {section_title}... \n") if not section_draft.next_section_markdown: return None, None else: logger.info(f"Checking section {section_title}... \n") check_result = await self._check_section( section_draft, ins_query, language ) logger.info(f"Checking section {section_title} completed... \n") summary = await self._generate_summary( check_result.next_section_markdown ) logger.info(f"Generating summary for section {section_title} completed... \n") return check_result, summary # 过滤出非空的section_drafts并记录它们的原始索引 non_empty_sections = [] for i, section_draft in enumerate(section_drafts): if section_draft and section_draft.next_section_markdown: non_empty_sections.append((i, section_draft)) checkouts_results = await asyncio.gather( *( check_section( section_draft, f" u are modifing the section num {j + 1}", report_plan.sections[i].title, ) for j, (i, section_draft) in enumerate(non_empty_sections) ) ) logger.info(f"Checkouts completed... \n") section_summaries = [] for j, (section_result, summary) in enumerate(checkouts_results): if section_result: # 使用原始索引来获取正确的section title original_index = non_empty_sections[j][0] report_draft.sections.append( ReportDraftSection( section_title=report_plan.sections[original_index].title, section_content=section_result.next_section_markdown, ) ) if summary: section_summaries.append(summary) if thoughts_callback: await thoughts_callback("\n **Generating final report...** \n") logger.info(f"Generating abstract... \n") await self.results_callback(f"Generating abstract... \n") abstract = await self._genrate_abstract(section_summaries, language) if language == "CH": report_plan.report_title = await self._translate_title_chinese( report_plan.report_title ) logger.info(f"Writing report from section drafts... \n") final_output = await write_report_from_section_drafts( query, abstract, report_plan.report_title, report_draft, ref, self.thoughts_callback, ) return final_output async def _generate_summary(self, sections: str) -> str: full_response = "" result = ResearchRunner.run_streamed( starting_agent=section_summary_agent, input=sections ) try: async for event in result.stream_events(): try: if event.type == "raw_response_event" and isinstance( event.data, ResponseTextDeltaEvent ): full_response += event.data.delta except Exception as e: logger.error(f"Error processing event: {e}") continue except ValidationError: pass except Exception as e: logger.error(f"Error processing generate summary event: {e}") pass final_result = result.final_output return final_result async def _translate_title_chinese(self, title: str) -> str: """Translate English title to Chinese""" input_str = f"LANGUAGE: Chinese\n\nTITLE: {title}" try: result = ResearchRunner.run( starting_agent=translate_title_chinese_agent, input=input_str, ) return result.final_output except ValidationError as e: logger.warning(f"Translation validation error: {e}") return title # Return original title if translation fails except Exception as e: logger.error(f"Translation error: {e}") return title # Return original title if translation fails async def _genrate_abstract(self, summarys: List[str], language: str = "EN") -> str: full_response = "" if language == "CH": language_str = "Chinese" else: language_str = "English" input_str = f"LANGUAGE: {language_str}\n\nSUMMARY: {str(summarys)}" result = ResearchRunner.run_streamed( starting_agent=abstract_agent, input=input_str, ) try: async for event in result.stream_events(): try: if event.type == "raw_response_event" and isinstance( event.data, ResponseTextDeltaEvent ): full_response += event.data.delta except Exception as e: logger.error(f"Error processing event: {e}") continue except ValidationError: pass final_result = result.final_output return final_result async def _log_message(self, message: str) -> None: """Log a message if verbose is True""" if self.verbose: await self.thoughts_callback(message) else: print(message) async def _check_section( self, section: LongWriterOutput, query: str = "", language: str = "EN" ) -> LongWriterOutput: if language == "CH": language_str = "Chinese" else: language_str = "English" section_str = section.next_section_markdown ins = f""" LANGUAGE: {language_str} PROCESS_REQUIRMENT: {query} SECTION: {section.next_section_markdown} """ try_num = 0 max_try_num = 3 full_response = "" if not section_str: return section while try_num < max_try_num: result = ResearchRunner.run_streamed( starting_agent=checkout_section_agent, input=ins ) try: async for event in result.stream_events(): try: if event.type == "raw_response_event" and isinstance( event.data, ResponseTextDeltaEvent ): full_response += event.data.delta except Exception as e: logger.error(f"Error processing event: {e}") continue final_result = result.final_output break except ValidationError: final_result = full_response break except Exception as e: logger.error(f"Error processing event in {try_num} times: {e}") try_num += 1 if try_num == max_try_num: return section section.next_section_markdown = final_result return section # 使用示例 async def example_usage(): """ 展示如何使用LiteratureResearchTool的示例 """ # 创建工具实例 # 定义回调函数 async def progress_callback(message): print(f"Progress: {message}") async def results_callback(token): print(token, end="", flush=True) # user_message = str( """Please write a comprehensive review on recent advances in CAR-T cell therapy, focusing on innovative target mining strategies to address core challenges in solid tumor treatment. The review should: (1) analyze key obstacles hindering CAR-T efficacy in solid tumors, including tumor heterogeneity, lack of tumor-specific antigens, and immunosuppressive microenvironments; (2) explore cutting-edge technologies such as single-cell RNA sequencing, spatial transcriptomics, and machine learning/AI in driving novel target discovery, emphasizing their roles in deciphering clonal evolution, predicting antigen immunogenicity, and integrating multi-omics data; (3) discuss engineering strategies (e.g., logic-gated CAR designs, affinity optimization) that link target selection to toxicity control, as well as target-informed combination therapies (e.g., with immune checkpoint inhibitors); (4) Link target profiles to combination approaches: Immune checkpoint inhibitors, Microenvironment modulators; (5) Future Directions: AI, Personalization, and Scalable Platforms outline future directions, including AI-powered target prediction, personalized neoantigen screening, and scalable manufacturing platforms. Maintain a cohesive narrative centered on target mining, incorporate tables where appropriate to compare technologies or summarize critical targets, and ensure academic rigor with logical progression from challenges to solutions and future perspectives.""" ) tool = LiteratureReviewTool( thoughts_callback=progress_callback, results_callback=results_callback, verbose=True, ) # 运行研究 report_plan = await tool.run( query=user_message, ) print(report_plan) # 运行研究 分段 async def example_tool(): async def collect_thoughts(thought): print(f"THOUGHT: {thought}") async def collect_results(result): # Only store first 100 chars to avoid overwhelming memory print(f"PARTIAL RESULT: {result[:100]}...") from dataclasses import dataclass from typing import Any, Callable, Optional from agents import RunContextWrapper @dataclass class InputCallbackTool: query: str thoughts_callback: Optional[Callable[[str], Any]] = None """callback of thinking .""" results_callback: Optional[Callable[[str], Any]] = None """callback of results""" @property def name(self): return "callback" @function_tool async def test_tool(wrapper: RunContextWrapper[InputCallbackTool]): """ a tool to generate a literature review """ tool = LiteratureReviewTool( verbose=True, thoughts_callback=wrapper.context.thoughts_callback, results_callback=wrapper.context.results_callback, ) response = await tool.run(wrapper.context.query) return response # 处理最后的相对导入 try: from .utils.llm_client import qianwen_plus_model except ImportError: from utils.llm_client import qianwen_plus_model INSTRUCTIONS = """ You are a research manager, managing a team of research agents. Given a research query, your job is to produce an initial outline of the report (section titles and key questions), as well as some background context. Each section will be assigned to a different researcher in your team who will then carry out research on the section. You will be given: - An initial research query Your task is to: use once of this tool to generate the review report return the full result of tool """ selected_model = qianwen_plus_model test_agent = ResearchAgent( name="testtool", instructions=INSTRUCTIONS, tools=[test_tool], model=selected_model, ) user_message = str( """Please write a comprehensive review on recent advances in CAR-T cell therapy, focusing on innovative target mining strategies to address core challenges in solid tumor treatment. The review should: (1) analyze key obstacles hindering CAR-T efficacy in solid tumors, including tumor heterogeneity, lack of tumor-specific antigens, and immunosuppressive microenvironments; (2) explore cutting-edge technologies such as single-cell RNA sequencing, spatial transcriptomics, and machine learning/AI in driving novel target discovery, emphasizing their roles in deciphering clonal evolution, predicting antigen immunogenicity, and integrating multi-omics data; (3) discuss engineering strategies (e.g., logic-gated CAR designs, affinity optimization) that link target selection to toxicity control, as well as target-informed combination therapies (e.g., with immune checkpoint inhibitors); (4) Link target profiles to combination approaches: Immune checkpoint inhibitors, Microenvironment modulators; (5) Future Directions: AI, Personalization, and Scalable Platforms outline future directions, including AI-powered target prediction, personalized neoantigen screening, and scalable manufacturing platforms. Maintain a cohesive narrative centered on target mining, incorporate tables where appropriate to compare technologies or summarize critical targets, and ensure academic rigor with logical progression from challenges to solutions and future perspectives.""" ) input = InputCallbackTool( query=user_message, thoughts_callback=collect_thoughts, results_callback=collect_results, ) result = await ResearchRunner.run(test_agent, user_message, context=input) # print(result) if __name__ == "__main__": asyncio.run(example_usage())