nova-api/api/providers/__main__.py

51 lines
1.2 KiB
Python
Raw Normal View History

2023-10-05 15:06:33 +02:00
import os
import sys
from rich import print
def remove_duplicate_keys(file):
with open(file, 'r', encoding='utf8') as f:
2023-10-06 16:42:42 +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)
try:
provider_name = sys.argv[1]
if provider_name == '--clear':
2023-10-06 09:37:16 +02:00
for file in os.listdir('secret'):
2023-10-05 15:06:33 +02:00
if file.endswith('.txt'):
2023-10-06 09:37:16 +02:00
remove_duplicate_keys(os.path.join('secret', file))
2023-10-05 15:06:33 +02:00
exit()
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('_'):
print(file_name.split('.')[0])
sys.exit(0)
try:
provider = __import__(provider_name)
except ModuleNotFoundError as exc:
print(f'Provider "{provider_name}" not found.')
print('Available providers:')
for file_name in os.listdir(os.path.dirname(__file__)):
if file_name.endswith('.py') and not file_name.startswith('_'):
print(file_name.split('.')[0])
sys.exit(1)
if len(sys.argv) > 2:
model = sys.argv[2]
else:
model = provider.MODELS[-1]
print(f'{provider_name} @ {model}')
comp = provider.chat_completion(model=model)
print(comp)