utils/api_client.py

"""
GLM API Client for AI Writer.
Handles communication with the GLM-5 API endpoint on Modal.
"""

import os
import requests
import json
from typing import Optional, Dict, Any


DEFAULT_API_URL = "https://api.us-west-2.modal.direct/v1/chat/completions"
DEFAULT_MODEL = "zai-org/GLM-5.1-FP8"
DEFAULT_TOKEN = os.environ.get("GLM_API_TOKEN", "modalresearch_-z6GDDZ_VYtv7RlIuppxp5Vll50nSaDLtAOO-A5OnrI")


class GLMClient:
    """Client for the GLM-5 API."""

    def __init__(
        self,
        api_url: str = DEFAULT_API_URL,
        model: str = DEFAULT_MODEL,
        token: str = DEFAULT_TOKEN,
    ):
        self.api_url = api_url
        self.model = model
        self.token = token

    def test_connection(self) -> Dict[str, Any]:
        """Test the API connection with a simple request."""
        try:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.token}",
            }
            payload = {
                "model": self.model,
                "messages": [
                    {"role": "user", "content": "Say 'Connection successful' in Russian."}
                ],
                "max_tokens": 50,
            }
            response = requests.post(
                self.api_url,
                headers=headers,
                json=payload,
                timeout=30,
            )
            if response.status_code == 200:
                data = response.json()
                content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
                return {
                    "success": True,
                    "message": f"API connection successful. Response: {content}",
                    "status_code": response.status_code,
                }
            else:
                return {
                    "success": False,
                    "message": f"API returned status {response.status_code}: {response.text[:500]}",
                    "status_code": response.status_code,
                }
        except requests.exceptions.Timeout:
            return {
                "success": False,
                "message": "Connection timed out after 30 seconds.",
                "status_code": None,
            }
        except requests.exceptions.ConnectionError:
            return {
                "success": False,
                "message": "Could not connect to the API endpoint. Check your internet connection.",
                "status_code": None,
            }
        except Exception as e:
            return {
                "success": False,
                "message": f"Error: {str(e)}",
                "status_code": None,
            }

    def generate(
        self,
        system_prompt: str,
        user_prompt: str,
        max_tokens: int = 8000,
        temperature: float = 0.7,
    ) -> Dict[str, Any]:
        """Generate text using the GLM API."""
        try:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.token}",
            }
            messages = []
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            messages.append({"role": "user", "content": user_prompt})

            payload = {
                "model": self.model,
                "messages": messages,
                "max_tokens": max_tokens,
                "temperature": temperature,
            }

            response = requests.post(
                self.api_url,
                headers=headers,
                json=payload,
                timeout=300,  # 5 minutes for long generation
            )

            if response.status_code == 200:
                data = response.json()
                content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
                usage = data.get("usage", {})
                return {
                    "success": True,
                    "content": content,
                    "usage": usage,
                    "status_code": response.status_code,
                }
            else:
                return {
                    "success": False,
                    "content": "",
                    "message": f"API returned status {response.status_code}: {response.text[:500]}",
                    "status_code": response.status_code,
                }
        except requests.exceptions.Timeout:
            return {
                "success": False,
                "content": "",
                "message": "Generation timed out after 5 minutes. Try reducing max tokens or shortening the prompt.",
            }
        except Exception as e:
            return {
                "success": False,
                "content": "",
                "message": f"Error during generation: {str(e)}",
            }

    def update_token(self, new_token: str):
        """Update the API token."""
        self.token = new_token

    def update_model(self, new_model: str):
        """Update the model name."""
        self.model = new_model