zai-cli 0.1.0__py3-none-any.whl

zai/core/cancellation.py

@@ -0,0 +1,67 @@
+"""Cooperative cancellation shared by providers, tools, and integrations."""
+from __future__ import annotations
+
+import threading
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from typing import Iterator
+
+
+class OperationCancelled(RuntimeError):
+    pass
+
+
+@dataclass
+class CancellationToken:
+    _event: threading.Event = field(default_factory=threading.Event)
+    reason: str = "operation cancelled"
+
+    def cancel(self, reason: str = "operation cancelled") -> None:
+        self.reason = reason
+        self._event.set()
+
+    @property
+    def cancelled(self) -> bool:
+        return self._event.is_set()
+
+    def raise_if_cancelled(self) -> None:
+        if self.cancelled:
+            raise OperationCancelled(self.reason)
+
+
+_current: ContextVar[CancellationToken | None] = ContextVar(
+    "zai_cancellation_token",
+    default=None,
+)
+
+
+def current_token() -> CancellationToken | None:
+    return _current.get()
+
+
+def raise_if_cancelled() -> None:
+    token = current_token()
+    if token:
+        token.raise_if_cancelled()
+
+
+def cancel_current(reason: str = "operation cancelled") -> bool:
+    token = current_token()
+    if not token:
+        return False
+    token.cancel(reason)
+    return True
+
+
+@contextmanager
+def operation(token: CancellationToken | None = None) -> Iterator[CancellationToken]:
+    active = token or CancellationToken()
+    marker = _current.set(active)
+    try:
+        yield active
+    except KeyboardInterrupt:
+        active.cancel("cancelled by user")
+        raise
+    finally:
+        _current.reset(marker)

zai/core/commands.py

@@ -0,0 +1,85 @@
+"""
+Slash commands system — inspired by Claude Code.
+Commands are .md files with YAML frontmatter in ~/.zai/commands/
+"""
+import os
+import re
+from pathlib import Path
+from .process import run_direct
+
+COMMANDS_DIR = Path.home() / ".zai" / "commands"
+BUILTIN_DIR = Path(__file__).parent.parent / "commands"
+
+
+def _parse_frontmatter(text: str) -> tuple[dict, str]:
+    """Parse YAML frontmatter from markdown file. Returns (meta, body)."""
+    if not text.startswith("---"):
+        return {}, text
+    end = text.find("\n---", 3)
+    if end == -1:
+        return {}, text
+    frontmatter = text[3:end].strip()
+    body = text[end + 4:].strip()
+    meta = {}
+    for line in frontmatter.splitlines():
+        if ":" in line:
+            k, v = line.split(":", 1)
+            meta[k.strip()] = v.strip()
+    return meta, body
+
+
+def _substitute_shell(prompt: str) -> str:
+    """Replace !`command` with command output (like Claude Code does)."""
+    pattern = r'!`([^`]+)`'
+    def run(m):
+        cmd = m.group(1)
+        result = run_direct(cmd, timeout=10)
+        if result.blocked_reason:
+            return f"(blocked command substitution: {result.blocked_reason})"
+        if result.cancelled:
+            return "(blocked command substitution: approval required)"
+        if result.returncode != 0:
+            return f"(failed: {cmd}: {result.output})"
+        return result.output
+    return re.sub(pattern, run, prompt)
+
+
+def load_commands() -> dict:
+    """Load all available slash commands. Returns {name: {meta, body, path}}."""
+    commands = {}
+
+    # Load built-in commands first
+    for search_dir in [BUILTIN_DIR, COMMANDS_DIR]:
+        if not search_dir.exists():
+            continue
+        for md_file in sorted(search_dir.glob("*.md")):
+            name = md_file.stem
+            text = md_file.read_text(encoding="utf-8", errors="ignore")
+            meta, body = _parse_frontmatter(text)
+            commands[name] = {
+                "meta": meta,
+                "body": body,
+                "path": str(md_file),
+                "description": meta.get("description", ""),
+            }
+
+    return commands
+
+
+def get_command_prompt(name: str, arguments: str = "") -> str | None:
+    """Get the prompt for a slash command, with substitutions applied."""
+    commands = load_commands()
+    if name not in commands:
+        return None
+    body = commands[name]["body"]
+    body = body.replace("$ARGUMENTS", arguments)
+    body = _substitute_shell(body)
+    return body
+
+
+def list_commands() -> list[dict]:
+    """List all commands with name and description."""
+    return [
+        {"name": name, "description": cmd["description"]}
+        for name, cmd in load_commands().items()
+    ]

zai/core/context.py

@@ -0,0 +1,299 @@
+"""Shared, model-aware conversation context management."""
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import replace
+
+from ..providers.base import Message
+
+DEFAULT_CONTEXT_WINDOW = 128_000
+DEFAULT_OUTPUT_RESERVE = 8_192
+MAX_TOOL_RESULT_CHARS = 16_000
+SUMMARY_MAX_CHARS = 12_000
+
+
+def estimate_text_tokens(text: str) -> int:
+    """Conservative provider-neutral token estimate."""
+    if not text:
+        return 0
+    words = len(re.findall(r"\S+", text))
+    return max(1, (len(text) + 2) // 3, int(words * 1.35))
+
+
+def estimate_message_tokens(message: Message) -> int:
+    overhead = 8
+    tool_payload = ""
+    if message.tool_calls:
+        tool_payload = json.dumps([
+            {
+                "id": call.id,
+                "name": call.name,
+                "arguments": call.arguments,
+            }
+            for call in message.tool_calls
+        ], ensure_ascii=False)
+    return overhead + estimate_text_tokens(message.content) + estimate_text_tokens(
+        tool_payload
+    )
+
+
+def bound_tool_result(content: str, max_chars: int = MAX_TOOL_RESULT_CHARS) -> str:
+    if max_chars <= 0:
+        return ""
+    if len(content) <= max_chars:
+        return content
+    head = max_chars * 2 // 3
+    tail = max_chars - head
+    removed = len(content) - max_chars
+    return (
+        content[:head]
+        + f"\n\n[... {removed} characters omitted from tool result ...]\n\n"
+        + content[-tail:]
+    )
+
+
+def _fit_message(message: Message, token_budget: int) -> Message | None:
+    if token_budget <= 8:
+        return None
+    if estimate_message_tokens(message) <= token_budget:
+        return message
+    low, high = 0, len(message.content)
+    best = ""
+    while low <= high:
+        middle = (low + high) // 2
+        candidate_content = bound_tool_result(message.content, middle)
+        candidate = replace(message, content=candidate_content)
+        if estimate_message_tokens(candidate) <= token_budget:
+            best = candidate_content
+            low = middle + 1
+        else:
+            high = middle - 1
+    return replace(message, content=best) if best or not message.content else None
+
+
+def _context_units(messages: list[Message]) -> list[list[Message]]:
+    """Keep assistant tool calls and their tool results indivisible."""
+    units: list[list[Message]] = []
+    index = 0
+    while index < len(messages):
+        message = messages[index]
+        unit = [message]
+        index += 1
+        if message.role == "assistant" and message.tool_calls:
+            call_ids = {call.id for call in message.tool_calls}
+            while (
+                index < len(messages)
+                and messages[index].role == "tool"
+                and messages[index].tool_call_id in call_ids
+            ):
+                unit.append(messages[index])
+                index += 1
+        units.append(unit)
+    return units
+
+
+def _unit_tokens(unit: list[Message]) -> int:
+    return sum(estimate_message_tokens(message) for message in unit)
+
+
+def _summary(messages: list[Message]) -> Message:
+    lines = [
+        "[Compacted earlier context]",
+        "The following facts were retained from older conversation turns:",
+    ]
+    remaining = SUMMARY_MAX_CHARS
+    for message in messages:
+        if message.role == "tool":
+            label = f"TOOL {message.tool_name or 'result'}"
+        else:
+            label = message.role.upper()
+        content = " ".join(message.content.split())
+        if not content and message.tool_calls:
+            content = "Requested tools: " + ", ".join(
+                call.name for call in message.tool_calls
+            )
+        if not content:
+            continue
+        excerpt = content[:800]
+        line = f"- {label}: {excerpt}"
+        if len(line) > remaining:
+            break
+        lines.append(line)
+        remaining -= len(line)
+    return Message(role="user", content="\n".join(lines), pinned=True)
+
+
+def compact_messages(
+    messages: list[Message],
+    max_tokens: int,
+    *,
+    reserve_tokens: int = DEFAULT_OUTPUT_RESERVE,
+) -> list[Message]:
+    """Fit messages to a context budget while preserving pinned and recent units."""
+    input_budget = max(128, max_tokens - reserve_tokens)
+    normalized = [
+        replace(
+            message,
+            content=(
+                bound_tool_result(message.content)
+                if message.role == "tool"
+                else message.content
+            ),
+        )
+        for message in messages
+    ]
+    if sum(estimate_message_tokens(message) for message in normalized) <= input_budget:
+        return normalized
+
+    units = _context_units(normalized)
+    pinned_units = [
+        unit for unit in units if any(message.pinned for message in unit)
+    ]
+    regular_units = [
+        unit for unit in units if not any(message.pinned for message in unit)
+    ]
+    summary_reserve = min(512, max(64, input_budget // 4))
+    selection_budget = max(64, input_budget - summary_reserve)
+    selected: list[list[Message]] = []
+    used = sum(_unit_tokens(unit) for unit in pinned_units)
+
+    for position, unit in enumerate(reversed(regular_units)):
+        cost = _unit_tokens(unit)
+        if position == 0:
+            # The latest user/tool turn must never disappear entirely.
+            selected.append(unit)
+            used += cost
+            continue
+        if used + cost > selection_budget:
+            continue
+        selected.append(unit)
+        used += cost
+    selected.reverse()
+
+    selected_ids = {id(message) for unit in selected for message in unit}
+    omitted = [
+        message
+        for unit in regular_units
+        for message in unit
+        if id(message) not in selected_ids
+    ]
+    result = [message for unit in pinned_units for message in unit]
+    if omitted and used < input_budget:
+        summary = _summary(omitted)
+        if estimate_message_tokens(summary) > summary_reserve:
+            summary = replace(
+                summary,
+                content=bound_tool_result(
+                    summary.content,
+                    max(128, summary_reserve * 3 - 32),
+                ),
+            )
+        if used + estimate_message_tokens(summary) <= input_budget:
+            result.append(summary)
+    result.extend(message for unit in selected for message in unit)
+
+    # A single pinned or recent message may itself exceed the model limit.
+    final: list[Message] = []
+    remaining = input_budget
+    for message in result:
+        cost = estimate_message_tokens(message)
+        if cost <= remaining:
+            final.append(message)
+            remaining -= cost
+            continue
+        fitted = _fit_message(message, remaining)
+        if fitted is not None:
+            final.append(fitted)
+        break
+    return final
+
+
+class ContextManager:
+    def __init__(
+        self,
+        max_tokens: int = DEFAULT_CONTEXT_WINDOW,
+        model: str | None = None,
+    ):
+        self.messages: list[Message] = []
+        self.max_tokens = max_tokens
+        self.model = model
+        if model:
+            self.set_model(model)
+
+    def set_model(self, model: str | None) -> None:
+        self.model = model
+        if model:
+            try:
+                from ..config import get_model_config
+
+                self.max_tokens = get_model_config(model)["context_window"]
+            except KeyError:
+                pass
+
+    def add(
+        self,
+        role: str,
+        content: str,
+        *,
+        pinned: bool = False,
+        **message_fields,
+    ) -> Message:
+        message = Message(
+            role=role,
+            content=content,
+            pinned=pinned,
+            **message_fields,
+        )
+        self.messages.append(message)
+        return message
+
+    def add_message(self, message: Message, *, pinned: bool | None = None) -> None:
+        self.messages.append(
+            replace(message, pinned=pinned)
+            if pinned is not None
+            else message
+        )
+
+    def replace_messages(self, messages: list[Message]) -> None:
+        self.messages[:] = messages
+
+    def is_near_limit(self, threshold: float = 0.85) -> bool:
+        return self.token_count() >= self.max_tokens * threshold
+
+    def compress(self, preferred: str | None = None) -> None:
+        if preferred:
+            self.set_model(preferred)
+        self.messages[:] = compact_messages(
+            self.messages,
+            self.max_tokens,
+        )
+
+    def prepared(
+        self,
+        *,
+        model: str | None = None,
+        reserve_tokens: int = DEFAULT_OUTPUT_RESERVE,
+    ) -> list[Message]:
+        if model:
+            self.set_model(model)
+        return compact_messages(
+            self.messages,
+            self.max_tokens,
+            reserve_tokens=reserve_tokens,
+        )
+
+    def update_limit(self, new_max: int) -> None:
+        self.max_tokens = new_max
+
+    def get_messages(self) -> list[Message]:
+        return self.messages
+
+    def token_count(self) -> int:
+        return sum(estimate_message_tokens(message) for message in self.messages)
+
+    def usage_bar(self) -> str:
+        ratio = self.token_count() / max(1, self.max_tokens)
+        filled = min(int(ratio * 30), 30)
+        bar = "#" * filled + "-" * (30 - filled)
+        return f"[{bar}] {int(ratio * 100)}%"

zai/core/errors.py

@@ -0,0 +1,125 @@
+from rich.console import Console
+from .runtime import print_exception
+
+console = Console()
+
+
+class ZaiError(Exception):
+    pass
+
+
+class NoAPIKeyError(ZaiError):
+    def __init__(self, provider: str):
+        self.provider = provider
+        super().__init__(f"No API key for {provider}. Run: zai setup")
+
+
+class ProviderError(ZaiError):
+    category = "provider"
+
+    def __init__(self, provider: str, message: str):
+        self.provider = provider
+        self.detail = message
+        super().__init__(f"{provider} {self.category} error: {message}")
+
+
+class AuthenticationError(ProviderError):
+    category = "authentication"
+
+
+class QuotaError(ProviderError):
+    category = "quota"
+
+
+class ModelNotFoundError(ProviderError):
+    category = "model-not-found"
+
+
+class ProviderTimeoutError(ProviderError):
+    category = "timeout"
+
+
+class MalformedResponseError(ProviderError):
+    category = "malformed-response"
+
+
+class AllModelsFailedError(ZaiError):
+    def __init__(self, last_error: str = ""):
+        message = "All models failed. Check API keys with `zai setup` or verify Ollama."
+        if last_error:
+            message += f" Last error: {last_error}"
+        super().__init__(message)
+
+
+class RateLimitError(QuotaError):
+    def __init__(self, provider: str):
+        super().__init__(provider, "rate limit hit")
+
+
+class NetworkError(ZaiError):
+    def __init__(self, msg: str = ""):
+        super().__init__(f"Network error: {msg or 'Check your internet connection.'}")
+
+
+def classify_provider_error(provider: str, error: Exception) -> ZaiError:
+    """Convert SDK/HTTP failures into stable user-facing categories."""
+    if isinstance(error, ZaiError):
+        return error
+    status = getattr(getattr(error, "response", None), "status_code", None)
+    text = str(error)
+    lowered = text.lower()
+    if status in {401, 403} or any(
+        marker in lowered for marker in ("api key", "unauthorized", "authentication")
+    ):
+        return AuthenticationError(provider, text)
+    if status == 429 or any(
+        marker in lowered for marker in ("quota", "rate limit", "resource exhausted")
+    ):
+        return QuotaError(provider, text)
+    if status == 404 or any(
+        marker in lowered for marker in ("model not found", "unknown model", "does not exist")
+    ):
+        return ModelNotFoundError(provider, text)
+    if any(marker in lowered for marker in ("timeout", "timed out", "deadline exceeded")):
+        return ProviderTimeoutError(provider, text)
+    if any(marker in lowered for marker in ("connect", "network", "dns", "connection")):
+        return NetworkError(f"{provider}: {text}")
+    if isinstance(error, (KeyError, TypeError, ValueError)):
+        return MalformedResponseError(provider, text)
+    return ProviderError(provider, text)
+
+
+class FileError(ZaiError):
+    pass
+
+
+def show_error(msg: str):
+    console.print(f"[red]Error:[/red] {msg}")
+
+
+def show_warning(msg: str):
+    console.print(f"[yellow]Warning:[/yellow] {msg}")
+
+
+def handle(func):
+    """Decorator — catches ZaiError and unexpected errors for CLI commands."""
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except NoAPIKeyError as e:
+            show_error(str(e))
+        except AllModelsFailedError as e:
+            show_error(str(e))
+        except NetworkError as e:
+            show_error(str(e))
+        except FileError as e:
+            show_error(str(e))
+        except ZaiError as e:
+            show_error(str(e))
+        except KeyboardInterrupt:
+            console.print("\n[dim]Cancelled.[/dim]")
+        except Exception as e:
+            print_exception(console, e)
+    wrapper.__name__ = func.__name__
+    wrapper.__doc__ = func.__doc__
+    return wrapper