zai-cli 0.1.0__py3-none-any.whl

zai/providers/openrouter.py

@@ -0,0 +1,63 @@
+import httpx
+from .base import BaseProvider, Message, Response
+from ..config import get_api_key
+from ..core.errors import (
+    NoAPIKeyError,
+    AuthenticationError,
+    RateLimitError,
+    NetworkError,
+    classify_provider_error,
+)
+
+
+class OpenRouterProvider(BaseProvider):
+    name = "openrouter"
+    model_id = "meta-llama/llama-3.3-70b-instruct:free"
+    context_window = 128000
+
+    def is_available(self) -> bool:
+        return bool(get_api_key("openrouter"))
+
+    def chat(
+        self,
+        messages: list[Message],
+        system: str = "",
+        tools: list[dict] | None = None,
+    ) -> Response:
+        key = get_api_key("openrouter")
+        if not key:
+            raise NoAPIKeyError("openrouter")
+
+        formatted = [{"role": "system", "content": system or "You are zai, a helpful AI assistant."}]
+        formatted += [{"role": m.role, "content": m.content} for m in messages]
+
+        for attempt in range(self.retries + 1):
+            try:
+                response = httpx.post(
+                    "https://openrouter.ai/api/v1/chat/completions",
+                    headers={"Authorization": f"Bearer {key}"},
+                    json={"model": self.model_id, "messages": formatted},
+                    timeout=self.timeout,
+                )
+                if response.status_code == 429:
+                    raise RateLimitError("openrouter")
+                if response.status_code == 401:
+                    raise AuthenticationError(
+                        "openrouter", "invalid API credentials"
+                    )
+                try:
+                    response.raise_for_status()
+                    data = response.json()
+                except Exception as error:
+                    raise classify_provider_error("openrouter", error) from error
+                return Response(
+                    content=data["choices"][0]["message"]["content"],
+                    model=self.model_id,
+                    tokens_used=data.get("usage", {}).get("total_tokens", 0),
+                )
+            except (httpx.ConnectError, httpx.TimeoutException) as error:
+                if attempt == self.retries:
+                    raise NetworkError(str(error))
+            except (RateLimitError, NoAPIKeyError):
+                raise
+        raise NetworkError("OpenRouter request failed")

zai/providers/qwen.py

@@ -0,0 +1,47 @@
+import httpx
+from .base import BaseProvider, Message, Response
+from ..config import get_api_key
+from ..core.errors import NetworkError, NoAPIKeyError, classify_provider_error
+
+
+class QwenProvider(BaseProvider):
+    name = "qwen"
+    model_id = "qwen-turbo"
+    context_window = 1000000
+
+    def is_available(self) -> bool:
+        return bool(get_api_key("qwen"))
+
+    def chat(
+        self,
+        messages: list[Message],
+        system: str = "",
+        tools: list[dict] | None = None,
+    ) -> Response:
+        key = get_api_key("qwen")
+        if not key:
+            raise NoAPIKeyError("qwen")
+        formatted = [{"role": "system", "content": system or "You are zai, a helpful AI assistant."}]
+        formatted += [{"role": m.role, "content": m.content} for m in messages]
+        for attempt in range(self.retries + 1):
+            try:
+                response = httpx.post(
+                    "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
+                    headers={"Authorization": f"Bearer {key}"},
+                    json={"model": self.model_id, "messages": formatted},
+                    timeout=self.timeout,
+                )
+                try:
+                    response.raise_for_status()
+                    data = response.json()
+                except Exception as error:
+                    raise classify_provider_error("qwen", error) from error
+                return Response(
+                    content=data["choices"][0]["message"]["content"],
+                    model=self.model_id,
+                    tokens_used=data.get("usage", {}).get("total_tokens", 0),
+                )
+            except (httpx.ConnectError, httpx.TimeoutException) as error:
+                if attempt == self.retries:
+                    raise NetworkError(str(error))
+        raise NetworkError("Qwen request failed")

zai/skills/registry.py

@@ -0,0 +1,52 @@
+from ..tools.files import read_file
+
+SKILLS = {
+    "review": {
+        "description": "Review code and suggest improvements",
+        "prompt": lambda ctx: f"Do a thorough code review of the following. Point out bugs, security issues, and improvements:\n\n{ctx}",
+    },
+    "test": {
+        "description": "Generate unit tests",
+        "prompt": lambda ctx: f"Write comprehensive unit tests for this code using pytest:\n\n{ctx}",
+    },
+    "commit": {
+        "description": "Generate git commit message",
+        "prompt": lambda ctx: f"Write a clear, concise git commit message for these changes (conventional commits format):\n\n{ctx}",
+    },
+    "docs": {
+        "description": "Generate documentation",
+        "prompt": lambda ctx: f"Write clear documentation (docstrings + README section) for this code:\n\n{ctx}",
+    },
+    "fix": {
+        "description": "Find and fix bugs",
+        "prompt": lambda ctx: f"Find all bugs in this code and provide the fixed version with explanation:\n\n{ctx}",
+    },
+    "explain": {
+        "description": "Explain code clearly",
+        "prompt": lambda ctx: f"Explain this code clearly, step by step, for someone learning:\n\n{ctx}",
+    },
+    "refactor": {
+        "description": "Refactor code for quality",
+        "prompt": lambda ctx: f"Refactor this code for better readability, performance, and best practices:\n\n{ctx}",
+    },
+    "summarize": {
+        "description": "Summarize text or file",
+        "prompt": lambda ctx: f"Summarize the following in clear bullet points:\n\n{ctx}",
+    },
+    "translate": {
+        "description": "Translate text",
+        "prompt": lambda ctx: f"Translate the following text to English (or specify target language):\n\n{ctx}",
+    },
+}
+
+
+def get_skill_prompt(skill_name: str, file_path: str = None, text: str = None) -> str | None:
+    skill = SKILLS.get(skill_name)
+    if not skill:
+        return None
+    if file_path:
+        content = read_file(file_path)
+        return skill["prompt"](content)
+    if text:
+        return skill["prompt"](text)
+    return skill["prompt"]("")

zai/tools/browser.py

