104 lines
2.6 KiB
Python
104 lines
2.6 KiB
Python
|
import aiohttp
|
||
|
|
import asyncio
|
||
|
|
import json
|
||
|
|
import gconfig
|
||
|
|
|
||
|
|
try:
|
||
|
|
from config import *
|
||
|
|
except:
|
||
|
|
gconfig.copy_config()
|
||
|
|
finally:
|
||
|
|
from config import *
|
||
|
|
|
||
|
|
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 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
|
||
|
|
):
|
||
|
|
print(url)
|
||
|
|
|
||
|
|
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
|