nova-api/api/providers/__main__.py
2023-10-16 21:34:54 +00:00

62 lines
1.7 KiB
Python

import os
import sys
import aiohttp
import asyncio
import importlib
from rich import print
def remove_duplicate_keys(file):
with open(file, 'r', encoding='utf8') as f:
unique_lines = set(f)
with open(file, 'w', encoding='utf8') as f:
f.writelines(unique_lines)
async def main():
try:
provider_name = sys.argv[1]
except IndexError:
print('List of available providers:')
for file_name in os.listdir(os.path.dirname(__file__)):
if file_name.endswith('.py') and not file_name.startswith('_'):
name = file_name.split('.')[0]
models = importlib.import_module(f'.{file_name.split(".")[0]}', 'providers').MODELS
print(f' {name} @ {", ".join(models)}')
sys.exit(0)
try:
provider = importlib.import_module(f'.{provider_name}', 'providers')
except ModuleNotFoundError as exc:
print(exc)
sys.exit(1)
if len(sys.argv) > 2:
model = sys.argv[2] # choose a specific model
else:
model = provider.MODELS[-1] # choose best model
print(f'{provider_name} @ {model}')
req = await provider.chat_completion(model=model, messages=[{'role': 'user', 'content': '1+1='}])
print(req)
# launch aiohttp
async with aiohttp.ClientSession() as session:
async with session.request(
method=req['method'],
url=req['url'],
headers=req['headers'],
json=req['payload'],
) as response:
try:
res_json = await response.json()
except aiohttp.ContentTypeError:
res_json = await response.text()
print(response.status, res_json)
asyncio.run(main())