Aura.py 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748
  1. from __future__ import annotations
  2. from aiohttp import ClientSession
  3. from ...typing import AsyncResult, Messages
  4. from ..base_provider import AsyncGeneratorProvider
  5. class Aura(AsyncGeneratorProvider):
  6. url = "https://openchat.team"
  7. working = False
  8. @classmethod
  9. async def create_async_generator(
  10. cls,
  11. model: str,
  12. messages: Messages,
  13. proxy: str = None,
  14. temperature: float = 0.5,
  15. max_tokens: int = 8192,
  16. webdriver = None,
  17. **kwargs
  18. ) -> AsyncResult:
  19. args = get_args_from_browser(cls.url, webdriver, proxy)
  20. async with ClientSession(**args) as session:
  21. new_messages = []
  22. system_message = []
  23. for message in messages:
  24. if message["role"] == "system":
  25. system_message.append(message["content"])
  26. else:
  27. new_messages.append(message)
  28. data = {
  29. "model": {
  30. "id": "openchat_3.6",
  31. "name": "OpenChat 3.6 (latest)",
  32. "maxLength": 24576,
  33. "tokenLimit": max_tokens
  34. },
  35. "messages": new_messages,
  36. "key": "",
  37. "prompt": "\n".join(system_message),
  38. "temperature": temperature
  39. }
  40. async with session.post(f"{cls.url}/api/chat", json=data, proxy=proxy) as response:
  41. response.raise_for_status()
  42. async for chunk in response.content.iter_any():
  43. yield chunk.decode(error="ignore")