""" Research Orchestrator for GAIA Agent Intelligent coordination of multiple research tools with result synthesis """ import os import logging from typing import Dict, List, Any, Optional, Union, Tuple from dataclasses import dataclass from datetime import datetime import json import re from .web_research_tool import EnhancedWebSearchTool, SearchQuery, SearchResult from .wikipedia_tool import WikipediaSpecializedTool, WikipediaArticle logger = logging.getLogger(__name__) @dataclass class ResearchQuery: """Structured research query with analysis metadata.""" original_question: str query_type: str # factual, biographical, historical, technical, numerical entities: List[str] # Named entities extracted from question time_constraints: Optional[Dict[str, Any]] = None domain_hints: Optional[List[str]] = None expected_answer_type: str = "text" # text, number, date, list confidence_threshold: float = 0.7 @dataclass class ResearchResult: """Comprehensive research result with confidence scoring.""" answer: str confidence: float sources: List[Dict[str, Any]] reasoning: str alternative_answers: List[str] verification_status: str # verified, partial, unverified search_strategy_used: str class ResearchOrchestrator: """ Intelligent research orchestrator that coordinates multiple tools. Features: - Query analysis and classification - Multi-tool coordination - Result synthesis and validation - Confidence scoring - Source verification - Fallback strategies Note: This orchestrator is designed to work WITH AGNO's orchestration, not replace it. It provides specialized research capabilities that AGNO tools can call when needed. """ def __init__(self, exa_api_key: Optional[str] = None): """Initialize the research orchestrator.""" self.web_search = EnhancedWebSearchTool(exa_api_key) self.wikipedia = WikipediaSpecializedTool() # Research strategies for different question types self.strategies = { 'factual': self._factual_research_strategy, 'biographical': self._biographical_research_strategy, 'historical': self._historical_research_strategy, 'technical': self._technical_research_strategy, 'numerical': self._numerical_research_strategy, 'discography': self._discography_research_strategy, 'featured_article': self._featured_article_research_strategy } logger.info("✅ Research Orchestrator initialized") def research(self, question: str, **kwargs) -> ResearchResult: """ Perform comprehensive research on a question. Args: question: The research question **kwargs: Additional parameters Returns: ResearchResult with comprehensive findings """ try: logger.info(f"🔬 Starting research: {question[:100]}...") # Analyze the query research_query = self._analyze_query(question, **kwargs) # Select and execute research strategy strategy = self.strategies.get( research_query.query_type, self._general_research_strategy ) result = strategy(research_query) logger.info(f"✅ Research completed with confidence: {result.confidence:.2f}") return result except Exception as e: logger.error(f"❌ Research error: {e}") return ResearchResult( answer="Research failed", confidence=0.0, sources=[], reasoning=f"Error during research: {str(e)}", alternative_answers=[], verification_status="unverified", search_strategy_used="error" ) def _analyze_query(self, question: str, **kwargs) -> ResearchQuery: """Analyze and classify the research query.""" question_lower = question.lower() # Determine query type query_type = "factual" # default if any(word in question_lower for word in ['album', 'song', 'discography', 'studio album']): query_type = "discography" elif any(word in question_lower for word in ['featured article', 'wikipedia featured']): query_type = "featured_article" elif any(word in question_lower for word in ['born', 'died', 'biography', 'life']): query_type = "biographical" elif any(word in question_lower for word in ['when', 'year', 'date', 'time']): query_type = "historical" elif any(word in question_lower for word in ['how many', 'count', 'number']): query_type = "numerical" elif any(word in question_lower for word in ['technical', 'algorithm', 'method']): query_type = "technical" # Extract entities (simplified) entities = self._extract_entities(question) # Extract time constraints time_constraints = self._extract_time_constraints(question) return ResearchQuery( original_question=question, query_type=query_type, entities=entities, time_constraints=time_constraints, expected_answer_type=kwargs.get('expected_answer_type', 'text'), confidence_threshold=kwargs.get('confidence_threshold', 0.7) ) def _extract_entities(self, question: str) -> List[str]: """Extract named entities from the question.""" # Simplified entity extraction # In production, you'd use spaCy or similar NLP library entities = [] # Look for quoted strings quoted_entities = re.findall(r'"([^"]*)"', question) entities.extend(quoted_entities) # Look for capitalized words (potential proper nouns) words = question.split() for word in words: if word[0].isupper() and len(word) > 2 and word not in ['The', 'A', 'An', 'In', 'On', 'At']: entities.append(word) return list(set(entities)) def _extract_time_constraints(self, question: str) -> Optional[Dict[str, Any]]: """Extract time-related constraints from the question.""" time_patterns = [ (r'(\d{4})-(\d{4})', 'year_range'), (r'between (\d{4}) and (\d{4})', 'year_range'), (r'in (\d{4})', 'specific_year'), (r'(\d{4})', 'year_mention'), (r'(January|February|March|April|May|June|July|August|September|October|November|December) (\d{4})', 'month_year') ] for pattern, constraint_type in time_patterns: match = re.search(pattern, question, re.IGNORECASE) if match: if constraint_type == 'year_range': return { 'type': 'range', 'start_year': int(match.group(1)), 'end_year': int(match.group(2)) } elif constraint_type == 'specific_year': return { 'type': 'specific', 'year': int(match.group(1)) } elif constraint_type == 'month_year': return { 'type': 'month_year', 'month': match.group(1), 'year': int(match.group(2)) } return None def _factual_research_strategy(self, query: ResearchQuery) -> ResearchResult: """Research strategy for factual questions.""" sources = [] answers = [] # Try web search first web_results = self.web_search.search( SearchQuery( query=query.original_question, query_type="factual", num_results=5 ) ) for result in web_results[:3]: sources.append({ 'type': 'web', 'title': result.title, 'url': result.url, 'score': result.score }) # Try to extract answer from content if result.content: potential_answer = self._extract_factual_answer(result.content, query.original_question) if potential_answer: answers.append(potential_answer) # Try Wikipedia if web search didn't yield good results if len(answers) < 2: wiki_results = self.wikipedia.search_articles(query.original_question, limit=3) for wiki_result in wiki_results: article = self.wikipedia.get_article(wiki_result.title, include_content=False) if article: sources.append({ 'type': 'wikipedia', 'title': article.title, 'url': article.url, 'score': 0.8 }) if article.summary: potential_answer = self._extract_factual_answer(article.summary, query.original_question) if potential_answer: answers.append(potential_answer) # Synthesize final answer final_answer, confidence = self._synthesize_answers(answers, query) return ResearchResult( answer=final_answer, confidence=confidence, sources=sources, reasoning=f"Used factual research strategy with {len(sources)} sources", alternative_answers=answers[1:] if len(answers) > 1 else [], verification_status="verified" if confidence > 0.8 else "partial", search_strategy_used="factual" ) def _discography_research_strategy(self, query: ResearchQuery) -> ResearchResult: """Research strategy for discography questions.""" sources = [] # Extract artist name from entities artist_name = None for entity in query.entities: if len(entity) > 3: # Likely an artist name artist_name = entity break if not artist_name: # Try to extract from question words = query.original_question.split() for i, word in enumerate(words): if word.lower() in ['albums', 'discography'] and i > 0: artist_name = words[i-1] break if not artist_name: return ResearchResult( answer="Could not identify artist name", confidence=0.1, sources=[], reasoning="Failed to extract artist name from question", alternative_answers=[], verification_status="unverified", search_strategy_used="discography" ) # Get discography information albums = self.wikipedia.extract_discography_info(artist_name, "studio") # Filter by time constraints if present if query.time_constraints and query.time_constraints.get('type') == 'range': start_year = query.time_constraints['start_year'] end_year = query.time_constraints['end_year'] albums = [album for album in albums if start_year <= album.get('year', 0) <= end_year] sources.append({ 'type': 'wikipedia_discography', 'artist': artist_name, 'albums_found': len(albums) }) # Format answer if albums: album_count = len(albums) answer = str(album_count) confidence = 0.9 if album_count > 0 else 0.3 else: answer = "0" confidence = 0.3 return ResearchResult( answer=answer, confidence=confidence, sources=sources, reasoning=f"Found {len(albums)} studio albums for {artist_name}", alternative_answers=[], verification_status="verified" if confidence > 0.7 else "partial", search_strategy_used="discography" ) def _featured_article_research_strategy(self, query: ResearchQuery) -> ResearchResult: """Research strategy for Wikipedia featured article questions.""" sources = [] # Extract date and topic from query date_str = None topic_keywords = [] if query.time_constraints: if query.time_constraints.get('type') == 'month_year': month = query.time_constraints['month'] year = query.time_constraints['year'] # Convert to date format (assuming mid-month) month_num = { 'january': 1, 'february': 2, 'march': 3, 'april': 4, 'may': 5, 'june': 6, 'july': 7, 'august': 8, 'september': 9, 'october': 10, 'november': 11, 'december': 12 }.get(month.lower(), 1) date_str = f"{year}-{month_num:02d}-15" # Extract topic keywords question_lower = query.original_question.lower() if 'dinosaur' in question_lower: topic_keywords = ['dinosaur', 'paleontology', 'fossil'] # Search for featured article if date_str and topic_keywords: featured_article = self.wikipedia.find_featured_article_by_date(date_str, topic_keywords) if featured_article: sources.append({ 'type': 'wikipedia_featured', 'date': date_str, 'article': featured_article }) return ResearchResult( answer=featured_article, confidence=0.9, sources=sources, reasoning=f"Found featured article for {date_str}: {featured_article}", alternative_answers=[], verification_status="verified", search_strategy_used="featured_article" ) return ResearchResult( answer="Featured article not found", confidence=0.1, sources=sources, reasoning="Could not locate featured article for specified criteria", alternative_answers=[], verification_status="unverified", search_strategy_used="featured_article" ) def _general_research_strategy(self, query: ResearchQuery) -> ResearchResult: """General research strategy for unclassified questions.""" return self._factual_research_strategy(query) def _biographical_research_strategy(self, query: ResearchQuery) -> ResearchResult: """Research strategy for biographical questions.""" return self._factual_research_strategy(query) def _historical_research_strategy(self, query: ResearchQuery) -> ResearchResult: """Research strategy for historical questions.""" return self._factual_research_strategy(query) def _technical_research_strategy(self, query: ResearchQuery) -> ResearchResult: """Research strategy for technical questions.""" return self._factual_research_strategy(query) def _numerical_research_strategy(self, query: ResearchQuery) -> ResearchResult: """Research strategy for numerical questions.""" return self._factual_research_strategy(query) def _extract_factual_answer(self, content: str, question: str) -> Optional[str]: """Extract a factual answer from content.""" # Simplified answer extraction sentences = content.split('.') question_words = set(question.lower().split()) best_sentence = None best_score = 0 for sentence in sentences: sentence = sentence.strip() if 10 < len(sentence) < 200: # Reasonable length sentence_words = set(sentence.lower().split()) overlap = len(question_words & sentence_words) if overlap > best_score: best_score = overlap best_sentence = sentence return best_sentence if best_score > 2 else None def _synthesize_answers(self, answers: List[str], query: ResearchQuery) -> Tuple[str, float]: """Synthesize multiple answers into a final answer with confidence.""" if not answers: return "No answer found", 0.0 # For now, return the first answer with confidence based on number of sources final_answer = answers[0] confidence = min(0.9, 0.3 + (len(answers) * 0.2)) return final_answer, confidence # AGNO Integration Methods def research_mercedes_sosa_albums(self, start_year: int = 2000, end_year: int = 2009) -> str: """ Specific method for Mercedes Sosa album research (GAIA question). This method can be called directly by AGNO tools. """ try: albums = self.wikipedia.search_mercedes_sosa_albums(start_year, end_year) return str(len(albums)) except Exception as e: logger.error(f"Mercedes Sosa research error: {e}") return "0" def research_featured_article(self, date: str, topic: str) -> str: """ Specific method for featured article research (GAIA question). This method can be called directly by AGNO tools. """ try: topic_keywords = [topic.lower()] if topic.lower() == 'dinosaur': topic_keywords = ['dinosaur', 'paleontology', 'fossil'] result = self.wikipedia.find_featured_article_by_date(date, topic_keywords) return result or "Not found" except Exception as e: logger.error(f"Featured article research error: {e}") return "Not found" def quick_factual_search(self, question: str) -> str: """ Quick factual search method for AGNO integration. Returns just the answer string for easy integration. """ try: result = self.research(question) return result.answer if result.confidence > 0.5 else "Not found" except Exception as e: logger.error(f"Quick search error: {e}") return "Error in search"