ha-text-ai/custom_components/ha_text_ai/api_client.py

"""
API Client for HA Text AI.

@license: CC BY-NC-SA 4.0 International
@author: SMKRV
@github: https://github.com/smkrv/ha-text-ai
@source: https://github.com/smkrv/ha-text-ai
"""
import logging
import asyncio
from typing import Any, Dict, List, Optional
from aiohttp import ClientSession, ClientTimeout
from async_timeout import timeout

from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from .const import (
    API_TIMEOUT,
    API_RETRY_COUNT,
    API_PROVIDER_ANTHROPIC,
    API_PROVIDER_DEEPSEEK,
    API_PROVIDER_OPENAI,
    API_PROVIDER_GEMINI,
    MIN_TEMPERATURE,
    MAX_TEMPERATURE,
    MIN_MAX_TOKENS,
    MAX_MAX_TOKENS,
)

_LOGGER = logging.getLogger(__name__)


class APIClient:
    """API Client for OpenAI and Anthropic."""

    def __init__(
        self,
        session: ClientSession,
        endpoint: str,
        headers: Dict[str, str],
        api_provider: str,
        model: str,
    ) -> None:
        """Initialize API client."""
        self.session = session
        self.endpoint = endpoint
        self.headers = headers
        self.api_provider = api_provider
        self.model = model
        self.timeout = ClientTimeout(total=API_TIMEOUT)

    def _validate_parameters(
        self,
        temperature: float,
        max_tokens: int,
    ) -> None:
        """Validate API parameters."""
        if not MIN_TEMPERATURE <= temperature <= MAX_TEMPERATURE:
            raise ValueError(
                f"Temperature must be between {MIN_TEMPERATURE} and {MAX_TEMPERATURE}"
            )
        if not MIN_MAX_TOKENS <= max_tokens <= MAX_MAX_TOKENS:
            raise ValueError(
                f"Max tokens must be between {MIN_MAX_TOKENS} and {MAX_MAX_TOKENS}"
            )

    async def _make_request(
        self,
        url: str,
        payload: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Make API request with retry logic."""
        _LOGGER.debug(f"API Request: URL={url}, Payload={payload}")
        for attempt in range(API_RETRY_COUNT):
            try:
                async with timeout(API_TIMEOUT):
                    async with self.session.post(
                        url,
                        json=payload,
                        headers=self.headers,
                        timeout=self.timeout,
                    ) as response:
                        _LOGGER.debug(f"Response status: {response.status}")
                        if response.status != 200:
                            error_data = await response.json()
                            _LOGGER.error(f"API error: {error_data}")
                            raise HomeAssistantError(f"API error: {error_data}")
                        return await response.json()
            except asyncio.TimeoutError:
                _LOGGER.warning(f"Timeout on attempt {attempt + 1}")
                if attempt == API_RETRY_COUNT - 1:
                    raise HomeAssistantError("API request timed out")
                await asyncio.sleep(1 * (attempt + 1))
            except Exception as e:
                _LOGGER.warning(f"API request failed on attempt {attempt + 1}: {str(e)}")
                if attempt == API_RETRY_COUNT - 1:
                    raise
                await asyncio.sleep(1 * (attempt + 1))

    async def create(
        self,
        model: str,
        messages: List[Dict[str, str]],
        temperature: float,
        max_tokens: int,
    ) -> Dict[str, Any]:
        """Create completion using appropriate API."""
        try:
            self._validate_parameters(temperature, max_tokens)

            if self.api_provider == API_PROVIDER_ANTHROPIC:
                return await self._create_anthropic_completion(
                    model, messages, temperature, max_tokens
                )
            elif self.api_provider == API_PROVIDER_DEEPSEEK:
                return await self._create_deepseek_completion(
                    model, messages, temperature, max_tokens
                )
            elif self.api_provider == API_PROVIDER_GEMINI:
                return await self._create_gemini_completion(
                    model, messages, temperature, max_tokens
                )
            else:
                return await self._create_openai_completion(
                    model, messages, temperature, max_tokens
                )
        except Exception as e:
            _LOGGER.error("API request failed: %s", str(e))
            raise HomeAssistantError(f"API request failed: {str(e)}")

    async def _create_deepseek_completion(
        self,
        model: str,
        messages: List[Dict[str, str]],
        temperature: float,
        max_tokens: int,
    ) -> Dict[str, Any]:
        """Create completion using DeepSeek API."""
        url = f"{self.endpoint}/chat/completions"
        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "stream": False
        }

        data = await self._make_request(url, payload)
        return {
            "choices": [
                {
                    "message": {"content": data["choices"][0]["message"]["content"]},
                }
            ],
            "usage": {
                "prompt_tokens": data["usage"]["prompt_tokens"],
                "completion_tokens": data["usage"]["completion_tokens"],
                "total_tokens": data["usage"]["total_tokens"],
            },
        }

    async def _create_openai_completion(
        self,
        model: str,
        messages: List[Dict[str, str]],
        temperature: float,
        max_tokens: int,
    ) -> Dict[str, Any]:
        """Create completion using OpenAI API."""
        url = f"{self.endpoint}/chat/completions"
        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }

        data = await self._make_request(url, payload)
        return {
            "choices": [
                {
                    "message": {"content": data["choices"][0]["message"]["content"]},
                }
            ],
            "usage": {
                "prompt_tokens": data["usage"]["prompt_tokens"],
                "completion_tokens": data["usage"]["completion_tokens"],
                "total_tokens": data["usage"]["total_tokens"],
            },
        }

    async def _create_anthropic_completion(
        self,
        model: str,
        messages: List[Dict[str, str]],
        temperature: float,
        max_tokens: int,
    ) -> Dict[str, Any]:
        """Create completion using Anthropic API."""
        url = f"{self.endpoint}/v1/messages"

        system_prompt = None
        filtered_messages = []
        for msg in messages:
            if msg['role'] == 'system':
                if system_prompt is None:
                    system_prompt = msg['content']
                else:
                    system_prompt += f" {msg['content']}"
            else:
                filtered_messages.append(msg)

        payload = {
            "model": model,
            "messages": filtered_messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
        }

        if system_prompt:
            payload["system"] = system_prompt

        data = await self._make_request(url, payload)
        return {
            "choices": [
                {
                    "message": {"content": data["content"][0]["text"]},
                }
            ],
            "usage": {
                "prompt_tokens": data["usage"]["input_tokens"],
                "completion_tokens": data["usage"]["output_tokens"],
                "total_tokens": data["usage"]["input_tokens"] + data["usage"]["output_tokens"],
            },
        }

    async def check_connection(self) -> bool:
        """Check API connection."""
        try:
            await self._make_request(self.endpoint, {"test": "connection"})
            return True
        except Exception as e:
            _LOGGER.error(f"Connection check failed: {str(e)}")
            return False

    async def _create_gemini_completion(
        self,
        model: str,
        messages: List[Dict[str, str]],
        temperature: float,
        max_tokens: int,
    ) -> Dict[str, Any]:
        """Create completion using Gemini API."""
        # Extract API key from headers (Bearer token)
        api_key = self.headers.get("Authorization", "").replace("Bearer ", "")
        url = f"{self.endpoint}/models/{model}:generateContent?key={api_key}"

        # Convert messages to Gemini format
        contents = []
        system_instruction = ""

        # Process messages and ensure proper role alternation
        current_role = None
        current_content = ""

        for msg in messages:
            if msg['role'] == 'system':
                system_instruction += msg['content'] + "\n"
            else:
                role = "user" if msg['role'] == 'user' else "model"

                # If same role as previous, combine content
                if role == current_role:
                    current_content += "\n" + msg['content']
                else:
                    # Add previous message if exists
                    if current_role is not None:
                        contents.append({
                            "role": current_role,
                            "parts": [{"text": current_content}]
                        })
                    # Start new message
                    current_role = role
                    current_content = msg['content']

        # Add the last message if exists
        if current_role is not None:
            contents.append({
                "role": current_role,
                "parts": [{"text": current_content}]
            })

        # Ensure contents starts with user message if not empty
        if contents and contents[0]["role"] != "user":
            # Add a placeholder user message if needed
            contents.insert(0, {
                "role": "user",
                "parts": [{"text": "I need your assistance."}]
            })

        # Ensure contents is not empty
        if not contents:
            contents.append({
                "role": "user",
                "parts": [{"text": "I need your assistance."}]
            })

        # Create payload with snake_case keys as required by Gemini API
        payload = {
            "contents": contents,
            "generation_config": {  # Changed from camelCase to snake_case
                "temperature": temperature,
                "max_output_tokens": max_tokens  # Changed from camelCase to snake_case
            }
        }

        if system_instruction:
            payload["system_instruction"] = {  # Changed from camelCase to snake_case
                "parts": [{"text": system_instruction.strip()}]
            }

        try:
            data = await self._make_request(url, payload)

            # Safely extract response data with fallbacks
            candidates = data.get("candidates", [])
            if not candidates:
                raise HomeAssistantError("Gemini API returned no candidates")

            content = candidates[0].get("content", {})
            parts = content.get("parts", [])
            if not parts:
                raise HomeAssistantError("Gemini API response contains no content parts")

            answer_text = parts[0].get("text", "")

            # Safely extract usage data
            usage = data.get("usageMetadata", {})
            prompt_tokens = usage.get("promptTokenCount", 0)
            completion_tokens = usage.get("candidatesTokenCount", 0)

            # Handle case where candidatesTokenCount might be a list
            if isinstance(completion_tokens, list):
                completion_tokens = sum(completion_tokens)

            total_tokens = usage.get("totalTokenCount", prompt_tokens + completion_tokens)

            return {
                "choices": [{
                    "message": {
                        "content": answer_text
                    }
                }],
                "usage": {
                    "prompt_tokens": prompt_tokens,
                    "completion_tokens": completion_tokens,
                    "total_tokens": total_tokens
                }
            }
        except Exception as e:
            _LOGGER.error(f"Gemini API error: {str(e)}")
            raise HomeAssistantError(f"Gemini API error: {str(e)}")

    async def shutdown(self) -> None:
        """Shutdown API client."""
        _LOGGER.debug("Shutting down API client")
        await self.session.close()