diff --git a/app/tool/__init__.py b/app/tool/__init__.py index b1d25e2..d5c36e5 100644 --- a/app/tool/__init__.py +++ b/app/tool/__init__.py @@ -2,6 +2,7 @@ from app.tool.base import BaseTool from app.tool.bash import Bash from app.tool.browser_use_tool import BrowserUseTool from app.tool.create_chat_completion import CreateChatCompletion +from app.tool.deep_research import DeepResearch from app.tool.planning import PlanningTool from app.tool.str_replace_editor import StrReplaceEditor from app.tool.terminate import Terminate @@ -13,6 +14,7 @@ __all__ = [ "BaseTool", "Bash", "BrowserUseTool", + "DeepResearch", "Terminate", "StrReplaceEditor", "WebSearch", diff --git a/app/tool/deep_research.py b/app/tool/deep_research.py index be6f093..f2deed0 100644 --- a/app/tool/deep_research.py +++ b/app/tool/deep_research.py @@ -4,7 +4,7 @@ import re import time from typing import List, Optional, Set -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, model_validator from app.exceptions import ToolError from app.llm import LLM @@ -100,9 +100,11 @@ class ResearchContext(BaseModel): ) -class ResearchSummary(BaseModel): +class ResearchSummary(ToolResult): """Comprehensive summary of deep research results.""" + model_config = ConfigDict(arbitrary_types_allowed=True) + query: str = Field(description="The original research query") insights: List[ResearchInsight] = Field( default_factory=list, description="Key insights discovered" @@ -114,8 +116,9 @@ class ResearchSummary(BaseModel): default=0, description="Maximum depth of research reached", ge=0 ) - def to_tool_result(self) -> ToolResult: - """Convert research summary to a formatted ToolResult.""" + @model_validator(mode="after") + def populate_output(self) -> "ResearchSummary": + """Populate the output field after validation.""" # Group and sort insights by relevance grouped_insights = { "Key Findings": [i for i in self.insights if i.relevance_score >= 0.8], @@ -143,7 +146,9 @@ class ResearchSummary(BaseModel): ] ) - return ToolResult(output="\n".join(sections)) + # Assign the formatted string to the 'output' field inherited from ToolResult + self.output = "\n".join(sections) + return self class DeepResearch(BaseTool): @@ -316,17 +321,23 @@ class DeepResearch(BaseTool): # 4. Continue research with follow-up queries if follow_up_queries and context.current_depth < context.max_depth: + tasks = [] # Create a list to hold the tasks for follow_up in follow_up_queries[:2]: # Limit branching factor if time.time() >= deadline: break - # Recursive research with reduced result count at deeper levels - await self._research_graph( + # Create a coroutine for the recursive research call + task = self._research_graph( context=context, query=follow_up, - results_count=max(1, results_count - 1), + results_count=max(1, results_count - 1), # Reduce result count deadline=deadline, ) + tasks.append(task) # Add the task to the list + + # Run all the created tasks concurrently + if tasks: + await asyncio.gather(*tasks) async def _search_web(self, query: str, results_count: int) -> List[SearchResult]: """Perform web search for the given query.""" @@ -523,4 +534,4 @@ if __name__ == "__main__": "What is deep learning", max_depth=1, results_per_search=2 ) ) - print(result.to_tool_result().output) + print(result) diff --git a/app/tool/web_search.py b/app/tool/web_search.py index fc94586..5a6be30 100644 --- a/app/tool/web_search.py +++ b/app/tool/web_search.py @@ -125,7 +125,7 @@ class WebContentFetcher: Extracted text content or None if fetching fails """ headers = { - "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36" + "WebSearch": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36" } try: