zai-cli 0.1.0__py3-none-any.whl

zai/cli/workflows.py

@@ -0,0 +1,137 @@
+from typing import Optional
+
+import typer
+from rich import print as rprint
+from rich.console import Console
+from rich.panel import Panel
+from rich.prompt import Confirm
+
+from ..core.fallback import chat_with_fallback
+from ..plugins import loader as plugin_loader
+from ..skills.registry import SKILLS, get_skill_prompt
+from ..tools.files import read_file
+from ..tools.git import (
+    commit,
+    get_diff,
+    get_log,
+    get_staged_diff,
+    get_status,
+)
+
+console = Console()
+
+
+def register_workflow_commands(app: typer.Typer, context, system_prompt: str) -> None:
+    """Register AI-assisted skill and Git workflows."""
+
+    @app.command()
+    def skill(
+        name: str = typer.Argument(
+            ...,
+            help="Skill: review, test, commit, docs, fix, explain, refactor, summarize",
+        ),
+        path: Optional[str] = typer.Argument(
+            None,
+            help="File to apply skill to",
+        ),
+        model: Optional[str] = typer.Option(None, "--model", "-m"),
+    ):
+        """Run a built-in skill: /review, /test, /commit, /docs, /fix, /explain..."""
+        plugin_loader.load_all()
+        plugin_skills = plugin_loader.get_all_skills()
+        all_skills = {**SKILLS, **plugin_skills}
+        if name not in all_skills:
+            console.print(f"[red]Unknown skill: {name}[/red]")
+            console.print(f"Available: {', '.join(all_skills.keys())}")
+            return
+
+        if name == "commit":
+            diff = get_staged_diff() or get_diff()
+            prompt = get_skill_prompt("commit", text=diff)
+        elif name in plugin_skills:
+            content = read_file(path) if path else ""
+            prompt = plugin_skills[name]["prompt"](content)
+        elif path:
+            prompt = get_skill_prompt(name, file_path=path)
+        else:
+            console.print(f"[red]Provide a file: zai skill {name} <file>[/red]")
+            return
+
+        context.add("user", prompt)
+        with console.status(f"[cyan]Running /{name}...[/cyan]"):
+            response, _ = chat_with_fallback(
+                context.get_messages(),
+                system=system_prompt,
+                preferred=model,
+            )
+        context.add("assistant", response.content)
+        rprint(Panel(
+            response.content,
+            title=f"[cyan]/{name}[/cyan]",
+            border_style="cyan",
+        ))
+
+    @app.command()
+    def git(
+        action: str = typer.Argument(
+            "status",
+            help="status, log, diff, commit, review",
+        ),
+        model: Optional[str] = typer.Option(None, "--model", "-m"),
+    ):
+        """Git operations with AI assistance."""
+        if action == "status":
+            console.print(Panel(
+                get_status() or "Clean.",
+                title="[cyan]Git Status[/cyan]",
+            ))
+        elif action == "log":
+            console.print(Panel(get_log(), title="[cyan]Git Log[/cyan]"))
+        elif action == "diff":
+            console.print(Panel(
+                get_diff() or "No changes.",
+                title="[cyan]Git Diff[/cyan]",
+            ))
+        elif action == "commit":
+            diff = get_staged_diff() or get_diff()
+            if not diff:
+                console.print("[yellow]No changes to commit.[/yellow]")
+                return
+            prompt = get_skill_prompt("commit", text=diff)
+            context.add("user", prompt)
+            with console.status("[cyan]Generating commit message...[/cyan]"):
+                response, _ = chat_with_fallback(
+                    context.get_messages(),
+                    system=system_prompt,
+                    preferred=model,
+                )
+            message = response.content.strip().strip("`").strip()
+            console.print(Panel(
+                message,
+                title="[cyan]Suggested Commit Message[/cyan]",
+            ))
+            if Confirm.ask("Use this message?"):
+                console.print(f"[green]{commit(message)}[/green]")
+        elif action == "review":
+            diff = get_diff()
+            if not diff:
+                console.print("[yellow]No changes to review.[/yellow]")
+                return
+            prompt = get_skill_prompt("review", text=diff)
+            context.add("user", prompt)
+            with console.status("[cyan]Reviewing changes...[/cyan]"):
+                response, _ = chat_with_fallback(
+                    context.get_messages(),
+                    system=system_prompt,
+                    preferred=model,
+                )
+            context.add("assistant", response.content)
+            rprint(Panel(
+                response.content,
+                title="[cyan]Code Review[/cyan]",
+                border_style="cyan",
+            ))
+        else:
+            console.print(
+                "[red]Usage: zai git status | log | diff | commit | review[/red]"
+            )

zai/commands/commit.md

