zai-cli 0.1.0__py3-none-any.whl

zai/providers/anthropic.py

@@ -0,0 +1,131 @@
+from .base import BaseProvider, Message, Response, ToolCall
+from ..config import get_api_key
+from ..core.errors import (
+    NoAPIKeyError,
+    AuthenticationError as ZaiAuthenticationError,
+    RateLimitError,
+    NetworkError,
+    classify_provider_error,
+)
+
+
+class AnthropicProvider(BaseProvider):
+    name = "anthropic"
+    model_key = "claude"
+    model_id = "claude-sonnet-4-6"
+    context_window = 200000
+    supports_streaming = True
+    supports_native_tools = True
+
+    def is_available(self) -> bool:
+        return bool(get_api_key("anthropic"))
+
+    def stream_chat(self, messages: list[Message], system: str = "") -> str:
+        from ..core.streaming import stream_anthropic
+        key = get_api_key("anthropic")
+        if not key:
+            raise NoAPIKeyError("anthropic")
+        return stream_anthropic(
+            messages, system or "You are zai, a helpful AI assistant.",
+            self.model_id, key, self.timeout, self.retries,
+        )
+
+    def chat(
+        self,
+        messages: list[Message],
+        system: str = "",
+        tools: list[dict] | None = None,
+    ) -> Response:
+        import anthropic
+
+        key = get_api_key("anthropic")
+        if not key:
+            raise NoAPIKeyError("anthropic")
+
+        try:
+            client = anthropic.Anthropic(
+                api_key=key,
+                timeout=self.timeout,
+                max_retries=self.retries,
+            )
+            formatted = []
+            for message in messages:
+                if message.role == "assistant" and message.tool_calls:
+                    content = []
+                    if message.content:
+                        content.append({"type": "text", "text": message.content})
+                    content.extend({
+                        "type": "tool_use",
+                        "id": call.id,
+                        "name": call.name,
+                        "input": call.arguments,
+                    } for call in message.tool_calls)
+                    formatted.append({"role": "assistant", "content": content})
+                elif message.role == "tool":
+                    result_block = {
+                        "type": "tool_result",
+                        "tool_use_id": message.tool_call_id,
+                        "content": message.content,
+                    }
+                    if (
+                        formatted
+                        and formatted[-1]["role"] == "user"
+                        and isinstance(formatted[-1]["content"], list)
+                        and all(
+                            block.get("type") == "tool_result"
+                            for block in formatted[-1]["content"]
+                        )
+                    ):
+                        formatted[-1]["content"].append(result_block)
+                    else:
+                        formatted.append({
+                            "role": "user",
+                            "content": [result_block],
+                        })
+                else:
+                    formatted.append({
+                        "role": message.role,
+                        "content": message.content,
+                    })
+            request = {
+                "model": self.model_id,
+                "max_tokens": 8096,
+                "system": system or "You are zai, a helpful AI assistant.",
+                "messages": formatted,
+            }
+            if tools:
+                request["tools"] = [{
+                    "name": tool["name"],
+                    "description": tool["description"],
+                    "input_schema": tool["parameters"],
+                } for tool in tools]
+            result = client.messages.create(
+                **request,
+            )
+            text = "".join(
+                block.text for block in result.content
+                if (
+                    getattr(block, "type", None) == "text"
+                    or not isinstance(getattr(block, "type", None), str)
+                )
+            )
+            tool_calls = [
+                ToolCall(id=block.id, name=block.name, arguments=block.input)
+                for block in result.content
+                if getattr(block, "type", "") == "tool_use"
+            ]
+            return Response(
+                content=text,
+                model=self.model_id,
+                tokens_used=result.usage.input_tokens + result.usage.output_tokens,
+                tool_calls=tool_calls,
+            )
+        except anthropic.RateLimitError:
+            raise RateLimitError("anthropic")
+        except anthropic.AuthenticationError:
+            raise ZaiAuthenticationError("anthropic", "invalid API credentials")
+        except Exception as e:
+            err = str(e).lower()
+            if "connect" in err or "timeout" in err or "network" in err:
+                raise NetworkError(str(e))
+            raise classify_provider_error("anthropic", e) from e

zai/providers/base.py

