aiecs 1.2.1__py3-none-any.whl → 1.3.1__py3-none-any.whl

aiecs/tools/search_tool/core.py

@@ -0,0 +1,773 @@
+"""
+Core SearchTool Implementation
+
+Enhanced Google Custom Search Tool with quality analysis, intent understanding,
+intelligent caching, and comprehensive metrics.
+"""
+
+import asyncio
+import logging
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field, ConfigDict
+
+from aiecs.tools.base_tool import BaseTool
+from aiecs.tools.tool_executor import cache_result_with_strategy
+from aiecs.config.config import get_settings
+
+# Import Google API with graceful fallback
+try:
+    from googleapiclient.discovery import build
+    from googleapiclient.errors import HttpError
+    from google.auth.exceptions import GoogleAuthError
+    from google.oauth2 import service_account
+    GOOGLE_API_AVAILABLE = True
+except ImportError:
+    GOOGLE_API_AVAILABLE = False
+    HttpError = Exception
+    GoogleAuthError = Exception
+
+# Import search tool components
+from .constants import *
+from .rate_limiter import RateLimiter, CircuitBreaker
+from .analyzers import ResultQualityAnalyzer, QueryIntentAnalyzer, ResultSummarizer
+from .deduplicator import ResultDeduplicator
+from .context import SearchContext
+from .cache import IntelligentCache
+from .metrics import EnhancedMetrics
+from .error_handler import AgentFriendlyErrorHandler
+from .schemas import (
+    SearchWebSchema,
+    SearchImagesSchema,
+    SearchNewsSchema,
+    SearchVideosSchema,
+    SearchPaginatedSchema,
+    SearchBatchSchema,
+    ValidateCredentialsSchema,
+    GetQuotaStatusSchema,
+    GetMetricsSchema,
+    GetMetricsReportSchema,
+    GetHealthScoreSchema,
+    GetSearchContextSchema,
+)
+
+
+class SearchTool(BaseTool):
+    """
+    Enhanced web search tool using Google Custom Search API.
+    
+    Provides intelligent search with:
+    - Quality scoring and ranking
+    - Query intent analysis
+    - Result deduplication
+    - Context-aware search
+    - Intelligent Redis caching
+    - Comprehensive metrics
+    - Agent-friendly error handling
+    """
+    
+    # Configuration schema
+    class Config(BaseModel):
+        """Configuration for the search tool"""
+        model_config = ConfigDict(env_prefix="SEARCH_TOOL_")
+        
+        google_api_key: Optional[str] = Field(
+            default=None,
+            description="Google API key for Custom Search"
+        )
+        google_cse_id: Optional[str] = Field(
+            default=None,
+            description="Custom Search Engine ID"
+        )
+        google_application_credentials: Optional[str] = Field(
+            default=None,
+            description="Path to service account JSON"
+        )
+        max_results_per_query: int = Field(
+            default=10,
+            description="Maximum results per single query"
+        )
+        cache_ttl: int = Field(
+            default=3600,
+            description="Default cache time-to-live in seconds"
+        )
+        rate_limit_requests: int = Field(
+            default=100,
+            description="Maximum requests per time window"
+        )
+        rate_limit_window: int = Field(
+            default=86400,
+            description="Time window for rate limiting in seconds"
+        )
+        circuit_breaker_threshold: int = Field(
+            default=5,
+            description="Failures before opening circuit"
+        )
+        circuit_breaker_timeout: int = Field(
+            default=60,
+            description="Timeout before trying half-open in seconds"
+        )
+        retry_attempts: int = Field(
+            default=3,
+            description="Number of retry attempts"
+        )
+        retry_backoff: float = Field(
+            default=2.0,
+            description="Exponential backoff factor"
+        )
+        timeout: int = Field(
+            default=30,
+            description="API request timeout in seconds"
+        )
+        user_agent: str = Field(
+            default="AIECS-SearchTool/2.0",
+            description="User agent string"
+        )
+        
+        # Enhanced features
+        enable_quality_analysis: bool = Field(
+            default=True,
+            description="Enable result quality analysis"
+        )
+        enable_intent_analysis: bool = Field(
+            default=True,
+            description="Enable query intent analysis"
+        )
+        enable_deduplication: bool = Field(
+            default=True,
+            description="Enable result deduplication"
+        )
+        enable_context_tracking: bool = Field(
+            default=True,
+            description="Enable search context tracking"
+        )
+        enable_intelligent_cache: bool = Field(
+            default=True,
+            description="Enable intelligent Redis caching"
+        )
+        similarity_threshold: float = Field(
+            default=0.85,
+            description="Similarity threshold for deduplication"
+        )
+        max_search_history: int = Field(
+            default=10,
+            description="Maximum search history to maintain"
+        )
+
+    # Operation schemas for input validation and documentation
+    SearchWebSchema = SearchWebSchema
+    SearchImagesSchema = SearchImagesSchema
+    SearchNewsSchema = SearchNewsSchema
+    SearchVideosSchema = SearchVideosSchema
+    SearchPaginatedSchema = SearchPaginatedSchema
+    SearchBatchSchema = SearchBatchSchema
+    ValidateCredentialsSchema = ValidateCredentialsSchema
+    GetQuotaStatusSchema = GetQuotaStatusSchema
+    GetMetricsSchema = GetMetricsSchema
+    GetMetricsReportSchema = GetMetricsReportSchema
+    GetHealthScoreSchema = GetHealthScoreSchema
+    GetSearchContextSchema = GetSearchContextSchema
+
+    # Tool metadata
+    description = "Comprehensive web search tool using Google Custom Search API."
+    category = "task"
+
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        """
+        Initialize SearchTool with enhanced capabilities.
+        
+        Args:
+            config: Optional configuration overrides
+            
+        Raises:
+            AuthenticationError: If Google API libraries not available
+            ValidationError: If configuration is invalid
+        """
+        super().__init__(config)
+        
+        if not GOOGLE_API_AVAILABLE:
+            raise AuthenticationError(
+                "Google API client libraries not available. "
+                "Install with: pip install google-api-python-client google-auth google-auth-httplib2"
+            )
+        
+        # Load settings
+        global_settings = get_settings()
+        
+        # Merge configuration
+        merged_config = {
+            'google_api_key': global_settings.google_api_key,
+            'google_cse_id': global_settings.google_cse_id,
+            'google_application_credentials': global_settings.google_application_credentials
+        }
+        if config:
+            merged_config.update(config)
+        
+        # Parse configuration
+        self.config = self.Config(**merged_config)
+        
+        # Initialize logger
+        self.logger = logging.getLogger(__name__)
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            handler.setFormatter(
+                logging.Formatter('%(asctime)s %(levelname)s [SearchTool] %(message)s')
+            )
+            self.logger.addHandler(handler)
+        self.logger.setLevel(logging.INFO)
+        
+        # Initialize API client
+        self._service = None
+        self._credentials = None
+        self._init_credentials()
+        
+        # Initialize core components
+        self.rate_limiter = RateLimiter(
+            self.config.rate_limit_requests,
+            self.config.rate_limit_window
+        )
+        
+        self.circuit_breaker = CircuitBreaker(
+            self.config.circuit_breaker_threshold,
+            self.config.circuit_breaker_timeout
+        )
+        
+        # Initialize enhanced components
+        self.quality_analyzer = ResultQualityAnalyzer() if self.config.enable_quality_analysis else None
+        self.intent_analyzer = QueryIntentAnalyzer() if self.config.enable_intent_analysis else None
+        self.deduplicator = ResultDeduplicator() if self.config.enable_deduplication else None
+        self.result_summarizer = ResultSummarizer() if self.config.enable_quality_analysis else None
+        self.search_context = SearchContext(self.config.max_search_history) if self.config.enable_context_tracking else None
+        self.error_handler = AgentFriendlyErrorHandler()
+        
+        # Initialize intelligent cache (Redis)
+        self.intelligent_cache = None
+        if self.config.enable_intelligent_cache:
+            try:
+                from aiecs.infrastructure.persistence import RedisClient
+                redis_client = RedisClient()
+                # Note: Redis client needs to be initialized asynchronously
+                self.intelligent_cache = IntelligentCache(redis_client, enabled=True)
+            except Exception as e:
+                self.logger.warning(f"Could not initialize Redis cache: {e}")
+                self.intelligent_cache = IntelligentCache(None, enabled=False)
+        
+        # Initialize enhanced metrics
+        self.metrics = EnhancedMetrics()
+        
+        self.logger.info("SearchTool initialized with enhanced capabilities")
+
+    def _create_search_ttl_strategy(self):
+        """
+        Create intelligent TTL strategy for search results.
+
+        This strategy calculates TTL based on:
+        1. Query intent type (from result metadata)
+        2. Result freshness score
+        3. Result quality score
+
+        Returns:
+            Callable: TTL strategy function compatible with cache_result_with_strategy
+        """
+        def calculate_search_ttl(result: Any, args: tuple, kwargs: dict) -> int:
+            """
+            Calculate intelligent TTL for search results.
+
+            Args:
+                result: Search result (dict with 'results' and '_metadata')
+                args: Positional arguments (not used)
+                kwargs: Keyword arguments containing 'query', etc.
+
+            Returns:
+                int: TTL in seconds
+            """
+            # Extract metadata from result
+            if not isinstance(result, dict):
+                return 3600  # Default 1 hour for non-dict results
+
+            metadata = result.get('_metadata', {})
+            intent_type = metadata.get('intent_type', 'GENERAL')
+            results_list = result.get('results', [])
+            query = kwargs.get('query', '')
+
+            # Use IntelligentCache logic if available
+            if hasattr(self, 'intelligent_cache') and self.intelligent_cache:
+                try:
+                    return self.intelligent_cache.calculate_ttl(
+                        query,
+                        intent_type,
+                        results_list
+                    )
+                except Exception as e:
+                    self.logger.warning(f"Failed to calculate intelligent TTL: {e}")
+
+            # Fallback: Use intent-based TTL
+            from .cache import IntelligentCache
+            ttl_strategies = IntelligentCache.TTL_STRATEGIES
+            base_ttl = ttl_strategies.get(intent_type, ttl_strategies.get('GENERAL', 3600))
+
+            # Adjust based on result count
+            if not results_list:
+                return base_ttl // 2  # Shorter TTL for empty results
+
+            return base_ttl
+
+        return calculate_search_ttl
+
+    def _init_credentials(self):
+        """Initialize Google API credentials"""
+        # Method 1: API Key
+        if self.config.google_api_key and self.config.google_cse_id:
+            try:
+                self._service = build(
+                    'customsearch',
+                    'v1',
+                    developerKey=self.config.google_api_key,
+                    cache_discovery=False
+                )
+                self.logger.info("Initialized with API key")
+                return
+            except Exception as e:
+                self.logger.warning(f"Failed to initialize with API key: {e}")
+        
+        # Method 2: Service Account
+        if self.config.google_application_credentials:
+            creds_path = self.config.google_application_credentials
+            if os.path.exists(creds_path):
+                try:
+                    credentials = service_account.Credentials.from_service_account_file(
+                        creds_path,
+                        scopes=['https://www.googleapis.com/auth/cse']
+                    )
+                    self._credentials = credentials
+                    self._service = build(
+                        'customsearch',
+                        'v1',
+                        credentials=credentials,
+                        cache_discovery=False
+                    )
+                    self.logger.info("Initialized with service account")
+                    return
+                except Exception as e:
+                    self.logger.warning(f"Failed to initialize with service account: {e}")
+        
+        raise AuthenticationError(
+            "No valid Google API credentials found. Set GOOGLE_API_KEY and GOOGLE_CSE_ID"
+        )
+
+    def _execute_search(
+        self,
+        query: str,
+        num_results: int = 10,
+        start_index: int = 1,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """Execute search with rate limiting and circuit breaker"""
+        # Check rate limit
+        self.rate_limiter.acquire()
+        
+        # Prepare parameters
+        search_params = {
+            'q': query,
+            'cx': self.config.google_cse_id,
+            'num': min(num_results, 10),
+            'start': start_index,
+            **kwargs
+        }
+        
+        # Execute with circuit breaker
+        def _do_search():
+            try:
+                result = self._service.cse().list(**search_params).execute()
+                return result
+            except HttpError as e:
+                if e.resp.status == 429:
+                    raise QuotaExceededError(f"API quota exceeded: {e}")
+                elif e.resp.status == 403:
+                    raise AuthenticationError(f"Authentication failed: {e}")
+                else:
+                    raise SearchAPIError(f"Search API error: {e}")
+            except Exception as e:
+                raise SearchAPIError(f"Unexpected error: {e}")
+        
+        return self.circuit_breaker.call(_do_search)
+
+    def _retry_with_backoff(self, func, *args, **kwargs) -> Any:
+        """Execute with exponential backoff retry"""
+        last_exception = None
+        
+        for attempt in range(self.config.retry_attempts):
+            try:
+                return func(*args, **kwargs)
+            except (RateLimitError, CircuitBreakerOpenError) as e:
+                # Don't retry these
+                raise e
+            except Exception as e:
+                last_exception = e
+                if attempt < self.config.retry_attempts - 1:
+                    wait_time = self.config.retry_backoff ** attempt
+                    self.logger.warning(
+                        f"Attempt {attempt + 1} failed: {e}. Retrying in {wait_time}s..."
+                    )
+                    time.sleep(wait_time)
+        
+        raise last_exception
+
+    def _parse_search_results(
+        self, 
+        raw_results: Dict[str, Any],
+        query: str = "",
+        enable_quality_analysis: bool = True
+    ) -> List[Dict[str, Any]]:
+        """Parse and enhance search results"""
+        items = raw_results.get('items', [])
+        results = []
+        
+        for position, item in enumerate(items, start=1):
+            result = {
+                'title': item.get('title', ''),
+                'link': item.get('link', ''),
+                'snippet': item.get('snippet', ''),
+                'displayLink': item.get('displayLink', ''),
+                'formattedUrl': item.get('formattedUrl', ''),
+            }
+            
+            # Add image metadata
+            if 'image' in item:
+                result['image'] = {
+                    'contextLink': item['image'].get('contextLink', ''),
+                    'height': item['image'].get('height', 0),
+                    'width': item['image'].get('width', 0),
+                    'byteSize': item['image'].get('byteSize', 0),
+                    'thumbnailLink': item['image'].get('thumbnailLink', '')
+                }
+            
+            # Add page metadata
+            if 'pagemap' in item:
+                result['metadata'] = item['pagemap']
+            
+            # Add quality analysis
+            if enable_quality_analysis and self.quality_analyzer and query:
+                quality_analysis = self.quality_analyzer.analyze_result_quality(
+                    result, query, position
+                )
+                result['_quality'] = quality_analysis
+                
+                # Add agent-friendly quality summary
+                result['_quality_summary'] = {
+                    'score': quality_analysis['quality_score'],
+                    'level': quality_analysis['credibility_level'],
+                    'is_authoritative': quality_analysis['authority_score'] > 0.8,
+                    'is_relevant': quality_analysis['relevance_score'] > 0.7,
+                    'is_fresh': quality_analysis['freshness_score'] > 0.7,
+                    'warnings_count': len(quality_analysis['warnings'])
+                }
+            
+            results.append(result)
+        
+        return results
+
+    # ========================================================================
+    # Core Search Methods
+    # ========================================================================
+
+    @cache_result_with_strategy(ttl_strategy=lambda self, result, args, kwargs:
+                                 self._create_search_ttl_strategy()(result, args, kwargs))
+    def search_web(
+        self,
+        query: str,
+        num_results: int = 10,
+        start_index: int = 1,
+        language: str = "en",
+        country: str = "us",
+        safe_search: str = "medium",
+        date_restrict: Optional[str] = None,
+        file_type: Optional[str] = None,
+        exclude_terms: Optional[str] = None,
+        auto_enhance: bool = True,
+        return_summary: bool = False
+    ) -> Dict[str, Any]:
+        """
+        Search the web with enhanced intelligence.
+        
+        Args:
+            query: Search query string
+            num_results: Number of results to return
+            start_index: Starting index for pagination
+            language: Language code
+            country: Country code
+            safe_search: Safe search level
+            date_restrict: Date restriction
+            file_type: File type filter
+            exclude_terms: Terms to exclude
+            auto_enhance: Enable automatic query enhancement
+            return_summary: Return summary metadata
+            
+        Returns:
+            List of search results (or dict with results and summary)
+        """
+        start_time = time.time()
+        intent_analysis = None
+        
+        try:
+            if not query or not query.strip():
+                raise ValidationError("Query cannot be empty")
+            
+            if num_results < 1 or num_results > 100:
+                raise ValidationError("num_results must be between 1 and 100")
+            
+            # Analyze query intent
+            enhanced_query = query
+            if auto_enhance and self.intent_analyzer:
+                intent_analysis = self.intent_analyzer.analyze_query_intent(query)
+                enhanced_query = intent_analysis['enhanced_query']
+                
+                # Merge suggested parameters
+                for param, value in intent_analysis['suggested_params'].items():
+                    if param == 'date_restrict' and not date_restrict:
+                        date_restrict = value
+                    elif param == 'file_type' and not file_type:
+                        file_type = value
+                    elif param == 'num_results':
+                        num_results = min(num_results, value)
+                
+                self.logger.info(
+                    f"Intent: {intent_analysis['intent_type']} "
+                    f"(confidence: {intent_analysis['confidence']:.2f})"
+                )
+            
+            # Note: Cache is now handled by @cache_result_with_strategy decorator
+            # No need for manual cache check here
+
+            # Prepare search parameters
+            search_params = {
+                'lr': f'lang_{language}',
+                'cr': f'country{country.upper()}',
+                'safe': safe_search,
+            }
+            
+            if date_restrict:
+                search_params['dateRestrict'] = date_restrict
+            
+            if file_type:
+                search_params['fileType'] = file_type
+            
+            if exclude_terms:
+                enhanced_query = f"{enhanced_query} -{exclude_terms}"
+            
+            # Execute search
+            raw_results = self._retry_with_backoff(
+                self._execute_search,
+                enhanced_query,
+                num_results,
+                start_index,
+                **search_params
+            )
+            
+            # Parse results
+            results = self._parse_search_results(
+                raw_results,
+                query=query,
+                enable_quality_analysis=self.config.enable_quality_analysis
+            )
+            
+            # Deduplicate
+            if self.deduplicator:
+                results = self.deduplicator.deduplicate_results(
+                    results,
+                    self.config.similarity_threshold
+                )
+            
+            # Add search metadata
+            if intent_analysis:
+                for result in results:
+                    result['_search_metadata'] = {
+                        'original_query': query,
+                        'enhanced_query': enhanced_query,
+                        'intent_type': intent_analysis['intent_type'],
+                        'intent_confidence': intent_analysis['confidence'],
+                        'suggestions': intent_analysis['suggestions']
+                    }
+            
+            # Update context
+            if self.search_context:
+                self.search_context.add_search(query, results)
+
+            # Note: Cache is now handled by @cache_result_with_strategy decorator
+            # The decorator will call _create_search_ttl_strategy() to calculate TTL
+
+            # Record metrics
+            response_time = (time.time() - start_time) * 1000
+            self.metrics.record_search(
+                query, 'web', results, response_time, cached=False
+            )
+
+            # Prepare result with metadata for TTL calculation
+            result_data = {
+                'results': results,
+                '_metadata': {
+                    'intent_type': intent_analysis['intent_type'] if intent_analysis else 'GENERAL',
+                    'query': query,
+                    'enhanced_query': enhanced_query,
+                    'timestamp': time.time(),
+                    'response_time_ms': response_time
+                }
+            }
+
+            # Generate summary if requested
+            if return_summary and self.result_summarizer:
+                summary = self.result_summarizer.generate_summary(results, query)
+                result_data['summary'] = summary
+
+            return result_data
+            
+        except Exception as e:
+            response_time = (time.time() - start_time) * 1000
+            self.metrics.record_search(
+                query, 'web', [], response_time, error=e
+            )
+            
+            # Format error for agent
+            error_info = self.error_handler.format_error_for_agent(
+                e,
+                {'circuit_breaker_timeout': self.config.circuit_breaker_timeout}
+            )
+            
+            self.logger.error(f"Search failed: {error_info['user_message']}")
+            raise
+
+    def search_images(
+        self,
+        query: str,
+        num_results: int = 10,
+        image_size: Optional[str] = None,
+        image_type: Optional[str] = None,
+        image_color_type: Optional[str] = None,
+        safe_search: str = "medium"
+    ) -> List[Dict[str, Any]]:
+        """Search for images"""
+        if not query or not query.strip():
+            raise ValidationError("Query cannot be empty")
+        
+        search_params = {
+            'searchType': 'image',
+            'safe': safe_search,
+        }
+        
+        if image_size:
+            search_params['imgSize'] = image_size
+        if image_type:
+            search_params['imgType'] = image_type
+        if image_color_type:
+            search_params['imgColorType'] = image_color_type
+        
+        raw_results = self._retry_with_backoff(
+            self._execute_search,
+            query,
+            num_results,
+            1,
+            **search_params
+        )
+        
+        return self._parse_search_results(raw_results, query=query)
+
+    def search_news(
+        self,
+        query: str,
+        num_results: int = 10,
+        start_index: int = 1,
+        language: str = "en",
+        date_restrict: Optional[str] = None,
+        sort_by: str = "date"
+    ) -> List[Dict[str, Any]]:
+        """Search for news articles"""
+        if not query or not query.strip():
+            raise ValidationError("Query cannot be empty")
+        
+        news_query = f"{query} news"
+        
+        search_params = {
+            'lr': f'lang_{language}',
+            'sort': sort_by if sort_by == 'date' else '',
+        }
+        
+        if date_restrict:
+            search_params['dateRestrict'] = date_restrict
+        
+        raw_results = self._retry_with_backoff(
+            self._execute_search,
+            news_query,
+            num_results,
+            start_index,
+            **search_params
+        )
+        
+        return self._parse_search_results(raw_results, query=query)
+
+    def search_videos(
+        self,
+        query: str,
+        num_results: int = 10,
+        start_index: int = 1,
+        language: str = "en",
+        safe_search: str = "medium"
+    ) -> List[Dict[str, Any]]:
+        """Search for videos"""
+        if not query or not query.strip():
+            raise ValidationError("Query cannot be empty")
+        
+        video_query = f"{query} filetype:mp4 OR filetype:webm OR filetype:mov"
+        
+        search_params = {
+            'lr': f'lang_{language}',
+            'safe': safe_search,
+        }
+        
+        raw_results = self._retry_with_backoff(
+            self._execute_search,
+            video_query,
+            num_results,
+            start_index,
+            **search_params
+        )
+        
+        return self._parse_search_results(raw_results, query=query)
+
+    # ========================================================================
+    # Utility Methods
+    # ========================================================================
+
+    def get_metrics(self) -> Dict[str, Any]:
+        """Get comprehensive metrics"""
+        return self.metrics.get_metrics()
+
+    def get_metrics_report(self) -> str:
+        """Get human-readable metrics report"""
+        return self.metrics.generate_report()
+
+    def get_health_score(self) -> float:
+        """Get system health score (0-1)"""
+        return self.metrics.get_health_score()
+
+    def get_quota_status(self) -> Dict[str, Any]:
+        """Get quota and rate limit status"""
+        return {
+            'remaining_quota': self.rate_limiter.get_remaining_quota(),
+            'max_requests': self.config.rate_limit_requests,
+            'time_window_seconds': self.config.rate_limit_window,
+            'circuit_breaker_state': self.circuit_breaker.get_state(),
+            'health_score': self.get_health_score()
+        }
+
+    def get_search_context(self) -> Optional[Dict[str, Any]]:
+        """Get search context information"""
+        if not self.search_context:
+            return None
+        
+        return {
+            'history': self.search_context.get_history(5),
+            'preferences': self.search_context.get_preferences()
+        }
+