@@ -0,0 +1,224 @@
+import ipaddress
+import socket
+from pathlib import Path
+from urllib.parse import urlsplit
+
+from ..core.errors import FileError, ZaiError
+from ..core.security import resolve_project_path
+
+MAX_DOCUMENT_BYTES = 2_000_000
+
+
+class BrowserError(ZaiError):
+    pass
+
+
+def _get_playwright():
+    try:
+        from playwright.sync_api import sync_playwright
+
+        return sync_playwright
+    except ImportError:
+        raise BrowserError(
+            "Playwright not installed. Run: pip install playwright && "
+            "playwright install chromium"
+        )
+
+
+def _validate_public_url(url: str) -> str:
+    """Allow only credential-free public HTTP(S) destinations."""
+    try:
+        parsed = urlsplit(url)
+        port = parsed.port
+    except ValueError as error:
+        raise BrowserError(f"Invalid URL: {error}")
+    if parsed.scheme.lower() not in {"http", "https"}:
+        raise BrowserError("Only http:// and https:// URLs are allowed")
+    if not parsed.hostname:
+        raise BrowserError("URL must include a hostname")
+    if parsed.username is not None or parsed.password is not None:
+        raise BrowserError("URLs containing credentials are not allowed")
+
+    hostname = parsed.hostname.rstrip(".").lower()
+    if hostname == "localhost" or hostname.endswith(".localhost"):
+        raise BrowserError("Local and private network URLs are blocked")
+    try:
+        addresses = {
+            item[4][0]
+            for item in socket.getaddrinfo(
+                hostname,
+                port or (443 if parsed.scheme.lower() == "https" else 80),
+                type=socket.SOCK_STREAM,
+            )
+        }
+    except socket.gaierror as error:
+        raise BrowserError(f"Cannot resolve URL hostname: {error}")
+    if not addresses:
+        raise BrowserError("URL hostname did not resolve")
+    for address in addresses:
+        try:
+            if not ipaddress.ip_address(address).is_global:
+                raise BrowserError("Local and private network URLs are blocked")
+        except ValueError:
+            raise BrowserError("URL resolved to an invalid network address")
+    return url
+
+
+def _safe_output_path(path: str, project_dir: str = ".") -> Path:
+    if Path(path).suffix.lower() not in {".png", ".jpg", ".jpeg"}:
+        raise BrowserError("Screenshot output must use .png, .jpg, or .jpeg")
+    try:
+        output = resolve_project_path(project_dir, path)
+    except FileError as error:
+        raise BrowserError(str(error))
+    output.parent.mkdir(parents=True, exist_ok=True)
+    return output
+
+
+def _prepare_page(browser, initial_url: str):
+    context = browser.new_context(
+        service_workers="block",
+        accept_downloads=False,
+    )
+    page = context.new_page()
+
+    def guard_request(route):
+        try:
+            _validate_public_url(route.request.url)
+        except BrowserError:
+            route.abort("blockedbyclient")
+            return
+        route.continue_()
+
+    page.route("http://**/*", guard_request)
+    page.route("https://**/*", guard_request)
+    response = page.goto(
+        _validate_public_url(initial_url),
+        timeout=30000,
+        wait_until="domcontentloaded",
+    )
+    if response is not None:
+        content_length = response.headers.get("content-length")
+        if content_length:
+            try:
+                if int(content_length) > MAX_DOCUMENT_BYTES:
+                    raise BrowserError("Webpage exceeded the 2 MB document limit")
+            except ValueError:
+                pass
+    _validate_public_url(page.url)
+    return context, page
+
+
+def scrape(url: str) -> str:
+    """Open a public URL and return bounded visible text content."""
+    sync_playwright = _get_playwright()
+    try:
+        with sync_playwright() as playwright:
+            browser = playwright.chromium.launch(headless=True)
+            try:
+                context, page = _prepare_page(browser, url)
+                try:
+                    text = page.evaluate("() => document.body.innerText")
+                    return text[:8000] if text else "No content found."
+                finally:
+                    context.close()
+            finally:
+                browser.close()
+    except BrowserError:
+        raise
+    except Exception as error:
+        raise BrowserError(f"Cannot scrape {url}: {error}")
+
+
+def screenshot(url: str, path: str = "screenshot.png", project_dir: str = ".") -> str:
+    """Take a screenshot of a public URL inside the current project."""
+    output = _safe_output_path(path, project_dir)
+    sync_playwright = _get_playwright()
+    try:
+        with sync_playwright() as playwright:
+            browser = playwright.chromium.launch(headless=True)
+            try:
+                context, page = _prepare_page(browser, url)
+                try:
+                    page.screenshot(path=str(output), full_page=True)
+                    return str(output)
+                finally:
+                    context.close()
+            finally:
+                browser.close()
+    except BrowserError:
+        raise
+    except Exception as error:
+        raise BrowserError(f"Screenshot failed: {error}")
+
+
+def click_and_get(url: str, selector: str) -> str:
+    """Click an element on a public page and return bounded content."""
+    sync_playwright = _get_playwright()
+    try:
+        with sync_playwright() as playwright:
+            browser = playwright.chromium.launch(headless=True)
+            try:
+                context, page = _prepare_page(browser, url)
+                try:
+                    page.click(selector, timeout=5000)
+                    page.wait_for_load_state("networkidle", timeout=10000)
+                    _validate_public_url(page.url)
+                    text = page.evaluate("() => document.body.innerText")
+                    return text[:8000] if text else "No content after click."
+                finally:
+                    context.close()
+            finally:
+                browser.close()
+    except BrowserError:
+        raise
+    except Exception as error:
+        raise BrowserError(f"Click failed: {error}")
+
+
+def fill_form(url: str, fields: dict[str, str], submit_selector: str = None) -> str:
+    """Fill a form on a public page and optionally submit it."""
+    sync_playwright = _get_playwright()
+    try:
+        with sync_playwright() as playwright:
+            browser = playwright.chromium.launch(headless=True)
+            try:
+                context, page = _prepare_page(browser, url)
+                try:
+                    for selector, value in fields.items():
+                        page.fill(selector, value)
+                    if submit_selector:
+                        page.click(submit_selector)
+                        page.wait_for_load_state("networkidle", timeout=10000)
+                        _validate_public_url(page.url)
+                    text = page.evaluate("() => document.body.innerText")
+                    return text[:5000] if text else "Done."
+                finally:
+                    context.close()
+            finally:
+                browser.close()
+    except BrowserError:
+        raise
+    except Exception as error:
+        raise BrowserError(f"Form fill failed: {error}")
+
+
+def run_js(url: str, script: str) -> str:
+    """Run JavaScript on a public page and return a bounded result."""
+    sync_playwright = _get_playwright()
+    try:
+        with sync_playwright() as playwright:
+            browser = playwright.chromium.launch(headless=True)
+            try:
+                context, page = _prepare_page(browser, url)
+                try:
+                    result = page.evaluate(script)
+                    return str(result)[:8000]
+                finally:
+                    context.close()
+            finally:
+                browser.close()
+    except BrowserError:
+        raise
+    except Exception as error:
+        raise BrowserError(f"JS execution failed: {error}")