@@ -0,0 +1,67 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class ToolCall:
+    id: str
+    name: str
+    arguments: dict[str, Any]
+
+
+@dataclass
+class Message:
+    role: str  # "user", "assistant", or "tool"
+    content: str = ""
+    tool_calls: list[ToolCall] = field(default_factory=list)
+    tool_call_id: str = ""
+    tool_name: str = ""
+    pinned: bool = False
+
+
+@dataclass
+class Response:
+    content: str
+    model: str
+    tokens_used: int = 0
+    tokens_remaining: int = -1
+    tool_calls: list[ToolCall] = field(default_factory=list)
+
+
+class BaseProvider(ABC):
+    name: str = ""
+    model_key: str = ""
+    model_id: str = ""
+    context_window: int = 128000
+    supports_streaming: bool = False
+    supports_native_tools: bool = False
+    timeout: float = 60
+    retries: int = 2
+
+    def __init__(self):
+        from ..config import get_model_config
+
+        model = get_model_config(self.model_key or self.name)
+        self.model_id = model["model_id"]
+        self.context_window = model["context_window"]
+        self.timeout = model["timeout"]
+        self.retries = model["retries"]
+
+    @abstractmethod
+    def chat(
+        self,
+        messages: list[Message],
+        system: str = "",
+        tools: list[dict] | None = None,
+    ) -> Response:
+        pass
+
+    def stream_chat(self, messages: list[Message], system: str = "") -> str:
+        """Stream response — override in providers that support it."""
+        response = self.chat(messages, system=system)
+        return response.content
+
+    @abstractmethod
+    def is_available(self) -> bool:
+        pass

zai/providers/cerebras.py

@@ -0,0 +1,57 @@
+from .base import BaseProvider, Message, Response
+from ..config import get_api_key
+from ..core.errors import (
+    NoAPIKeyError,
+    AuthenticationError,
+    RateLimitError,
+    NetworkError,
+    classify_provider_error,
+)
+
+
+class CerebrasProvider(BaseProvider):
+    name = "cerebras"
+    model_id = "llama-3.3-70b"
+    context_window = 128000
+
+    def is_available(self) -> bool:
+        return bool(get_api_key("cerebras"))
+
+    def chat(
+        self,
+        messages: list[Message],
+        system: str = "",
+        tools: list[dict] | None = None,
+    ) -> Response:
+        from cerebras.cloud.sdk import Cerebras
+
+        key = get_api_key("cerebras")
+        if not key:
+            raise NoAPIKeyError("cerebras")
+
+        try:
+            client = Cerebras(
+                api_key=key,
+                timeout=self.timeout,
+                max_retries=self.retries,
+            )
+            formatted = [{"role": "system", "content": system or "You are zai, a helpful AI assistant."}]
+            formatted += [{"role": m.role, "content": m.content} for m in messages]
+            result = client.chat.completions.create(
+                model=self.model_id,
+                messages=formatted,
+            )
+            return Response(
+                content=result.choices[0].message.content,
+                model=self.model_id,
+                tokens_used=result.usage.total_tokens if result.usage else 0,
+            )
+        except Exception as e:
+            err = str(e).lower()
+            if "429" in err or "rate" in err or "quota" in err:
+                raise RateLimitError("cerebras")
+            if "401" in err or "api key" in err or "auth" in err:
+                raise AuthenticationError("cerebras", str(e))
+            if "connect" in err or "timeout" in err:
+                raise NetworkError(str(e))
+            raise classify_provider_error("cerebras", e) from e

zai/providers/gemini.py

