llm.py 3.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. # coding=utf-8
  2. """
  3. @project: MaxKB
  4. @Author:虎
  5. @file: llm.py
  6. @date:2024/7/11 18:32
  7. @desc:
  8. """
  9. from typing import Dict
  10. from django.utils.translation import gettext_lazy as _, gettext
  11. from langchain_core.messages import HumanMessage
  12. from common import forms
  13. from common.exception.app_exception import AppApiException
  14. from common.forms import BaseForm, TooltipLabel
  15. from models_provider.base_model_provider import BaseModelCredential, ValidCode
  16. from common.utils.logger import maxkb_logger
  17. class RegoloLLMModelParams(BaseForm):
  18. temperature = forms.SliderField(TooltipLabel(_('Temperature'),
  19. _('Higher values make the output more random, while lower values make it more focused and deterministic')),
  20. required=True, default_value=0.7,
  21. _min=0.1,
  22. _max=1.0,
  23. _step=0.01,
  24. precision=2)
  25. max_tokens = forms.SliderField(
  26. TooltipLabel(_('Output the maximum Tokens'),
  27. _('Specify the maximum number of tokens that the model can generate')),
  28. required=True, default_value=800,
  29. _min=1,
  30. _max=100000,
  31. _step=1,
  32. precision=0)
  33. class RegoloLLMModelCredential(BaseForm, BaseModelCredential):
  34. def is_valid(self, model_type: str, model_name, model_credential: Dict[str, object], model_params, provider,
  35. raise_exception=False):
  36. model_type_list = provider.get_model_type_list()
  37. if not any(list(filter(lambda mt: mt.get('value') == model_type, model_type_list))):
  38. raise AppApiException(ValidCode.valid_error.value,
  39. gettext('{model_type} Model type is not supported').format(model_type=model_type))
  40. for key in ['api_key', 'api_base']:
  41. if key not in model_credential:
  42. if raise_exception:
  43. raise AppApiException(ValidCode.valid_error.value, gettext('{key} is required').format(key=key))
  44. else:
  45. return False
  46. try:
  47. model = provider.get_model(model_type, model_name, model_credential, **model_params)
  48. model.invoke([HumanMessage(content=gettext('Hello'))])
  49. except Exception as e:
  50. maxkb_logger.error(f'Exception: {e}', exc_info=True)
  51. if isinstance(e, AppApiException):
  52. raise e
  53. if raise_exception:
  54. raise AppApiException(ValidCode.valid_error.value,
  55. gettext(
  56. 'Verification failed, please check whether the parameters are correct: {error}').format(
  57. error=str(e)))
  58. else:
  59. return False
  60. return True
  61. def encryption_dict(self, model: Dict[str, object]):
  62. return {**model, 'api_key': super().encryption(model.get('api_key', ''))}
  63. api_base = forms.TextInputField('API URL', required=True, default_value='https://api.regolo.ai/v1')
  64. api_key = forms.PasswordInputField('API Key', required=True)
  65. def get_model_params_setting_form(self, model_name):
  66. return RegoloLLMModelParams()