zai/tools/code_runner.py

@@ -0,0 +1,49 @@
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+
+def run_python(code: str, timeout: int = 30) -> str:
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+        f.write(code)
+        tmp = f.name
+    try:
+        result = subprocess.run(
+            [sys.executable, tmp],
+            capture_output=True, text=True, timeout=timeout
+        )
+        out = result.stdout.strip()
+        err = result.stderr.strip()
+        if err:
+            return f"Output:\n{out}\n\nErrors:\n{err}" if out else f"Error:\n{err}"
+        return out or "No output."
+    except subprocess.TimeoutExpired:
+        return f"Error: Code timed out after {timeout} seconds."
+    except Exception as e:
+        return f"Error: {e}"
+    finally:
+        Path(tmp).unlink(missing_ok=True)
+
+
+def run_file(path: str, timeout: int = 30) -> str:
+    from ..core.errors import FileError
+    target = Path(path)
+    if not target.exists():
+        raise FileError(f"File not found: {path}")
+    if not target.is_file():
+        raise FileError(f"Not a file: {path}")
+    try:
+        result = subprocess.run(
+            [sys.executable, str(target)],
+            capture_output=True, text=True, timeout=timeout
+        )
+        out = result.stdout.strip()
+        err = result.stderr.strip()
+        if err:
+            return f"Output:\n{out}\n\nErrors:\n{err}" if out else f"Error:\n{err}"
+        return out or "No output."
+    except subprocess.TimeoutExpired:
+        return f"Error: Timed out after {timeout}s."
+    except Exception as e:
+        return f"Error: {e}"

zai/tools/files.py