@@ -0,0 +1,119 @@
+from .base import BaseProvider, Message, Response, ToolCall
+from ..config import MODELS, get_api_key
+from ..core.errors import (
+    NoAPIKeyError,
+    AuthenticationError,
+    RateLimitError,
+    NetworkError,
+    classify_provider_error,
+)
+
+
+class GeminiProvider(BaseProvider):
+    name = "gemini"
+    model_id = MODELS["gemini"]["model_id"]
+    context_window = MODELS["gemini"]["context_window"]
+    supports_streaming = True
+    supports_native_tools = True
+
+    def is_available(self) -> bool:
+        return bool(get_api_key("gemini"))
+
+    def stream_chat(self, messages: list[Message], system: str = "") -> str:
+        from ..core.streaming import stream_gemini
+        key = get_api_key("gemini")
+        if not key:
+            raise NoAPIKeyError("gemini")
+        return stream_gemini(
+            messages, system or "You are zai, a helpful AI assistant.",
+            self.model_id, key, self.timeout,
+        )
+
+    def chat(
+        self,
+        messages: list[Message],
+        system: str = "",
+        tools: list[dict] | None = None,
+    ) -> Response:
+        key = get_api_key("gemini")
+        if not key:
+            raise NoAPIKeyError("gemini")
+
+        from google import genai
+        from google.genai import types
+
+        try:
+            client = genai.Client(
+                api_key=key,
+                http_options=types.HttpOptions(
+                    timeout=int(self.timeout * 1000),
+                ),
+            )
+            contents = []
+            for message in messages:
+                if message.role == "assistant" and message.tool_calls:
+                    parts = []
+                    if message.content:
+                        parts.append(types.Part.from_text(text=message.content))
+                    parts.extend(
+                        types.Part.from_function_call(
+                            name=call.name,
+                            args=call.arguments,
+                        )
+                        for call in message.tool_calls
+                    )
+                    contents.append(types.Content(role="model", parts=parts))
+                elif message.role == "tool":
+                    contents.append(types.Content(
+                        role="user",
+                        parts=[types.Part.from_function_response(
+                            name=message.tool_name,
+                            response={"result": message.content},
+                        )],
+                    ))
+                else:
+                    contents.append(types.Content(
+                        role=message.role if message.role == "user" else "model",
+                        parts=[types.Part.from_text(text=message.content)],
+                    ))
+            config_args = {
+                "system_instruction": system or "You are zai, a helpful AI assistant.",
+            }
+            if tools:
+                declarations = [
+                    types.FunctionDeclaration(
+                        name=tool["name"],
+                        description=tool["description"],
+                        parameters_json_schema=tool["parameters"],
+                    )
+                    for tool in tools
+                ]
+                config_args["tools"] = [types.Tool(function_declarations=declarations)]
+            result = client.models.generate_content(
+                model=self.model_id,
+                contents=contents,
+                config=types.GenerateContentConfig(**config_args),
+            )
+            usage = getattr(result, "usage_metadata", None)
+            tool_calls = []
+            for index, call in enumerate(getattr(result, "function_calls", None) or []):
+                tool_calls.append(ToolCall(
+                    id=f"gemini-{index}-{call.name}",
+                    name=call.name,
+                    arguments=dict(call.args or {}),
+                ))
+            return Response(
+                content=result.text or "",
+                model=self.model_id,
+                tokens_used=getattr(usage, "total_token_count", 0) or 0,
+                tool_calls=tool_calls,
+            )
+        except Exception as e:
+            err = str(e).lower()
+            if "quota" in err or "429" in err or "rate" in err:
+                raise RateLimitError("gemini")
+            if "api key" in err or "invalid" in err or "401" in err:
+                raise AuthenticationError("gemini", str(e))
+            if "network" in err or "connect" in err or "timeout" in err:
+                raise NetworkError(str(e))
+            raise classify_provider_error("gemini", e) from e

zai/providers/groq.py

