DarkAI.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. from __future__ import annotations
  2. import json
  3. from aiohttp import ClientSession
  4. from ..typing import AsyncResult, Messages
  5. from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
  6. from .helper import format_prompt
  7. class DarkAI(AsyncGeneratorProvider, ProviderModelMixin):
  8. url = "https://darkai.foundation/chat"
  9. api_endpoint = "https://darkai.foundation/chat"
  10. working = True
  11. supports_stream = True
  12. supports_system_message = True
  13. supports_message_history = True
  14. default_model = 'llama-3-405b'
  15. models = [
  16. 'gpt-4o', # Uncensored
  17. 'gpt-3.5-turbo', # Uncensored
  18. 'llama-3-70b', # Uncensored
  19. default_model,
  20. ]
  21. model_aliases = {
  22. "llama-3.1-70b": "llama-3-70b",
  23. "llama-3.1-405b": "llama-3-405b",
  24. }
  25. @classmethod
  26. def get_model(cls, model: str) -> str:
  27. if model in cls.models:
  28. return model
  29. elif model in cls.model_aliases:
  30. return cls.model_aliases[model]
  31. else:
  32. return cls.default_model
  33. @classmethod
  34. async def create_async_generator(
  35. cls,
  36. model: str,
  37. messages: Messages,
  38. proxy: str = None,
  39. **kwargs
  40. ) -> AsyncResult:
  41. model = cls.get_model(model)
  42. headers = {
  43. "accept": "text/event-stream",
  44. "content-type": "application/json",
  45. "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36"
  46. }
  47. async with ClientSession(headers=headers) as session:
  48. prompt = format_prompt(messages)
  49. data = {
  50. "query": prompt,
  51. "model": model,
  52. }
  53. async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
  54. response.raise_for_status()
  55. full_text = ""
  56. async for chunk in response.content:
  57. if chunk:
  58. try:
  59. chunk_str = chunk.decode().strip()
  60. if chunk_str.startswith('data: '):
  61. chunk_data = json.loads(chunk_str[6:])
  62. if chunk_data['event'] == 'text-chunk':
  63. full_text += chunk_data['data']['text']
  64. elif chunk_data['event'] == 'stream-end':
  65. if full_text:
  66. yield full_text.strip()
  67. return
  68. except json.JSONDecodeError:
  69. pass
  70. except Exception:
  71. pass
  72. if full_text:
  73. yield full_text.strip()