ThebApi.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. from __future__ import annotations
  2. from ...typing import CreateResult, Messages
  3. from ..helper import filter_none
  4. from ..template import OpenaiTemplate
  5. models = {
  6. "theb-ai": "TheB.AI",
  7. "gpt-3.5-turbo": "GPT-3.5",
  8. "gpt-4-turbo": "GPT-4 Turbo",
  9. "gpt-4": "GPT-4",
  10. "claude-3.5-sonnet": "Claude",
  11. "llama-2-7b-chat": "Llama 2 7B",
  12. "llama-2-13b-chat": "Llama 2 13B",
  13. "llama-2-70b-chat": "Llama 2 70B",
  14. "code-llama-7b": "Code Llama 7B",
  15. "code-llama-13b": "Code Llama 13B",
  16. "code-llama-34b": "Code Llama 34B",
  17. "qwen-2-72b": "Qwen"
  18. }
  19. class ThebApi(OpenaiTemplate):
  20. label = "TheB.AI API"
  21. url = "https://theb.ai"
  22. login_url = "https://beta.theb.ai/home"
  23. api_base = "https://api.theb.ai/v1"
  24. working = True
  25. needs_auth = True
  26. default_model = "theb-ai"
  27. fallback_models = list(models)
  28. @classmethod
  29. def create_async_generator(
  30. cls,
  31. model: str,
  32. messages: Messages,
  33. temperature: float = None,
  34. top_p: float = None,
  35. **kwargs
  36. ) -> CreateResult:
  37. system_message = "\n".join([message["content"] for message in messages if message["role"] == "system"])
  38. messages = [message for message in messages if message["role"] != "system"]
  39. data = {
  40. "model_params": filter_none(
  41. system_prompt=system_message,
  42. temperature=temperature,
  43. top_p=top_p,
  44. )
  45. }
  46. return super().create_async_generator(model, messages, extra_data=data, **kwargs)