api_client.py

import json
import os
from typing import Any, Dict, List, Optional, Tuple

import requests


def _strip_quotes(value: str) -> str:
    v = value.strip()
    if (v.startswith('"') and v.endswith('"')) or (v.startswith("'") and v.endswith("'")):
        return v[1:-1]
    return v


def load_env_file(env_path: str = 'api.ENV') -> Dict[str, str]:
    env: Dict[str, str] = {}
    if not os.path.exists(env_path):
        raise FileNotFoundError(f'ENV file not found: {env_path}')
    with open(env_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            if '=' not in line:
                continue
            key, value = line.split('=', 1)
            env[key.strip()] = _strip_quotes(value.strip())
    return env


def normalize_chat_completions_url(api_url: str) -> str:
    if not api_url:
        return 'https://api.openai.com/v1/chat/completions'
    lower = api_url.rstrip('/')
    if lower.endswith('/v1/chat/completions'):
        return lower
    if lower.endswith('/chat/completions'):
        return lower
    if lower.endswith('/v1'):
        return lower + '/chat/completions'
    if lower.endswith('/v1/chat'):
        return lower + '/completions'
    # Treat as base URL
    return lower + '/v1/chat/completions'


class ApiClient:
    def __init__(self, api_url: str, api_key: str, default_model: str):
        self.chat_url = normalize_chat_completions_url(api_url)
        self.api_key = api_key
        self.default_model = default_model
        self.session = requests.Session()
        self.headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json',
        }

    def chat_completion(
        self,
        messages: List[Dict[str, str]],
        model: Optional[str] = None,
        max_tokens: int = 15000,
        temperature: float = 0.5,
        top_p: float = 1.0,
        extra_body: Optional[Dict[str, Any]] = None,
        timeout: int = 300,
    ) -> Tuple[str, Dict[str, Any]]:
        payload: Dict[str, Any] = {
            'model': model or self.default_model,
            'messages': messages,
            'max_tokens': max_tokens,
            'temperature': temperature,
            'top_p': top_p,
            'stream': False,
        }
        if extra_body:
            payload.update(extra_body)

        try:
            resp = self.session.post(self.chat_url, headers=self.headers, json=payload, timeout=timeout)
        except requests.RequestException as e:
            print('[API ERROR] Request failed:')
            print(str(e))
            raise

        if resp.status_code >= 400:
            print('[API ERROR] HTTP', resp.status_code)
            try:
                print(resp.text)
            except Exception:
                pass
            resp.raise_for_status()

        try:
            data = resp.json()
        except Exception:
            print('[API ERROR] Non-JSON response:')
            print(resp.text)
            raise

        text = ''
        try:
            # Standard OpenAI format
            if 'choices' in data and data['choices']:
                choice = data['choices'][0]
                if 'message' in choice and 'content' in choice['message']:
                    text = choice['message']['content'] or ''
                elif 'text' in choice:
                    text = choice['text'] or ''
        except Exception:
            pass

        if not text:
            # Fallback: try common fields
            text = data.get('content') or data.get('message', {}).get('content', '') or ''

        return text, data


def create_client_from_env(env_path: str = 'api.ENV') -> ApiClient:
    env = load_env_file(env_path)
    api_key = env.get('API_KEY', '')
    api_url = env.get('API_URL', '')
    model = env.get('DEFAULT_MODEL', '')
    if not api_key:
        raise ValueError('API_KEY is missing in api.ENV')
    if not api_url:
        print('[WARN] API_URL is empty, defaulting to OpenAI public endpoint')
    if not model:
        raise ValueError('DEFAULT_MODEL is missing in api.ENV')
    return ApiClient(api_url=api_url, api_key=api_key, default_model=model)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='Quick API test using api.ENV config')
    parser.add_argument('--prompt', type=str, default='Say hello in one sentence', help='User prompt')
    parser.add_argument('--system', type=str, default='You are a helpful assistant.', help='System prompt')
    parser.add_argument('--max_tokens', type=int, default=256)
    parser.add_argument('--temperature', type=float, default=0.2)
    args = parser.parse_args()

    client = create_client_from_env('api.ENV')
    messages = [
        {'role': 'system', 'content': args.system},
        {'role': 'user', 'content': args.prompt},
    ]
    text, raw = client.chat_completion(messages, max_tokens=args.max_tokens, temperature=args.temperature)
    print('--- RESPONSE START ---')
    print(text)
    print('--- RESPONSE END ---')