two version of R2R are hereHEADmaster

1 files changed, 697 insertions, 0 deletions

diff --git a/.venv/lib/python3.12/site-packages/core/agent/research.py b/.venv/lib/python3.12/site-packages/core/agent/research.py
new file mode 100644
index 00000000..6ea35783
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/core/agent/research.py
@@ -0,0 +1,697 @@
+import logging
+import os
+import subprocess
+import sys
+import tempfile
+from copy import copy
+from typing import Any, Callable, Optional
+
+from core.base import AppConfig
+from core.base.abstractions import GenerationConfig, Message, SearchSettings
+from core.base.agent import Tool
+from core.base.providers import DatabaseProvider
+from core.providers import (
+    AnthropicCompletionProvider,
+    LiteLLMCompletionProvider,
+    OpenAICompletionProvider,
+    R2RCompletionProvider,
+)
+from core.utils import extract_citations
+
+from ..base.agent.agent import RAGAgentConfig  # type: ignore
+
+# Import the RAG agents we'll leverage
+from .rag import (  # type: ignore
+    R2RRAGAgent,
+    R2RStreamingRAGAgent,
+    R2RXMLToolsRAGAgent,
+    R2RXMLToolsStreamingRAGAgent,
+    RAGAgentMixin,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ResearchAgentMixin(RAGAgentMixin):
+    """
+    A mixin that extends RAGAgentMixin to add research capabilities to any R2R agent.
+
+    This mixin provides all RAG capabilities plus additional research tools:
+    - A RAG tool for knowledge retrieval (which leverages the underlying RAG capabilities)
+    - A Python execution tool for code execution and computation
+    - A reasoning tool for complex problem solving
+    - A critique tool for analyzing conversation history
+    """
+
+    def __init__(
+        self,
+        *args,
+        app_config: AppConfig,
+        search_settings: SearchSettings,
+        knowledge_search_method: Callable,
+        content_method: Callable,
+        file_search_method: Callable,
+        max_tool_context_length=10_000,
+        **kwargs,
+    ):
+        # Store the app configuration needed for research tools
+        self.app_config = app_config
+
+        # Call the parent RAGAgentMixin's __init__ with explicitly passed parameters
+        super().__init__(
+            *args,
+            search_settings=search_settings,
+            knowledge_search_method=knowledge_search_method,
+            content_method=content_method,
+            file_search_method=file_search_method,
+            max_tool_context_length=max_tool_context_length,
+            **kwargs,
+        )
+
+        # Register our research-specific tools
+        self._register_research_tools()
+
+    def _register_research_tools(self):
+        """
+        Register research-specific tools to the agent.
+        This is called by the mixin's __init__ after the parent class initialization.
+        """
+        # Add our research tools to whatever tools are already registered
+        research_tools = []
+        for tool_name in set(self.config.research_tools):
+            if tool_name == "rag":
+                research_tools.append(self.rag_tool())
+            elif tool_name == "reasoning":
+                research_tools.append(self.reasoning_tool())
+            elif tool_name == "critique":
+                research_tools.append(self.critique_tool())
+            elif tool_name == "python_executor":
+                research_tools.append(self.python_execution_tool())
+            else:
+                logger.warning(f"Unknown research tool: {tool_name}")
+                raise ValueError(f"Unknown research tool: {tool_name}")
+
+        logger.debug(f"Registered research tools: {research_tools}")
+        self.tools = research_tools
+
+    def rag_tool(self) -> Tool:
+        """Tool that provides access to the RAG agent's search capabilities."""
+        return Tool(
+            name="rag",
+            description=(
+                "Search for information using RAG (Retrieval-Augmented Generation). "
+                "This tool searches across relevant sources and returns comprehensive information. "
+                "Use this tool when you need to find specific information on any topic. Be sure to pose your query as a comprehensive query."
+            ),
+            results_function=self._rag,
+            llm_format_function=self._format_search_results,
+            parameters={
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "The search query to find information.",
+                    }
+                },
+                "required": ["query"],
+            },
+        )
+
+    def reasoning_tool(self) -> Tool:
+        """Tool that provides access to a strong reasoning model."""
+        return Tool(
+            name="reasoning",
+            description=(
+                "A dedicated reasoning system that excels at solving complex problems through step-by-step analysis. "
+                "This tool connects to a separate AI system optimized for deep analytical thinking.\n\n"
+                "USAGE GUIDELINES:\n"
+                "1. Formulate your request as a complete, standalone question to a reasoning expert.\n"
+                "2. Clearly state the problem/question at the beginning.\n"
+                "3. Provide all relevant context, data, and constraints.\n\n"
+                "IMPORTANT: This system has no memory of previous interactions or context from your conversation.\n\n"
+                "STRENGTHS: Mathematical reasoning, logical analysis, evaluating complex scenarios, "
+                "solving multi-step problems, and identifying potential errors in reasoning."
+            ),
+            results_function=self._reason,
+            llm_format_function=self._format_search_results,
+            parameters={
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "A complete, standalone question with all necessary context, appropriate for a dedicated reasoning system.",
+                    }
+                },
+                "required": ["query"],
+            },
+        )
+
+    def critique_tool(self) -> Tool:
+        """Tool that provides critical analysis of the reasoning done so far in the conversation."""
+        return Tool(
+            name="critique",
+            description=(
+                "Analyzes the conversation history to identify potential flaws, biases, and alternative "
+                "approaches to the reasoning presented so far.\n\n"
+                "Use this tool to get a second opinion on your reasoning, find overlooked considerations, "
+                "identify biases or fallacies, explore alternative hypotheses, and improve the robustness "
+                "of your conclusions."
+            ),
+            results_function=self._critique,
+            llm_format_function=self._format_search_results,
+            parameters={
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "A specific aspect of the reasoning you want critiqued, or leave empty for a general critique.",
+                    },
+                    "focus_areas": {
+                        "type": "array",
+                        "items": {"type": "string"},
+                        "description": "Optional specific areas to focus the critique (e.g., ['logical fallacies', 'methodology'])",
+                    },
+                },
+                "required": ["query"],
+            },
+        )
+
+    def python_execution_tool(self) -> Tool:
+        """Tool that provides Python code execution capabilities."""
+        return Tool(
+            name="python_executor",
+            description=(
+                "Executes Python code and returns the results, output, and any errors. "
+                "Use this tool for complex calculations, statistical operations, or algorithmic implementations.\n\n"
+                "The execution environment includes common libraries such as numpy, pandas, sympy, scipy, statsmodels, biopython, etc.\n\n"
+                "USAGE:\n"
+                "1. Send complete, executable Python code as a string.\n"
+                "2. Use print statements for output you want to see.\n"
+                "3. Assign to the 'result' variable for values you want to return.\n"
+                "4. Do not use input() or plotting (matplotlib). Output is text-based."
+            ),
+            results_function=self._execute_python_with_process_timeout,
+            llm_format_function=self._format_python_results,
+            parameters={
+                "type": "object",
+                "properties": {
+                    "code": {
+                        "type": "string",
+                        "description": "Python code to execute.",
+                    }
+                },
+                "required": ["code"],
+            },
+        )
+
+    async def _rag(
+        self,
+        query: str,
+        *args,
+        **kwargs,
+    ) -> dict[str, Any]:
+        """Execute a search using an internal RAG agent."""
+        # Create a copy of the current configuration for the RAG agent
+        config_copy = copy(self.config)
+        config_copy.max_iterations = 10  # Could be configurable
+        config_copy.rag_tools = [
+            "web_search",
+            "web_scrape",
+        ]  # HACK HACK TODO - Fix.
+
+        # Create a generation config for the RAG agent
+        generation_config = GenerationConfig(
+            model=self.app_config.quality_llm,
+            max_tokens_to_sample=16000,
+        )
+
+        # Create a new RAG agent - we'll use the non-streaming variant for consistent results
+        rag_agent = R2RRAGAgent(
+            database_provider=self.database_provider,
+            llm_provider=self.llm_provider,
+            config=config_copy,
+            search_settings=self.search_settings,
+            rag_generation_config=generation_config,
+            knowledge_search_method=self.knowledge_search_method,
+            content_method=self.content_method,
+            file_search_method=self.file_search_method,
+            max_tool_context_length=self.max_tool_context_length,
+        )
+
+        # Run the RAG agent with the query
+        user_message = Message(role="user", content=query)
+        response = await rag_agent.arun(messages=[user_message])
+
+        # Get the content from the response
+        structured_content = response[-1].get("structured_content")
+        if structured_content:
+            possible_text = structured_content[-1].get("text")
+            content = response[-1].get("content") or possible_text
+        else:
+            content = response[-1].get("content")
+
+        # Extract citations and transfer search results from RAG agent to research agent
+        short_ids = extract_citations(content)
+        if short_ids:
+            logger.info(f"Found citations in RAG response: {short_ids}")
+
+            for short_id in short_ids:
+                result = rag_agent.search_results_collector.find_by_short_id(
+                    short_id
+                )
+                if result:
+                    self.search_results_collector.add_result(result)
+
+            # Log confirmation for successful transfer
+            logger.info(
+                "Transferred search results from RAG agent to research agent for citations"
+            )
+        return content
+
+    async def _reason(
+        self,
+        query: str,
+        *args,
+        **kwargs,
+    ) -> dict[str, Any]:
+        """Execute a reasoning query using a specialized reasoning LLM."""
+        msg_list = await self.conversation.get_messages()
+
+        # Create a specialized generation config for reasoning
+        gen_cfg = self.get_generation_config(msg_list[-1], stream=False)
+        gen_cfg.model = self.app_config.reasoning_llm
+        gen_cfg.top_p = None
+        gen_cfg.temperature = 0.1
+        gen_cfg.max_tokens_to_sample = 64000
+        gen_cfg.stream = False
+        gen_cfg.tools = None
+        gen_cfg.functions = None
+        gen_cfg.reasoning_effort = "high"
+        gen_cfg.add_generation_kwargs = None
+
+        # Call the LLM with the reasoning request
+        response = await self.llm_provider.aget_completion(
+            [{"role": "user", "content": query}], gen_cfg
+        )
+        return response.choices[0].message.content
+
+    async def _critique(
+        self,
+        query: str,
+        focus_areas: Optional[list] = None,
+        *args,
+        **kwargs,
+    ) -> dict[str, Any]:
+        """Critique the conversation history."""
+        msg_list = await self.conversation.get_messages()
+        if not focus_areas:
+            focus_areas = []
+        # Build the critique prompt
+        critique_prompt = (
+            "You are a critical reasoning expert. Your task is to analyze the following conversation "
+            "and critique the reasoning. Look for:\n"
+            "1. Logical fallacies or inconsistencies\n"
+            "2. Cognitive biases\n"
+            "3. Overlooked questions or considerations\n"
+            "4. Alternative approaches\n"
+            "5. Improvements in rigor\n\n"
+        )
+
+        if focus_areas:
+            critique_prompt += f"Focus areas: {', '.join(focus_areas)}\n\n"
+
+        if query.strip():
+            critique_prompt += f"Specific question: {query}\n\n"
+
+        critique_prompt += (
+            "Structure your critique:\n"
+            "1. Summary\n"
+            "2. Key strengths\n"
+            "3. Potential issues\n"
+            "4. Alternatives\n"
+            "5. Recommendations\n\n"
+        )
+
+        # Add the conversation history to the prompt
+        conversation_text = "\n--- CONVERSATION HISTORY ---\n\n"
+        for msg in msg_list:
+            role = msg.get("role", "")
+            content = msg.get("content", "")
+            if content and role in ["user", "assistant", "system"]:
+                conversation_text += f"{role.upper()}: {content}\n\n"
+
+        final_prompt = critique_prompt + conversation_text
+
+        # Use the reasoning tool to process the critique
+        return await self._reason(final_prompt, *args, **kwargs)
+
+    async def _execute_python_with_process_timeout(
+        self, code: str, timeout: int = 10, *args, **kwargs
+    ) -> dict[str, Any]:
+        """
+        Executes Python code in a separate subprocess with a timeout.
+        This provides isolation and prevents re-importing the current agent module.
+
+        Parameters:
+          code (str): Python code to execute.
+          timeout (int): Timeout in seconds (default: 10).
+
+        Returns:
+          dict[str, Any]: Dictionary containing stdout, stderr, return code, etc.
+        """
+        # Write user code to a temporary file
+        with tempfile.NamedTemporaryFile(
+            mode="w", suffix=".py", delete=False
+        ) as tmp_file:
+            tmp_file.write(code)
+            script_path = tmp_file.name
+
+        try:
+            # Run the script in a fresh subprocess
+            result = subprocess.run(
+                [sys.executable, script_path],
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+            )
+
+            return {
+                "result": None,  # We'll parse from stdout if needed
+                "stdout": result.stdout,
+                "stderr": result.stderr,
+                "error": (
+                    None
+                    if result.returncode == 0
+                    else {
+                        "type": "SubprocessError",
+                        "message": f"Process exited with code {result.returncode}",
+                        "traceback": "",
+                    }
+                ),
+                "locals": {},  # No direct local var capture in a separate process
+                "success": (result.returncode == 0),
+                "timed_out": False,
+                "timeout": timeout,
+            }
+        except subprocess.TimeoutExpired as e:
+            return {
+                "result": None,
+                "stdout": e.output or "",
+                "stderr": e.stderr or "",
+                "error": {
+                    "type": "TimeoutError",
+                    "message": f"Execution exceeded {timeout} second limit.",
+                    "traceback": "",
+                },
+                "locals": {},
+                "success": False,
+                "timed_out": True,
+                "timeout": timeout,
+            }
+        finally:
+            # Clean up the temp file
+            if os.path.exists(script_path):
+                os.remove(script_path)
+
+    def _format_python_results(self, results: dict[str, Any]) -> str:
+        """Format Python execution results for display."""
+        output = []
+
+        # Timeout notification
+        if results.get("timed_out", False):
+            output.append(
+                f"⚠️ **Execution Timeout**: Code exceeded the {results.get('timeout', 10)} second limit."
+            )
+            output.append("")
+
+        # Stdout
+        if results.get("stdout"):
+            output.append("## Output:")
+            output.append("```")
+            output.append(results["stdout"].rstrip())
+            output.append("```")
+            output.append("")
+
+        # If there's a 'result' variable to display
+        if results.get("result") is not None:
+            output.append("## Result:")
+            output.append("```")
+            output.append(str(results["result"]))
+            output.append("```")
+            output.append("")
+
+        # Error info
+        if not results.get("success", True):
+            output.append("## Error:")
+            output.append("```")
+            stderr_out = results.get("stderr", "").rstrip()
+            if stderr_out:
+                output.append(stderr_out)
+
+            err_obj = results.get("error")
+            if err_obj and err_obj.get("message"):
+                output.append(err_obj["message"])
+            output.append("```")
+
+        # Return formatted output
+        return (
+            "\n".join(output)
+            if output
+            else "Code executed with no output or result."
+        )
+
+    def _format_search_results(self, results) -> str:
+        """Simple pass-through formatting for RAG search results."""
+        return results
+
+
+class R2RResearchAgent(ResearchAgentMixin, R2RRAGAgent):
+    """
+    A non-streaming research agent that uses the standard R2R agent as its base.
+
+    This agent combines research capabilities with the non-streaming RAG agent,
+    providing tools for deep research through tool-based interaction.
+    """
+
+    def __init__(
+        self,
+        app_config: AppConfig,
+        database_provider: DatabaseProvider,
+        llm_provider: (
+            AnthropicCompletionProvider
+            | LiteLLMCompletionProvider
+            | OpenAICompletionProvider
+            | R2RCompletionProvider
+        ),
+        config: RAGAgentConfig,
+        search_settings: SearchSettings,
+        rag_generation_config: GenerationConfig,
+        knowledge_search_method: Callable,
+        content_method: Callable,
+        file_search_method: Callable,
+        max_tool_context_length: int = 20_000,
+    ):
+        # Set a higher max iterations for research tasks
+        config.max_iterations = config.max_iterations or 15
+
+        # Initialize the RAG agent first
+        R2RRAGAgent.__init__(
+            self,
+            database_provider=database_provider,
+            llm_provider=llm_provider,
+            config=config,
+            search_settings=search_settings,
+            rag_generation_config=rag_generation_config,
+            knowledge_search_method=knowledge_search_method,
+            content_method=content_method,
+            file_search_method=file_search_method,
+            max_tool_context_length=max_tool_context_length,
+        )
+
+        # Then initialize the ResearchAgentMixin
+        ResearchAgentMixin.__init__(
+            self,
+            app_config=app_config,
+            database_provider=database_provider,
+            llm_provider=llm_provider,
+            config=config,
+            search_settings=search_settings,
+            rag_generation_config=rag_generation_config,
+            max_tool_context_length=max_tool_context_length,
+            knowledge_search_method=knowledge_search_method,
+            file_search_method=file_search_method,
+            content_method=content_method,
+        )
+
+
+class R2RStreamingResearchAgent(ResearchAgentMixin, R2RStreamingRAGAgent):
+    """
+    A streaming research agent that uses the streaming RAG agent as its base.
+
+    This agent combines research capabilities with streaming text generation,
+    providing real-time responses while still offering research tools.
+    """
+
+    def __init__(
+        self,
+        app_config: AppConfig,
+        database_provider: DatabaseProvider,
+        llm_provider: (
+            AnthropicCompletionProvider
+            | LiteLLMCompletionProvider
+            | OpenAICompletionProvider
+            | R2RCompletionProvider
+        ),
+        config: RAGAgentConfig,
+        search_settings: SearchSettings,
+        rag_generation_config: GenerationConfig,
+        knowledge_search_method: Callable,
+        content_method: Callable,
+        file_search_method: Callable,
+        max_tool_context_length: int = 10_000,
+    ):
+        # Force streaming on
+        config.stream = True
+        config.max_iterations = config.max_iterations or 15
+
+        # Initialize the streaming RAG agent first
+        R2RStreamingRAGAgent.__init__(
+            self,
+            database_provider=database_provider,
+            llm_provider=llm_provider,
+            config=config,
+            search_settings=search_settings,
+            rag_generation_config=rag_generation_config,
+            knowledge_search_method=knowledge_search_method,
+            content_method=content_method,
+            file_search_method=file_search_method,
+            max_tool_context_length=max_tool_context_length,
+        )
+
+        # Then initialize the ResearchAgentMixin
+        ResearchAgentMixin.__init__(
+            self,
+            app_config=app_config,
+            database_provider=database_provider,
+            llm_provider=llm_provider,
+            config=config,
+            search_settings=search_settings,
+            rag_generation_config=rag_generation_config,
+            max_tool_context_length=max_tool_context_length,
+            knowledge_search_method=knowledge_search_method,
+            content_method=content_method,
+            file_search_method=file_search_method,
+        )
+
+
+class R2RXMLToolsResearchAgent(ResearchAgentMixin, R2RXMLToolsRAGAgent):
+    """
+    A non-streaming research agent that uses XML tool formatting.
+
+    This agent combines research capabilities with the XML-based tool calling format,
+    which might be more appropriate for certain LLM providers.
+    """
+
+    def __init__(
+        self,
+        app_config: AppConfig,
+        database_provider: DatabaseProvider,
+        llm_provider: (
+            AnthropicCompletionProvider
+            | LiteLLMCompletionProvider
+            | OpenAICompletionProvider
+            | R2RCompletionProvider
+        ),
+        config: RAGAgentConfig,
+        search_settings: SearchSettings,
+        rag_generation_config: GenerationConfig,
+        knowledge_search_method: Callable,
+        content_method: Callable,
+        file_search_method: Callable,
+        max_tool_context_length: int = 20_000,
+    ):
+        # Set higher max iterations
+        config.max_iterations = config.max_iterations or 15
+
+        # Initialize the XML Tools RAG agent first
+        R2RXMLToolsRAGAgent.__init__(
+            self,
+            database_provider=database_provider,
+            llm_provider=llm_provider,
+            config=config,
+            search_settings=search_settings,
+            rag_generation_config=rag_generation_config,
+            knowledge_search_method=knowledge_search_method,
+            content_method=content_method,
+            file_search_method=file_search_method,
+            max_tool_context_length=max_tool_context_length,
+        )
+
+        # Then initialize the ResearchAgentMixin
+        ResearchAgentMixin.__init__(
+            self,
+            app_config=app_config,
+            search_settings=search_settings,
+            knowledge_search_method=knowledge_search_method,
+            content_method=content_method,
+            file_search_method=file_search_method,
+            max_tool_context_length=max_tool_context_length,
+        )
+
+
+class R2RXMLToolsStreamingResearchAgent(
+    ResearchAgentMixin, R2RXMLToolsStreamingRAGAgent
+):
+    """
+    A streaming research agent that uses XML tool formatting.
+
+    This agent combines research capabilities with streaming and XML-based tool calling,
+    providing real-time responses in a format suitable for certain LLM providers.
+    """
+
+    def __init__(
+        self,
+        app_config: AppConfig,
+        database_provider: DatabaseProvider,
+        llm_provider: (
+            AnthropicCompletionProvider
+            | LiteLLMCompletionProvider
+            | OpenAICompletionProvider
+            | R2RCompletionProvider
+        ),
+        config: RAGAgentConfig,
+        search_settings: SearchSettings,
+        rag_generation_config: GenerationConfig,
+        knowledge_search_method: Callable,
+        content_method: Callable,
+        file_search_method: Callable,
+        max_tool_context_length: int = 10_000,
+    ):
+        # Force streaming on
+        config.stream = True
+        config.max_iterations = config.max_iterations or 15
+
+        # Initialize the XML Tools Streaming RAG agent first
+        R2RXMLToolsStreamingRAGAgent.__init__(
+            self,
+            database_provider=database_provider,
+            llm_provider=llm_provider,
+            config=config,
+            search_settings=search_settings,
+            rag_generation_config=rag_generation_config,
+            knowledge_search_method=knowledge_search_method,
+            content_method=content_method,
+            file_search_method=file_search_method,
+            max_tool_context_length=max_tool_context_length,
+        )
+
+        # Then initialize the ResearchAgentMixin
+        ResearchAgentMixin.__init__(
+            self,
+            app_config=app_config,
+            search_settings=search_settings,
+            knowledge_search_method=knowledge_search_method,
+            content_method=content_method,
+            file_search_method=file_search_method,
+            max_tool_context_length=max_tool_context_length,
+        )