Spaces:
Sleeping
Sleeping
| """ | |
| 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 |