@@ -0,0 +1,116 @@
+import json
+
+from .base import BaseProvider, Message, Response, ToolCall
+from ..config import MODELS, get_api_key
+from ..core.errors import (
+    NoAPIKeyError,
+    AuthenticationError as ZaiAuthenticationError,
+    RateLimitError,
+    NetworkError,
+    classify_provider_error,
+)
+
+
+class GroqProvider(BaseProvider):
+    name = "groq"
+    model_id = MODELS["groq"]["model_id"]
+    context_window = MODELS["groq"]["context_window"]
+    supports_streaming = True
+    supports_native_tools = True
+
+    def is_available(self) -> bool:
+        return bool(get_api_key("groq"))
+
+    def stream_chat(self, messages: list[Message], system: str = "") -> str:
+        from ..core.streaming import stream_groq
+        key = get_api_key("groq")
+        if not key:
+            raise NoAPIKeyError("groq")
+        return stream_groq(
+            messages, system or "You are zai, a helpful AI assistant.",
+            self.model_id, key, self.timeout, self.retries,
+        )
+
+    def chat(
+        self,
+        messages: list[Message],
+        system: str = "",
+        tools: list[dict] | None = None,
+    ) -> Response:
+        from groq import Groq, RateLimitError as GroqRateLimit, AuthenticationError
+
+        key = get_api_key("groq")
+        if not key:
+            raise NoAPIKeyError("groq")
+
+        try:
+            client = Groq(
+                api_key=key,
+                timeout=self.timeout,
+                max_retries=self.retries,
+            )
+            formatted = [{"role": "system", "content": system or "You are zai, a helpful AI assistant."}]
+            for message in messages:
+                if message.role == "assistant" and message.tool_calls:
+                    formatted.append({
+                        "role": "assistant",
+                        "content": message.content or None,
+                        "tool_calls": [{
+                            "id": call.id,
+                            "type": "function",
+                            "function": {
+                                "name": call.name,
+                                "arguments": json.dumps(call.arguments),
+                            },
+                        } for call in message.tool_calls],
+                    })
+                elif message.role == "tool":
+                    formatted.append({
+                        "role": "tool",
+                        "tool_call_id": message.tool_call_id,
+                        "name": message.tool_name,
+                        "content": message.content,
+                    })
+                else:
+                    formatted.append({
+                        "role": message.role,
+                        "content": message.content,
+                    })
+            request = {"model": self.model_id, "messages": formatted}
+            if tools:
+                request["tools"] = [{
+                    "type": "function",
+                    "function": {
+                        "name": tool["name"],
+                        "description": tool["description"],
+                        "parameters": tool["parameters"],
+                    },
+                } for tool in tools]
+            result = client.chat.completions.create(**request)
+            message = result.choices[0].message
+            tool_calls = []
+            for call in message.tool_calls or []:
+                try:
+                    arguments = json.loads(call.function.arguments)
+                except (TypeError, json.JSONDecodeError):
+                    arguments = {}
+                tool_calls.append(ToolCall(
+                    id=call.id,
+                    name=call.function.name,
+                    arguments=arguments,
+                ))
+            return Response(
+                content=message.content or "",
+                model=self.model_id,
+                tokens_used=result.usage.total_tokens,
+                tool_calls=tool_calls,
+            )
+        except GroqRateLimit:
+            raise RateLimitError("groq")
+        except AuthenticationError:
+            raise ZaiAuthenticationError("groq", "invalid API credentials")
+        except Exception as e:
+            err = str(e).lower()
+            if "connect" in err or "timeout" in err or "network" in err:
+                raise NetworkError(str(e))
+            raise classify_provider_error("groq", e) from e

zai/providers/ollama.py

@@ -0,0 +1,62 @@
+from .base import BaseProvider, Message, Response
+from ..core.errors import NetworkError
+
+OLLAMA_URL = "http://localhost:11434"
+
+
+class OllamaProvider(BaseProvider):
+    name = "ollama"
+    model_id = "llama3.2"
+    context_window = 128000
+    supports_streaming = True
+
+    def is_available(self) -> bool:
+        import httpx
+        try:
+            r = httpx.get(f"{OLLAMA_URL}/api/tags", timeout=2)
+            return r.status_code == 200
+        except Exception:
+            return False
+
+    def stream_chat(self, messages: list[Message], system: str = "") -> str:
+        from ..core.streaming import stream_ollama
+        return stream_ollama(
+            messages, system or "You are zai, a helpful AI assistant.",
+            self.model_id, self.timeout,
+        )
+
+    def chat(
+        self,
+        messages: list[Message],
+        system: str = "",
+        tools: list[dict] | None = None,
+    ) -> Response:
+        import httpx
+        import json
+
+        formatted = [{"role": "system", "content": system or "You are zai, a helpful AI assistant."}]
+        formatted += [{"role": m.role, "content": m.content} for m in messages]
+
+        for attempt in range(self.retries + 1):
+            try:
+                r = httpx.post(
+                    f"{OLLAMA_URL}/api/chat",
+                    json={"model": self.model_id, "messages": formatted, "stream": False},
+                    timeout=self.timeout,
+                )
+                r.raise_for_status()
+                data = r.json()
+                content = data.get("message", {}).get("content", "")
+                return Response(
+                    content=content,
+                    model=self.model_id,
+                    tokens_used=data.get("eval_count", 0),
+                )
+            except (httpx.ConnectError, httpx.TimeoutException) as error:
+                if attempt == self.retries:
+                    if isinstance(error, httpx.ConnectError):
+                        raise NetworkError(
+                            "Ollama not running. Start it with: ollama serve"
+                        )
+                    raise NetworkError(str(error))
+        raise NetworkError("Ollama request failed")

