mcp-agents/deep_research_integration.py

"""
Integration with Deep Research workflow.

This module provides functions to integrate the MCP agents into the Deep Research workflow.
It hooks into the research question generation and hypothesis enhancement steps.
"""
import asyncio
import json
import logging
from typing import List, Dict, Any, Optional

from integration import process_research_questions, enhance_hypotheses, detect_research_domain, ResearchDomain

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("deep-research-integration")

async def enhance_research_questions(
    query: str,
    questions: List[str],
    domain_context: Optional[str] = None
) -> Dict[str, Any]:
    """
    Enhance research questions using the appropriate specialized agent.
    
    Args:
        query (str): The original research query.
        questions (List[str]): List of research questions to enhance.
        domain_context (str, optional): Additional domain context to provide.
        
    Returns:
        Dict[str, Any]: Enhanced research questions with explanations and context.
    """
    # Detect the research domain based on the query
    domain = detect_research_domain(query)
    logger.info(f"Detected research domain: {domain.value}")
    
    # Process the research questions using the appropriate agent
    enhanced_questions = await process_research_questions(questions, domain, domain_context)
    
    return {
        "domain": domain.value,
        "enhanced_questions": enhanced_questions
    }

async def enhance_research_hypotheses(
    query: str,
    hypotheses: List[Dict[str, Any]],
    research_goal: str
) -> Dict[str, Any]:
    """
    Enhance research hypotheses using the appropriate specialized agent.
    
    Args:
        query (str): The original research query.
        hypotheses (List[Dict[str, Any]]): List of hypotheses to enhance.
        research_goal (str): The research goal.
        
    Returns:
        Dict[str, Any]: Enhanced hypotheses with explanations and context.
    """
    # Detect the research domain based on the query
    domain = detect_research_domain(query)
    logger.info(f"Detected research domain: {domain.value}")
    
    # Enhance the hypotheses using the appropriate agent
    enhanced_hypotheses = await enhance_hypotheses(hypotheses, research_goal, domain)
    
    return {
        "domain": domain.value,
        "enhanced_hypotheses": enhanced_hypotheses
    }

# Function to be called from the Deep Research workflow after question generation
def post_question_generation_hook(query: str, questions: List[str]) -> Dict[str, Any]:
    """
    Hook to be called after question generation in the Deep Research workflow.
    
    Args:
        query (str): The original research query.
        questions (List[str]): The generated research questions.
        
    Returns:
        Dict[str, Any]: Enhanced research questions with explanations and context.
    """
    return asyncio.run(enhance_research_questions(query, questions))

# Function to be called from the AI Co-Scientist workflow after hypothesis generation
def post_hypothesis_generation_hook(query: str, hypotheses: List[Dict[str, Any]], research_goal: str) -> Dict[str, Any]:
    """
    Hook to be called after hypothesis generation in the AI Co-Scientist workflow.
    
    Args:
        query (str): The original research query.
        hypotheses (List[Dict[str, Any]]): The generated hypotheses.
        research_goal (str): The research goal.
        
    Returns:
        Dict[str, Any]: Enhanced hypotheses with explanations and context.
    """
    return asyncio.run(enhance_research_hypotheses(query, hypotheses, research_goal))

if __name__ == "__main__":
    # Example usage
    def main():
        # Example research query
        query = "Investigating the impact of climate change on biodiversity and ecosystem services"
        
        # Example research questions
        questions = [
            "What are the effects of climate change on biodiversity in tropical rainforests?",
            "How do changes in temperature affect species distribution in marine ecosystems?",
            "What are the economic implications of biodiversity loss due to climate change?"
        ]
        
        # Example research goal
        research_goal = "Understand the complex relationships between climate change, biodiversity loss, and ecosystem services to develop effective conservation strategies."
        
        # Example hypotheses
        hypotheses = [
            {
                "id": "H1",
                "title": "Temperature Effects on Species Distribution",
                "text": "Rising temperatures will cause a poleward shift in species distribution, leading to increased competition and potential extinction of species unable to adapt or migrate."
            },
            {
                "id": "H2",
                "title": "Ecosystem Service Degradation",
                "text": "Climate-induced biodiversity loss will significantly reduce ecosystem services such as carbon sequestration, water purification, and pollination, resulting in cascading economic impacts."
            }
        ]
        
        # Enhance research questions
        enhanced_questions = post_question_generation_hook(query, questions)
        print("Enhanced Questions:")
        print(json.dumps(enhanced_questions, indent=2))
        
        # Enhance hypotheses
        enhanced_hypotheses = post_hypothesis_generation_hook(query, hypotheses, research_goal)
        print("\nEnhanced Hypotheses:")
        print(json.dumps(enhanced_hypotheses, indent=2))
    
    main()