krons 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

krons/agent/third_party/gemini_models.py

@@ -0,0 +1,508 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio  # Required for subprocess management (create_subprocess_exec)
+import codecs
+import contextlib
+import inspect
+import json  # Required for JSONDecoder.raw_decode() (streaming JSON parsing)
+import logging
+import shutil
+import warnings
+from collections.abc import AsyncIterator, Callable
+from dataclasses import dataclass
+from dataclasses import field as datafield
+from pathlib import Path
+from textwrap import shorten
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field, model_validator
+
+from krons.utils import json_dump
+
+HAS_GEMINI_CLI = False
+GEMINI_CLI = None
+
+if (g := (shutil.which("gemini") or "gemini")) and shutil.which(g):
+    HAS_GEMINI_CLI = True
+    GEMINI_CLI = g
+
+logging.basicConfig(level=logging.INFO)
+log = logging.getLogger("gemini-cli")
+
+__all__ = (
+    "GeminiChunk",
+    "GeminiCodeRequest",
+    "GeminiSession",
+    "stream_gemini_cli",
+)
+
+
+class GeminiCodeRequest(BaseModel):
+    """Request model for Gemini CLI execution."""
+
+    # -- conversational bits -------------------------------------------------
+    prompt: str = Field(description="The prompt for Gemini CLI")
+    system_prompt: str | None = None
+
+    # -- repo / workspace ----------------------------------------------------
+    repo: Path = Field(default_factory=Path.cwd, exclude=True)
+    ws: str | None = None  # sub-directory under repo
+    include_directories: list[str] = Field(default_factory=list)
+
+    # -- runtime & safety ----------------------------------------------------
+    model: str | None = Field(
+        default="gemini-2.5-pro",
+        description="Gemini model to use (gemini-2.5-pro, gemini-2.5-flash, gemini-3-pro, etc.)",
+    )
+    yolo: bool = Field(
+        default=False,
+        description="Auto-approve all actions without confirmation (--yolo flag)",
+    )
+    approval_mode: Literal["suggest", "auto_edit", "full_auto"] | None = None
+    debug: bool = False
+    sandbox: bool = Field(
+        default=True,
+        description="Run in sandbox mode for safety",
+    )
+
+    # -- MCP integration -----------------------------------------------------
+    mcp_tools: list[str] = Field(default_factory=list)
+
+    # -- internal use --------------------------------------------------------
+    verbose_output: bool = Field(default=False)
+    cli_include_summary: bool = Field(default=False)
+
+    @model_validator(mode="before")
+    @classmethod
+    def _validate_message_prompt(cls, data):
+        """Convert messages format to prompt if needed."""
+        if data.get("prompt"):
+            return data
+
+        if not (msg := data.get("messages")):
+            raise ValueError("messages or prompt required")
+
+        # Extract prompt from messages
+        prompts = []
+        for message in msg:
+            if message["role"] != "system":
+                content = message["content"]
+                if isinstance(content, (dict, list)):
+                    prompts.append(json_dump(content))
+                else:
+                    prompts.append(content)
+            elif message["role"] == "system" and not data.get("system_prompt"):
+                data["system_prompt"] = message["content"]
+
+        data["prompt"] = "\n".join(prompts)
+        return data
+
+    @model_validator(mode="after")
+    def _warn_dangerous_settings(self):
+        """Emit security warnings for dangerous CLI settings."""
+        if self.yolo:
+            warnings.warn(
+                "GeminiCodeRequest: yolo=True enables auto-approval of ALL actions "
+                "without confirmation. This bypasses safety prompts and may allow "
+                "unintended file modifications, command execution, or data access. "
+                "Only use in trusted, isolated environments.",
+                UserWarning,
+                stacklevel=4,
+            )
+
+        if not self.sandbox:
+            warnings.warn(
+                "GeminiCodeRequest: sandbox=False disables sandbox protection. "
+                "The Gemini CLI will have unrestricted access to the file system "
+                "and can execute arbitrary commands. This significantly increases "
+                "security risk. Only disable sandbox in controlled environments.",
+                UserWarning,
+                stacklevel=4,
+            )
+
+        return self
+
+    def cwd(self) -> Path:
+        """Get working directory, validating workspace path."""
+        if not self.ws:
+            return self.repo
+
+        ws_path = Path(self.ws)
+
+        if ws_path.is_absolute():
+            raise ValueError(
+                f"Workspace path must be relative, got absolute: {self.ws}"
+            )
+
+        if ".." in ws_path.parts:
+            raise ValueError(
+                f"Directory traversal detected in workspace path: {self.ws}"
+            )
+
+        repo_resolved = self.repo.resolve()
+        result = (self.repo / ws_path).resolve()
+
+        try:
+            result.relative_to(repo_resolved)
+        except ValueError:
+            raise ValueError(
+                f"Workspace path escapes repository bounds. "
+                f"Repository: {repo_resolved}, Workspace: {result}"
+            )
+
+        return result
+
+    def as_cmd_args(self) -> list[str]:
+        """Build argument list for the Gemini CLI."""
+        args: list[str] = ["-p", self.prompt, "--output-format", "stream-json"]
+
+        if self.model:
+            args += ["-m", self.model]
+
+        if self.yolo:
+            args.append("--yolo")
+
+        if self.approval_mode:
+            args += ["--approval-mode", self.approval_mode]
+
+        if self.debug:
+            args.append("--debug")
+
+        if not self.sandbox:
+            args.append("--no-sandbox")
+
+        for directory in self.include_directories:
+            args += ["--include-directories", directory]
+
+        return args
+
+
+@dataclass
+class GeminiChunk:
+    """Low-level wrapper around every JSON object from the CLI."""
+
+    raw: dict[str, Any]
+    type: str
+    # convenience views
+    text: str | None = None
+    tool_use: dict[str, Any] | None = None
+    tool_result: dict[str, Any] | None = None
+
+
+@dataclass
+class GeminiSession:
+    """Aggregated view of a whole CLI conversation."""
+
+    session_id: str | None = None
+    model: str | None = None
+
+    # chronological log
+    chunks: list[GeminiChunk] = datafield(default_factory=list)
+
+    # materialized views
+    messages: list[dict[str, Any]] = datafield(default_factory=list)
+    tool_uses: list[dict[str, Any]] = datafield(default_factory=list)
+    tool_results: list[dict[str, Any]] = datafield(default_factory=list)
+
+    # final summary
+    result: str = ""
+    usage: dict[str, Any] = datafield(default_factory=dict)
+    total_cost_usd: float | None = None
+    num_turns: int | None = None
+    duration_ms: int | None = None
+    is_error: bool = False
+    summary: dict | None = None
+
+    def populate_summary(self) -> None:
+        self.summary = _extract_summary(self)
+
+
+def _extract_summary(session: GeminiSession) -> dict[str, Any]:
+    """Extract summary from session data."""
+    tool_counts: dict[str, int] = {}
+    tool_details: list[dict[str, Any]] = []
+    file_operations: dict[str, list[str]] = {"reads": [], "writes": [], "edits": []}
+    key_actions = []
+
+    for tool_use in session.tool_uses:
+        tool_name = tool_use.get("name", "unknown")
+        tool_input = tool_use.get("input", {})
+        tool_id = tool_use.get("id", "")
+
+        tool_counts[tool_name] = tool_counts.get(tool_name, 0) + 1
+        tool_details.append({"tool": tool_name, "id": tool_id, "input": tool_input})
+
+        # Categorize by tool type
+        if tool_name in ["read_file", "Read"]:
+            file_path = tool_input.get("path", tool_input.get("file_path", "unknown"))
+            file_operations["reads"].append(file_path)
+            key_actions.append(f"Read {file_path}")
+
+        elif tool_name in ["write_file", "Write"]:
+            file_path = tool_input.get("path", tool_input.get("file_path", "unknown"))
+            file_operations["writes"].append(file_path)
+            key_actions.append(f"Wrote {file_path}")
+
+        elif tool_name in ["edit_file", "Edit"]:
+            file_path = tool_input.get("path", tool_input.get("file_path", "unknown"))
+            file_operations["edits"].append(file_path)
+            key_actions.append(f"Edited {file_path}")
+
+        elif tool_name in ["run_shell_command", "shell", "Bash"]:
+            command = tool_input.get("command", "")
+            command_summary = command[:50] + "..." if len(command) > 50 else command
+            key_actions.append(f"Ran: {command_summary}")
+
+        elif tool_name.startswith("mcp_"):
+            operation = tool_name.replace("mcp_", "")
+            key_actions.append(f"MCP {operation}")
+
+        else:
+            key_actions.append(f"Used {tool_name}")
+
+    key_actions = (
+        list(dict.fromkeys(key_actions)) if key_actions else ["No specific actions"]
+    )
+
+    for op_type in file_operations:
+        file_operations[op_type] = list(dict.fromkeys(file_operations[op_type]))
+
+    result_summary = (
+        (session.result[:200] + "...") if len(session.result) > 200 else session.result
+    )
+
+    return {
+        "tool_counts": tool_counts,
+        "tool_details": tool_details,
+        "file_operations": file_operations,
+        "key_actions": key_actions,
+        "total_tool_calls": sum(tool_counts.values()),
+        "result_summary": result_summary,
+        "usage_stats": {
+            "total_cost_usd": session.total_cost_usd,
+            "num_turns": session.num_turns,
+            "duration_ms": session.duration_ms,
+            **session.usage,
+        },
+    }
+
+
+async def _ndjson_from_cli(request: GeminiCodeRequest):
+    """
+    Yields each JSON object emitted by the Gemini CLI.
+
+    Robust against UTF-8 splits and uses json.JSONDecoder.raw_decode.
+    """
+    if GEMINI_CLI is None:
+        raise RuntimeError("Gemini CLI not found. Please install the gemini CLI tool.")
+
+    workspace = request.cwd()
+    workspace.mkdir(parents=True, exist_ok=True)
+
+    proc = await asyncio.create_subprocess_exec(
+        GEMINI_CLI,
+        *request.as_cmd_args(),
+        cwd=str(workspace),
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE,
+    )
+
+    decoder = codecs.getincrementaldecoder("utf-8")()
+    json_decoder = json.JSONDecoder()
+    buffer: str = ""
+
+    if proc.stdout is None:
+        raise RuntimeError("Failed to capture stdout from Gemini CLI")
+
+    try:
+        while True:
+            chunk = await proc.stdout.read(4096)
+            if not chunk:
+                break
+
+            buffer += decoder.decode(chunk)
+
+            while buffer:
+                buffer = buffer.lstrip()
+                if not buffer:
+                    break
+                try:
+                    obj, idx = json_decoder.raw_decode(buffer)
+                    yield obj
+                    buffer = buffer[idx:]
+                except json.JSONDecodeError:
+                    break
+
+        buffer += decoder.decode(b"", final=True)
+        buffer = buffer.strip()
+        if buffer:
+            try:
+                obj, idx = json_decoder.raw_decode(buffer)
+                yield obj
+            except json.JSONDecodeError:
+                log.error("Skipped unrecoverable JSON tail: %.120s...", buffer)
+
+        if await proc.wait() != 0:
+            err = ""
+            if proc.stderr is not None:
+                err = (await proc.stderr.read()).decode().strip()
+            raise RuntimeError(err or "Gemini CLI exited non-zero")
+
+    finally:
+        with contextlib.suppress(ProcessLookupError):
+            proc.terminate()
+        await proc.wait()
+
+
+async def stream_gemini_cli_events(request: GeminiCodeRequest):
+    """Stream events from Gemini CLI."""
+    if not GEMINI_CLI:
+        raise RuntimeError("Gemini CLI not found (npm i -g @google/gemini-cli)")
+    async for obj in _ndjson_from_cli(request):
+        yield obj
+    yield {"type": "done"}
+
+
+async def _maybe_await(func, *args, **kw):
+    """Call func which may be sync or async."""
+    res = func(*args, **kw) if func else None
+    if inspect.iscoroutine(res):
+        await res
+
+
+def _pp_text(text: str) -> None:
+    print(f"\n> Gemini:\n{text}\n")
+
+
+def _pp_tool_use(tu: dict[str, Any]) -> None:
+    preview = shorten(str(tu.get("input", {})).replace("\n", " "), 130)
+    print(f"- Tool Use - {tu.get('name', 'unknown')}: {preview}")
+
+
+def _pp_tool_result(tr: dict[str, Any]) -> None:
+    body_preview = shorten(str(tr.get("content", "")).replace("\n", " "), 130)
+    status = "ERR" if tr.get("is_error") else "OK"
+    print(f"- Tool Result - {status}: {body_preview}")
+
+
+def _pp_final(sess: GeminiSession) -> None:
+    usage = sess.usage or {}
+    print(
+        f"\n### Session complete\n"
+        f"**Result:** {sess.result or ''}\n"
+        f"- turns: {sess.num_turns}\n"
+        f"- duration: {sess.duration_ms} ms\n"
+        f"- tokens: {usage.get('input_tokens', 0)}/{usage.get('output_tokens', 0)}"
+    )
+
+
+async def stream_gemini_cli(
+    request: GeminiCodeRequest,
+    session: GeminiSession | None = None,
+    *,
+    on_text: Callable[[str], None] | None = None,
+    on_tool_use: Callable[[dict[str, Any]], None] | None = None,
+    on_tool_result: Callable[[dict[str, Any]], None] | None = None,
+    on_final: Callable[[GeminiSession], None] | None = None,
+) -> AsyncIterator[GeminiChunk | dict | GeminiSession]:
+    """
+    Consume the ND-JSON stream from Gemini CLI and return a populated GeminiSession.
+    """
+    if session is None:
+        session = GeminiSession()
+
+    stream = stream_gemini_cli_events(request)
+
+    async for obj in stream:
+        typ = obj.get("type", "unknown")
+        chunk = GeminiChunk(raw=obj, type=typ)
+        session.chunks.append(chunk)
+
+        # Handle different event types based on Gemini CLI output format
+        if typ in ("system", "init"):
+            session.session_id = obj.get("session_id", obj.get("id"))
+            session.model = obj.get("model")
+            yield obj
+
+        elif typ in ("message", "assistant"):
+            msg = obj.get("message", obj)
+            session.messages.append(msg)
+
+            content = msg.get("content", "")
+            if isinstance(content, str):
+                chunk.text = content
+                await _maybe_await(on_text, content)
+                if request.verbose_output:
+                    _pp_text(content)
+            elif isinstance(content, list):
+                for blk in content:
+                    if isinstance(blk, dict):
+                        btype = blk.get("type")
+                        if btype == "text":
+                            text = blk.get("text", "")
+                            chunk.text = text
+                            await _maybe_await(on_text, text)
+                            if request.verbose_output:
+                                _pp_text(text)
+                        elif btype == "tool_use":
+                            tu = {
+                                "id": blk.get("id", ""),
+                                "name": blk.get("name", ""),
+                                "input": blk.get("input", {}),
+                            }
+                            chunk.tool_use = tu
+                            session.tool_uses.append(tu)
+                            await _maybe_await(on_tool_use, tu)
+                            if request.verbose_output:
+                                _pp_tool_use(tu)
+            yield chunk
+
+        elif typ in ("tool_call", "tool_use"):
+            tu = {
+                "id": obj.get("id", obj.get("tool_use_id", "")),
+                "name": obj.get("name", obj.get("tool_name", "")),
+                "input": obj.get("input", obj.get("args", {})),
+            }
+            chunk.tool_use = tu
+            session.tool_uses.append(tu)
+            await _maybe_await(on_tool_use, tu)
+            if request.verbose_output:
+                _pp_tool_use(tu)
+            yield chunk
+
+        elif typ == "tool_result":
+            tr = {
+                "tool_use_id": obj.get("tool_use_id", obj.get("id", "")),
+                "content": obj.get("content", obj.get("result", "")),
+                "is_error": obj.get("is_error", False),
+            }
+            chunk.tool_result = tr
+            session.tool_results.append(tr)
+            await _maybe_await(on_tool_result, tr)
+            if request.verbose_output:
+                _pp_tool_result(tr)
+            yield chunk
+
+        elif typ in ("result", "response"):
+            session.result = obj.get("result", obj.get("response", "")).strip()
+            session.usage = obj.get("usage", obj.get("stats", {}))
+            session.total_cost_usd = obj.get("total_cost_usd", obj.get("cost"))
+            session.num_turns = obj.get("num_turns", obj.get("turns"))
+            session.duration_ms = obj.get("duration_ms", obj.get("duration"))
+            session.is_error = obj.get("is_error", obj.get("error") is not None)
+
+        elif typ == "error":
+            session.is_error = True
+            session.result = obj.get("message", obj.get("error", "Unknown error"))
+
+        elif typ == "done":
+            break
+
+    await _maybe_await(on_final, session)
+    if request.verbose_output:
+        _pp_final(session)
+
+    yield session