zai-cli 0.1.0__py3-none-any.whl

zai/core/streaming.py

@@ -0,0 +1,157 @@
+from rich.console import Console
+from rich.live import Live
+from rich.text import Text
+
+console = Console()
+
+
+def stream_groq(
+    messages: list, system: str, model_id: str, api_key: str,
+    timeout: float = 60, retries: int = 2,
+):
+    from groq import Groq
+    client = Groq(api_key=api_key, timeout=timeout, max_retries=retries)
+    formatted = [{"role": "system", "content": system}]
+    formatted += [{"role": m.role, "content": m.content} for m in messages]
+    stream = client.chat.completions.create(
+        model=model_id,
+        messages=formatted,
+        stream=True,
+    )
+    full = ""
+    with Live(Text(""), refresh_per_second=15, console=console) as live:
+        for chunk in stream:
+            delta = chunk.choices[0].delta.content or ""
+            full += delta
+            live.update(Text(full))
+    return full
+
+
+def stream_gemini(
+    messages: list, system: str, model_id: str, api_key: str,
+    timeout: float = 60,
+):
+    from google import genai
+    from google.genai import types
+
+    client = genai.Client(
+        api_key=api_key,
+        http_options=types.HttpOptions(timeout=int(timeout * 1000)),
+    )
+    contents = [
+        types.Content(
+            role=m.role if m.role == "user" else "model",
+            parts=[types.Part.from_text(text=m.content)],
+        )
+        for m in messages
+    ]
+    full = ""
+    with Live(Text(""), refresh_per_second=15, console=console) as live:
+        for chunk in client.models.generate_content_stream(
+            model=model_id,
+            contents=contents,
+            config=types.GenerateContentConfig(system_instruction=system),
+        ):
+            full += chunk.text or ""
+            live.update(Text(full))
+    return full
+
+
+def stream_anthropic(
+    messages: list, system: str, model_id: str, api_key: str,
+    timeout: float = 60, retries: int = 2,
+):
+    import anthropic
+    client = anthropic.Anthropic(
+        api_key=api_key, timeout=timeout, max_retries=retries,
+    )
+    formatted = [{"role": m.role, "content": m.content} for m in messages]
+    full = ""
+    with client.messages.stream(
+        model=model_id,
+        max_tokens=8096,
+        system=system,
+        messages=formatted,
+    ) as stream:
+        with Live(Text(""), refresh_per_second=15, console=console) as live:
+            for delta in stream.text_stream:
+                full += delta
+                live.update(Text(full))
+    return full
+
+
+def stream_openai(
+    messages: list, system: str, model_id: str, api_key: str,
+    timeout: float = 60, retries: int = 2,
+):
+    from openai import OpenAI
+    client = OpenAI(api_key=api_key, timeout=timeout, max_retries=retries)
+    formatted = [{"role": "system", "content": system}]
+    formatted += [{"role": m.role, "content": m.content} for m in messages]
+    full = ""
+    stream = client.chat.completions.create(
+        model=model_id,
+        messages=formatted,
+        stream=True,
+    )
+    with Live(Text(""), refresh_per_second=15, console=console) as live:
+        for chunk in stream:
+            delta = chunk.choices[0].delta.content or ""
+            full += delta
+            live.update(Text(full))
+    return full
+
+
+def stream_ollama(
+    messages: list, system: str, model_id: str, timeout: float = 120,
+):
+    import httpx
+    import json
+    formatted = [{"role": "system", "content": system}]
+    formatted += [{"role": m.role, "content": m.content} for m in messages]
+    full = ""
+    with httpx.stream(
+        "POST",
+        "http://localhost:11434/api/chat",
+        json={"model": model_id, "messages": formatted, "stream": True},
+        timeout=timeout,
+    ) as response:
+        response.raise_for_status()
+        with Live(Text(""), refresh_per_second=15, console=console) as live:
+            for line in response.iter_lines():
+                if line:
+                    try:
+                        data = json.loads(line)
+                        delta = data.get("message", {}).get("content", "")
+                        full += delta
+                        live.update(Text(full))
+                    except Exception:
+                        pass
+    return full
+
+
+def stream_openrouter(messages: list, system: str, model_id: str, api_key: str):
+    import httpx
+    formatted = [{"role": "system", "content": system}]
+    formatted += [{"role": m.role, "content": m.content} for m in messages]
+    full = ""
+    with httpx.stream(
+        "POST",
+        "https://openrouter.ai/api/v1/chat/completions",
+        headers={"Authorization": f"Bearer {api_key}"},
+        json={"model": model_id, "messages": formatted, "stream": True},
+        timeout=60,
+    ) as response:
+        response.raise_for_status()
+        with Live(Text(""), refresh_per_second=15, console=console) as live:
+            for line in response.iter_lines():
+                if line.startswith("data: ") and line != "data: [DONE]":
+                    import json
+                    try:
+                        data = json.loads(line[6:])
+                        delta = data["choices"][0]["delta"].get("content", "")
+                        full += delta
+                        live.update(Text(full))
+                    except Exception:
+                        pass
+    return full

