spark_provider.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. import json
  2. import logging
  3. from json import JSONDecodeError
  4. from typing import Type
  5. from flask import current_app
  6. from langchain.schema import HumanMessage
  7. from core.helper import encrypter
  8. from core.model_providers.models.base import BaseProviderModel
  9. from core.model_providers.models.entity.model_params import ModelKwargsRules, KwargRule, ModelType
  10. from core.model_providers.models.llm.spark_model import SparkModel
  11. from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
  12. from core.third_party.langchain.llms.spark import ChatSpark
  13. from core.third_party.spark.spark_llm import SparkError
  14. from models.provider import ProviderType, ProviderQuotaType
  15. class SparkProvider(BaseModelProvider):
  16. @property
  17. def provider_name(self):
  18. """
  19. Returns the name of a provider.
  20. """
  21. return 'spark'
  22. def _get_fixed_model_list(self, model_type: ModelType) -> list[dict]:
  23. if model_type == ModelType.TEXT_GENERATION:
  24. return [
  25. {
  26. 'id': 'spark',
  27. 'name': '星火认知大模型',
  28. }
  29. ]
  30. else:
  31. return []
  32. def get_model_class(self, model_type: ModelType) -> Type[BaseProviderModel]:
  33. """
  34. Returns the model class.
  35. :param model_type:
  36. :return:
  37. """
  38. if model_type == ModelType.TEXT_GENERATION:
  39. model_class = SparkModel
  40. else:
  41. raise NotImplementedError
  42. return model_class
  43. def get_model_parameter_rules(self, model_name: str, model_type: ModelType) -> ModelKwargsRules:
  44. """
  45. get model parameter rules.
  46. :param model_name:
  47. :param model_type:
  48. :return:
  49. """
  50. return ModelKwargsRules(
  51. temperature=KwargRule[float](min=0, max=1, default=0.5),
  52. top_p=KwargRule[float](enabled=False),
  53. presence_penalty=KwargRule[float](enabled=False),
  54. frequency_penalty=KwargRule[float](enabled=False),
  55. max_tokens=KwargRule[int](min=10, max=4096, default=2048),
  56. )
  57. @classmethod
  58. def is_provider_credentials_valid_or_raise(cls, credentials: dict):
  59. """
  60. Validates the given credentials.
  61. """
  62. if 'app_id' not in credentials:
  63. raise CredentialsValidateFailedError('Spark app_id must be provided.')
  64. if 'api_key' not in credentials:
  65. raise CredentialsValidateFailedError('Spark api_key must be provided.')
  66. if 'api_secret' not in credentials:
  67. raise CredentialsValidateFailedError('Spark api_secret must be provided.')
  68. try:
  69. credential_kwargs = {
  70. 'app_id': credentials['app_id'],
  71. 'api_key': credentials['api_key'],
  72. 'api_secret': credentials['api_secret'],
  73. }
  74. chat_llm = ChatSpark(
  75. max_tokens=10,
  76. temperature=0.01,
  77. **credential_kwargs
  78. )
  79. messages = [
  80. HumanMessage(
  81. content="ping"
  82. )
  83. ]
  84. chat_llm(messages)
  85. except SparkError as ex:
  86. raise CredentialsValidateFailedError(str(ex))
  87. except Exception as ex:
  88. logging.exception('Spark config validation failed')
  89. raise ex
  90. @classmethod
  91. def encrypt_provider_credentials(cls, tenant_id: str, credentials: dict) -> dict:
  92. credentials['api_key'] = encrypter.encrypt_token(tenant_id, credentials['api_key'])
  93. credentials['api_secret'] = encrypter.encrypt_token(tenant_id, credentials['api_secret'])
  94. return credentials
  95. def get_provider_credentials(self, obfuscated: bool = False) -> dict:
  96. if self.provider.provider_type == ProviderType.CUSTOM.value \
  97. or (self.provider.provider_type == ProviderType.SYSTEM.value
  98. and self.provider.quota_type == ProviderQuotaType.FREE.value):
  99. try:
  100. credentials = json.loads(self.provider.encrypted_config)
  101. except JSONDecodeError:
  102. credentials = {
  103. 'app_id': None,
  104. 'api_key': None,
  105. 'api_secret': None,
  106. }
  107. if credentials['api_key']:
  108. credentials['api_key'] = encrypter.decrypt_token(
  109. self.provider.tenant_id,
  110. credentials['api_key']
  111. )
  112. if obfuscated:
  113. credentials['api_key'] = encrypter.obfuscated_token(credentials['api_key'])
  114. if credentials['api_secret']:
  115. credentials['api_secret'] = encrypter.decrypt_token(
  116. self.provider.tenant_id,
  117. credentials['api_secret']
  118. )
  119. if obfuscated:
  120. credentials['api_secret'] = encrypter.obfuscated_token(credentials['api_secret'])
  121. return credentials
  122. else:
  123. return {
  124. 'app_id': None,
  125. 'api_key': None,
  126. 'api_secret': None,
  127. }
  128. def should_deduct_quota(self):
  129. return True
  130. @classmethod
  131. def is_model_credentials_valid_or_raise(cls, model_name: str, model_type: ModelType, credentials: dict):
  132. """
  133. check model credentials valid.
  134. :param model_name:
  135. :param model_type:
  136. :param credentials:
  137. """
  138. return
  139. @classmethod
  140. def encrypt_model_credentials(cls, tenant_id: str, model_name: str, model_type: ModelType,
  141. credentials: dict) -> dict:
  142. """
  143. encrypt model credentials for save.
  144. :param tenant_id:
  145. :param model_name:
  146. :param model_type:
  147. :param credentials:
  148. :return:
  149. """
  150. return {}
  151. def get_model_credentials(self, model_name: str, model_type: ModelType, obfuscated: bool = False) -> dict:
  152. """
  153. get credentials for llm use.
  154. :param model_name:
  155. :param model_type:
  156. :param obfuscated:
  157. :return:
  158. """
  159. return self.get_provider_credentials(obfuscated)