@shirlytaylor73/smart-search 0.2.0-beta.1

package/src/smart_search/providers/zhipu.py

@@ -0,0 +1,143 @@
+import json
+import time
+from typing import Any
+
+import httpx
+from tenacity import AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential
+
+from .base import BaseSearchProvider
+from ..config import config
+from ..logger import log_info
+
+
+RETRYABLE_STATUS_CODES = {408, 500, 502, 503, 504}
+
+
+def _is_retryable_exception(exc) -> bool:
+    if isinstance(exc, (httpx.TimeoutException, httpx.NetworkError, httpx.ConnectError)):
+        return True
+    if isinstance(exc, httpx.HTTPStatusError):
+        return exc.response.status_code in RETRYABLE_STATUS_CODES
+    return False
+
+
+def _normalize_result(item: dict[str, Any]) -> dict[str, Any]:
+    return {
+        "title": item.get("title") or "",
+        "url": item.get("link") or item.get("url") or "",
+        "description": item.get("content") or "",
+        "provider": "zhipu",
+        "source": item.get("media") or "",
+        "published_date": item.get("publish_date") or "",
+        "icon": item.get("icon") or "",
+        "refer": item.get("refer") or "",
+    }
+
+
+def _error_payload(exc: Exception) -> dict[str, Any]:
+    if isinstance(exc, httpx.HTTPStatusError):
+        status_code = exc.response.status_code
+        if status_code == 429:
+            error_type = "rate_limited"
+        elif status_code in {401, 403}:
+            error_type = "auth_error"
+        else:
+            error_type = "network_error"
+        return {"error_type": error_type, "error": f"HTTP {status_code}: {exc.response.reason_phrase}"}
+    if isinstance(exc, httpx.TimeoutException):
+        return {"error_type": "timeout", "error": "request timed out"}
+    if isinstance(exc, httpx.RequestError):
+        return {"error_type": "network_error", "error": str(exc)}
+    return {"error_type": "runtime_error", "error": str(exc)}
+
+
+class ZhipuWebSearchProvider(BaseSearchProvider):
+    def __init__(
+        self,
+        api_url: str,
+        api_key: str,
+        search_engine: str = "search_std",
+        timeout: float = 30.0,
+    ):
+        super().__init__(api_url.rstrip("/"), api_key)
+        self.search_engine = search_engine
+        self.timeout = timeout
+
+    def get_provider_name(self) -> str:
+        return "Zhipu Web Search"
+
+    async def search(
+        self,
+        query: str,
+        count: int = 10,
+        search_engine: str | None = None,
+        search_intent: bool = True,
+        search_domain_filter: str = "",
+        search_recency_filter: str = "noLimit",
+        content_size: str = "medium",
+        user_id: str = "",
+        ctx=None,
+    ) -> str:
+        endpoint = f"{self.api_url}/paas/v4/web_search"
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        }
+        payload: dict[str, Any] = {
+            "search_query": query[:70],
+            "search_engine": search_engine or self.search_engine,
+            "search_intent": search_intent,
+            "count": count,
+            "search_recency_filter": search_recency_filter,
+            "content_size": content_size,
+        }
+        if search_domain_filter:
+            payload["search_domain_filter"] = search_domain_filter
+        if user_id:
+            payload["user_id"] = user_id
+
+        await log_info(ctx, f"Zhipu search: {query}", config.debug_enabled)
+        start_time = time.time()
+        try:
+            data = await self._request_with_retry(endpoint, headers, payload)
+            elapsed_ms = round((time.time() - start_time) * 1000, 2)
+            results = [_normalize_result(item) for item in data.get("search_result", []) or []]
+            output = {
+                "ok": True,
+                "query": query,
+                "provider": "zhipu",
+                "search_engine": payload["search_engine"],
+                "results": results,
+                "total": len(results),
+                "search_intent": data.get("search_intent", []),
+                "request_id": data.get("request_id", ""),
+                "elapsed_ms": elapsed_ms,
+            }
+        except Exception as e:
+            elapsed_ms = round((time.time() - start_time) * 1000, 2)
+            error = _error_payload(e)
+            output = {
+                "ok": False,
+                "query": query,
+                "provider": "zhipu",
+                "error_type": error["error_type"],
+                "error": error["error"],
+                "elapsed_ms": elapsed_ms,
+            }
+        return json.dumps(output, ensure_ascii=False, indent=2)
+
+    async def _request_with_retry(self, endpoint: str, headers: dict, payload: dict) -> dict[str, Any]:
+        timeout = httpx.Timeout(connect=6.0, read=self.timeout, write=10.0, pool=None)
+        async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
+            async for attempt in AsyncRetrying(
+                stop=stop_after_attempt(config.retry_max_attempts + 1),
+                wait=wait_random_exponential(multiplier=config.retry_multiplier, max=config.retry_max_wait),
+                retry=retry_if_exception(_is_retryable_exception),
+                reraise=True,
+            ):
+                with attempt:
+                    response = await client.post(endpoint, headers=headers, json=payload)
+                    response.raise_for_status()
+                    return response.json()
+        return {}

