agent-maxim/agent/runtime/tools/web_search.py

"""Web search tool using Bing with snippets."""
from __future__ import annotations

import re
import requests
from bs4 import BeautifulSoup
from datetime import datetime
from core.state import PlanStep, StepResult, StepStatus


def _enhance_search_query(query: str) -> str:
    """Enhance search query for better AI/news results.
    
    Args:
        query: Original search query (may contain instruction text)
        
    Returns:
        Enhanced search query with AI/news focus
    """
    # Clean up query - remove instruction text
    query = query.strip()
    
    # Remove common instruction prefixes
    instruction_patterns = [
        r'^Task:\s*',
        r'^Please complete the task above\.?\s*',
        r'^Respond with a concise result\.?\s*',
        r'^IMPORTANT INSTRUCTIONS.*?(?=\n|$)',
        r'^Use web search.*?(?=\n|$)',
    ]
    
    for pattern in instruction_patterns:
        query = re.sub(pattern, '', query, flags=re.IGNORECASE | re.DOTALL)
    
    # Take only first line or first 100 characters if multi-line
    if '\n' in query:
        query = query.split('\n')[0].strip()
    elif len(query) > 100:
        query = query[:100].strip()
    
    # Remove any remaining instruction-like text
    query = re.sub(r'(Please|complete|task|result|instruction)s?\b.*$', '', query, flags=re.IGNORECASE)
    query = query.strip()
    
    # If query is too short after cleaning, use generic AI news
    if len(query) < 5:
        query = "AI artificial intelligence news"
    
    current_year = datetime.now().year
    query_lower = query.lower()
    
    # Check if query is about news/AI/OpenAI
    news_keywords = ["новости", "news", "openai", "ai", "artificial intelligence", 
                    "ml", "machine learning", "gpt", "llm", "chatgpt", "claude", 
                    "developments", "latest", "updates", "breakthrough"]
    
    is_news_query = any(keyword in query_lower for keyword in news_keywords)
    
    if is_news_query:
        # Add AI/news specific terms only if not already present
        if not any(term in query_lower for term in ["latest", "news", "current"]):
            query = f"{query} latest"
            
        if not any(term in query_lower for term in ["2026", str(current_year)]):
            query = f"{query} {current_year}"
        
        # Add tech site restrictions for better results
        tech_sites = ["site:techcrunch.com", "site:theverge.com", "site:arstechnica.com",
                     "site:wired.com", "site:openai.com", "site:anthropic.com"]
        
        # Use up to 2 tech sites to avoid overly restrictive queries
        site_filter = " OR ".join(tech_sites[:2])
        query = f"{query} {site_filter}"
        
        # Add AI-specific terms if not already present
        ai_terms = ["artificial intelligence", "machine learning", "AI technology"]
        if not any(term in query_lower for term in ai_terms):
            query = f"{query} artificial intelligence"
    else:
        # For non-news queries, just ensure basic quality
        # Add relevant context if missing
        if len(query) < 10:
            query = f"{query} information"
            
        # Exclude common low-quality results
        exclude_terms = ["tasks", "calendar", "todo", "reminder", "schedule", "app"]
        exclusions = []
        for term in exclude_terms:
            if term not in query_lower:
                exclusions.append(f"-{term}")
        
        if exclusions:
            query = f"{query} {' '.join(exclusions)}"
    
    return query


def web_search(step: PlanStep) -> StepResult:
    """Search the web using Bing with informative snippets.
    
    Args:
        step: PlanStep with payload containing 'query'
        
    Returns:
        StepResult with search results (title + URL + snippet)
    """
    query = step.payload.get("query", "").strip()
    
    if not query:
        return StepResult(
            step_id=step.step_id,
            status=StepStatus.FAILED,
            error="payload missing 'query'"
        )
    
    # Enhance query for better AI/news results
    enhanced_query = _enhance_search_query(query)
    
    try:
        url = "https://www.bing.com/search"
        params = {"q": enhanced_query, "setlang": "en", "cc": "US"}
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
            "Accept-Language": "en-US,en;q=0.9"
        }
        
        resp = requests.get(url, params=params, headers=headers, timeout=10)
        soup = BeautifulSoup(resp.text, "html.parser")
        
        results = []
        items = soup.select("li.b_algo")[:5]
        
        for i, item in enumerate(items, 1):
            title_tag = item.select_one("h2, h3")
            link_tag = item.select_one("a[href]")
            snippet_tag = item.select_one("[data-bm], p")
            
            if not title_tag or not link_tag:
                continue
            
            title = title_tag.get_text(strip=True)
            link = link_tag.get("href", "")
            
            # Decode Bing redirect URLs to get actual URLs
            if "bing.com/ck/a" in link:
                try:
                    match = re.search(r'[?&]u=([a-zA-Z0-9_-]+)', link)
                    if match:
                        encoded_url = match.group(1)
                        # Remove 'a1' prefix if present
                        if encoded_url.startswith('a1'):
                            encoded_url = encoded_url[2:]
                        # Add proper padding and decode
                        padding = 4 - (len(encoded_url) % 4)
                        if padding != 4:
                            encoded_url += '=' * padding
                        import base64
                        link = base64.b64decode(encoded_url).decode('utf-8')
                except:
                    pass  # Keep original URL if decoding fails
            
            snippet = ""
            if snippet_tag:
                snippet = snippet_tag.get_text(strip=True)
                snippet = re.sub(r"\s+", " ", snippet)
                snippet = snippet[:200]
            
            results.append(f"{i}.\nTitle: {title}\nURL: {link}\nSnippet: {snippet}\n")
        
        if not results:
            return StepResult(
                step_id=step.step_id,
                status=StepStatus.DONE,
                output_text=f"No results found for query: {enhanced_query}"
            )
        
        return StepResult(
            step_id=step.step_id,
            status=StepStatus.DONE,
            output_text="\n".join(results),
            metadata={"source": "bing", "results_count": len(results)}
        )
        
    except Exception as e:
        return StepResult(
            step_id=step.step_id,
            status=StepStatus.FAILED,
            error=f"search error: {str(e)}"
        )