nova-api/api/chat_balancing.py

31 lines
765 B
Python

import random
import asyncio
import chat_providers
provider_modules = [
# chat_providers.twa,
chat_providers.quantum,
# chat_providers.churchless,
chat_providers.closed
]
async def balance(payload: dict) -> dict:
providers_available = []
for provider_module in provider_modules:
if payload['stream'] and not provider_module.STREAMING:
continue
if payload['model'] not in provider_module.MODELS:
continue
providers_available.append(provider_module)
provider = random.choice(providers_available)
return provider.chat_completion(**payload)
if __name__ == '__main__':
req = asyncio.run(balance(payload={'model': 'gpt-3.5-turbo', 'stream': True}))
print(req['url'])