sdk/python/m5_client.py

"""
M5 Python SDK
"""
import httpx
from typing import List, Dict, Optional, AsyncGenerator
from dataclasses import dataclass


@dataclass
class Message:
    role: str
    content: str


@dataclass
class ChatChoice:
    index: int
    message: Dict
    finish_reason: str


@dataclass
class Usage:
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


@dataclass
class ChatCompletion:
    id: str
    object: str
    created: int
    model: str
    choices: List[ChatChoice]
    usage: Optional[Usage] = None


class M5Client:
    """M5 API Client for Python"""
    
    def __init__(self, api_key: str, base_url: str = "https://m5.hf.space"):
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
    
    async def chat_completion(
        self,
        model: str,
        messages: List[Message],
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        stream: bool = False,
        ollama_url: Optional[str] = None
    ) -> ChatCompletion:
        """Create a chat completion."""
        payload = {
            "model": model,
            "messages": [{"role": m.role, "content": m.content} for m in messages],
            "temperature": temperature,
            "stream": stream
        }
        if max_tokens:
            payload["max_tokens"] = max_tokens
        if ollama_url:
            payload["ollama_url"] = ollama_url
        
        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{self.base_url}/v1/chat/completions",
                headers=self.headers,
                json=payload
            )
            response.raise_for_status()
            data = response.json()
            
            return ChatCompletion(
                id=data["id"],
                object=data["object"],
                created=data["created"],
                model=data["model"],
                choices=[
                    ChatChoice(
                        index=c["index"],
                        message=c["message"],
                        finish_reason=c["finish_reason"]
                    )
                    for c in data["choices"]
                ],
                usage=Usage(**data["usage"]) if "usage" in data else None
            )
    
    async def chat_completion_stream(
        self,
        model: str,
        messages: List[Message],
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        ollama_url: Optional[str] = None
    ) -> AsyncGenerator[str, None]:
        """Stream a chat completion."""
        payload = {
            "model": model,
            "messages": [{"role": m.role, "content": m.content} for m in messages],
            "temperature": temperature,
            "stream": True
        }
        if max_tokens:
            payload["max_tokens"] = max_tokens
        if ollama_url:
            payload["ollama_url"] = ollama_url
        
        async with httpx.AsyncClient() as client:
            async with client.stream(
                "POST",
                f"{self.base_url}/v1/chat/completions",
                headers=self.headers,
                json=payload
            ) as response:
                response.raise_for_status()
                async for line in response.aiter_lines():
                    if line.startswith("data: "):
                        data = line[6:]
                        if data == "[DONE]":
                            break
                        try:
                            import json
                            chunk = json.loads(data)
                            if "choices" in chunk and chunk["choices"]:
                                delta = chunk["choices"][0].get("delta", {})
                                if "content" in delta:
                                    yield delta["content"]
                        except:
                            pass
    
    async def list_models(self) -> List[Dict]:
        """List available models."""
        async with httpx.AsyncClient() as client:
            response = await client.get(
                f"{self.base_url}/v1/models",
                headers=self.headers
            )
            response.raise_for_status()
            data = response.json()
            return data.get("data", [])


# Convenience classes
class Chat:
    """Chat completions interface"""
    def __init__(self, client: M5Client):
        self.client = client
    
    async def create(self, **kwargs) -> ChatCompletion:
        return await self.client.chat_completion(**kwargs)
    
    async def create_stream(self, **kwargs) -> AsyncGenerator[str, None]:
        async for chunk in self.client.chat_completion_stream(**kwargs):
            yield chunk


class M5:
    """Main M5 client"""
    def __init__(self, api_key: str, base_url: str = "https://m5.hf.space"):
        self.client = M5Client(api_key, base_url)
        self.chat = Chat(self.client)