zai-cli 0.1.0__py3-none-any.whl

zai/tools/vision.py

@@ -0,0 +1,128 @@
+import base64
+from pathlib import Path
+
+from PIL import Image, UnidentifiedImageError
+
+from ..core.errors import ZaiError, NoAPIKeyError
+from ..config import get_api_key, get_model_config
+
+SUPPORTED = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"}
+MAX_IMAGE_BYTES = 10 * 1024 * 1024
+MAX_IMAGE_PIXELS = 40_000_000
+
+
+class VisionError(ZaiError):
+    pass
+
+
+def _validate_image(path: str) -> tuple[Path, str]:
+    p = Path(path)
+    if not p.exists():
+        raise VisionError(f"Image not found: {path}")
+    if not p.is_file():
+        raise VisionError(f"Not an image file: {path}")
+    suffix = p.suffix.lower()
+    if suffix not in SUPPORTED:
+        allowed = ", ".join(sorted(SUPPORTED))
+        raise VisionError(f"Unsupported format: {suffix}. Use: {allowed}")
+    try:
+        size = p.stat().st_size
+    except OSError as error:
+        raise VisionError(f"Cannot inspect image: {error}") from error
+    if size > MAX_IMAGE_BYTES:
+        raise VisionError("Image exceeds the 10 MB size limit")
+    try:
+        with Image.open(p) as image:
+            width, height = image.size
+            if width <= 0 or height <= 0 or width * height > MAX_IMAGE_PIXELS:
+                raise VisionError("Image exceeds the 40 megapixel limit")
+            image.verify()
+    except VisionError:
+        raise
+    except (UnidentifiedImageError, OSError, ValueError) as error:
+        raise VisionError(f"Invalid or corrupt image: {error}") from error
+    mime = {
+        ".png": "image/png",
+        ".jpg": "image/jpeg",
+        ".jpeg": "image/jpeg",
+        ".gif": "image/gif",
+        ".webp": "image/webp",
+        ".bmp": "image/bmp",
+    }[suffix]
+    return p, mime
+
+
+def _encode_image(path: str) -> tuple[str, str]:
+    p, mime = _validate_image(path)
+    try:
+        data = base64.b64encode(p.read_bytes()).decode()
+    except OSError as error:
+        raise VisionError(f"Cannot read image: {error}") from error
+    return data, mime
+
+
+def analyze_with_gemini(image_path: str, prompt: str = "Describe this image in detail.") -> str:
+    key = get_api_key("gemini")
+    if not key:
+        raise NoAPIKeyError("gemini")
+    from google import genai
+    from google.genai import types
+
+    data, mime = _encode_image(image_path)
+    client = genai.Client(api_key=key)
+    try:
+        result = client.models.generate_content(
+            model=get_model_config("gemini")["model_id"],
+            contents=[
+                prompt,
+                types.Part.from_bytes(
+                    data=base64.b64decode(data),
+                    mime_type=mime,
+                ),
+            ],
+        )
+        return result.text or ""
+    except Exception as e:
+        raise VisionError(f"Gemini vision failed: {e}")
+
+
+def analyze_with_groq(image_path: str, prompt: str = "Describe this image in detail.") -> str:
+    key = get_api_key("groq")
+    if not key:
+        raise NoAPIKeyError("groq")
+    from groq import Groq
+    data, mime = _encode_image(image_path)
+    client = Groq(api_key=key)
+    try:
+        result = client.chat.completions.create(
+            model="llama-4-scout-17b-16e-instruct",
+            messages=[{
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                    {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{data}"}},
+                ],
+            }],
+        )
+        return result.choices[0].message.content or ""
+    except Exception as e:
+        raise VisionError(f"Groq vision failed: {e}")
+
+
+def analyze_image(image_path: str, prompt: str = "Describe this image in detail.") -> tuple[str, str]:
+    """Try Gemini first, then Groq for vision."""
+    _validate_image(image_path)
+    failures = []
+    if get_api_key("gemini"):
+        try:
+            return analyze_with_gemini(image_path, prompt), "gemini"
+        except (VisionError, NoAPIKeyError) as error:
+            failures.append(str(error))
+    if get_api_key("groq"):
+        try:
+            return analyze_with_groq(image_path, prompt), "groq"
+        except (VisionError, NoAPIKeyError) as error:
+            failures.append(str(error))
+    if failures:
+        raise VisionError("Vision analysis failed: " + "; ".join(failures))
+    raise VisionError("No vision-capable model available. Add Gemini or Groq API key: zai setup")

zai/ui/input.py

