Upstage.py 3.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  1. from __future__ import annotations
  2. from aiohttp import ClientSession
  3. import json
  4. from ..typing import AsyncResult, Messages
  5. from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
  6. from .helper import format_prompt
  7. class Upstage(AsyncGeneratorProvider, ProviderModelMixin):
  8. url = "https://console.upstage.ai/playground/chat"
  9. api_endpoint = "https://ap-northeast-2.apistage.ai/v1/web/demo/chat/completions"
  10. working = True
  11. default_model = 'solar-pro'
  12. models = [
  13. 'upstage/solar-1-mini-chat',
  14. 'upstage/solar-1-mini-chat-ja',
  15. 'solar-pro',
  16. ]
  17. model_aliases = {
  18. "solar-mini": "upstage/solar-1-mini-chat",
  19. "solar-mini": "upstage/solar-1-mini-chat-ja",
  20. }
  21. @classmethod
  22. def get_model(cls, model: str) -> str:
  23. if model in cls.models:
  24. return model
  25. elif model in cls.model_aliases:
  26. return cls.model_aliases[model]
  27. else:
  28. return cls.default_model
  29. @classmethod
  30. async def create_async_generator(
  31. cls,
  32. model: str,
  33. messages: Messages,
  34. proxy: str = None,
  35. **kwargs
  36. ) -> AsyncResult:
  37. model = cls.get_model(model)
  38. headers = {
  39. "accept": "*/*",
  40. "accept-language": "en-US,en;q=0.9",
  41. "cache-control": "no-cache",
  42. "content-type": "application/json",
  43. "dnt": "1",
  44. "origin": "https://console.upstage.ai",
  45. "pragma": "no-cache",
  46. "priority": "u=1, i",
  47. "referer": "https://console.upstage.ai/",
  48. "sec-ch-ua": '"Not?A_Brand";v="99", "Chromium";v="130"',
  49. "sec-ch-ua-mobile": "?0",
  50. "sec-ch-ua-platform": '"Linux"',
  51. "sec-fetch-dest": "empty",
  52. "sec-fetch-mode": "cors",
  53. "sec-fetch-site": "cross-site",
  54. "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36"
  55. }
  56. async with ClientSession(headers=headers) as session:
  57. data = {
  58. "stream": True,
  59. "messages": [{"role": "user", "content": format_prompt(messages)}],
  60. "model": model
  61. }
  62. async with session.post(f"{cls.api_endpoint}", json=data, proxy=proxy) as response:
  63. response.raise_for_status()
  64. response_text = ""
  65. async for line in response.content:
  66. if line:
  67. line = line.decode('utf-8').strip()
  68. if line.startswith("data: ") and line != "data: [DONE]":
  69. try:
  70. data = json.loads(line[6:])
  71. content = data['choices'][0]['delta'].get('content', '')
  72. if content:
  73. response_text += content
  74. yield content
  75. except json.JSONDecodeError:
  76. continue
  77. if line == "data: [DONE]":
  78. break