Spaces:
Paused
Paused
File size: 5,101 Bytes
d958e80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """
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)
|