api_client.py

"""
Unified API client abstraction for OpenAI-compatible and Anthropic APIs.
"""

import json
import re
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, field, asdict
from typing import Any, Optional


@dataclass
class UsageInfo:
    """Token usage information from an API response."""

    prompt_tokens: Optional[int] = None
    completion_tokens: Optional[int] = None
    total_tokens: Optional[int] = None

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass
class APIResponse:
    """Normalized API response."""

    content: str
    usage: UsageInfo
    model: str
    elapsed_seconds: float
    raw: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "content": self.content,
            "usage": self.usage.to_dict(),
            "model": self.model,
            "elapsed_seconds": self.elapsed_seconds,
        }


class APIClient:
    """Abstract base for API clients."""

    def __init__(
        self,
        api_key: str,
        model: str,
        base_url: Optional[str] = None,
        temperature: float = 0.0,
        top_p: float = 1.0,
        max_tokens: int = 4096,
        timeout: float = 120.0,
    ):
        self.api_key = api_key
        self.model = model
        self.base_url = base_url
        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        self.timeout = timeout

    def chat(self, messages: list[dict], **kwargs) -> APIResponse:
        """Send a chat completion request."""
        raise NotImplementedError

    def count_tokens(self, text: str) -> int:
        """Count tokens for a given text using local tokenizer if available."""
        raise NotImplementedError

    def close(self):
        """Clean up resources."""
        pass