zai/core/tool_schema.py

@@ -0,0 +1,133 @@
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+
+
+class StrictArguments(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class PathArguments(StrictArguments):
+    path: str = Field(min_length=1)
+
+
+class WriteFileArguments(PathArguments):
+    content: str
+
+
+class EditFileArguments(PathArguments):
+    search: str = Field(min_length=1)
+    replace: str
+
+
+class ListFilesArguments(StrictArguments):
+    path: str = "."
+
+
+class RenamePathArguments(StrictArguments):
+    source: str = Field(min_length=1)
+    destination: str = Field(min_length=1)
+
+
+class RunCommandArguments(StrictArguments):
+    command: str = Field(min_length=1)
+
+
+class MCPCallArguments(StrictArguments):
+    server: str = Field(min_length=1)
+    tool: str = Field(min_length=1)
+    arguments: dict[str, Any] = Field(default_factory=dict)
+
+
+class PluginCallArguments(StrictArguments):
+    plugin: str = Field(min_length=1)
+    tool: str = Field(min_length=1)
+    arguments: dict[str, Any] = Field(default_factory=dict)
+
+
+class ToolCallEnvelope(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    name: Literal[
+        "write_file",
+        "read_file",
+        "edit_file",
+        "create_folder",
+        "rename_path",
+        "list_files",
+        "run_command",
+        "mcp_call",
+        "plugin_call",
+    ]
+    arguments: dict[str, Any] = Field(default_factory=dict)
+
+
+ARGUMENT_MODELS = {
+    "write_file": WriteFileArguments,
+    "read_file": PathArguments,
+    "edit_file": EditFileArguments,
+    "create_folder": PathArguments,
+    "rename_path": RenamePathArguments,
+    "list_files": ListFilesArguments,
+    "run_command": RunCommandArguments,
+    "mcp_call": MCPCallArguments,
+    "plugin_call": PluginCallArguments,
+}
+
+TOOL_DESCRIPTIONS = {
+    "write_file": "Create a new text file. Use edit_file for an existing file.",
+    "read_file": "Read a text file inside the current project.",
+    "edit_file": "Replace one exact text occurrence in an existing file.",
+    "create_folder": "Create a directory inside the current project.",
+    "rename_path": "Rename or move a file or directory inside the current project.",
+    "list_files": "List files in a project directory.",
+    "run_command": "Run an approved direct command in the current project.",
+    "mcp_call": "Call a tool exposed by a connected MCP server.",
+    "plugin_call": "Call a tool exposed by an enabled zai plugin.",
+}
+
+
+def get_tool_definitions() -> list[dict[str, Any]]:
+    """Return provider-neutral JSON Schema function definitions."""
+    definitions = []
+    for name, model in ARGUMENT_MODELS.items():
+        schema = model.model_json_schema()
+        schema.pop("title", None)
+        schema["additionalProperties"] = False
+        # Native providers are most reliable when every property is explicit.
+        schema["required"] = list(schema.get("properties", {}).keys())
+        definitions.append({
+            "name": name,
+            "description": TOOL_DESCRIPTIONS[name],
+            "parameters": schema,
+        })
+    return definitions
+
+
+def legacy_tool_instructions() -> str:
+    """Instructions for providers that do not expose native function calling."""
+    lines = [
+        "Your API does not support native tools here. Emit each call as JSON:",
+        '<tool_call>{"name":"tool_name","arguments":{...}}</tool_call>',
+        "Valid tools:",
+    ]
+    for tool in get_tool_definitions():
+        properties = ", ".join(tool["parameters"].get("properties", {}))
+        lines.append(f"- {tool['name']}: {properties}")
+    lines.append("JSON must use double quotes and must not contain trailing commas.")
+    return "\n".join(lines)
+
+
+def validate_tool_call(data: Any) -> tuple[str, dict[str, Any]]:
+    """Validate a structured tool call and return normalized arguments."""
+    envelope = ToolCallEnvelope.model_validate(data)
+    arguments = ARGUMENT_MODELS[envelope.name].model_validate(envelope.arguments)
+    return envelope.name, arguments.model_dump()
+
+
+def format_validation_error(error: ValidationError) -> str:
+    issues = []
+    for item in error.errors(include_url=False):
+        location = ".".join(str(part) for part in item["loc"])
+        issues.append(f"{location}: {item['msg']}")
+    return "; ".join(issues)