@@ -0,0 +1,24 @@
+---
+description: Create a git commit with auto-generated message
+allowed-tools: run_command
+---
+
+## Context
+
+- Git status: !`git status`
+- Staged changes: !`git diff --cached`
+- All changes: !`git diff HEAD`
+- Recent commits: !`git log --oneline -5`
+- Current branch: !`git branch --show-current`
+
+## Task
+
+Based on the changes above, create a descriptive git commit message and commit the changes.
+
+1. Stage all relevant changes with:
+   <tool_call>{"name":"run_command","arguments":{"command":"git add -A"}}</tool_call>
+2. Create a commit with a clear, concise message describing WHAT changed and WHY
+3. Use conventional commit format: `type(scope): description` (e.g. feat, fix, refactor, docs)
+4. Confirm what was committed
+
+$ARGUMENTS

zai/commands/explain.md

@@ -0,0 +1,17 @@
+---
+description: Explain what a file or function does
+allowed-tools: read_file
+---
+
+## Task
+
+$ARGUMENTS
+
+Read the file and explain clearly:
+1. What this code does (plain English)
+2. How it works step by step
+3. Key functions/classes and their purpose
+4. Any important patterns or design decisions
+
+Use this tool call to read it first:
+<tool_call>{"name":"read_file","arguments":{"path":"$ARGUMENTS"}}</tool_call>

zai/commands/feature.md

@@ -0,0 +1,34 @@
+---
+description: Build a new feature with guided development workflow
+argument-hint: describe the feature you want to build
+---
+
+# Feature Development
+
+## Phase 1: Understand
+
+Feature request: $ARGUMENTS
+
+First, list files in current folder:
+<tool_call>{"name":"list_files","arguments":{"path":"."}}</tool_call>
+
+## Phase 2: Plan
+
+Based on existing code, create a plan:
+1. What files need to be created/modified?
+2. What is the simplest implementation?
+3. Any dependencies needed?
+
+## Phase 3: Implement
+
+Build the feature step by step:
+- Create new files with structured `write_file` tool calls
+- Read existing files with `read_file` before modifying
+- Run commands with `run_command` to test
+
+## Phase 4: Confirm
+
+After implementation, confirm:
+- What was created/changed
+- How to test it
+- Next steps

zai/commands/fix.md

@@ -0,0 +1,14 @@
+---
+description: Find and fix errors or bugs in a file
+allowed-tools: read_file, write_file
+---
+
+## Task
+
+$ARGUMENTS
+
+1. Use this tool call to read the file:
+   <tool_call>{"name":"read_file","arguments":{"path":"$ARGUMENTS"}}</tool_call>
+2. Identify all bugs, errors, and issues
+3. Fix them using `edit_file`, or `write_file` when replacing the complete file
+4. Explain what was fixed and why

zai/commands/review.md

@@ -0,0 +1,22 @@
+---
+description: Review code for bugs, security issues, and improvements
+---
+
+## Context
+
+- Current changes: !`git diff HEAD`
+- Files changed: !`git diff --name-only HEAD`
+
+## Task
+
+Review the code changes above and provide:
+
+1. **Bugs** — any logic errors, null pointer issues, off-by-one errors
+2. **Security** — SQL injection, XSS, insecure operations, hardcoded secrets
+3. **Performance** — slow loops, unnecessary DB calls, memory leaks
+4. **Code quality** — duplication, naming, complexity, missing error handling
+5. **Suggestions** — concrete improvements with example code
+
+Be specific — mention exact line numbers or function names where possible.
+
+$ARGUMENTS

zai/config.py