zai/providers/openai.py

@@ -0,0 +1,124 @@
+import json
+
+from .base import BaseProvider, Message, Response, ToolCall
+from ..config import get_api_key
+from ..core.errors import (
+    NoAPIKeyError,
+    AuthenticationError as ZaiAuthenticationError,
+    RateLimitError,
+    NetworkError,
+    classify_provider_error,
+)
+
+
+class OpenAIProvider(BaseProvider):
+    name = "openai"
+    model_key = "gpt4o"
+    model_id = "gpt-4o-mini"
+    context_window = 128000
+    supports_streaming = True
+    supports_native_tools = True
+
+    def is_available(self) -> bool:
+        return bool(get_api_key("openai"))
+
+    def stream_chat(self, messages: list[Message], system: str = "") -> str:
+        from ..core.streaming import stream_openai
+        key = get_api_key("openai")
+        if not key:
+            raise NoAPIKeyError("openai")
+        return stream_openai(
+            messages, system or "You are zai, a helpful AI assistant.",
+            self.model_id, key, self.timeout, self.retries,
+        )
+
+    def chat(
+        self,
+        messages: list[Message],
+        system: str = "",
+        tools: list[dict] | None = None,
+    ) -> Response:
+        from openai import OpenAI, RateLimitError as OAIRateLimit, AuthenticationError
+
+        key = get_api_key("openai")
+        if not key:
+            raise NoAPIKeyError("openai")
+
+        try:
+            client = OpenAI(
+                api_key=key,
+                timeout=self.timeout,
+                max_retries=self.retries,
+            )
+            formatted = [{
+                "role": "system",
+                "content": system or "You are zai, a helpful AI assistant.",
+            }]
+            for message in messages:
+                if message.role == "assistant" and message.tool_calls:
+                    formatted.append({
+                        "role": "assistant",
+                        "content": message.content or None,
+                        "tool_calls": [{
+                            "id": call.id,
+                            "type": "function",
+                            "function": {
+                                "name": call.name,
+                                "arguments": json.dumps(call.arguments),
+                            },
+                        } for call in message.tool_calls],
+                    })
+                elif message.role == "tool":
+                    formatted.append({
+                        "role": "tool",
+                        "tool_call_id": message.tool_call_id,
+                        "content": message.content,
+                    })
+                else:
+                    formatted.append({"role": message.role, "content": message.content})
+
+            request = {
+                "model": self.model_id,
+                "messages": formatted,
+            }
+            if tools:
+                request["tools"] = [{
+                    "type": "function",
+                    "function": {
+                        "name": tool["name"],
+                        "description": tool["description"],
+                        "parameters": tool["parameters"],
+                        "strict": tool["name"] not in {"mcp_call", "plugin_call"},
+                    },
+                } for tool in tools]
+                request["parallel_tool_calls"] = True
+            result = client.chat.completions.create(
+                **request,
+            )
+            message = result.choices[0].message
+            tool_calls = []
+            for call in message.tool_calls or []:
+                try:
+                    arguments = json.loads(call.function.arguments)
+                except (TypeError, json.JSONDecodeError):
+                    arguments = {}
+                tool_calls.append(ToolCall(
+                    id=call.id,
+                    name=call.function.name,
+                    arguments=arguments,
+                ))
+            return Response(
+                content=message.content or "",
+                model=self.model_id,
+                tokens_used=result.usage.total_tokens if result.usage else 0,
+                tool_calls=tool_calls,
+            )
+        except OAIRateLimit:
+            raise RateLimitError("openai")
+        except AuthenticationError:
+            raise ZaiAuthenticationError("openai", "invalid API credentials")
+        except Exception as e:
+            err = str(e).lower()
+            if "connect" in err or "timeout" in err or "network" in err:
+                raise NetworkError(str(e))
+            raise classify_provider_error("openai", e) from e