base.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481
  1. import json
  2. import os
  3. import re
  4. from abc import abstractmethod
  5. from typing import List, Optional, Any, Union, Tuple
  6. import decimal
  7. from langchain.callbacks.manager import Callbacks
  8. from langchain.memory.chat_memory import BaseChatMemory
  9. from langchain.schema import LLMResult, SystemMessage, AIMessage, HumanMessage, BaseMessage, ChatGeneration
  10. from core.callback_handler.std_out_callback_handler import DifyStreamingStdOutCallbackHandler, DifyStdOutCallbackHandler
  11. from core.helper import moderation
  12. from core.model_providers.models.base import BaseProviderModel
  13. from core.model_providers.models.entity.message import PromptMessage, MessageType, LLMRunResult, to_prompt_messages
  14. from core.model_providers.models.entity.model_params import ModelType, ModelKwargs, ModelMode, ModelKwargsRules
  15. from core.model_providers.providers.base import BaseModelProvider
  16. from core.prompt.prompt_builder import PromptBuilder
  17. from core.prompt.prompt_template import JinjaPromptTemplate
  18. from core.third_party.langchain.llms.fake import FakeLLM
  19. import logging
  20. logger = logging.getLogger(__name__)
  21. class BaseLLM(BaseProviderModel):
  22. model_mode: ModelMode = ModelMode.COMPLETION
  23. name: str
  24. model_kwargs: ModelKwargs
  25. credentials: dict
  26. streaming: bool = False
  27. type: ModelType = ModelType.TEXT_GENERATION
  28. deduct_quota: bool = True
  29. def __init__(self, model_provider: BaseModelProvider,
  30. name: str,
  31. model_kwargs: ModelKwargs,
  32. streaming: bool = False,
  33. callbacks: Callbacks = None):
  34. self.name = name
  35. self.model_rules = model_provider.get_model_parameter_rules(name, self.type)
  36. self.model_kwargs = model_kwargs if model_kwargs else ModelKwargs(
  37. max_tokens=None,
  38. temperature=None,
  39. top_p=None,
  40. presence_penalty=None,
  41. frequency_penalty=None
  42. )
  43. self.credentials = model_provider.get_model_credentials(
  44. model_name=name,
  45. model_type=self.type
  46. )
  47. self.streaming = streaming
  48. if streaming:
  49. default_callback = DifyStreamingStdOutCallbackHandler()
  50. else:
  51. default_callback = DifyStdOutCallbackHandler()
  52. if not callbacks:
  53. callbacks = [default_callback]
  54. else:
  55. callbacks.append(default_callback)
  56. self.callbacks = callbacks
  57. client = self._init_client()
  58. super().__init__(model_provider, client)
  59. @abstractmethod
  60. def _init_client(self) -> Any:
  61. raise NotImplementedError
  62. @property
  63. def base_model_name(self) -> str:
  64. """
  65. get llm base model name
  66. :return: str
  67. """
  68. return self.name
  69. @property
  70. def price_config(self) -> dict:
  71. def get_or_default():
  72. default_price_config = {
  73. 'prompt': decimal.Decimal('0'),
  74. 'completion': decimal.Decimal('0'),
  75. 'unit': decimal.Decimal('0'),
  76. 'currency': 'USD'
  77. }
  78. rules = self.model_provider.get_rules()
  79. price_config = rules['price_config'][
  80. self.base_model_name] if 'price_config' in rules else default_price_config
  81. price_config = {
  82. 'prompt': decimal.Decimal(price_config['prompt']),
  83. 'completion': decimal.Decimal(price_config['completion']),
  84. 'unit': decimal.Decimal(price_config['unit']),
  85. 'currency': price_config['currency']
  86. }
  87. return price_config
  88. self._price_config = self._price_config if hasattr(self, '_price_config') else get_or_default()
  89. logger.debug(f"model: {self.name} price_config: {self._price_config}")
  90. return self._price_config
  91. def run(self, messages: List[PromptMessage],
  92. stop: Optional[List[str]] = None,
  93. callbacks: Callbacks = None,
  94. **kwargs) -> LLMRunResult:
  95. """
  96. run predict by prompt messages and stop words.
  97. :param messages:
  98. :param stop:
  99. :param callbacks:
  100. :return:
  101. """
  102. moderation_result = moderation.check_moderation(
  103. self.model_provider,
  104. "\n".join([message.content for message in messages])
  105. )
  106. if not moderation_result:
  107. kwargs['fake_response'] = "I apologize for any confusion, " \
  108. "but I'm an AI assistant to be helpful, harmless, and honest."
  109. if self.deduct_quota:
  110. self.model_provider.check_quota_over_limit()
  111. if not callbacks:
  112. callbacks = self.callbacks
  113. else:
  114. callbacks.extend(self.callbacks)
  115. if 'fake_response' in kwargs and kwargs['fake_response']:
  116. prompts = self._get_prompt_from_messages(messages, ModelMode.CHAT)
  117. fake_llm = FakeLLM(
  118. response=kwargs['fake_response'],
  119. num_token_func=self.get_num_tokens,
  120. streaming=self.streaming,
  121. callbacks=callbacks
  122. )
  123. result = fake_llm.generate([prompts])
  124. else:
  125. try:
  126. result = self._run(
  127. messages=messages,
  128. stop=stop,
  129. callbacks=callbacks if not (self.streaming and not self.support_streaming) else None,
  130. **kwargs
  131. )
  132. except Exception as ex:
  133. raise self.handle_exceptions(ex)
  134. if isinstance(result.generations[0][0], ChatGeneration):
  135. completion_content = result.generations[0][0].message.content
  136. else:
  137. completion_content = result.generations[0][0].text
  138. if self.streaming and not self.support_streaming:
  139. # use FakeLLM to simulate streaming when current model not support streaming but streaming is True
  140. prompts = self._get_prompt_from_messages(messages, ModelMode.CHAT)
  141. fake_llm = FakeLLM(
  142. response=completion_content,
  143. num_token_func=self.get_num_tokens,
  144. streaming=self.streaming,
  145. callbacks=callbacks
  146. )
  147. fake_llm.generate([prompts])
  148. if result.llm_output and result.llm_output['token_usage']:
  149. prompt_tokens = result.llm_output['token_usage']['prompt_tokens']
  150. completion_tokens = result.llm_output['token_usage']['completion_tokens']
  151. total_tokens = result.llm_output['token_usage']['total_tokens']
  152. else:
  153. prompt_tokens = self.get_num_tokens(messages)
  154. completion_tokens = self.get_num_tokens(
  155. [PromptMessage(content=completion_content, type=MessageType.ASSISTANT)])
  156. total_tokens = prompt_tokens + completion_tokens
  157. self.model_provider.update_last_used()
  158. if self.deduct_quota:
  159. self.model_provider.deduct_quota(total_tokens)
  160. return LLMRunResult(
  161. content=completion_content,
  162. prompt_tokens=prompt_tokens,
  163. completion_tokens=completion_tokens
  164. )
  165. @abstractmethod
  166. def _run(self, messages: List[PromptMessage],
  167. stop: Optional[List[str]] = None,
  168. callbacks: Callbacks = None,
  169. **kwargs) -> LLMResult:
  170. """
  171. run predict by prompt messages and stop words.
  172. :param messages:
  173. :param stop:
  174. :param callbacks:
  175. :return:
  176. """
  177. raise NotImplementedError
  178. @abstractmethod
  179. def get_num_tokens(self, messages: List[PromptMessage]) -> int:
  180. """
  181. get num tokens of prompt messages.
  182. :param messages:
  183. :return:
  184. """
  185. raise NotImplementedError
  186. def calc_tokens_price(self, tokens: int, message_type: MessageType) -> decimal.Decimal:
  187. """
  188. calc tokens total price.
  189. :param tokens:
  190. :param message_type:
  191. :return:
  192. """
  193. if message_type == MessageType.HUMAN or message_type == MessageType.SYSTEM:
  194. unit_price = self.price_config['prompt']
  195. else:
  196. unit_price = self.price_config['completion']
  197. unit = self.get_price_unit(message_type)
  198. total_price = tokens * unit_price * unit
  199. total_price = total_price.quantize(decimal.Decimal('0.0000001'), rounding=decimal.ROUND_HALF_UP)
  200. logging.debug(f"tokens={tokens}, unit_price={unit_price}, unit={unit}, total_price:{total_price}")
  201. return total_price
  202. def get_tokens_unit_price(self, message_type: MessageType) -> decimal.Decimal:
  203. """
  204. get token price.
  205. :param message_type:
  206. :return: decimal.Decimal('0.0001')
  207. """
  208. if message_type == MessageType.HUMAN or message_type == MessageType.SYSTEM:
  209. unit_price = self.price_config['prompt']
  210. else:
  211. unit_price = self.price_config['completion']
  212. unit_price = unit_price.quantize(decimal.Decimal('0.0001'), rounding=decimal.ROUND_HALF_UP)
  213. logging.debug(f"unit_price={unit_price}")
  214. return unit_price
  215. def get_price_unit(self, message_type: MessageType) -> decimal.Decimal:
  216. """
  217. get price unit.
  218. :param message_type:
  219. :return: decimal.Decimal('0.000001')
  220. """
  221. if message_type == MessageType.HUMAN or message_type == MessageType.SYSTEM:
  222. price_unit = self.price_config['unit']
  223. else:
  224. price_unit = self.price_config['unit']
  225. price_unit = price_unit.quantize(decimal.Decimal('0.000001'), rounding=decimal.ROUND_HALF_UP)
  226. logging.debug(f"price_unit={price_unit}")
  227. return price_unit
  228. def get_currency(self) -> str:
  229. """
  230. get token currency.
  231. :return: get from price config, default 'USD'
  232. """
  233. currency = self.price_config['currency']
  234. return currency
  235. def get_model_kwargs(self):
  236. return self.model_kwargs
  237. def set_model_kwargs(self, model_kwargs: ModelKwargs):
  238. self.model_kwargs = model_kwargs
  239. self._set_model_kwargs(model_kwargs)
  240. @abstractmethod
  241. def _set_model_kwargs(self, model_kwargs: ModelKwargs):
  242. raise NotImplementedError
  243. @abstractmethod
  244. def handle_exceptions(self, ex: Exception) -> Exception:
  245. """
  246. Handle llm run exceptions.
  247. :param ex:
  248. :return:
  249. """
  250. raise NotImplementedError
  251. def add_callbacks(self, callbacks: Callbacks):
  252. """
  253. Add callbacks to client.
  254. :param callbacks:
  255. :return:
  256. """
  257. if not self.client.callbacks:
  258. self.client.callbacks = callbacks
  259. else:
  260. self.client.callbacks.extend(callbacks)
  261. @property
  262. def support_streaming(self):
  263. return False
  264. def get_prompt(self, mode: str,
  265. pre_prompt: str, inputs: dict,
  266. query: str,
  267. context: Optional[str],
  268. memory: Optional[BaseChatMemory]) -> \
  269. Tuple[List[PromptMessage], Optional[List[str]]]:
  270. prompt_rules = self._read_prompt_rules_from_file(self.prompt_file_name(mode))
  271. prompt, stops = self._get_prompt_and_stop(prompt_rules, pre_prompt, inputs, query, context, memory)
  272. return [PromptMessage(content=prompt)], stops
  273. def prompt_file_name(self, mode: str) -> str:
  274. if mode == 'completion':
  275. return 'common_completion'
  276. else:
  277. return 'common_chat'
  278. def _get_prompt_and_stop(self, prompt_rules: dict, pre_prompt: str, inputs: dict,
  279. query: str,
  280. context: Optional[str],
  281. memory: Optional[BaseChatMemory]) -> Tuple[str, Optional[list]]:
  282. context_prompt_content = ''
  283. if context and 'context_prompt' in prompt_rules:
  284. prompt_template = JinjaPromptTemplate.from_template(template=prompt_rules['context_prompt'])
  285. context_prompt_content = prompt_template.format(
  286. context=context
  287. )
  288. pre_prompt_content = ''
  289. if pre_prompt:
  290. prompt_template = JinjaPromptTemplate.from_template(template=pre_prompt)
  291. prompt_inputs = {k: inputs[k] for k in prompt_template.input_variables if k in inputs}
  292. pre_prompt_content = prompt_template.format(
  293. **prompt_inputs
  294. )
  295. prompt = ''
  296. for order in prompt_rules['system_prompt_orders']:
  297. if order == 'context_prompt':
  298. prompt += context_prompt_content
  299. elif order == 'pre_prompt':
  300. prompt += pre_prompt_content
  301. query_prompt = prompt_rules['query_prompt'] if 'query_prompt' in prompt_rules else '{{query}}'
  302. if memory and 'histories_prompt' in prompt_rules:
  303. # append chat histories
  304. tmp_human_message = PromptBuilder.to_human_message(
  305. prompt_content=prompt + query_prompt,
  306. inputs={
  307. 'query': query
  308. }
  309. )
  310. if self.model_rules.max_tokens.max:
  311. curr_message_tokens = self.get_num_tokens(to_prompt_messages([tmp_human_message]))
  312. max_tokens = self.model_kwargs.max_tokens
  313. rest_tokens = self.model_rules.max_tokens.max - max_tokens - curr_message_tokens
  314. rest_tokens = max(rest_tokens, 0)
  315. else:
  316. rest_tokens = 2000
  317. memory.human_prefix = prompt_rules['human_prefix'] if 'human_prefix' in prompt_rules else 'Human'
  318. memory.ai_prefix = prompt_rules['assistant_prefix'] if 'assistant_prefix' in prompt_rules else 'Assistant'
  319. histories = self._get_history_messages_from_memory(memory, rest_tokens)
  320. prompt_template = JinjaPromptTemplate.from_template(template=prompt_rules['histories_prompt'])
  321. histories_prompt_content = prompt_template.format(
  322. histories=histories
  323. )
  324. prompt = ''
  325. for order in prompt_rules['system_prompt_orders']:
  326. if order == 'context_prompt':
  327. prompt += context_prompt_content
  328. elif order == 'pre_prompt':
  329. prompt += (pre_prompt_content + '\n') if pre_prompt_content else ''
  330. elif order == 'histories_prompt':
  331. prompt += histories_prompt_content
  332. prompt_template = JinjaPromptTemplate.from_template(template=query_prompt)
  333. query_prompt_content = prompt_template.format(
  334. query=query
  335. )
  336. prompt += query_prompt_content
  337. prompt = re.sub(r'<\|.*?\|>', '', prompt)
  338. stops = prompt_rules.get('stops')
  339. if stops is not None and len(stops) == 0:
  340. stops = None
  341. return prompt, stops
  342. def _read_prompt_rules_from_file(self, prompt_name: str) -> dict:
  343. # Get the absolute path of the subdirectory
  344. prompt_path = os.path.join(
  345. os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))),
  346. 'prompt/generate_prompts')
  347. json_file_path = os.path.join(prompt_path, f'{prompt_name}.json')
  348. # Open the JSON file and read its content
  349. with open(json_file_path, 'r') as json_file:
  350. return json.load(json_file)
  351. def _get_history_messages_from_memory(self, memory: BaseChatMemory,
  352. max_token_limit: int) -> str:
  353. """Get memory messages."""
  354. memory.max_token_limit = max_token_limit
  355. memory_key = memory.memory_variables[0]
  356. external_context = memory.load_memory_variables({})
  357. return external_context[memory_key]
  358. def _get_prompt_from_messages(self, messages: List[PromptMessage],
  359. model_mode: Optional[ModelMode] = None) -> Union[str | List[BaseMessage]]:
  360. if not model_mode:
  361. model_mode = self.model_mode
  362. if model_mode == ModelMode.COMPLETION:
  363. if len(messages) == 0:
  364. return ''
  365. return messages[0].content
  366. else:
  367. if len(messages) == 0:
  368. return []
  369. chat_messages = []
  370. for message in messages:
  371. if message.type == MessageType.HUMAN:
  372. chat_messages.append(HumanMessage(content=message.content))
  373. elif message.type == MessageType.ASSISTANT:
  374. chat_messages.append(AIMessage(content=message.content))
  375. elif message.type == MessageType.SYSTEM:
  376. chat_messages.append(SystemMessage(content=message.content))
  377. return chat_messages
  378. def _to_model_kwargs_input(self, model_rules: ModelKwargsRules, model_kwargs: ModelKwargs) -> dict:
  379. """
  380. convert model kwargs to provider model kwargs.
  381. :param model_rules:
  382. :param model_kwargs:
  383. :return:
  384. """
  385. model_kwargs_input = {}
  386. for key, value in model_kwargs.dict().items():
  387. rule = getattr(model_rules, key)
  388. if not rule.enabled:
  389. continue
  390. if rule.alias:
  391. key = rule.alias
  392. if rule.default is not None and value is None:
  393. value = rule.default
  394. if rule.min is not None:
  395. value = max(value, rule.min)
  396. if rule.max is not None:
  397. value = min(value, rule.max)
  398. model_kwargs_input[key] = value
  399. return model_kwargs_input