@@ -0,0 +1,53 @@
+from pathlib import Path
+from ..core.errors import FileError
+from ..core.storage import atomic_write_text
+
+
+def read_file(path: str) -> str:
+    p = Path(path)
+    if not p.exists():
+        raise FileError(f"File not found: {path}")
+    if not p.is_file():
+        raise FileError(f"Not a file: {path}")
+    try:
+        return p.read_text(encoding="utf-8", errors="ignore")
+    except PermissionError:
+        raise FileError(f"Permission denied: {path}")
+    except Exception as e:
+        raise FileError(f"Cannot read {path}: {e}")
+
+
+def write_file(path: str, content: str):
+    p = Path(path)
+    try:
+        p.parent.mkdir(parents=True, exist_ok=True)
+        atomic_write_text(p, content, mode=0o644, lock=False)
+    except PermissionError:
+        raise FileError(f"Permission denied: {path}")
+    except Exception as e:
+        raise FileError(f"Cannot write {path}: {e}")
+
+
+def edit_file(path: str, old: str, new: str) -> bool:
+    p = Path(path)
+    if not p.exists():
+        raise FileError(f"File not found: {path}")
+    content = p.read_text(encoding="utf-8")
+    if old not in content:
+        return False
+    atomic_write_text(
+        p,
+        content.replace(old, new, 1),
+        mode=0o644,
+        lock=False,
+    )
+    return True
+
+
+def list_files(directory: str = ".", pattern: str = "*") -> list[str]:
+    d = Path(directory)
+    if not d.exists():
+        raise FileError(f"Directory not found: {directory}")
+    if not d.is_dir():
+        raise FileError(f"Not a directory: {directory}")
+    return [str(p) for p in d.rglob(pattern) if p.is_file()]

zai/tools/git.py

@@ -0,0 +1,38 @@
+import subprocess
+
+
+def run_git(args: list[str]) -> str:
+    try:
+        result = subprocess.run(
+            ["git"] + args,
+            capture_output=True, text=True, timeout=30
+        )
+        return result.stdout.strip() or result.stderr.strip()
+    except FileNotFoundError:
+        return "Error: git not found. Please install git."
+    except Exception as e:
+        return f"Error: {e}"
+
+
+def get_diff() -> str:
+    return run_git(["diff", "HEAD"])
+
+
+def get_staged_diff() -> str:
+    return run_git(["diff", "--staged"])
+
+
+def get_status() -> str:
+    return run_git(["status", "--short"])
+
+
+def get_log(n: int = 10) -> str:
+    return run_git(["log", f"-{n}", "--oneline"])
+
+
+def commit(message: str) -> str:
+    return run_git(["commit", "-m", message])
+
+
+def get_branch() -> str:
+    return run_git(["branch", "--show-current"])

zai/tools/search.py