class OpenAIClient(APIClient):
    """Client for OpenAI-compatible APIs."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._client = None
        try:
            import openai
        except ModuleNotFoundError:
            self._openai_module = None
        else:
            self._openai_module = openai
            client_kwargs = {"api_key": self.api_key, "timeout": self.timeout}
            if self.base_url:
                client_kwargs["base_url"] = self.base_url
            if self.base_url and "azure.com" in self.base_url.lower():
                client_kwargs["api_version"] = kwargs.get("api_version", "2024-02-15-preview")
            self._client = openai.OpenAI(**client_kwargs)

    def _chat_with_urllib(self, messages: list[dict], **kwargs) -> APIResponse:
        start = time.time()
        root = (self.base_url or "https://api.openai.com/v1").rstrip("/")
        url = root if root.endswith("/chat/completions") else root + "/chat/completions"
        payload = {
            "model": kwargs.get("model", self.model),
            "messages": messages,
            "temperature": kwargs.get("temperature", self.temperature),
            "top_p": kwargs.get("top_p", self.top_p),
            "max_tokens": kwargs.get("max_tokens", self.max_tokens),
        }
        body = json.dumps(payload).encode("utf-8")
        request = urllib.request.Request(
            url,
            data=body,
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "Accept": "application/json",
            },
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                raw = json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            error_body = exc.read().decode("utf-8", errors="replace")[:800]
            raise RuntimeError(f"OpenAI-compatible request failed with HTTP {exc.code}: {error_body}") from exc
        except urllib.error.URLError as exc:
            raise RuntimeError(f"OpenAI-compatible request failed: {exc.reason}") from exc

        elapsed = time.time() - start
        choices = raw.get("choices") or []
        message = choices[0].get("message", {}) if choices else {}
        content = message.get("content") or ""
        usage = raw.get("usage") or {}
        usage_data = UsageInfo(
            prompt_tokens=usage.get("prompt_tokens"),
            completion_tokens=usage.get("completion_tokens"),
            total_tokens=usage.get("total_tokens"),
        )
        return APIResponse(
            content=content,
            usage=usage_data,
            model=raw.get("model", self.model),
            elapsed_seconds=elapsed,
            raw=raw,
        )

    def chat(self, messages: list[dict], **kwargs) -> APIResponse:
        if self._client is None:
            return self._chat_with_urllib(messages, **kwargs)

        start = time.time()
        response = self._client.chat.completions.create(
            model=kwargs.get("model", self.model),
            messages=messages,
            temperature=kwargs.get("temperature", self.temperature),
            top_p=kwargs.get("top_p", self.top_p),
            max_tokens=kwargs.get("max_tokens", self.max_tokens),
        )
        elapsed = time.time() - start

        content = response.choices[0].message.content or ""
        usage_data = UsageInfo(
            prompt_tokens=response.usage.prompt_tokens if response.usage else None,
            completion_tokens=response.usage.completion_tokens if response.usage else None,
            total_tokens=response.usage.total_tokens if response.usage else None,
        )

        return APIResponse(
            content=content,
            usage=usage_data,
            model=getattr(response, "model", self.model),
            elapsed_seconds=elapsed,
            raw=response.model_dump() if hasattr(response, "model_dump") else {},
        )

    def count_tokens(self, text: str, model: Optional[str] = None) -> int:
        """Count tokens using tiktoken when available, otherwise estimate."""
        try:
            import tiktoken
        except ModuleNotFoundError:
            return max(1, len(text) // 4)

        model_name = model or self.model
        try:
            encoding = tiktoken.encoding_for_model(model_name)
        except KeyError:
            encoding = tiktoken.get_encoding("cl100k_base")
        return len(encoding.encode(text))

    def close(self):
        pass

class AnthropicClient(APIClient):
    """Client for Anthropic API."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._client = None
        try:
            import anthropic
        except ModuleNotFoundError:
            self._anthropic_module = None
        else:
            self._anthropic_module = anthropic
            client_kwargs = {"api_key": self.api_key}
            if self.base_url:
                client_kwargs["base_url"] = self.base_url
            self._client = anthropic.Anthropic(**client_kwargs)

    def _split_system_messages(self, messages: list[dict]) -> tuple[Optional[str], list[dict]]:
        system = None
        msgs = []
        for m in messages:
            if m.get("role") == "system":
                system = m.get("content", "")
            else:
                msgs.append({"role": m["role"], "content": m["content"]})
        return system, msgs

    def _chat_with_urllib(self, messages: list[dict], **kwargs) -> APIResponse:
        start = time.time()
        system, msgs = self._split_system_messages(messages)
        root = (self.base_url or "https://api.anthropic.com").rstrip("/")
        url = root if root.endswith("/v1/messages") else root + "/v1/messages"
        payload = {
            "model": kwargs.get("model", self.model),
            "max_tokens": kwargs.get("max_tokens", self.max_tokens),
            "messages": msgs,
            "temperature": kwargs.get("temperature", self.temperature),
        }
        if not kwargs.get("disable_top_p", False):
            payload["top_p"] = kwargs.get("top_p", self.top_p)
        if system:
            payload["system"] = system

        request = urllib.request.Request(
            url,
            data=json.dumps(payload).encode("utf-8"),
            headers={
                "x-api-key": self.api_key,
                "anthropic-version": "2023-06-01",
                "Content-Type": "application/json",
                "Accept": "application/json",
            },
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                raw = json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            error_body = exc.read().decode("utf-8", errors="replace")[:800]
            raise RuntimeError(f"Anthropic request failed with HTTP {exc.code}: {error_body}") from exc
        except urllib.error.URLError as exc:
            raise RuntimeError(f"Anthropic request failed: {exc.reason}") from exc

        elapsed = time.time() - start
        content = "".join(
            block.get("text", "")
            for block in raw.get("content", [])
            if isinstance(block, dict)
        )
        usage = raw.get("usage") or {}
        input_tokens = usage.get("input_tokens")
        output_tokens = usage.get("output_tokens")
        usage_data = UsageInfo(
            prompt_tokens=input_tokens,
            completion_tokens=output_tokens,
            total_tokens=(input_tokens + output_tokens) if input_tokens is not None and output_tokens is not None else None,
        )
        return APIResponse(
            content=content,
            usage=usage_data,
            model=raw.get("model", self.model),
            elapsed_seconds=elapsed,
            raw=raw,
        )

    def chat(self, messages: list[dict], **kwargs) -> APIResponse:
        if self._client is None:
            return self._chat_with_urllib(messages, **kwargs)

        start = time.time()
        system, msgs = self._split_system_messages(messages)

        create_kwargs = {
            "model": kwargs.get("model", self.model),
            "max_tokens": kwargs.get("max_tokens", self.max_tokens),
            "messages": msgs,
            "temperature": kwargs.get("temperature", self.temperature),
        }

        if not kwargs.get("disable_top_p", False):
            create_kwargs["top_p"] = kwargs.get("top_p", self.top_p)

        if system:
            create_kwargs["system"] = system

        response = self._client.messages.create(**create_kwargs)
        elapsed = time.time() - start

        content = "".join(block.text for block in response.content if hasattr(block, "text"))

        usage_data = UsageInfo(
            prompt_tokens=response.usage.input_tokens if response.usage else None,
            completion_tokens=response.usage.output_tokens if response.usage else None,
            total_tokens=(response.usage.input_tokens + response.usage.output_tokens)
            if response.usage
            else None,
        )

        return APIResponse(
            content=content,
            usage=usage_data,
            model=getattr(response, "model", self.model),
            elapsed_seconds=elapsed,
            raw={},
        )

    def count_tokens(self, text: str) -> int:
        """Use Anthropic token counting when available, otherwise estimate."""
        if self._client is not None:
            try:
                response = self._client.count_tokens(text)
                return response.input_tokens
            except Exception:
                pass
        return max(1, len(text) // 4)

    def close(self):
        pass

def create_client(
    api_type: str,
    api_key: str,
    model: str,
    base_url: Optional[str] = None,
    temperature: float = 0.0,
    top_p: float = 1.0,
    max_tokens: int = 4096,
    timeout: float = 120.0,
) -> APIClient:
    """Factory function to create the appropriate API client."""
    if api_type == "openai":
        return OpenAIClient(
            api_key=api_key,
            model=model,
            base_url=base_url,
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
            timeout=timeout,
        )
    elif api_type == "anthropic":
        return AnthropicClient(
            api_key=api_key,
            model=model,
            base_url=base_url,
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
            timeout=timeout,
        )
    else:
        raise ValueError(f"Unsupported API type: {api_type}. Use 'openai' or 'anthropic'.")