src/hermes/tools/web_search.py

"""Web search tool implementation."""

from __future__ import annotations

import logging
from html.parser import HTMLParser
from typing import Any

import httpx

from hermes.tools.base.tool import BaseTool, ToolSchema

logger = logging.getLogger(__name__)


class DDGLiteParser(HTMLParser):
    """HTML Parser for DuckDuckGo Lite search results."""

    def __init__(self) -> None:
        super().__init__()
        self.results: list[dict[str, str]] = []
        self.current_result: dict[str, str] = {}
        self.in_result_link = False
        self.in_result_snippet = False

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        attrs_dict = {k: v for k, v in attrs if v is not None}
        class_name = attrs_dict.get("class", "")
        if tag == "a" and "result-link" in class_name:
            self.in_result_link = True
            self.current_result["url"] = attrs_dict.get("href", "")
        elif tag == "td" and "result-snippet" in class_name:
            self.in_result_snippet = True

    def handle_data(self, data: str) -> None:
        if self.in_result_link:
            self.current_result["title"] = self.current_result.get("title", "") + data
        elif self.in_result_snippet:
            self.current_result["description"] = self.current_result.get("description", "") + data

    def handle_endtag(self, tag: str) -> None:
        if tag == "a" and self.in_result_link:
            self.in_result_link = False
        elif tag == "td" and self.in_result_snippet:
            self.in_result_snippet = False
            title = self.current_result.get("title", "").strip()
            url = self.current_result.get("url", "").strip()
            desc = self.current_result.get("description", "").strip()

            # Filter out advertisements and help pages
            if title and url and "company/ads-by-microsoft" not in url and "duckduckgo-help-pages" not in url:
                self.results.append({
                    "title": title,
                    "url": url,
                    "description": desc
                })
            self.current_result = {}


class WebSearchTool(BaseTool):
    """Tool for searching the web."""

    def __init__(self, api_key: str | None = None) -> None:
        super().__init__()
        self.api_key = api_key
        self.base_url = "https://api.search.brave.com/res/v1/web/search"

    def _define_schema(self) -> ToolSchema:
        return ToolSchema(
            name="search_web",
            description="Search the web for information on a given query",
            parameters={
                "query": {
                    "type": "string",
                    "description": "The search query",
                },
                "max_results": {
                    "type": "integer",
                    "description": "Maximum number of results to return",
                    "default": 10,
                },
                "language": {
                    "type": "string",
                    "description": "Language for search results",
                    "default": "en",
                },
            },
            required=["query"],
            category="research",
            tags=["search", "web", "research"],
        )

    async def execute(self, **kwargs: Any) -> dict[str, Any]:
        """Execute web search."""
        query = kwargs["query"]
        max_results = kwargs.get("max_results", 10)
        language = kwargs.get("language", "en")

        if not self.api_key:
            return await self._fallback_search(query, max_results)

        try:
            async with httpx.AsyncClient() as client:
                response = await client.get(
                    self.base_url,
                    headers={"X-Subscription-Token": self.api_key, "Accept": "application/json"},
                    params={"q": query, "count": max_results, "search_lang": language},
                    timeout=30.0,
                )
                response.raise_for_status()
                data = response.json()

                results = []
                for item in data.get("web", {}).get("results", [])[:max_results]:
                    results.append(
                        {
                            "title": item.get("title", ""),
                            "url": item.get("url", ""),
                            "description": item.get("description", ""),
                            "age": item.get("age", ""),
                        }
                    )

                return {"query": query, "results": results, "total": len(results)}

        except Exception as e:
            logger.error(f"Web search failed: {e}")
            return await self._fallback_search(query, max_results)

    async def _fallback_search(self, query: str, max_results: int) -> dict[str, Any]:
        """Fallback search using DuckDuckGo Lite."""
        try:
            async with httpx.AsyncClient() as client:
                headers = {
                    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
                    "Content-Type": "application/x-www-form-urlencoded"
                }
                response = await client.post(
                    "https://lite.duckduckgo.com/lite/",
                    data={"q": query},
                    headers=headers,
                    timeout=30.0,
                )
                response.raise_for_status()

                parser = DDGLiteParser()
                parser.feed(response.text)
                results = parser.results[:max_results]
                return {"query": query, "results": results, "total": len(results)}
        except Exception as e:
            logger.error(f"Fallback search failed: {e}")
            return {"query": query, "results": [], "total": 0, "error": str(e)}