@@ -0,0 +1,157 @@
+"""Bounded web-search provider adapters."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from html.parser import HTMLParser
+from urllib.parse import parse_qs, unquote, urlsplit
+
+import httpx
+
+from ..core.errors import NetworkError
+
+MAX_SEARCH_RESPONSE_BYTES = 1_000_000
+MAX_RESULTS = 8
+DUCKDUCKGO_HTML_URL = "https://html.duckduckgo.com/html/"
+
+
+@dataclass(frozen=True)
+class SearchResult:
+    title: str
+    url: str
+    snippet: str = ""
+
+
+class SearchProvider:
+    """Interface implemented by general web-search backends."""
+
+    name = "search"
+
+    def search(self, query: str, *, limit: int = MAX_RESULTS) -> list[SearchResult]:
+        raise NotImplementedError
+
+
+class _DuckDuckGoHTMLParser(HTMLParser):
+    def __init__(self) -> None:
+        super().__init__(convert_charrefs=True)
+        self.results: list[SearchResult] = []
+        self._title_parts: list[str] | None = None
+        self._snippet_parts: list[str] | None = None
+        self._url = ""
+
+    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
+        attributes = dict(attrs)
+        classes = set((attributes.get("class") or "").split())
+        if tag == "a" and "result__a" in classes:
+            self._url = _unwrap_duckduckgo_url(attributes.get("href") or "")
+            self._title_parts = []
+        elif tag in {"a", "div"} and "result__snippet" in classes:
+            self._snippet_parts = []
+
+    def handle_data(self, data: str) -> None:
+        if self._title_parts is not None:
+            self._title_parts.append(data)
+        if self._snippet_parts is not None:
+            self._snippet_parts.append(data)
+
+    def handle_endtag(self, tag: str) -> None:
+        if tag == "a" and self._title_parts is not None:
+            title = " ".join("".join(self._title_parts).split())
+            if title and self._url:
+                self.results.append(SearchResult(title=title, url=self._url))
+            self._title_parts = None
+            self._url = ""
+        if tag in {"a", "div"} and self._snippet_parts is not None:
+            snippet = " ".join("".join(self._snippet_parts).split())
+            if snippet and self.results:
+                previous = self.results[-1]
+                self.results[-1] = SearchResult(
+                    title=previous.title,
+                    url=previous.url,
+                    snippet=snippet,
+                )
+            self._snippet_parts = None
+
+
+def _unwrap_duckduckgo_url(url: str) -> str:
+    parsed = urlsplit(url)
+    redirected = parse_qs(parsed.query).get("uddg")
+    return unquote(redirected[0]) if redirected else url
+
+
+def _read_bounded_response(
+    client: httpx.Client,
+    url: str,
+    *,
+    data: dict[str, str],
+) -> str:
+    with client.stream("POST", url, data=data) as response:
+        response.raise_for_status()
+        content_length = response.headers.get("content-length")
+        if content_length and int(content_length) > MAX_SEARCH_RESPONSE_BYTES:
+            raise NetworkError("Search response exceeded the 1 MB limit.")
+        body = bytearray()
+        for chunk in response.iter_bytes():
+            body.extend(chunk)
+            if len(body) > MAX_SEARCH_RESPONSE_BYTES:
+                raise NetworkError("Search response exceeded the 1 MB limit.")
+        return body.decode(response.encoding or "utf-8", errors="replace")
+
+
+class DuckDuckGoHTMLSearchProvider(SearchProvider):
+    """General web results from DuckDuckGo's HTML search endpoint."""
+
+    name = "duckduckgo"
+
+    def search(self, query: str, *, limit: int = MAX_RESULTS) -> list[SearchResult]:
+        query = " ".join(query.split())
+        if not query:
+            raise NetworkError("Search query cannot be empty.")
+        headers = {
+            "User-Agent": "zai-cli/0.1 (+https://github.com/zai-cli/zai)",
+            "Accept": "text/html,application/xhtml+xml",
+        }
+        try:
+            with httpx.Client(
+                headers=headers,
+                timeout=httpx.Timeout(10.0),
+                follow_redirects=True,
+            ) as client:
+                html = _read_bounded_response(
+                    client,
+                    DUCKDUCKGO_HTML_URL,
+                    data={"q": query},
+                )
+        except NetworkError:
+            raise
+        except httpx.ConnectError as error:
+            raise NetworkError("Cannot connect. Check your internet.") from error
+        except httpx.TimeoutException as error:
+            raise NetworkError("Search timed out.") from error
+        except httpx.HTTPError as error:
+            raise NetworkError(f"Search provider error: {error}") from error
+
+        parser = _DuckDuckGoHTMLParser()
+        parser.feed(html)
+        return parser.results[: max(1, min(limit, MAX_RESULTS))]
+
+
+def _format_results(query: str, results: list[SearchResult]) -> str:
+    if not results:
+        return f"No results found for: {query}"
+    lines = []
+    for index, result in enumerate(results, 1):
+        lines.append(f"{index}. {result.title}\n   {result.url}")
+        if result.snippet:
+            lines.append(f"   {result.snippet}")
+    return "\n".join(lines)
+
+
+def web_search(query: str, provider: SearchProvider | None = None) -> str:
+    """Search through an injectable provider and return bounded plain text."""
+    backend = provider or DuckDuckGoHTMLSearchProvider()
+    try:
+        return _format_results(query, backend.search(query))
+    except NetworkError:
+        raise
+    except Exception as error:
+        raise NetworkError(f"Search failed: {error}") from error