ata-coder 2.4.2__py3-none-any.whl

ata_coder/tools/web.py

@@ -0,0 +1,361 @@
+"""Web search, web fetch, and HTML text extraction — mixin for ToolExecutor."""
+import asyncio
+import html
+import html.parser
+import logging
+import re
+
+import httpx
+
+from .result import ToolResult  # noqa: E402 — circular-safe, ToolResult is defined before mixin
+
+logger = logging.getLogger(__name__)
+
+
+class WebToolsMixin:
+    """Web search and fetch capabilities for ToolExecutor."""
+
+    # Internal HTTP client (lazy-init, shared across web tools)
+    _http: httpx.Client | None = None
+
+    # ── Web tools ──────────────────────────────────────────────────────────
+
+    async def _run_in_thread(self, func, *args, **kwargs):
+        """Run a sync function in a thread pool to avoid blocking the event loop."""
+        from functools import partial
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, partial(func, *args, **kwargs))
+
+    @property
+    def http(self) -> httpx.Client:
+        if self._http is None:
+            self._http = httpx.Client(
+                timeout=httpx.Timeout(30.0),
+                follow_redirects=True,
+                headers={
+                    "User-Agent": (
+                        "ATA-Coder/2.0 (AI Coding Assistant; "
+                        "+https://github.com/ata-coder/ata-coder)"
+                    ),
+                    "Accept": "text/html,application/xhtml+xml,*/*",
+                    "Accept-Language": "en-US,zh-CN;q=0.9",
+                },
+            )
+        return self._http
+
+    async def _tool_web_search(
+        self,
+        query: str,
+        max_results: int = 10,
+    ) -> ToolResult:
+        """Search the web with tiered fallback: Bing → Baidu → Google.
+
+        All three use web scraping (no API key required).
+        Set ATA_CODER_SEARCH_BACKEND to force a single backend:
+          "bing" / "baidu" / "google" / "duckduckgo"
+        """
+        import os
+        max_results = min(max(max_results, 1), 20)
+        forced = os.environ.get("ATA_CODER_SEARCH_BACKEND", "")
+
+        # Whitelist valid backend names
+        _VALID_BACKENDS = {"bing", "baidu", "google", "duckduckgo"}
+        if forced and forced.lower() not in _VALID_BACKENDS:
+            logger.warning("Unknown ATA_CODER_SEARCH_BACKEND=%r — ignoring, using fallback chain", forced)
+            forced = ""
+
+        errors: list[str] = []
+
+        # Build fallback chain: respect forced backend, otherwise tiered
+        if forced:
+            chain = [(forced, getattr(self, f"_search_{forced}", None))]
+        else:
+            chain = [
+                ("Bing",   self._search_bing),
+                ("Baidu",  self._search_baidu),
+                ("Google", self._search_google),
+            ]
+
+        for name, searcher in chain:
+            if searcher is None:
+                errors.append(f"{name}: unsupported backend")
+                continue
+            # Real-time progress: tell the user which backend we're trying
+            if self._stream_cb:
+                self._stream_cb("web_search", f"🔍 Searching {name}...\n")
+            try:
+                # Run sync search in thread pool to avoid blocking event loop
+                results = await self._run_in_thread(searcher, query)
+                if results:
+                    if self._stream_cb:
+                        self._stream_cb("web_search", f"✓ {name}: {len(results)} results\n")
+                    return self._format_search_results(query, results, max_results, name)
+                if self._stream_cb:
+                    self._stream_cb("web_search", f"✗ {name}: no results\n")
+                errors.append(f"{name} returned no results")
+            except httpx.TimeoutException:
+                if self._stream_cb:
+                    self._stream_cb("web_search", f"✗ {name}: timed out\n")
+                errors.append(f"{name} timed out")
+            except httpx.HTTPStatusError as e:
+                if self._stream_cb:
+                    self._stream_cb("web_search", f"✗ {name}: HTTP {e.response.status_code}\n")
+                errors.append(f"{name} HTTP {e.response.status_code}")
+            except Exception as e:
+                if self._stream_cb:
+                    self._stream_cb("web_search", f"✗ {name}: {e}\n")
+                errors.append(f"{name}: {e}")
+
+        return ToolResult(
+            success=False, output="",
+            error=f"Search failed: {'; '.join(errors)}"
+        )
+
+    def _search_bing(self, query: str) -> list[dict[str, str]]:
+        """Search Bing (web scraping, no API key)."""
+        import urllib.parse
+        url = f"https://www.bing.com/search?q={urllib.parse.quote(query)}&setlang=en"
+        headers = {
+            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
+                          "(KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
+            "Accept-Language": "en-US,en;q=0.9",
+        }
+        resp = self.http.get(url, headers=headers)
+        resp.raise_for_status()
+
+        results: list[dict[str, str]] = []
+        # Bing results are in <li class="b_algo"> blocks
+        blocks = re.findall(
+            r'<li[^>]*class="[^"]*b_algo[^"]*"[^>]*>(.*?)</li>',
+            resp.text, re.DOTALL | re.IGNORECASE,
+        )
+        for block in blocks:
+            # Title + link in <h2><a href="...">title</a></h2>
+            m = re.search(r'<a[^>]*href="([^"]*)"[^>]*>(.*?)</a>', block, re.DOTALL)
+            if not m:
+                continue
+            href = html.unescape(m.group(1).strip())
+            title = re.sub(r'<[^>]+>', '', m.group(2)).strip()
+            if not title or not href.startswith("http"):
+                continue
+            # Snippet in <p> or <div class="b_caption">
+            snippet = ""
+            sm = re.search(
+                r'<(?:p|div)[^>]*class="[^"]*(?:b_caption|b_lineclamp)[^"]*"[^>]*>(.*?)</(?:p|div)>',
+                block, re.DOTALL | re.IGNORECASE,
+            )
+            if sm:
+                snippet = re.sub(r'<[^>]+>', '', sm.group(1)).strip()
+            snippet = html.unescape(snippet)
+            results.append({"title": title, "url": href, "snippet": snippet})
+
+        return results
+
+    def _search_baidu(self, query: str) -> list[dict[str, str]]:
+        """Search Baidu (web scraping, no API key)."""
+        import urllib.parse
+        url = f"https://www.baidu.com/s?wd={urllib.parse.quote(query)}&ie=utf-8"
+        headers = {
+            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
+                          "(KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
+            "Accept-Language": "zh-CN,zh;q=0.9",
+        }
+        resp = self.http.get(url, headers=headers)
+        resp.raise_for_status()
+
+        results: list[dict[str, str]] = []
+        # Baidu results: <div class="result c-container"> or <div class="c-container">
+        blocks = re.findall(
+            r'<div[^>]*class="[^"]*(?:result|c-container)[^"]*"[^>]*>(.*?)</div>\s*(?=<div[^>]*class="[^"]*(?:result|c-container)|$)',
+            resp.text, re.DOTALL | re.IGNORECASE,
+        )
+        if not blocks:
+            # Fallback: match h3 titles with links
+            blocks = re.findall(
+                r'<div[^>]*class="[^"]*c-container[^"]*"[^>]*>(.*?)</div>',
+                resp.text, re.DOTALL | re.IGNORECASE,
+            )
+
+        for block in blocks:
+            m = re.search(r'<a[^>]*href="([^"]*)"[^>]*>(.*?)</a>', block, re.DOTALL)
+            if not m:
+                continue
+            href = html.unescape(m.group(1).strip())
+            title = re.sub(r'<[^>]+>', '', m.group(2)).strip()
+            if not title or not href.startswith("http"):
+                continue
+            snippet = ""
+            sm = re.search(
+                r'<(?:span|div|p)[^>]*class="[^"]*(?:content-right_[^"]*|c-abstract|content)[^"]*"[^>]*>(.*?)</(?:span|div|p)>',
+                block, re.DOTALL | re.IGNORECASE,
+            )
+            if sm:
+                snippet = re.sub(r'<[^>]+>', '', sm.group(1)).strip()
+            snippet = html.unescape(snippet)
+            results.append({"title": title, "url": href, "snippet": snippet})
+
+        return results
+
+    def _search_google(self, query: str) -> list[dict[str, str]]:
+        """Search Google (web scraping, no API key)."""
+        import urllib.parse
+        url = f"https://www.google.com/search?q={urllib.parse.quote(query)}&hl=en"
+        headers = {
+            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
+                          "(KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
+            "Accept-Language": "en-US,en;q=0.9",
+        }
+        resp = self.http.get(url, headers=headers)
+        resp.raise_for_status()
+
+        results: list[dict[str, str]] = []
+        # Google results are in <div class="g"> or <div data-sokoban-container>
+        blocks = re.findall(
+            r'<(?:div|li)[^>]*\b(?:class="g\b|data-sokoban-container)[^>]*>(.*?)</(?:div|li)>',
+            resp.text, re.DOTALL | re.IGNORECASE,
+        )
+        for block in blocks:
+            # Title + link: <h3>...<a href="...">title</a></h3>
+            m = re.search(r'<a[^>]*href="(/url\?q=|)([^"&]*)"[^>]*>(.*?)</a>', block, re.DOTALL)
+            if not m:
+                continue
+            href = html.unescape(m.group(2).strip())
+            if not href.startswith("http"):
+                href = "https://www.google.com" + m.group(1) + m.group(2)
+            title = re.sub(r'<[^>]+>', '', m.group(3)).strip()
+            if not title:
+                continue
+            # Snippet: <span class="aCOpRe"> or various other classes
+            snippet = ""
+            sm = re.search(
+                r'<(?:span|div)[^>]*\b(?:class="[^"]*(?:\baCOpRe\b|st\b)[^"]*")[^>]*>(.*?)</(?:span|div)>',
+                block, re.DOTALL | re.IGNORECASE,
+            )
+            if sm:
+                snippet = re.sub(r'<[^>]+>', '', sm.group(1)).strip()
+            snippet = html.unescape(snippet)
+            results.append({"title": title, "url": href, "snippet": snippet})
+
+        return results
+
+    @staticmethod
+    def _format_search_results(
+        query: str, results: list[dict[str, str]], max_results: int, source: str
+    ) -> ToolResult:
+        out = [f"Search results for: {query}  (via {source})\n"]
+        for i, r in enumerate(results[:max_results], 1):
+            out.append(f"{i}. **{html.unescape(r['title'])}**")
+            out.append(f"   {r['url']}")
+            if r.get("snippet"):
+                out.append(f"   {html.unescape(r['snippet'])}")
+            out.append("")
+        return ToolResult(success=True, output="\n".join(out))
+
+    @staticmethod
+    def _parse_ddg_lite(html_text: str) -> list[dict[str, str]]:
+        """Extract search results from DuckDuckGo Lite HTML."""
+        results: list[dict[str, str]] = []
+
+        # DDG Lite: results are in <a> tags with class="result-link"
+        # and snippets in <td class="result-snippet">
+        link_pattern = re.compile(
+            r'<a[^>]*href="([^"]*)"[^>]*class="[^"]*result-link[^"]*"[^>]*>(.*?)</a>',
+            re.DOTALL | re.IGNORECASE,
+        )
+        snippet_pattern = re.compile(
+            r'<td[^>]*class="[^"]*result-snippet[^"]*"[^>]*>(.*?)</td>',
+            re.DOTALL | re.IGNORECASE,
+        )
+
+        links = link_pattern.findall(html_text)
+        snippets = snippet_pattern.findall(html_text)
+
+        for i, (href, title) in enumerate(links):
+            href = html.unescape(href.strip())
+            title = re.sub(r'<[^>]+>', '', title).strip()
+            if not title:
+                continue
+
+            # Pick corresponding snippet
+            snippet = ""
+            if i < len(snippets):
+                snippet = re.sub(r'<[^>]+>', '', snippets[i])
+                snippet = html.unescape(snippet.strip())
+
+            results.append({
+                "title": title,
+                "url": href,
+                "snippet": snippet[:300],
+            })
+
+        return results
+
+    async def _tool_web_fetch(self, url: str) -> ToolResult:
+        """Fetch a URL and extract its text content."""
+        if not url.startswith(("http://", "https://")):
+            return ToolResult(
+                success=False, output="",
+                error="Invalid URL: must start with http:// or https://"
+            )
+
+        if self._stream_cb:
+            self._stream_cb("web_fetch", f"🌐 Fetching {url}...\n")
+
+        def _do_fetch():
+            return self.http.get(url)
+
+        try:
+            resp = await self._run_in_thread(_do_fetch)
+            resp.raise_for_status()
+        except httpx.TimeoutException:
+            if self._stream_cb:
+                self._stream_cb("web_fetch", f"✗ Timeout: {url}\n")
+            return ToolResult(
+                success=False, output="",
+                error=f"Request timed out: {url}"
+            )
+        except httpx.HTTPStatusError as e:
+            if self._stream_cb:
+                self._stream_cb("web_fetch", f"✗ HTTP {e.response.status_code}: {url}\n")
+            return ToolResult(
+                success=False, output="",
+                error=f"HTTP {e.response.status_code} for {url}"
+            )
+        except Exception as e:
+            if self._stream_cb:
+                self._stream_cb("web_fetch", f"✗ Failed: {url} — {e}\n")
+            return ToolResult(
+                success=False, output="",
+                error=f"Fetch failed: {e}"
+            )
+
+        content_type = resp.headers.get("content-type", "")
+        if "text/html" not in content_type and "text/plain" not in content_type:
+            return ToolResult(
+                success=False, output="",
+                error=f"Cannot process content type: {content_type}. Only text/html and text/plain are supported."
+            )
+
+        if self._stream_cb:
+            size_kb = len(resp.text) // 1024
+            self._stream_cb("web_fetch", f"✓ Downloaded {size_kb}KB, extracting text...\n")
+
+        text = self._extract_text(resp.text, url)
+
+        if self._stream_cb:
+            self._stream_cb("web_fetch", f"✓ Extracted {len(text):,} chars\n")
+
+        # Truncate
+        MAX_CHARS = 15_000
+        if len(text) > MAX_CHARS:
+            text = text[:MAX_CHARS] + (
+                f"\n\n... [truncated {len(text) - MAX_CHARS:,} "
+                f"chars from {url}]"
+            )
+
+        return ToolResult(
+            success=True,
+            output=f"Content from: {url}\n\n{text}",
+        )
+
+