vtx-coding-agent 0.1.1__py3-none-any.whl

vtx/llm/base.py

@@ -0,0 +1,217 @@
+import os
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from dataclasses import dataclass, field
+from ipaddress import ip_address
+from typing import ClassVar, Literal
+from urllib.parse import urlparse
+
+import httpx
+
+from vtx import config as vtx_config
+
+from ..core.types import Message, StreamPart, ToolDefinition, Usage
+
+DEFAULT_THINKING_LEVELS: list[str] = ["none", "minimal", "low", "medium", "high", "xhigh"]
+LOCAL_API_KEY_PLACEHOLDER = "vtx-local"
+AuthMode = Literal["auto", "required", "none"]
+
+ENV_API_KEY_MAP: dict[str, str] = {
+    "openai": "OPENAI_API_KEY",
+    "anthropic": "ANTHROPIC_API_KEY",
+    "deepseek": "DEEPSEEK_API_KEY",
+    "zhipu": "ZAI_API_KEY",
+    "airouter": "AIROUTER_API_KEY",
+    "opencode": "OPENCODE_API_KEY",
+    "kilo": "KILO_API_KEY",
+    "tokenrouter": "TOKENROUTER_API_KEY",
+    "openrouter": "OPENROUTER_API_KEY",
+}
+
+
+def get_env_api_key(provider: str) -> str | None:
+    env_var = ENV_API_KEY_MAP.get(provider)
+    return os.environ.get(env_var) if env_var else None
+
+
+def is_local_base_url(base_url: str | None) -> bool:
+    if not base_url:
+        return False
+
+    parsed = urlparse(base_url if "://" in base_url else f"https://{base_url}")
+    hostname = parsed.hostname
+    if hostname is None:
+        return False
+
+    normalized = hostname.lower()
+    if normalized in {"localhost", "127.0.0.1", "0.0.0.0", "::1"}:
+        return True
+    if normalized.endswith(".local"):
+        return True
+
+    try:
+        addr = ip_address(normalized)
+    except ValueError:
+        return False
+
+    return addr.is_loopback or addr.is_private or addr.is_link_local
+
+
+def make_http_client() -> httpx.AsyncClient | None:
+    # Returns None when verify is required so the SDK uses its own default client.
+    if not vtx_config.llm.tls.insecure_skip_verify:
+        return None
+    return httpx.AsyncClient(
+        verify=False, timeout=httpx.Timeout(vtx_config.llm.request_timeout_seconds)
+    )
+
+
+def resolve_api_key(
+    explicit_api_key: str | None,
+    *,
+    env_vars: list[str] | tuple[str, ...] = (),
+    base_url: str | None = None,
+    auth_mode: AuthMode = "required",
+) -> str | None:
+    if explicit_api_key:
+        return explicit_api_key
+
+    for env_var in env_vars:
+        value = os.environ.get(env_var)
+        if value:
+            return value
+
+    if auth_mode == "none":
+        return LOCAL_API_KEY_PLACEHOLDER
+    if auth_mode == "auto" and is_local_base_url(base_url):
+        return LOCAL_API_KEY_PLACEHOLDER
+
+    return None
+
+
+@dataclass
+class ProviderConfig:
+    api_key: str | None = None
+    base_url: str | None = None
+    model: str = ""
+    max_tokens: int = 8192
+    temperature: float | None = None
+    thinking_level: str = "high"
+    provider: str | None = None
+    session_id: str | None = None
+    openai_compat_auth_mode: AuthMode = "auto"
+    anthropic_compat_auth_mode: AuthMode = "auto"
+    default_headers: dict[str, str] = field(default_factory=dict)
+
+
+class LLMStream(AsyncIterator["StreamPart"]):
+    """
+    Async iterator over stream parts with access to final usage/metadata.
+
+    Usage:
+        stream = await provider.stream(messages, tools)
+        async for part in stream:
+            match part:
+                case TextPart(text=t):
+                    print(t, end="")
+                case ThinkPart(think=t):
+                    print(f"[thinking] {t}")
+                case ToolCallStart(id=id, name=name):
+                    print(f"Tool call: {name}")
+                ...
+
+        # After iteration, access final stats
+        print(f"Usage: {stream.usage}")
+    """
+
+    def __init__(self) -> None:
+        self._iterator: AsyncIterator[StreamPart] | None = None
+        self._usage: Usage | None = None
+        self._id: str | None = None
+
+    def set_iterator(self, iterator: AsyncIterator[StreamPart]) -> None:
+        self._iterator = iterator
+
+    def __aiter__(self) -> AsyncIterator[StreamPart]:
+        return self
+
+    async def __anext__(self) -> StreamPart:
+        if self._iterator is None:
+            raise StopAsyncIteration
+        return await self._iterator.__anext__()
+
+    async def aclose(self) -> None:
+        if self._iterator is None:
+            return
+        close = getattr(self._iterator, "aclose", None)
+        if close is not None:
+            await close()
+
+    @property
+    def usage(self) -> Usage | None:
+        return self._usage
+
+    @property
+    def id(self) -> str | None:
+        return self._id
+
+
+class BaseProvider(ABC):
+    name: str
+    thinking_levels: ClassVar[list[str]] = DEFAULT_THINKING_LEVELS
+
+    def __init__(self, config: ProviderConfig):
+        self.config = config
+
+    @property
+    def thinking_level(self) -> str:
+        return self.config.thinking_level
+
+    def set_thinking_level(self, level: str) -> None:
+        if level not in self.thinking_levels:
+            raise ValueError(
+                f"Invalid thinking level '{level}' for {self.name}. "
+                f"Valid levels: {self.thinking_levels}"
+            )
+        self.config.thinking_level = level
+
+    def cycle_thinking_level(self) -> str:
+        levels = self.thinking_levels
+        current_idx = (
+            levels.index(self.config.thinking_level) if self.config.thinking_level in levels else 0
+        )
+        next_idx = (current_idx + 1) % len(levels)
+        new_level = levels[next_idx]
+        self.config.thinking_level = new_level
+        return new_level
+
+    async def stream(
+        self,
+        messages: list[Message],
+        *,
+        system_prompt: str | None = None,
+        tools: list[ToolDefinition] | None = None,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+    ) -> LLMStream:
+        return await self._stream_impl(
+            messages,
+            system_prompt=system_prompt,
+            tools=tools,
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )
+
+    @abstractmethod
+    async def _stream_impl(
+        self,
+        messages: list[Message],
+        *,
+        system_prompt: str | None = None,
+        tools: list[ToolDefinition] | None = None,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+    ) -> LLMStream: ...
+
+    @abstractmethod
+    def should_retry_for_error(self, error: Exception) -> bool: ...