package/src/smart_search/providers/zhipu_mcp.py

@@ -0,0 +1,230 @@
+import json
+import re
+import time
+from typing import Any
+
+import httpx
+
+
+def _elapsed_ms(start: float) -> float:
+    return round((time.time() - start) * 1000, 2)
+
+
+def _error_payload(exc: Exception) -> dict[str, str]:
+    if isinstance(exc, httpx.HTTPStatusError):
+        status_code = exc.response.status_code
+        if status_code in {401, 403}:
+            error_type = "auth_error"
+        elif status_code == 429:
+            error_type = "rate_limited"
+        else:
+            error_type = "network_error"
+        body = (exc.response.text or exc.response.reason_phrase or "")[:300]
+        return {"error_type": error_type, "error": f"HTTP {status_code}: {body}"}
+    if isinstance(exc, httpx.TimeoutException):
+        return {"error_type": "timeout", "error": "request timed out"}
+    if isinstance(exc, httpx.RequestError):
+        return {"error_type": "network_error", "error": str(exc)}
+    return {"error_type": "runtime_error", "error": str(exc)}
+
+
+def _mask_secret(text: str, secret: str) -> str:
+    return text.replace(secret, "***") if secret else text
+
+
+def _extract_text(result: dict[str, Any]) -> str:
+    content = result.get("content") or []
+    if isinstance(content, list):
+        parts = []
+        for item in content:
+            if isinstance(item, dict) and isinstance(item.get("text"), str):
+                parts.append(item["text"])
+        return "\n".join(parts).strip()
+    if isinstance(content, str):
+        return content.strip()
+    return ""
+
+
+def _parse_sse_or_json(response: httpx.Response) -> dict[str, Any]:
+    content_type = response.headers.get("content-type", "")
+    if "json" in content_type:
+        return response.json()
+    data_lines = []
+    for line in response.text.splitlines():
+        line = line.strip()
+        if line.startswith("data:"):
+            data_lines.append(line.removeprefix("data:").strip())
+    for line in reversed(data_lines):
+        if not line or line == "[DONE]":
+            continue
+        try:
+            return json.loads(line)
+        except json.JSONDecodeError:
+            continue
+    return response.json()
+
+
+def _parse_markdown_results(text: str, provider: str) -> list[dict[str, str]]:
+    results: list[dict[str, str]] = []
+    current: dict[str, str] | None = None
+    for line in text.splitlines():
+        heading = re.match(r"^#{2,4}\s+(?:\d+[.)]\s*)?(.+?)\s*$", line.strip())
+        if heading:
+            if current:
+                results.append(current)
+            current = {"title": heading.group(1).strip(), "url": "", "description": "", "provider": provider}
+            continue
+        url_match = re.search(r"https?://[^\s)>\]]+", line)
+        if url_match:
+            if current is None:
+                current = {"title": url_match.group(0), "url": "", "description": "", "provider": provider}
+            current["url"] = current.get("url") or url_match.group(0).rstrip(".,")
+            continue
+        if current is not None and line.strip() and not line.startswith("#"):
+            current["description"] = (current.get("description", "") + " " + line.strip()).strip()
+    if current:
+        results.append(current)
+    if results:
+        return results
+    urls = re.findall(r"https?://[^\s)>\]]+", text)
+    return [{"title": url, "url": url, "description": "", "provider": provider} for url in dict.fromkeys(urls)]
+
+
+def _content_error(text: str) -> tuple[str, str] | None:
+    stripped = (text or "").strip()
+    if not stripped:
+        return None
+    if stripped.lower().startswith("mcp error"):
+        lowered = stripped.lower()
+        error_type = "auth_error" if "-401" in stripped or "api key" in lowered else "provider_error"
+        return error_type, stripped
+    decoded: Any = stripped
+    for _ in range(2):
+        if not isinstance(decoded, str):
+            break
+        try:
+            decoded = json.loads(decoded)
+        except json.JSONDecodeError:
+            break
+    if isinstance(decoded, dict) and decoded.get("error"):
+        return "provider_error", str(decoded.get("error"))
+    return None
+
+
+class ZhipuMCPProvider:
+    def __init__(self, api_url: str, api_key: str, timeout: float = 30.0, provider_id: str = "zhipu-mcp"):
+        self.api_url = api_url.rstrip("/")
+        self.api_key = api_key
+        self.timeout = timeout
+        self.provider_id = provider_id
+
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> str:
+        start = time.time()
+        if not self.api_key:
+            return json.dumps(
+                {
+                    "ok": False,
+                    "provider": self.provider_id,
+                    "tool": name,
+                    "error_type": "config_error",
+                    "error": "ZHIPU_MCP_API_KEY is not configured.",
+                    "elapsed_ms": _elapsed_ms(start),
+                },
+                ensure_ascii=False,
+                indent=2,
+            )
+
+        payload = {
+            "jsonrpc": "2.0",
+            "id": 1,
+            "method": "tools/call",
+            "params": {"name": name, "arguments": arguments},
+        }
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "Accept": "application/json, text/event-stream",
+        }
+
+        try:
+            timeout = httpx.Timeout(connect=6.0, read=self.timeout, write=10.0, pool=None)
+            async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
+                response = await client.post(self.api_url, headers=headers, json=payload)
+                response.raise_for_status()
+                data = _parse_sse_or_json(response)
+            output = self._normalize_response(name, arguments, data, start)
+        except Exception as e:
+            error = _error_payload(e)
+            output = {
+                "ok": False,
+                "provider": self.provider_id,
+                "tool": name,
+                "error_type": error["error_type"],
+                "error": _mask_secret(error["error"], self.api_key),
+                "elapsed_ms": _elapsed_ms(start),
+            }
+        return json.dumps(output, ensure_ascii=False, indent=2)
+
+    def _normalize_response(self, name: str, arguments: dict[str, Any], data: dict[str, Any], start: float) -> dict[str, Any]:
+        if "error" in data:
+            error = data.get("error") or {}
+            message = error.get("message") if isinstance(error, dict) else str(error)
+            return {
+                "ok": False,
+                "provider": self.provider_id,
+                "tool": name,
+                "error_type": "provider_error",
+                "error": message or "Zhipu MCP JSON-RPC error",
+                "elapsed_ms": _elapsed_ms(start),
+            }
+
+        result = data.get("result") or {}
+        text = _extract_text(result)
+        content_error = _content_error(text)
+        is_error = bool(result.get("isError")) or bool(content_error)
+        output: dict[str, Any] = {
+            "ok": not is_error,
+            "provider": self.provider_id,
+            "tool": name,
+            "content": text,
+            "raw_content": text,
+            "elapsed_ms": _elapsed_ms(start),
+        }
+        for key in ("query", "url", "repo", "path", "ref"):
+            if arguments.get(key):
+                output[key] = arguments[key]
+        if name == "webReader":
+            output["url"] = arguments.get("url", "")
+        else:
+            results = [] if is_error else _parse_markdown_results(text, self.provider_id)
+            output["results"] = results
+            output["total"] = len(results)
+        if is_error:
+            output["error_type"] = content_error[0] if content_error else "provider_error"
+            output["error"] = content_error[1] if content_error else (text or "Zhipu MCP tool returned isError=true")
+        return output
+
+    async def web_search(self, query: str, count: int = 5) -> str:
+        del count
+        return await self.call_tool("web_search_prime", {"search_query": query})
+
+    async def web_reader(self, url: str) -> str:
+        return await self.call_tool("webReader", {"url": url})
+
+    async def search_doc(self, repo: str, query: str, max_results: int = 5) -> str:
+        del max_results
+        return await self.call_tool("search_doc", {"repo_name": repo, "query": query})
+
+    async def get_repo_structure(self, repo: str, path: str = "", ref: str = "") -> str:
+        arguments = {"repo_name": repo}
+        if path:
+            arguments["dir_path"] = path
+        if ref:
+            arguments["ref"] = ref
+        return await self.call_tool("get_repo_structure", arguments)
+
+    async def read_file(self, repo: str, path: str, ref: str = "") -> str:
+        arguments = {"repo_name": repo, "file_path": path}
+        if ref:
+            arguments["ref"] = ref
+        return await self.call_tool("read_file", arguments)