@@ -0,0 +1,307 @@
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+from .core.storage import atomic_write_json, atomic_write_text, read_json, file_lock
+
+ZAI_DIR = Path.home() / ".zai"
+ENV_FILE = ZAI_DIR / ".env"
+MEMORY_FILE = ZAI_DIR / "memory.json"
+CONFIG_FILE = ZAI_DIR / "config.json"
+
+ZAI_DIR.mkdir(exist_ok=True)
+
+load_dotenv(ENV_FILE)
+
+DEFAULT_CONFIG = {
+    "config_version": 3,
+    "default_model": "gemini",
+    "fallback_order": ["gemini", "groq", "cerebras", "openrouter", "qwen", "claude", "gpt4o", "ollama"],
+    "model_overrides": {},
+    "models": {},
+    "model_settings": {},
+    "show_token_count": True,
+    "auto_fallback": True,
+}
+
+MODELS = {
+    "gemini": {
+        "name": "Gemini 3.5 Flash",
+        "provider": "gemini",
+        "model_id": "gemini-3.5-flash",
+        "context_window": 1000000,
+        "free": True,
+    },
+    "groq": {
+        "name": "GPT-OSS 120B (Groq)",
+        "provider": "groq",
+        "model_id": "openai/gpt-oss-120b",
+        "context_window": 128000,
+        "free": True,
+    },
+    "cerebras": {
+        "name": "Llama 3.3 70B (Cerebras)",
+        "provider": "cerebras",
+        "model_id": "llama-3.3-70b",
+        "context_window": 128000,
+        "free": True,
+    },
+    "openrouter": {
+        "name": "Llama 3.3 70B (OpenRouter)",
+        "provider": "openrouter",
+        "model_id": "meta-llama/llama-3.3-70b-instruct:free",
+        "context_window": 128000,
+        "free": True,
+    },
+    "qwen": {
+        "name": "Qwen Turbo",
+        "provider": "qwen",
+        "model_id": "qwen-turbo",
+        "context_window": 1000000,
+        "free": False,
+    },
+    "claude": {
+        "name": "Claude Sonnet",
+        "provider": "anthropic",
+        "model_id": "claude-sonnet-4-6",
+        "context_window": 200000,
+        "free": False,
+    },
+    "gpt4o": {
+        "name": "GPT-4o Mini",
+        "provider": "openai",
+        "model_id": "gpt-4o-mini",
+        "context_window": 128000,
+        "free": False,
+    },
+    "ollama": {
+        "name": "Ollama (local)",
+        "provider": "ollama",
+        "model_id": "llama3.2",
+        "context_window": 128000,
+        "free": True,
+    },
+}
+
+MODEL_PROVIDERS = {
+    data["provider"] for data in MODELS.values()
+}
+
+
+def _valid_custom_model(name: str, data: object) -> bool:
+    return (
+        isinstance(name, str)
+        and bool(name.strip())
+        and isinstance(data, dict)
+        and data.get("provider") in MODEL_PROVIDERS
+        and isinstance(data.get("model_id"), str)
+        and bool(data["model_id"].strip())
+        and isinstance(data.get("context_window"), int)
+        and data["context_window"] >= 1024
+        and isinstance(data.get("timeout", 60), (int, float))
+        and 1 <= data.get("timeout", 60) <= 600
+        and isinstance(data.get("retries", 2), int)
+        and 0 <= data.get("retries", 2) <= 10
+    )
+
+
+def get_models(config: dict | None = None) -> dict[str, dict]:
+    """Return built-in and user-defined model aliases."""
+    if config is None:
+        config = load_config()
+    models = {name: dict(data) for name, data in MODELS.items()}
+    for name, data in config.get("models", {}).items():
+        if _valid_custom_model(name, data):
+            models[name] = {
+                "name": data.get("name") or name,
+                "provider": data["provider"],
+                "model_id": data["model_id"].strip(),
+                "context_window": data["context_window"],
+                "free": bool(data.get("free", False)),
+                "timeout": data.get("timeout", 60),
+                "retries": data.get("retries", 2),
+            }
+    return models
+
+
+def validate_config(config: dict) -> list[str]:
+    """Return human-readable configuration problems without changing the input."""
+    issues = []
+    available_models = get_models(config)
+    default_model = config.get("default_model")
+    if default_model not in available_models:
+        issues.append(f"Unknown default model: {default_model!r}")
+
+    fallback_order = config.get("fallback_order")
+    if not isinstance(fallback_order, list):
+        issues.append("fallback_order must be a list")
+    else:
+        unknown = [name for name in fallback_order if name not in available_models]
+        if unknown:
+            issues.append(f"Unknown fallback models: {', '.join(map(str, unknown))}")
+
+    overrides = config.get("model_overrides", {})
+    if not isinstance(overrides, dict):
+        issues.append("model_overrides must be an object")
+    else:
+        for name, model_id in overrides.items():
+            if name not in available_models:
+                issues.append(f"Unknown model override: {name}")
+            elif not isinstance(model_id, str) or not model_id.strip():
+                issues.append(f"Invalid model ID override for {name}")
+    custom_models = config.get("models", {})
+    if not isinstance(custom_models, dict):
+        issues.append("models must be an object")
+    else:
+        for name, data in custom_models.items():
+            if name in MODELS:
+                issues.append(f"Custom model cannot replace built-in model: {name}")
+            elif not _valid_custom_model(name, data):
+                issues.append(f"Invalid custom model: {name}")
+    model_settings = config.get("model_settings", {})
+    if not isinstance(model_settings, dict):
+        issues.append("model_settings must be an object")
+    else:
+        for name, settings in model_settings.items():
+            if name not in available_models or not isinstance(settings, dict):
+                issues.append(f"Invalid model settings: {name}")
+                continue
+            timeout = settings.get("timeout", 60)
+            retries = settings.get("retries", 2)
+            if (
+                not isinstance(timeout, (int, float))
+                or not 1 <= timeout <= 600
+                or not isinstance(retries, int)
+                or not 0 <= retries <= 10
+            ):
+                issues.append(f"Invalid model settings: {name}")
+    return issues
+
+
+def _normalize_config(config: dict) -> dict:
+    normalized = {**DEFAULT_CONFIG, **config}
+    custom_models = normalized.get("models")
+    if not isinstance(custom_models, dict):
+        custom_models = {}
+    normalized["models"] = {
+        name.strip(): {
+            "name": data.get("name") or name.strip(),
+            "provider": data["provider"],
+            "model_id": data["model_id"].strip(),
+            "context_window": data["context_window"],
+            "free": bool(data.get("free", False)),
+            "timeout": data.get("timeout", 60),
+            "retries": data.get("retries", 2),
+        }
+        for name, data in custom_models.items()
+        if name not in MODELS and _valid_custom_model(name, data)
+    }
+    available_models = get_models(normalized)
+    model_settings = normalized.get("model_settings")
+    if not isinstance(model_settings, dict):
+        model_settings = {}
+    normalized["model_settings"] = {
+        name: {
+            "timeout": settings.get("timeout", 60),
+            "retries": settings.get("retries", 2),
+        }
+        for name, settings in model_settings.items()
+        if (
+            name in available_models
+            and isinstance(settings, dict)
+            and isinstance(settings.get("timeout", 60), (int, float))
+            and 1 <= settings.get("timeout", 60) <= 600
+            and isinstance(settings.get("retries", 2), int)
+            and 0 <= settings.get("retries", 2) <= 10
+        )
+    }
+    if normalized["default_model"] not in available_models:
+        normalized["default_model"] = DEFAULT_CONFIG["default_model"]
+
+    fallback = normalized.get("fallback_order")
+    if not isinstance(fallback, list):
+        fallback = DEFAULT_CONFIG["fallback_order"]
+    normalized["fallback_order"] = [
+        name for name in fallback if name in available_models
+    ]
+
+    overrides = normalized.get("model_overrides")
+    if not isinstance(overrides, dict):
+        overrides = {}
+    normalized["model_overrides"] = {
+        name: model_id.strip()
+        for name, model_id in overrides.items()
+        if name in available_models and isinstance(model_id, str) and model_id.strip()
+    }
+    normalized["config_version"] = DEFAULT_CONFIG["config_version"]
+    return normalized
+
+
+def load_config() -> dict:
+    saved = read_json(CONFIG_FILE, {}, expected_type=dict)
+    return _normalize_config(saved)
+
+
+def save_config(config: dict):
+    atomic_write_json(CONFIG_FILE, _normalize_config(config))
+
+
+def get_model_config(name: str, config: dict | None = None) -> dict:
+    """Return effective model metadata, including user/environment overrides."""
+    config = config or load_config()
+    models = get_models(config)
+    if name not in models:
+        raise KeyError(f"Unknown model: {name}")
+    effective = dict(models[name])
+    effective.setdefault("timeout", 60)
+    effective.setdefault("retries", 2)
+    effective.update(config.get("model_settings", {}).get(name, {}))
+    configured = config.get("model_overrides", {}).get(name)
+    environment = os.getenv(f"ZAI_{name.upper()}_MODEL")
+    effective["model_id"] = environment or configured or effective["model_id"]
+    return effective
+
+
+def get_api_key(provider: str) -> str | None:
+    key_map = {
+        "gemini": "GEMINI_API_KEY",
+        "groq": "GROQ_API_KEY",
+        "cerebras": "CEREBRAS_API_KEY",
+        "openrouter": "OPENROUTER_API_KEY",
+        "qwen": "QWEN_API_KEY",
+        "anthropic": "ANTHROPIC_API_KEY",
+        "openai": "OPENAI_API_KEY",
+    }
+    return os.getenv(key_map.get(provider, ""))
+
+
+def save_api_key(provider: str, key: str):
+    key_map = {
+        "gemini": "GEMINI_API_KEY",
+        "groq": "GROQ_API_KEY",
+        "cerebras": "CEREBRAS_API_KEY",
+        "openrouter": "OPENROUTER_API_KEY",
+        "qwen": "QWEN_API_KEY",
+        "anthropic": "ANTHROPIC_API_KEY",
+        "openai": "OPENAI_API_KEY",
+    }
+    env_key = key_map.get(provider)
+    if not env_key:
+        return
+    with file_lock(ENV_FILE):
+        lines = []
+        if ENV_FILE.exists():
+            lines = ENV_FILE.read_text(encoding="utf-8").splitlines()
+        updated = False
+        for i, line in enumerate(lines):
+            if line.startswith(f"{env_key}="):
+                lines[i] = f"{env_key}={key}"
+                updated = True
+                break
+        if not updated:
+            lines.append(f"{env_key}={key}")
+        atomic_write_text(
+            ENV_FILE,
+            "\n".join(lines) + "\n",
+            lock=False,
+        )
+    os.environ[env_key] = key