xinference_provider.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. import json
  2. from typing import Type
  3. from langchain.llms import Xinference
  4. from core.helper import encrypter
  5. from core.model_providers.models.embedding.xinference_embedding import XinferenceEmbedding
  6. from core.model_providers.models.entity.model_params import KwargRule, ModelKwargsRules, ModelType
  7. from core.model_providers.models.llm.xinference_model import XinferenceModel
  8. from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
  9. from core.model_providers.models.base import BaseProviderModel
  10. from models.provider import ProviderType
  11. class XinferenceProvider(BaseModelProvider):
  12. @property
  13. def provider_name(self):
  14. """
  15. Returns the name of a provider.
  16. """
  17. return 'xinference'
  18. def _get_fixed_model_list(self, model_type: ModelType) -> list[dict]:
  19. return []
  20. def get_model_class(self, model_type: ModelType) -> Type[BaseProviderModel]:
  21. """
  22. Returns the model class.
  23. :param model_type:
  24. :return:
  25. """
  26. if model_type == ModelType.TEXT_GENERATION:
  27. model_class = XinferenceModel
  28. elif model_type == ModelType.EMBEDDINGS:
  29. model_class = XinferenceEmbedding
  30. else:
  31. raise NotImplementedError
  32. return model_class
  33. def get_model_parameter_rules(self, model_name: str, model_type: ModelType) -> ModelKwargsRules:
  34. """
  35. get model parameter rules.
  36. :param model_name:
  37. :param model_type:
  38. :return:
  39. """
  40. return ModelKwargsRules(
  41. temperature=KwargRule[float](min=0, max=2, default=1),
  42. top_p=KwargRule[float](min=0, max=1, default=0.7),
  43. presence_penalty=KwargRule[float](min=-2, max=2, default=0),
  44. frequency_penalty=KwargRule[float](min=-2, max=2, default=0),
  45. max_tokens=KwargRule[int](min=10, max=4000, default=256),
  46. )
  47. @classmethod
  48. def is_model_credentials_valid_or_raise(cls, model_name: str, model_type: ModelType, credentials: dict):
  49. """
  50. check model credentials valid.
  51. :param model_name:
  52. :param model_type:
  53. :param credentials:
  54. """
  55. if 'server_url' not in credentials:
  56. raise CredentialsValidateFailedError('Xinference Server URL must be provided.')
  57. if 'model_uid' not in credentials:
  58. raise CredentialsValidateFailedError('Xinference Model UID must be provided.')
  59. try:
  60. credential_kwargs = {
  61. 'server_url': credentials['server_url'],
  62. 'model_uid': credentials['model_uid'],
  63. }
  64. llm = Xinference(
  65. **credential_kwargs
  66. )
  67. llm("ping", generate_config={'max_tokens': 10})
  68. except Exception as ex:
  69. raise CredentialsValidateFailedError(str(ex))
  70. @classmethod
  71. def encrypt_model_credentials(cls, tenant_id: str, model_name: str, model_type: ModelType,
  72. credentials: dict) -> dict:
  73. """
  74. encrypt model credentials for save.
  75. :param tenant_id:
  76. :param model_name:
  77. :param model_type:
  78. :param credentials:
  79. :return:
  80. """
  81. credentials['server_url'] = encrypter.encrypt_token(tenant_id, credentials['server_url'])
  82. return credentials
  83. def get_model_credentials(self, model_name: str, model_type: ModelType, obfuscated: bool = False) -> dict:
  84. """
  85. get credentials for llm use.
  86. :param model_name:
  87. :param model_type:
  88. :param obfuscated:
  89. :return:
  90. """
  91. if self.provider.provider_type != ProviderType.CUSTOM.value:
  92. raise NotImplementedError
  93. provider_model = self._get_provider_model(model_name, model_type)
  94. if not provider_model.encrypted_config:
  95. return {
  96. 'server_url': None,
  97. 'model_uid': None,
  98. }
  99. credentials = json.loads(provider_model.encrypted_config)
  100. if credentials['server_url']:
  101. credentials['server_url'] = encrypter.decrypt_token(
  102. self.provider.tenant_id,
  103. credentials['server_url']
  104. )
  105. if obfuscated:
  106. credentials['server_url'] = encrypter.obfuscated_token(credentials['server_url'])
  107. return credentials
  108. @classmethod
  109. def is_provider_credentials_valid_or_raise(cls, credentials: dict):
  110. return
  111. @classmethod
  112. def encrypt_provider_credentials(cls, tenant_id: str, credentials: dict) -> dict:
  113. return {}
  114. def get_provider_credentials(self, obfuscated: bool = False) -> dict:
  115. return {}