125 lines
3.4 KiB
Python
125 lines
3.4 KiB
Python
import aiohttp
|
|
import asyncio
|
|
import json
|
|
import re
|
|
import gconfig
|
|
|
|
try:
|
|
from config import *
|
|
except:
|
|
gconfig.copy_config()
|
|
finally:
|
|
from config import *
|
|
|
|
|
|
def parse_error_body(status: int, body: str) -> str:
|
|
try:
|
|
data = json.loads(body)
|
|
if isinstance(data, dict):
|
|
err = data.get("error")
|
|
if isinstance(err, dict):
|
|
return f"{status} {err.get('message') or err.get('status') or err}"
|
|
if err:
|
|
return f"{status} {err}"
|
|
return f"{status} {data.get('message') or body[:200]}"
|
|
return f"{status} {body[:200]}"
|
|
except Exception:
|
|
pass
|
|
clean = re.sub(r"<[^>]+>", " ", body).strip()
|
|
return f"{status} {clean[:200]}"
|
|
|
|
|
|
async def make_request(url, headers, payload):
|
|
async with aiohttp.ClientSession(headers=headers) as session:
|
|
async with session.post(url, json=payload) as response:
|
|
if response.status != 200:
|
|
error_body = await response.text()
|
|
raise Exception(parse_error_body(response.status, error_body))
|
|
if payload.get("stream") == True:
|
|
async for line in response.content:
|
|
line = line.decode("utf-8").strip()
|
|
if line.startswith("data:"):
|
|
data = line[5:].strip()
|
|
if data == "DONE":
|
|
break
|
|
else:
|
|
# data = await response.text()
|
|
yield data
|
|
else:
|
|
response_data = b""
|
|
async for chunk in response.content.iter_chunks():
|
|
response_data += chunk[0]
|
|
yield response_data.decode()
|
|
|
|
async def completions(
|
|
url,
|
|
api_key,
|
|
model,
|
|
messages,
|
|
stream=False,
|
|
max_tokens=MAX_TOKENS,
|
|
extra=None
|
|
):
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
|
|
payload = {
|
|
"messages": messages,
|
|
"model": model,
|
|
"stream": stream,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
|
|
if extra:
|
|
payload = payload | extra
|
|
|
|
if stream:
|
|
async for data in make_request(url+"/chat/completions", headers, payload):
|
|
try:
|
|
yield json.loads(data)
|
|
except:
|
|
yield data
|
|
else:
|
|
async for data in make_request(url+"/chat/completions", headers, payload):
|
|
try:
|
|
yield json.loads(data)
|
|
except:
|
|
yield data
|
|
|
|
async def responses(
|
|
url,
|
|
api_key,
|
|
model,
|
|
input,
|
|
stream=False,
|
|
max_tokens=MAX_TOKENS,
|
|
extra=None
|
|
):
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
payload = {
|
|
"model": model,
|
|
"input": input,
|
|
"stream": stream,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
if extra:
|
|
payload = payload | extra
|
|
|
|
if stream:
|
|
async for data in make_request(url+"/chat/completions", headers, payload):
|
|
try:
|
|
yield json.loads(data)
|
|
except:
|
|
yield data
|
|
else:
|
|
async for data in make_request(url+"/chat/completions", headers, payload):
|
|
try:
|
|
yield json.loads(data)
|
|
except:
|
|
yield data
|