nova-api/api/providers/__main__.py

57 lines
1.5 KiB
Python
Raw Normal View History

2023-10-05 15:06:33 +02:00
import os
import sys
2023-10-06 23:05:38 +02:00
import aiohttp
import asyncio
import importlib
2023-10-06 09:37:16 +02:00
import aiofiles.os
2023-10-05 15:06:33 +02:00
from rich import print
def remove_duplicate_keys(file):
with open(file, 'r', encoding='utf8') as f:
2023-10-06 09:37:16 +02:00
unique_lines = set(f)
2023-10-05 15:06:33 +02:00
with open(file, 'w', encoding='utf8') as f:
f.writelines(unique_lines)
2023-10-06 23:05:38 +02:00
async def main():
try:
provider_name = sys.argv[1]
except IndexError:
print('List of available providers:')
2023-10-06 09:37:16 +02:00
for file_name in await aiofiles.os.listdir(os.path.dirname(__file__)):
2023-10-06 23:05:38 +02:00
if file_name.endswith('.py') and not file_name.startswith('_'):
print(file_name.split('.')[0])
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:
res_json = await response.json()
print(response.status, res_json)
asyncio.run(main())