vtx/llm/context_length.py

@@ -0,0 +1,150 @@
+"""Dynamic context length manager for models.
+
+Fetches model context/output limits from the models.dev API and
+caches them. Provides lookup by model ID with fuzzy matching fallback.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import threading
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+
+logger = logging.getLogger(__name__)
+
+MODELS_DEV_API_URL = "https://models.dev/api.json"
+CACHE_FILE = "models_dev_limits.json"
+DEFAULT_CONTEXT_LENGTH = 128 * 1024  # 131072
+DEFAULT_OUTPUT_TOKENS = 16 * 1024  # 16384
+CACHE_TTL_SECONDS = 24 * 60 * 60  # 24 hours
+
+
+@dataclass
+class TokenLimits:
+    context: int
+    output: int
+    supports_reasoning: bool = False
+    supports_vision: bool = False
+    supports_tools: bool = False
+    supports_audio: bool = False
+
+
+class ContextLengthManager:
+    def __init__(self) -> None:
+        self._limits: dict[str, TokenLimits] = {}
+        self._loaded = False
+        self._lock = threading.Lock()
+
+    def _get_cache_path(self) -> Path:
+        from ..config import get_config_dir
+
+        return get_config_dir() / CACHE_FILE
+
+    def _load_from_cache(self) -> bool:
+        path = self._get_cache_path()
+        try:
+            import time
+
+            if path.exists() and (path.stat().st_mtime + CACHE_TTL_SECONDS) > time.time():
+                data = json.loads(path.read_text(encoding="utf-8"))
+                self._parse_limits(data)
+                return True
+        except Exception:
+            pass
+        return False
+
+    def _save_to_cache(self, data: dict[str, Any]) -> None:
+        try:
+            path = self._get_cache_path()
+            path.parent.mkdir(parents=True, exist_ok=True)
+            path.write_text(json.dumps(data, indent=2), encoding="utf-8")
+        except Exception as exc:
+            logger.debug("Failed to save models.dev cache: %s", exc)
+
+    def _fetch_and_parse(self) -> None:
+        try:
+            req = Request(MODELS_DEV_API_URL, headers={"User-Agent": "vtx/1.0"})
+            with urlopen(req, timeout=10) as resp:
+                data = json.loads(resp.read().decode("utf-8"))
+                self._parse_limits(data)
+                self._save_to_cache(data)
+                logger.info("Loaded model limits from models.dev")
+        except (URLError, Exception) as exc:
+            logger.debug("Failed to fetch model limits: %s", exc)
+
+    def _parse_limits(self, data: dict[str, Any]) -> None:
+        for _provider_name, provider_data in data.items():
+            if not isinstance(provider_data, dict):
+                continue
+            models = provider_data.get("models", {})
+            if not models:
+                continue
+            for model_id, model_info in models.items():
+                limit = model_info.get("limit", {})
+                if not limit:
+                    continue
+                context = limit.get("context", 0)
+                output = limit.get("output", DEFAULT_OUTPUT_TOKENS)
+                if context > 0:
+                    modalities = model_info.get("modalities", {})
+                    input_mods = modalities.get("input", [])
+                    output_mods = modalities.get("output", [])
+                    self._limits[model_id] = TokenLimits(
+                        context=context,
+                        output=output,
+                        supports_reasoning=bool(model_info.get("reasoning", False)),
+                        supports_vision="image" in input_mods,
+                        supports_tools=bool(model_info.get("tool_call", False)),
+                        supports_audio="audio" in input_mods or "audio" in output_mods,
+                    )
+
+    def ensure_loaded(self) -> None:
+        if self._loaded:
+            return
+        with self._lock:
+            if self._loaded:
+                return
+            if not self._load_from_cache():
+                self._fetch_and_parse()
+            self._loaded = True
+
+    def get_limits(self, model: str) -> TokenLimits:
+        self.ensure_loaded()
+
+        if model in self._limits:
+            return self._limits[model]
+
+        model_lower = model.lower()
+        for model_id, limits in self._limits.items():
+            if model_lower in model_id.lower() or model_id.lower() in model_lower:
+                return limits
+
+        return TokenLimits(context=DEFAULT_CONTEXT_LENGTH, output=DEFAULT_OUTPUT_TOKENS)
+
+    def register(self, model: str, context: int, output: int | None = None, **kwargs: Any) -> None:
+        self._limits[model] = TokenLimits(
+            context=context, output=output or DEFAULT_OUTPUT_TOKENS, **kwargs
+        )
+
+    def get_context_length(self, model: str) -> int:
+        return self.get_limits(model).context
+
+    def get_max_output(self, model: str) -> int:
+        return self.get_limits(model).output
+
+    def supports_reasoning(self, model: str) -> bool:
+        return self.get_limits(model).supports_reasoning
+
+    def supports_vision(self, model: str) -> bool:
+        return self.get_limits(model).supports_vision
+
+    def supports_tools(self, model: str) -> bool:
+        return self.get_limits(model).supports_tools
+
+
+context_length_manager = ContextLengthManager()