@@ -0,0 +1,199 @@
+"""Interactive terminal input with history, completion, and multiline editing."""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Callable, Iterable
+
+from rich.prompt import Prompt
+
+from ..config import ZAI_DIR
+
+INPUT_HISTORY_FILE = ZAI_DIR / "input_history"
+PATH_COMMANDS = {
+    "/fix",
+    "/explain",
+    "/review",
+    "/file",
+}
+
+
+def _path_candidates(fragment: str, cwd: str) -> list[tuple[str, str]]:
+    """Return (insert_text, display_text) project-relative path candidates."""
+    normalized = fragment.replace("\\", "/")
+    parent_text, separator, leaf = normalized.rpartition("/")
+    search_root = Path(cwd) / parent_text if separator else Path(cwd)
+    if not search_root.is_dir():
+        return []
+    candidates = []
+    try:
+        children = sorted(
+            search_root.iterdir(),
+            key=lambda path: (not path.is_dir(), path.name.lower()),
+        )
+    except OSError:
+        return []
+    for child in children:
+        if child.name.startswith(".") or not child.name.lower().startswith(leaf.lower()):
+            continue
+        relative = child.relative_to(Path(cwd)).as_posix()
+        insert = relative + ("/" if child.is_dir() else "")
+        display = insert
+        candidates.append((insert, display))
+    return candidates[:100]
+
+
+def completion_candidates(
+    text: str,
+    cwd: str,
+    commands: Iterable[str],
+    models: Iterable[str] = (),
+    sessions: Iterable[str] = (),
+) -> list[tuple[str, str, str]]:
+    """Return (insert_text, display, metadata) completion candidates."""
+    before = text.lstrip()
+    if not before.startswith("/"):
+        return []
+
+    parts = before.split(maxsplit=1)
+    command_token = parts[0]
+    if len(parts) == 1 and not before.endswith(" "):
+        return [
+            (f"/{name}", f"/{name}", "command")
+            for name in sorted(set(commands))
+            if f"/{name}".startswith(command_token.lower())
+        ]
+
+    argument = parts[1] if len(parts) > 1 else ""
+    command = command_token.lower()
+    if command == "/model":
+        options = ["list", *models]
+        return [
+            (option, option, "model")
+            for option in options
+            if option.lower().startswith(argument.lower())
+        ]
+    if command == "/session":
+        subcommands = ["list", "save", "load", "search", "rename", "delete"]
+        if " " not in argument:
+            return [
+                (option, option, "session action")
+                for option in subcommands
+                if option.startswith(argument.lower())
+            ]
+        action, value = argument.split(" ", 1)
+        if action in {"load", "rename", "delete"}:
+            session_value = value.split(" ", 1)[0]
+            return [
+                (name, name, "saved session")
+                for name in sessions
+                if name.lower().startswith(session_value.lower())
+            ]
+    if command == "/resume":
+        return [
+            (name, name, "saved session")
+            for name in sessions
+            if name.lower().startswith(argument.lower())
+        ]
+    if command in PATH_COMMANDS:
+        return [
+            (insert, display, "directory" if insert.endswith("/") else "file")
+            for insert, display in _path_candidates(argument, cwd)
+        ]
+    return []
+
+
+class InteractiveInput:
+    """Reusable input session. Enter submits; Alt+Enter/Ctrl+J inserts newline."""
+
+    def __init__(
+        self,
+        cwd: str,
+        command_provider: Callable[[], Iterable[str]],
+        model_provider: Callable[[], Iterable[str]] = lambda: (),
+        session_provider: Callable[[], Iterable[str]] = lambda: (),
+        history_file: str | Path = INPUT_HISTORY_FILE,
+        input=None,
+        output=None,
+    ):
+        self.cwd = cwd
+        self.command_provider = command_provider
+        self.model_provider = model_provider
+        self.session_provider = session_provider
+        self.history_file = Path(history_file)
+        self.input = input
+        self.output = output
+        self._session = self._build_session()
+
+    def _build_session(self):
+        try:
+            from prompt_toolkit import PromptSession
+            from prompt_toolkit.completion import Completer, Completion
+            from prompt_toolkit.history import FileHistory
+            from prompt_toolkit.key_binding import KeyBindings
+        except ImportError:
+            return None
+
+        owner = self
+
+        class ZaiCompleter(Completer):
+            def get_completions(self, document, complete_event):
+                before = document.text_before_cursor
+                candidates = completion_candidates(
+                    before,
+                    owner.cwd,
+                    owner.command_provider(),
+                    owner.model_provider(),
+                    owner.session_provider(),
+                )
+                current = before.split()[-1] if before.split() else before
+                for insert, display, metadata in candidates:
+                    yield Completion(
+                        insert,
+                        start_position=-len(current),
+                        display=display,
+                        display_meta=metadata,
+                    )
+
+        bindings = KeyBindings()
+
+        @bindings.add("enter")
+        def submit(event):
+            event.current_buffer.validate_and_handle()
+
+        @bindings.add("escape", "enter")
+        @bindings.add("c-j")
+        def newline(event):
+            event.current_buffer.insert_text("\n")
+
+        self.history_file.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            return PromptSession(
+                history=FileHistory(str(self.history_file)),
+                completer=ZaiCompleter(),
+                complete_while_typing=True,
+                multiline=True,
+                key_bindings=bindings,
+                enable_history_search=True,
+                input=self.input,
+                output=self.output,
+            )
+        except Exception:
+            # Redirected stdin/stdout and some Windows hosts have no console
+            # screen buffer. Rich's basic prompt remains usable there.
+            return None
+
+    @property
+    def enhanced(self) -> bool:
+        return self._session is not None
+
+    def prompt(self) -> str:
+        if self._session is None:
+            return Prompt.ask("[cyan]>[/cyan]")
+        return self._session.prompt(
+            "> ",
+            bottom_toolbar=(
+                "Enter submit  |  Alt+Enter/Ctrl+J newline  |  "
+                "Tab complete  |  Ctrl+C cancel input"
+            ),
+        )