research_engine.py

import os
import json
import logging
import time
from typing import List, Dict, Any, Optional, Tuple, Union

from crewai import Crew
from crewai.agent import Agent
from crewai.task import Task

from agents import create_researcher_agent, create_analyst_agent, create_writer_agent
from tasks import (
    create_query_refinement_task, 
    create_search_task, 
    create_content_scraping_task, 
    create_content_analysis_task, 
    create_response_writing_task
)
from utils import is_valid_query, format_research_results, extract_citations

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class ResearchEngine:
    """
    Main engine for web research using CrewAI.
    Orchestrates agents and tasks to provide comprehensive research results.
    """
    
    def __init__(self, llm=None, verbose=False):
        """
        Initialize the research engine.
        
        Args:
            llm: The language model to use for agents
            verbose: Whether to log detailed information
        """
        self.llm = llm
        self.verbose = verbose
        
        # Initialize agents
        logger.info("Initializing agents...")
        self.researcher = create_researcher_agent(llm=llm, verbose=verbose)
        self.analyst = create_analyst_agent(llm=llm, verbose=verbose)
        self.writer = create_writer_agent(llm=llm, verbose=verbose)
        
        # Chat history for maintaining conversation context
        self.chat_history = []
        
        logger.info("Research engine initialized with agents")
    
    def _validate_api_keys(self):
        """Validates that required API keys are present"""
        missing_keys = []
        
        if not os.getenv("BRAVE_API_KEY"):
            missing_keys.append("BRAVE_API_KEY")
            
        if not os.getenv("TAVILY_API_KEY"):
            missing_keys.append("TAVILY_API_KEY")
            
        if not os.getenv("OPENAI_API_KEY") and not self.llm:
            missing_keys.append("OPENAI_API_KEY or custom LLM")
            
        if missing_keys:
            logger.warning(f"Missing API keys: {', '.join(missing_keys)}")
            if "TAVILY_API_KEY" in missing_keys:
                logger.warning("Tavily API key is missing - search functionality may be limited")
            if "BRAVE_API_KEY" in missing_keys:
                logger.warning("Brave API key is missing - search functionality may be limited")
            
            # Only raise error if all search API keys are missing
            if "BRAVE_API_KEY" in missing_keys and "TAVILY_API_KEY" in missing_keys:
                raise ValueError(f"Missing required API keys: {', '.join(missing_keys)}")
        else:
            logger.info("All required API keys are present")
    
    def research(self, query: str, output_file=None) -> Dict[str, Any]:
        """
        Perform research on the given query.
        
        Args:
            query: The research query
            output_file: Optional file to save the research results
            
        Returns:
            Research results
        """
        logger.info(f"Research initiated for query: {query}")
        start_time = time.time()  # Initialize the start_time for tracking processing time
        
        try:
            self._validate_api_keys()
            logger.info(f"Starting research for query: {query}")
            
            # Add the query to chat history
            self.chat_history.append({"role": "user", "content": query})
            
            # Step 1: Initialize the crew
            logger.info("Initializing research crew...")
            crew = Crew(
                agents=[self.researcher],
                tasks=[create_query_refinement_task(self.researcher, query)],
                verbose=self.verbose,  # Use the instance's verbose setting
                process="sequential"
            )
            
            # Step 2: Start the research process
            logger.info("Starting research process...")
            refinement_result = crew.kickoff(inputs={"query": query})
            logger.info(f"Query refinement completed with result type: {type(refinement_result)}")
            logger.debug(f"Refinement result: {refinement_result}")
            
            # Extract the refined query
            refined_query = None
            try:
                logger.info(f"Attempting to extract refined query from result type: {type(refinement_result)}")
                
                # Handle CrewOutput object (new CrewAI format)
                if hasattr(refinement_result, '__class__') and refinement_result.__class__.__name__ == 'CrewOutput':
                    logger.info("Processing CrewOutput format refinement result")
                    
                    # Try to access raw attribute first (contains the raw output)
                    if hasattr(refinement_result, 'raw'):
                        refined_query = self._extract_query_from_string(refinement_result.raw)
                        logger.info(f"Extracted from CrewOutput.raw: {refined_query}")
                    
                    # Try to access as dictionary
                    elif hasattr(refinement_result, 'to_dict'):
                        crew_dict = refinement_result.to_dict()
                        logger.info(f"CrewOutput converted to dict: {crew_dict}")
                        
                        if 'result' in crew_dict:
                            refined_query = self._extract_query_from_string(crew_dict['result'])
                            logger.info(f"Extracted from CrewOutput dict result: {refined_query}")
                    
                    # Try string representation as last resort
                    else:
                        crew_str = str(refinement_result)
                        refined_query = self._extract_query_from_string(crew_str)
                        logger.info(f"Extracted from CrewOutput string representation: {refined_query}")
                
                # First try to access it as a dictionary (new CrewAI format)
                elif isinstance(refinement_result, dict):
                    logger.info("Processing dictionary format refinement result")
                    if "query" in refinement_result:
                        refined_query = refinement_result["query"]
                    elif "refined_query" in refinement_result:
                        refined_query = refinement_result["refined_query"]
                    elif "result" in refinement_result and isinstance(refinement_result["result"], str):
                        # Try to extract from nested result field
                        json_str = refinement_result["result"]
                        refined_query = self._extract_query_from_string(json_str)
                
                # Then try to access it as a string (old CrewAI format)
                elif isinstance(refinement_result, str):
                    logger.info("Processing string format refinement result")
                    refined_query = self._extract_query_from_string(refinement_result)
                
                else:
                    logger.warning(f"Unexpected refinement result format: {type(refinement_result)}")
                    # Try to extract information by examining the structure
                    try:
                        # Try to access common attributes
                        if hasattr(refinement_result, "result"):
                            result_str = str(getattr(refinement_result, "result"))
                            refined_query = self._extract_query_from_string(result_str)
                            logger.info(f"Extracted from .result attribute: {refined_query}")
                        elif hasattr(refinement_result, "task_output"):
                            task_output = getattr(refinement_result, "task_output")
                            refined_query = self._extract_query_from_string(str(task_output))
                            logger.info(f"Extracted from .task_output attribute: {refined_query}")
                        else:
                            # Last resort: convert to string and extract
                            refined_query = self._extract_query_from_string(str(refinement_result))
                            logger.info(f"Extracted from string representation: {refined_query}")
                    except Exception as attr_error:
                        logger.exception(f"Error trying to extract attributes: {attr_error}")
                        refined_query = query  # Fall back to original query
                    
                    logger.debug(f"Refinement result: {refinement_result}")
            except Exception as e:
                logger.exception(f"Error extracting refined query: {e}")
                refined_query = query  # Fall back to original query on error
            
            if not refined_query or refined_query.strip() == "":
                logger.warning("Refined query is empty, using original query")
                refined_query = query
                
            logger.info(f"Refined query: {refined_query}")
            
            # Step 3: Create tasks for research process
            logger.info("Creating research tasks...")
            search_task = create_search_task(self.researcher, refined_query)
            
            scrape_task = create_content_scraping_task(self.analyst, search_task)
            
            analyze_task = create_content_analysis_task(self.analyst, refined_query, scrape_task)
            
            write_task = create_response_writing_task(self.writer, refined_query, analyze_task)
            
            # Step 4: Create a new crew for the research tasks
            logger.info("Initializing main research crew...")
            research_crew = Crew(
                agents=[self.researcher, self.analyst, self.writer],
                tasks=[search_task, scrape_task, analyze_task, write_task],
                verbose=self.verbose,  # Use the instance's verbose setting
                process="sequential"
            )
            
            # Step 5: Start the research process
            logger.info("Starting main research process...")
            result = research_crew.kickoff()
            logger.info(f"Research completed with result type: {type(result)}")
            logger.debug(f"Research result: {result}")
            
            # Extract the result
            final_result = {"query": query, "refined_query": refined_query}
            
            # Handle different result types
            if isinstance(result, dict) and "result" in result:
                final_result["result"] = result["result"]
            # Handle CrewOutput object (new CrewAI format)
            elif hasattr(result, '__class__') and result.__class__.__name__ == 'CrewOutput':
                logger.info("Processing CrewOutput format result")
                
                # Try to access raw attribute first (contains the raw output)
                if hasattr(result, 'raw'):
                    final_result["result"] = result.raw
                    logger.info("Extracted result from CrewOutput.raw")
                
                # Try to access as dictionary
                elif hasattr(result, 'to_dict'):
                    crew_dict = result.to_dict()
                    if 'result' in crew_dict:
                        final_result["result"] = crew_dict['result']
                        logger.info("Extracted result from CrewOutput dict")
                
                # Use string representation as last resort
                else:
                    final_result["result"] = str(result)
                    logger.info("Used string representation of CrewOutput")
            else:
                # For any other type, use the string representation
                final_result["result"] = str(result)
                logger.info(f"Used string representation for result type: {type(result)}")
                
            logger.info("Research process completed successfully")
            
            # Save to file if requested
            if output_file:
                with open(output_file, 'w', encoding='utf-8') as f:
                    json.dump(final_result, f, ensure_ascii=False, indent=2)
                    
            # Extract citations for easy access (if possible from the final string)
            citations = extract_citations(final_result["result"])
            
            # Calculate total processing time
            processing_time = time.time() - start_time
            logger.info(f"Research completed successfully in {processing_time:.2f} seconds")
            
            return {
                "result": final_result["result"],
                "success": True,
                "refined_query": refined_query,
                "citations": citations,
                "processing_time": processing_time
            }
        except Exception as e:
            logger.exception(f"Error in research process: {e}")
            return {
                "result": f"I encountered an error while researching your query: {str(e)}",
                "success": False,
                "reason": "research_error",
                "error": str(e)
            }
    
    def chat(self, message: str) -> str:
        """
        Handle a chat message, which could be a research query or a follow-up question.
        
        Args:
            message: The user's message
            
        Returns:
            The assistant's response
        """
        # Treat all messages as new research queries for simplicity
        try:
            research_result = self.research(message)
            return research_result["result"]
        except Exception as e:
            logger.exception(f"Error during research for message: {message}")
            return f"I encountered an error while processing your request: {str(e)}"
    
    def clear_history(self):
        """Clear the chat history"""
        self.chat_history = []

    def _extract_query_from_string(self, text: str) -> str:
        """
        Extract refined query from text string, handling various formats including JSON embedded in strings.
        
        Args:
            text: The text to extract the query from
            
        Returns:
            The extracted query or None if not found
        """
        if not text:
            return None
            
        # Log the input for debugging
        logger.debug(f"Extracting query from: {text[:200]}...")
            
        # Try to parse as JSON first
        try:
            # Check if the entire string is valid JSON
            json_data = json.loads(text)
            
            # Check for known keys in the parsed JSON
            if isinstance(json_data, dict):
                if "refined_query" in json_data:
                    return json_data["refined_query"]
                elif "query" in json_data:
                    return json_data["query"]
                elif "result" in json_data and isinstance(json_data["result"], str):
                    # Try to recursively extract from nested result
                    return self._extract_query_from_string(json_data["result"])
        except json.JSONDecodeError:
            # Not valid JSON, continue with string parsing
            pass
            
        # Look for JSON blocks in the string
        try:
            import re
            # Match both markdown JSON blocks and regular JSON objects
            json_pattern = r'```(?:json)?\s*({[^`]*})```|({[\s\S]*})'
            json_matches = re.findall(json_pattern, text, re.DOTALL)
            
            for json_match in json_matches:
                # Handle tuple result from findall with multiple capture groups
                json_str = next((s for s in json_match if s), '')
                try:
                    json_data = json.loads(json_str)
                    if isinstance(json_data, dict):
                        if "refined_query" in json_data:
                            return json_data["refined_query"]
                        elif "query" in json_data:
                            return json_data["query"]
                except Exception:
                    continue
        except Exception as e:
            logger.debug(f"Error parsing JSON blocks: {e}")
            
        # Check for common patterns in CrewAI output format
        patterns = [
            r'refined query[:\s]+([^\n]+)',
            r'query[:\s]+([^\n]+)',
            r'search(?:ed)? for[:\s]+[\'"]([^\'"]+)[\'"]',
            r'search(?:ing)? for[:\s]+[\'"]([^\'"]+)[\'"]',
            r'research(?:ing)? (?:about|on)[:\s]+[\'"]([^\'"]+)[\'"]',
            r'query is[:\s]+[\'"]([^\'"]+)[\'"]'
        ]
        
        for pattern in patterns:
            try:
                match = re.search(pattern, text.lower())
                if match:
                    return match.group(1).strip()
            except Exception as e:
                logger.debug(f"Error matching pattern {pattern}: {e}")
                
        # Fall back to string parsing methods
        if "refined query:" in text.lower():
            return text.split("refined query:", 1)[1].strip()
        elif "query:" in text.lower():
            return text.split("query:", 1)[1].strip()
            
        # If all else fails, return the whole string
        return text