uncloseai.com/languages/python/aiohttp/uncloseai.py

#!/usr/bin/env python3
"""
UncloseAI - Async Python Client (aiohttp)
A Python async client library for OpenAI-compatible APIs with streaming support
Compatible with vLLM, Ollama, and OpenAI-compatible endpoints
"""

import asyncio
import aiohttp
import json
import os
from typing import List, Dict, Optional, AsyncIterator


class UncloseAI:
    """Async client for OpenAI-compatible API endpoints with streaming support"""

    def __init__(
        self,
        model_endpoints: Optional[List[str]] = None,
        tts_endpoints: Optional[List[str]] = None,
        api_key: Optional[str] = None,
        timeout: float = 30.0
    ):
        """
        Initialize UncloseAI async client

        Args:
            model_endpoints: List of model endpoint URLs (defaults to MODEL_ENDPOINT_* env vars)
            tts_endpoints: List of TTS endpoint URLs (defaults to TTS_ENDPOINT_* env vars)
            api_key: Optional API key for authentication
            timeout: Request timeout in seconds
        """
        self.timeout = timeout
        self.api_key = api_key
        self.models: List[Dict] = []
        self.tts_endpoints: List[str] = []
        self._initialized = False
        self._model_endpoints = model_endpoints or self._discover_env_endpoints("MODEL_ENDPOINT")
        self._tts_endpoints = tts_endpoints or self._discover_env_endpoints("TTS_ENDPOINT")

    def _discover_env_endpoints(self, prefix: str) -> List[str]:
        """Discover endpoints from environment variables like PREFIX_1, PREFIX_2, ..."""
        endpoints = []
        for i in range(1, 10000):
            endpoint = os.getenv(f"{prefix}_{i}")
            if not endpoint:
                break
            endpoints.append(endpoint)
        return endpoints

    async def _ensure_initialized(self):
        """Ensure client is initialized with model discovery"""
        if self._initialized:
            return

        async with aiohttp.ClientSession() as session:
            for endpoint in self._model_endpoints:
                await self._discover_models_from_endpoint(session, endpoint)

        self.tts_endpoints = self._tts_endpoints
        self._initialized = True

    async def _discover_models_from_endpoint(self, session: aiohttp.ClientSession, endpoint: str) -> None:
        """Discover available models from an endpoint"""
        try:
            headers = {}
            if self.api_key:
                headers["Authorization"] = f"Bearer {self.api_key}"

            async with session.get(
                f"{endpoint}/models",
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=10)
            ) as response:
                if response.status == 200:
                    data = await response.json()
                    for model in data.get("data", []):
                        model_id = model["id"]

                        # Filter out modelperm-* and chatcmpl-* entries
                        if model_id.startswith("modelperm-") or model_id.startswith("chatcmpl-"):
                            continue

                        self.models.append({
                            "id": model_id,
                            "endpoint": endpoint,
                            "max_tokens": model.get("max_model_len", 8192)
                        })
        except Exception:
            # Silently skip failed endpoints
            pass

    async def list_models(self) -> List[Dict]:
        """Return list of discovered models with their metadata"""
        await self._ensure_initialized()
        return self.models.copy()

    async def chat(
        self,
        messages: List[Dict[str, str]],
        model: Optional[str] = None,
        max_tokens: int = 100,
        temperature: float = 0.7,
        **kwargs
    ) -> Dict:
        """
        Non-streaming chat completion

        Args:
            messages: List of message dicts with 'role' and 'content'
            model: Model ID (defaults to first available model)
            max_tokens: Maximum tokens in response
            temperature: Sampling temperature
            **kwargs: Additional parameters to pass to the API

        Returns:
            Response dict with 'choices' containing the completion
        """
        await self._ensure_initialized()
        model_info = self._get_model_info(model)

        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        payload = {
            "model": model_info["id"],
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "stream": False,
            **kwargs
        }

        async with aiohttp.ClientSession() as session:
            async with session.post(
                f"{model_info['endpoint']}/chat/completions",
                headers=headers,
                json=payload,
                timeout=aiohttp.ClientTimeout(total=self.timeout)
            ) as response:
                response.raise_for_status()
                return await response.json()

    async def chat_stream(
        self,
        messages: List[Dict[str, str]],
        model: Optional[str] = None,
        max_tokens: int = 500,
        temperature: float = 0.7,
        **kwargs
    ) -> AsyncIterator[Dict]:
        """
        Streaming chat completion using Server-Sent Events

        Args:
            messages: List of message dicts with 'role' and 'content'
            model: Model ID (defaults to first available model)
            max_tokens: Maximum tokens in response
            temperature: Sampling temperature
            **kwargs: Additional parameters to pass to the API

        Yields:
            Chunk dicts with 'choices' containing delta content
        """
        await self._ensure_initialized()
        model_info = self._get_model_info(model)

        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        payload = {
            "model": model_info["id"],
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "stream": True,
            **kwargs
        }

        async with aiohttp.ClientSession() as session:
            async with session.post(
                f"{model_info['endpoint']}/chat/completions",
                headers=headers,
                json=payload,
                timeout=aiohttp.ClientTimeout(total=self.timeout)
            ) as response:
                response.raise_for_status()

                async for line in response.content:
                    line_str = line.decode('utf-8').strip()

                    if not line_str:
                        continue

                    # SSE format: "data: {...}"
                    if line_str.startswith('data: '):
                        data = line_str[6:]

                        # Check for stream termination
                        if data.strip() == '[DONE]':
                            break

                        try:
                            chunk = json.loads(data)
                            yield chunk
                        except json.JSONDecodeError:
                            continue

    async def tts(
        self,
        text: str,
        voice: str = "alloy",
        model: str = "tts-1",
        response_format: str = "mp3"
    ) -> bytes:
        """
        Generate speech from text

        Args:
            text: Input text to convert to speech
            voice: Voice name (alloy, echo, fable, onyx, nova, shimmer)
            model: TTS model (tts-1 or tts-1-hd)
            response_format: Audio format (mp3, opus, aac, flac)

        Returns:
            Audio data as bytes
        """
        await self._ensure_initialized()

        if not self.tts_endpoints:
            raise ValueError("No TTS endpoints available")

        endpoint = self.tts_endpoints[0]

        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        payload = {
            "model": model,
            "voice": voice,
            "input": text,
            "response_format": response_format
        }

        async with aiohttp.ClientSession() as session:
            async with session.post(
                f"{endpoint}/audio/speech",
                headers=headers,
                json=payload,
                timeout=aiohttp.ClientTimeout(total=self.timeout)
            ) as response:
                response.raise_for_status()
                return await response.read()

    def _get_model_info(self, model: Optional[str] = None) -> Dict:
        """Get model info by ID or return first available model"""
        if not self.models:
            raise ValueError("No models available. Check endpoint configuration.")

        if model is None:
            return self.models[0]

        for m in self.models:
            if m["id"] == model:
                return m

        raise ValueError(f"Model '{model}' not found in discovered models")


