Blame for tools.py

"""
tools.py

This module contains all available tools for the AI assistant.
Each tool is a function that can be called by the pipeline.

Tools:
- web_search: Google Serper API-based web search
"""

import os
import json
from typing import Dict, Any, Optional, List

# =============================================================================
# AVAILABLE TOOLS REGISTRY
# =============================================================================

# Dictionary mapping tool names to their functions
AVAILABLE_TOOLS = {
    "web_search": None  # Will be set after function definition
}

# =============================================================================
# SEARCH CONFIGURATION
# =============================================================================

WEB_SEARCH_MAX_QUERY_CHARS = int(os.getenv("WEB_SEARCH_MAX_QUERY_CHARS", "200"))
SEARCH_MAX_RESULTS = int(os.getenv("SEARCH_MAX_RESULTS", "1"))
SEARCH_WORKERS = int(os.getenv("SEARCH_WORKERS", "6"))

try:
    from search import run_search
except ImportError as e:
    raise ImportError("Could not import run_search from search.py: " + str(e))


# =============================================================================
# TOOL IMPLEMENTATIONS
# =============================================================================

def _format_search_results(structured_results: List[Dict[str, Any]]) -> str:
    """
    Format search results from search.py into readable output.
    Excludes unnecessary fields and returns full structured data.
    
    Args:
        structured_results: List of structured result dictionaries from search.py
        
    Returns:
        JSON formatted string with filtered search results
    """
    if not structured_results:
        print("⚠️ No results found from search")
        return json.dumps({"results": [], "count": 0})
    
    # Fields to exclude from each result
    exclude_fields = {
        "publish_date",
        "author",
        "fetch_time",
        "extraction_status",
        "confidence_score",
        "content_type",
        "cleaned_content",
        "first_paragraph",
        "quality_signals"
    }
    
    # Filter results and remove excluded fields
    filtered_results = []
    for result in structured_results[:max(1, SEARCH_MAX_RESULTS)]:
        filtered_result = {k: v for k, v in result.items() if k not in exclude_fields}
        filtered_results.append(filtered_result)
    
    # Format as JSON output
    output = {
        "query_results": filtered_results,
        "total_count": len(filtered_results)
    }
    
    result_json = json.dumps(output, indent=2, ensure_ascii=False)
    print(f"✅ Got {len(filtered_results)} results from search pipeline")
    print(f"📈 Total chars in results from search pipeline after formated( from tools.py): {len(structured_results)}")
    return result_json


def web_search(query: str) -> str:
    """
    Live web search tool using search.py pipeline.

    Args:
        query: The search query string
        
    Returns:
        Formatted string with search results containing:
        - Position (for each result)
        - Title
        - Link
        - Snippet
    """
    try:
        if not query or not isinstance(query, str):
            print("⚠️ Invalid query provided")
            return "No results found."

        q = query.strip()
        if not q:
            print("⚠️ Empty query string")
            return "No results found."

        # Truncate query to max chars
        q = q[:WEB_SEARCH_MAX_QUERY_CHARS]

        print(f"🔎 Web search initiated for: '{q}' (max results={SEARCH_MAX_RESULTS})")
        
        # Run search pipeline from search.py
        structured_results, stats = run_search(q, max_results=SEARCH_MAX_RESULTS, workers=SEARCH_WORKERS)
        
        print(f"   Pipeline: {stats['search_engine']['total']} total results, "
              f"{stats['cleaner']['successful']} successfully extracted")
        print(f"📈 Total chars in results from search pipeline Before formated ( from tools.py ): {len(structured_results)}")
        
        return _format_search_results(structured_results)

    except Exception as e:
        error_msg = f"Web search failed: {str(e)}"
        print(f"❌ {error_msg}")
        return f"Search failed: {str(e)}"


def execute_tool(tool_name: str, args: Dict[str, Any]) -> Optional[str]:
    """
    Execute a tool by name with the given arguments.
    
    Args:
        tool_name: Name of the tool to execute
        args: Dictionary of arguments for the tool
        
    Returns:
        Tool result as string, or None if tool not found or execution fails
    """
    try:
        # Validate tool name
        if not tool_name:
            print("No tool name provided")
            return None
        
        # Check if tool exists
        if tool_name not in AVAILABLE_TOOLS:
            print(f"Unknown tool: '{tool_name}'")
            return None
        
        # Get the tool function
        tool_func = AVAILABLE_TOOLS.get(tool_name)
        
        if tool_func is None:
            print(f"Tool '{tool_name}' not implemented")
            return None
        
        print(f"Executing tool: '{tool_name}' with args: {args}")
        
        # Execute the tool
        result = tool_func(**args)
        
        print(f"Tool '{tool_name}' executed successfully")
        return result
        
    except Exception as e:
        print(f"Error executing tool '{tool_name}': {str(e)}")
        return None


# =============================================================================
# INITIALIZE TOOL REGISTRY
# =============================================================================

# Register the web_search tool
AVAILABLE_TOOLS["web_search"] = web_search
AVAILABLE_TOOLS["websearch"] = web_search


def get_available_tools() -> list:
    """
    Get a list of available tool names.
    
    Returns:
        List of tool name strings
    """
    return list(AVAILABLE_TOOLS.keys())


def is_tool_available(tool_name: str) -> bool:
    """
    Check if a tool is available.
    
    Args:
        tool_name: Name of the tool to check
        
    Returns:
        True if tool exists, False otherwise
    """
    return tool_name in AVAILABLE_TOOLS and AVAILABLE_TOOLS[tool_name] is not None