# Demo usage when run as script
async def main():
    print("=== UncloseAI Python Async Client (aiohttp) ===\n")

    # Initialize client (auto-discovers from environment)
    client = UncloseAI()

    models = await client.list_models()
    if not models:
        print("ERROR: No models discovered. Set environment variables:")
        print("  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.")
        return

    print(f"Discovered {len(models)} model(s)")
    for model in models:
        print(f"  - {model['id']} (max_tokens: {model['max_tokens']})")
    print()

    # Non-streaming chat example
    print("=== Non-Streaming Chat ===")
    response = await client.chat(
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant."},
            {"role": "user", "content": "Explain quantum computing in one sentence."}
        ],
        max_tokens=100
    )
    print(f"Model: {response['model']}")
    print(f"Response: {response['choices'][0]['message']['content']}\n")

    # Streaming chat example
    print("=== Streaming Chat ===")
    model_id = models[1]["id"] if len(models) > 1 else None
    print(f"Model: {model_id or models[0]['id']}")
    print("Response: ", end="", flush=True)

    async for chunk in client.chat_stream(
        messages=[
            {"role": "system", "content": "You are a coding assistant."},
            {"role": "user", "content": "Write an async Python function to fetch multiple URLs"}
        ],
        model=model_id,
        max_tokens=200
    ):
        if chunk.get("choices") and len(chunk["choices"]) > 0:
            delta = chunk["choices"][0].get("delta", {})
            content = delta.get("content", "")
            if content:
                print(content, end="", flush=True)

    print("\n")

    # TTS example
    if client.tts_endpoints:
        print("=== TTS Speech Generation ===")
        audio_data = await client.tts(
            text="Hello from UncloseAI Python async client with aiohttp! This demonstrates text to speech with streaming support.",
            voice="alloy"
        )

        with open("speech.mp3", "wb") as f:
            f.write(audio_data)

        print(f"✓ Speech file created: speech.mp3 ({len(audio_data)} bytes)\n")

    print("=== Examples Complete ===")


if __name__ == "__main__":
    asyncio.run(main())